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