XlsDataSetWriter.java

  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. import java.io.IOException;
  23. import java.io.OutputStream;
  24. import java.math.BigDecimal;
  25. import java.util.Date;
  26. import java.util.HashMap;
  27. import java.util.Map;

  28. import org.apache.poi.hssf.usermodel.HSSFWorkbook;
  29. import org.apache.poi.ss.usermodel.Cell;
  30. import org.apache.poi.ss.usermodel.CellStyle;
  31. import org.apache.poi.ss.usermodel.CellType;
  32. import org.apache.poi.ss.usermodel.DataFormat;
  33. import org.apache.poi.ss.usermodel.Row;
  34. import org.apache.poi.ss.usermodel.Sheet;
  35. import org.apache.poi.ss.usermodel.Workbook;
  36. import org.dbunit.dataset.Column;
  37. import org.dbunit.dataset.DataSetException;
  38. import org.dbunit.dataset.IDataSet;
  39. import org.dbunit.dataset.ITable;
  40. import org.dbunit.dataset.ITableIterator;
  41. import org.dbunit.dataset.ITableMetaData;
  42. import org.dbunit.dataset.datatype.DataType;
  43. import org.slf4j.Logger;
  44. import org.slf4j.LoggerFactory;

  45. /**
  46.  * Writes an {@link IDataSet} to an XLS file or OutputStream.
  47.  *
  48.  * @author gommma (gommma AT users.sourceforge.net)
  49.  * @author Last changed by: $Author$
  50.  * @version $Revision$ $Date$
  51.  * @since 2.4.0
  52.  */
  53. public class XlsDataSetWriter
  54. {
  55.     private static final Logger logger = LoggerFactory.getLogger(XlsDataSetWriter.class);

  56.     public static final String ZEROS = "0000000000000000000000000000000000000000000000000000";

  57.     /**
  58.      * A special format pattern used to create a custom {@link DataFormat} which
  59.      * marks {@link Date} values that are stored via POI to an XLS file.
  60.      * Note that it might produce problems if a normal numeric value uses this format
  61.      * pattern incidentally.
  62.      */
  63.     public static final String DATE_FORMAT_AS_NUMBER_DBUNIT = "####################";

  64.     /**
  65.      * Instead of recreating a new style object for each numeric cell, which
  66.      * will cause the code to hit the POI limit of 4000 styles pretty quickly,
  67.      * only create one per format and reuse the same style for all cells with
  68.      * the same format.
  69.      */
  70.     private static final Map<Workbook, Map> cellStyleMap = new HashMap<Workbook, Map>();

  71.     private CellStyle dateCellStyle;

  72.     /**
  73.      * Write the specified dataset to the specified Excel document.
  74.      */
  75.     public void write(IDataSet dataSet, OutputStream out)
  76.             throws IOException, DataSetException
  77.     {
  78.         logger.debug("write(dataSet={}, out={}) - start", dataSet, out);

  79.         Workbook workbook = createWorkbook();

  80.         this.dateCellStyle = createDateCellStyle(workbook);
  81.        
  82.         int index = 0;
  83.         ITableIterator iterator = dataSet.iterator();
  84.         while(iterator.next())
  85.         {
  86.             // create the table i.e. sheet
  87.             ITable table = iterator.getTable();
  88.             ITableMetaData metaData = table.getTableMetaData();
  89.             Sheet sheet = workbook.createSheet(metaData.getTableName());

  90.             // write table metadata i.e. first row in sheet
  91.             workbook.setSheetName(index, metaData.getTableName());

  92.             Row headerRow = sheet.createRow(0);
  93.             Column[] columns = metaData.getColumns();
  94.             for (int j = 0; j < columns.length; j++)
  95.             {
  96.                 Column column = columns[j];
  97.                 Cell cell = headerRow.createCell(j);
  98.                 cell.setCellValue(column.getColumnName());
  99.             }
  100.            
  101.             // write table data
  102.             for (int j = 0; j < table.getRowCount(); j++)
  103.             {
  104.                 Row row = sheet.createRow(j + 1);
  105.                 for (int k = 0; k < columns.length; k++)
  106.                 {
  107.                     Column column = columns[k];
  108.                     Object value = table.getValue(j, column.getColumnName());
  109.                     if (value != null)
  110.                     {
  111.                         Cell cell = row.createCell(k);
  112.                         if(value instanceof Date){
  113.                             setDateCell(cell, (Date)value, workbook);
  114.                         }
  115.                         else if(value instanceof BigDecimal){
  116.                             setNumericCell(cell, (BigDecimal)value, workbook);
  117.                         }
  118.                         else if(value instanceof Long){
  119.                             setDateCell(cell, new Date( ((Long)value).longValue()), workbook);
  120.                         }
  121.                         else {
  122.                             cell.setCellValue(DataType.asString(value));
  123.                         }
  124.                     }
  125.                 }
  126.             }

  127.             index++;
  128.         }

  129.         // write xls document
  130.         workbook.write(out);
  131.         out.flush();
  132.     }
  133.    
  134.     protected static CellStyle createDateCellStyle(Workbook workbook) {
  135.         DataFormat format = workbook.createDataFormat();
  136.         short dateFormatCode = format.getFormat(DATE_FORMAT_AS_NUMBER_DBUNIT);
  137.         return getCellStyle(workbook, dateFormatCode);
  138.     }

  139.     protected static CellStyle getCellStyle(Workbook workbook, short formatCode)
  140.     {
  141.         Map<Short, CellStyle> map = findWorkbookCellStyleMap(workbook);
  142.         CellStyle cellStyle = findCellStyle(workbook, formatCode, map);

  143.         return cellStyle;
  144.     }

  145.     protected static Map<Short, CellStyle> findWorkbookCellStyleMap(
  146.             Workbook workbook)
  147.     {
  148.         Map<Short, CellStyle> map = cellStyleMap.get(workbook);
  149.         if (map == null)
  150.         {
  151.             map = new HashMap<Short, CellStyle>();
  152.             cellStyleMap.put(workbook, map);
  153.         }

  154.         return map;
  155.     }

  156.     protected static CellStyle findCellStyle(Workbook workbook,
  157.             Short formatCode, Map<Short, CellStyle> map)
  158.     {
  159.         CellStyle cellStyle = map.get(formatCode);
  160.         if (cellStyle == null)
  161.         {
  162.             cellStyle = workbook.createCellStyle();
  163.             cellStyle.setDataFormat(formatCode);
  164.             map.put(formatCode, cellStyle);
  165.         }

  166.         return cellStyle;
  167.     }

  168.     protected void setDateCell(Cell cell, Date value, Workbook workbook)
  169.     {
  170. //        double excelDateValue = HSSFDateUtil.getExcelDate(value);
  171. //        cell.setCellValue(excelDateValue);
  172. //        cell.setCellType(HSSFCell.CELL_TYPE_NUMERIC);

  173.         long timeMillis = value.getTime();
  174.         cell.setCellValue( (double)timeMillis );
  175.         cell.setCellType(CellType.NUMERIC);
  176.         cell.setCellStyle(this.dateCellStyle);
  177.        
  178. //      System.out.println(HSSFDataFormat.getBuiltinFormats());
  179.         // TODO Find out correct cell styles for date objects
  180. //        HSSFCellStyle cellStyleDate = workbook.createCellStyle();
  181. //        cellStyleDate.setDataFormat(HSSFDataFormat.getBuiltinFormat("m/d/yy h:mm"));
  182. //
  183. //        HSSFCellStyle cellStyleDateTimeWithSeconds = workbook.createCellStyle();
  184. //        cellStyleDateTimeWithSeconds.setDataFormat(HSSFDataFormat.getBuiltinFormat("h:mm:ss"));
  185. //
  186. //        HSSFDataFormat dataFormat = workbook.createDataFormat();
  187. //        HSSFCellStyle cellStyle = workbook.createCellStyle();
  188. //        cellStyle.setDataFormat(dataFormat.getFormat("dd/mm/yyyy hh:mm:ss"));
  189. //
  190. //        SimpleDateFormat formatter = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss.SSS");
  191. //        SimpleDateFormat formatter2 = new SimpleDateFormat("dd/MM/yyyy");
  192. //        SimpleDateFormat formatter3 = new SimpleDateFormat("HH:mm:ss.SSS");
  193. //
  194. //        
  195. //        Date dateValue = (Date)value;
  196. //        Calendar cal = null;
  197. //        
  198. //        // If it is a date value that does not have seconds
  199. //        if(dateValue.getTime() % 60000 == 0){
  200. ////            cellStyle = cellStyleDate;
  201. //            cal=Calendar.getInstance();
  202. //            cal.setTimeInMillis(dateValue.getTime());
  203. //
  204. //            cell.setCellValue(cal);
  205. //            cell.setCellStyle(cellStyle);
  206. ////            cell.setCellValue(cal);
  207. //        }
  208. //        else {
  209. ////            HSSFDataFormatter formatter = new HSSFDataFormatter();
  210. //            
  211. //            // If we have seconds assume that it is only h:mm:ss without date
  212. //            // TODO Clean implementation where user can control date formats would be nice
  213. ////            double dateDouble = dateValue.getTime() % (24*60*60*1000);
  214. //            cal = get1900Cal(dateValue);
  215. //            
  216. //            String formatted = formatter3.format(dateValue);
  217. //            //TODO Format ...
  218. ////            cellStyle = cellStyleDateTimeWithSeconds;
  219. //            System.out.println("date formatted:"+formatted);
  220. ////            HSSFRichTextString s = new HSSFRichTextString(formatted);
  221. ////            cell.setCellValue(s);
  222. //            cell.setCellType(HSSFCell.CELL_TYPE_NUMERIC);
  223. //            cell.setCellValue((double)dateValue.getTime());
  224. //            cell.setCellStyle(cellStyleDateTimeWithSeconds);
  225. //        }

  226.     }

  227.     protected void setNumericCell(Cell cell, BigDecimal value, Workbook workbook)
  228.     {
  229.         if(logger.isDebugEnabled())
  230.             logger.debug("setNumericCell(cell={}, value={}, workbook={}) - start",
  231.                 new Object[] {cell, value, workbook} );

  232.         cell.setCellValue( ((BigDecimal)value).doubleValue() );

  233.         DataFormat df = workbook.createDataFormat();
  234.         int scale = ((BigDecimal)value).scale();
  235.         short format;
  236.         if(scale <= 0){
  237.             format = df.getFormat("####");
  238.         }
  239.         else {
  240.             String zeros = createZeros(((BigDecimal)value).scale());
  241.             format = df.getFormat("####." + zeros);
  242.         }
  243.         if(logger.isDebugEnabled())
  244.             logger.debug("Using format '{}' for value '{}'.", String.valueOf(format), value);

  245.         CellStyle cellStyleNumber = getCellStyle(workbook, format);
  246.         cell.setCellStyle(cellStyleNumber);
  247.     }

  248. //    public static Date get1900(Date date) {
  249. //        Calendar cal = Calendar.getInstance();
  250. //        cal.setTimeInMillis(date.getTime() % (24*60*60*1000));
  251. //        cal.set(1900, 0, 1); // 1.1.1900
  252. //        return cal.getTime();
  253. //    }
  254. //
  255. //    public static Calendar get1900Cal(Date date) {
  256. //        Calendar cal = Calendar.getInstance();
  257. //        cal.clear();
  258. ////        long hoursInMillis = date.getTime() % (24*60*60*1000);
  259. ////        long smallerThanDays = date.getTime() % (24*60*60*1000);
  260. ////        cal.setTimeInMillis(date.getTime() % (24*60*60*1000));
  261. //        cal.set(Calendar.SECOND, (int) (date.getTime() % (24*60*60*1000)) / (1000) );
  262. //        cal.set(Calendar.MINUTE, (int) (date.getTime() % (24*60*60*1000)) / (1000*60) );
  263. //        cal.set(Calendar.HOUR, (int) (date.getTime() % (24*60*60*1000)) / (1000*60*60) );
  264. ////        cal.set(1900, 0, 1); // 1.1.1900
  265. //        System.out.println(cal.isSet(Calendar.DATE));
  266. //        return cal;
  267. //    }

  268.     private static String createZeros(int count) {
  269.         return ZEROS.substring(0, count);
  270.     }
  271.    
  272.     protected Workbook createWorkbook() {
  273.         return new HSSFWorkbook();
  274.     }
  275. }