View Javadoc
1   /*
2    *
3    * The DbUnit Database Testing Framework
4    * Copyright (C)2002-2008, DbUnit.org
5    *
6    * This library is free software; you can redistribute it and/or
7    * modify it under the terms of the GNU Lesser General Public
8    * License as published by the Free Software Foundation; either
9    * version 2.1 of the License, or (at your option) any later version.
10   *
11   * This library is distributed in the hope that it will be useful,
12   * but WITHOUT ANY WARRANTY; without even the implied warranty of
13   * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
14   * Lesser General Public License for more details.
15   *
16   * You should have received a copy of the GNU Lesser General Public
17   * License along with this library; if not, write to the Free Software
18   * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
19   *
20   */
21  package org.dbunit.dataset.excel;
22  
23  import java.io.IOException;
24  import java.io.OutputStream;
25  import java.math.BigDecimal;
26  import java.util.Date;
27  import java.util.HashMap;
28  import java.util.Map;
29  
30  import org.apache.poi.hssf.usermodel.HSSFWorkbook;
31  import org.apache.poi.ss.usermodel.Cell;
32  import org.apache.poi.ss.usermodel.CellStyle;
33  import org.apache.poi.ss.usermodel.CellType;
34  import org.apache.poi.ss.usermodel.DataFormat;
35  import org.apache.poi.ss.usermodel.Row;
36  import org.apache.poi.ss.usermodel.Sheet;
37  import org.apache.poi.ss.usermodel.Workbook;
38  import org.dbunit.dataset.Column;
39  import org.dbunit.dataset.DataSetException;
40  import org.dbunit.dataset.IDataSet;
41  import org.dbunit.dataset.ITable;
42  import org.dbunit.dataset.ITableIterator;
43  import org.dbunit.dataset.ITableMetaData;
44  import org.dbunit.dataset.datatype.DataType;
45  import org.slf4j.Logger;
46  import org.slf4j.LoggerFactory;
47  
48  /**
49   * Writes an {@link IDataSet} to an XLS file or OutputStream.
50   * 
51   * @author gommma (gommma AT users.sourceforge.net)
52   * @author Last changed by: $Author$
53   * @version $Revision$ $Date$
54   * @since 2.4.0
55   */
56  public class XlsDataSetWriter 
57  {
58      private static final Logger logger = LoggerFactory.getLogger(XlsDataSetWriter.class);
59  
60      public static final String ZEROS = "0000000000000000000000000000000000000000000000000000";
61  
62      /**
63       * A special format pattern used to create a custom {@link DataFormat} which
64       * marks {@link Date} values that are stored via POI to an XLS file.
65       * Note that it might produce problems if a normal numeric value uses this format
66       * pattern incidentally.
67       */
68      public static final String DATE_FORMAT_AS_NUMBER_DBUNIT = "####################";
69  
70      /**
71       * Instead of recreating a new style object for each numeric cell, which
72       * will cause the code to hit the POI limit of 4000 styles pretty quickly,
73       * only create one per format and reuse the same style for all cells with
74       * the same format.
75       */
76      private static final Map<Workbook, Map> cellStyleMap = new HashMap<Workbook, Map>();
77  
78      private CellStyle dateCellStyle;
79  
80      /**
81       * Write the specified dataset to the specified Excel document.
82       */
83      public void write(IDataSet dataSet, OutputStream out)
84              throws IOException, DataSetException
85      {
86          logger.debug("write(dataSet={}, out={}) - start", dataSet, out);
87  
88          Workbook workbook = createWorkbook();
89  
90          this.dateCellStyle = createDateCellStyle(workbook);
91          
92          int index = 0;
93          ITableIterator iterator = dataSet.iterator();
94          while(iterator.next())
95          {
96              // create the table i.e. sheet
97              ITable table = iterator.getTable();
98              ITableMetaData metaData = table.getTableMetaData();
99              Sheet sheet = workbook.createSheet(metaData.getTableName());
100 
101             // write table metadata i.e. first row in sheet
102             workbook.setSheetName(index, metaData.getTableName());
103 
104             Row headerRow = sheet.createRow(0);
105             Column[] columns = metaData.getColumns();
106             for (int j = 0; j < columns.length; j++)
107             {
108                 Column column = columns[j];
109                 Cell cell = headerRow.createCell(j);
110                 cell.setCellValue(column.getColumnName());
111             }
112             
113             // write table data
114             for (int j = 0; j < table.getRowCount(); j++)
115             {
116                 Row row = sheet.createRow(j + 1);
117                 for (int k = 0; k < columns.length; k++)
118                 {
119                     Column column = columns[k];
120                     Object value = table.getValue(j, column.getColumnName());
121                     if (value != null)
122                     {
123                         Cell cell = row.createCell(k);
124                         if(value instanceof Date){
125                             setDateCell(cell, (Date)value, workbook);
126                         }
127                         else if(value instanceof BigDecimal){
128                             setNumericCell(cell, (BigDecimal)value, workbook);
129                         }
130                         else if(value instanceof Long){
131                             setDateCell(cell, new Date( ((Long)value).longValue()), workbook);
132                         }
133                         else {
134                             cell.setCellValue(DataType.asString(value));
135                         }
136                     }
137                 }
138             }
139 
140             index++;
141         }
142 
143         // write xls document
144         workbook.write(out);
145         out.flush();
146     }
147     
148     protected static CellStyle createDateCellStyle(Workbook workbook) {
149         DataFormat format = workbook.createDataFormat();
150         short dateFormatCode = format.getFormat(DATE_FORMAT_AS_NUMBER_DBUNIT);
151         return getCellStyle(workbook, dateFormatCode);
152     }
153 
154     protected static CellStyle getCellStyle(Workbook workbook, short formatCode)
155     {
156         Map<Short, CellStyle> map = findWorkbookCellStyleMap(workbook);
157         CellStyle cellStyle = findCellStyle(workbook, formatCode, map);
158 
159         return cellStyle;
160     }
161 
162     protected static Map<Short, CellStyle> findWorkbookCellStyleMap(
163             Workbook workbook)
164     {
165         Map<Short, CellStyle> map = cellStyleMap.get(workbook);
166         if (map == null)
167         {
168             map = new HashMap<Short, CellStyle>();
169             cellStyleMap.put(workbook, map);
170         }
171 
172         return map;
173     }
174 
175     protected static CellStyle findCellStyle(Workbook workbook,
176             Short formatCode, Map<Short, CellStyle> map)
177     {
178         CellStyle cellStyle = map.get(formatCode);
179         if (cellStyle == null)
180         {
181             cellStyle = workbook.createCellStyle();
182             cellStyle.setDataFormat(formatCode);
183             map.put(formatCode, cellStyle);
184         }
185 
186         return cellStyle;
187     }
188 
189     protected void setDateCell(Cell cell, Date value, Workbook workbook) 
190     {
191 //        double excelDateValue = HSSFDateUtil.getExcelDate(value);
192 //        cell.setCellValue(excelDateValue);
193 //        cell.setCellType(HSSFCell.CELL_TYPE_NUMERIC);
194 
195         long timeMillis = value.getTime();
196         cell.setCellValue( (double)timeMillis );
197         cell.setCellType(CellType.NUMERIC);
198         cell.setCellStyle(this.dateCellStyle);
199         
200 //      System.out.println(HSSFDataFormat.getBuiltinFormats());
201         // TODO Find out correct cell styles for date objects
202 //        HSSFCellStyle cellStyleDate = workbook.createCellStyle();
203 //        cellStyleDate.setDataFormat(HSSFDataFormat.getBuiltinFormat("m/d/yy h:mm"));
204 //
205 //        HSSFCellStyle cellStyleDateTimeWithSeconds = workbook.createCellStyle();
206 //        cellStyleDateTimeWithSeconds.setDataFormat(HSSFDataFormat.getBuiltinFormat("h:mm:ss"));
207 //
208 //        HSSFDataFormat dataFormat = workbook.createDataFormat();
209 //        HSSFCellStyle cellStyle = workbook.createCellStyle();
210 //        cellStyle.setDataFormat(dataFormat.getFormat("dd/mm/yyyy hh:mm:ss"));
211 //
212 //        SimpleDateFormat formatter = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss.SSS");
213 //        SimpleDateFormat formatter2 = new SimpleDateFormat("dd/MM/yyyy");
214 //        SimpleDateFormat formatter3 = new SimpleDateFormat("HH:mm:ss.SSS");
215 //
216 //        
217 //        Date dateValue = (Date)value;
218 //        Calendar cal = null;
219 //        
220 //        // If it is a date value that does not have seconds
221 //        if(dateValue.getTime() % 60000 == 0){
222 ////            cellStyle = cellStyleDate;
223 //            cal=Calendar.getInstance();
224 //            cal.setTimeInMillis(dateValue.getTime());
225 //
226 //            cell.setCellValue(cal);
227 //            cell.setCellStyle(cellStyle);
228 ////            cell.setCellValue(cal);
229 //        }
230 //        else {
231 ////            HSSFDataFormatter formatter = new HSSFDataFormatter();
232 //            
233 //            // If we have seconds assume that it is only h:mm:ss without date
234 //            // TODO Clean implementation where user can control date formats would be nice
235 ////            double dateDouble = dateValue.getTime() % (24*60*60*1000);
236 //            cal = get1900Cal(dateValue);
237 //            
238 //            String formatted = formatter3.format(dateValue);
239 //            //TODO Format ...
240 ////            cellStyle = cellStyleDateTimeWithSeconds;
241 //            System.out.println("date formatted:"+formatted);
242 ////            HSSFRichTextString s = new HSSFRichTextString(formatted);
243 ////            cell.setCellValue(s);
244 //            cell.setCellType(HSSFCell.CELL_TYPE_NUMERIC);
245 //            cell.setCellValue((double)dateValue.getTime());
246 //            cell.setCellStyle(cellStyleDateTimeWithSeconds);
247 //        }
248 
249     }
250 
251     protected void setNumericCell(Cell cell, BigDecimal value, Workbook workbook)
252     {
253         if(logger.isDebugEnabled())
254             logger.debug("setNumericCell(cell={}, value={}, workbook={}) - start", 
255                 new Object[] {cell, value, workbook} );
256 
257         cell.setCellValue( ((BigDecimal)value).doubleValue() );
258 
259         DataFormat df = workbook.createDataFormat();
260         int scale = ((BigDecimal)value).scale();
261         short format;
262         if(scale <= 0){
263             format = df.getFormat("####");
264         }
265         else {
266             String zeros = createZeros(((BigDecimal)value).scale());
267             format = df.getFormat("####." + zeros);
268         }
269         if(logger.isDebugEnabled())
270             logger.debug("Using format '{}' for value '{}'.", String.valueOf(format), value);
271 
272         CellStyle cellStyleNumber = getCellStyle(workbook, format);
273         cell.setCellStyle(cellStyleNumber);
274     }
275 
276 //    public static Date get1900(Date date) {
277 //        Calendar cal = Calendar.getInstance();
278 //        cal.setTimeInMillis(date.getTime() % (24*60*60*1000));
279 //        cal.set(1900, 0, 1); // 1.1.1900
280 //        return cal.getTime();
281 //    }
282 //
283 //    public static Calendar get1900Cal(Date date) {
284 //        Calendar cal = Calendar.getInstance();
285 //        cal.clear();
286 ////        long hoursInMillis = date.getTime() % (24*60*60*1000);
287 ////        long smallerThanDays = date.getTime() % (24*60*60*1000);
288 ////        cal.setTimeInMillis(date.getTime() % (24*60*60*1000));
289 //        cal.set(Calendar.SECOND, (int) (date.getTime() % (24*60*60*1000)) / (1000) );
290 //        cal.set(Calendar.MINUTE, (int) (date.getTime() % (24*60*60*1000)) / (1000*60) );
291 //        cal.set(Calendar.HOUR, (int) (date.getTime() % (24*60*60*1000)) / (1000*60*60) );
292 ////        cal.set(1900, 0, 1); // 1.1.1900
293 //        System.out.println(cal.isSet(Calendar.DATE));
294 //        return cal;
295 //    }
296 
297     private static String createZeros(int count) {
298         return ZEROS.substring(0, count);
299     }
300     
301     protected Workbook createWorkbook() {
302         return new HSSFWorkbook();
303     }
304 }