View Javadoc
1   /*
2    *
3    * The DbUnit Database Testing Framework
4    * Copyright (C)2002-2004, 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.csv;
22  
23  import org.slf4j.Logger;
24  import org.slf4j.LoggerFactory;
25  
26  import java.io.File;
27  import java.io.FileWriter;
28  import java.io.IOException;
29  import java.io.PrintWriter;
30  import java.io.Writer;
31  import java.util.Iterator;
32  import java.util.LinkedList;
33  import java.util.List;
34  
35  import org.dbunit.dataset.Column;
36  import org.dbunit.dataset.DataSetException;
37  import org.dbunit.dataset.IDataSet;
38  import org.dbunit.dataset.ITable;
39  import org.dbunit.dataset.ITableMetaData;
40  import org.dbunit.dataset.datatype.DataType;
41  import org.dbunit.dataset.datatype.TypeCastException;
42  import org.dbunit.dataset.stream.DataSetProducerAdapter;
43  import org.dbunit.dataset.stream.IDataSetConsumer;
44  
45  /**
46   * @author fede
47   * @since 24-set-2003 15.27.05
48   * @version $Revision$
49   */
50  public class CsvDataSetWriter implements IDataSetConsumer {
51  
52      /**
53       * Logger for this class
54       */
55      private static final Logger logger = LoggerFactory.getLogger(CsvDataSetWriter.class);
56  
57      /**
58       * todo: customizable separators (field, lines), manage the writers opened for each table
59       */
60  
61      public static final String NULL = "null";
62      private static final String NONE = "none";
63      private static final String FIELD_SEPARATOR = ", ";
64      private static final String QUOTE = "\"";
65      private static final String ESCAPE = "\\";
66  
67      private Writer writer;
68      private ITableMetaData _activeMetaData;
69      private String theDirectory;
70      private static char testExport;
71      /** list of tables */
72      private List tableList;
73  
74      public CsvDataSetWriter(String theDirectory) {
75          setTheDirectory(theDirectory);
76      }
77  
78      public CsvDataSetWriter(File theDirectory) {
79          setTheDirectory(theDirectory.getAbsolutePath());
80      }
81  
82      public void write(IDataSet dataSet) throws DataSetException {
83          logger.debug("write(dataSet={}) - start", dataSet);
84  
85          DataSetProducerAdapter provider = new DataSetProducerAdapter(dataSet);
86          provider.setConsumer(this);
87          provider.produce();
88      }
89  
90      public void startDataSet() throws DataSetException {
91          logger.debug("startDataSet() - start");
92  
93          try {
94          	tableList = new LinkedList();
95              new File(getTheDirectory()).mkdirs();
96          } catch (Exception e) {
97              throw new DataSetException("Error while creating the destination directory '" + getTheDirectory() + "'", e);
98          }
99      }
100 
101     public void endDataSet() throws DataSetException {
102         logger.debug("endDataSet() - start");
103 
104     	// write out table ordering file
105     	File orderingFile = new File(getTheDirectory(), CsvDataSet.TABLE_ORDERING_FILE);
106     	
107     	PrintWriter pw = null;
108     	try {
109 			pw = new PrintWriter(new FileWriter(orderingFile));
110 			for (Iterator fileNames = tableList.iterator(); fileNames.hasNext();) {
111 				String file = (String) fileNames.next();
112 				pw.println(file);
113 			}
114 		} 
115     	catch (IOException e) {
116 			throw new DataSetException("problems writing the table ordering file", e);
117 		}
118     	finally {
119     	    if(pw != null){
120     	        pw.close();
121     	    }
122     	}
123     }
124 
125     public void startTable(ITableMetaData metaData) throws DataSetException {
126         logger.debug("startTable(metaData={}) - start", metaData);
127 
128         try {
129             _activeMetaData = metaData;
130             String tableName = _activeMetaData.getTableName();
131             setWriter(new FileWriter(getTheDirectory() + File.separator + tableName + ".csv"));
132             writeColumnNames();
133             getWriter().write(System.getProperty("line.separator"));
134         } catch (IOException e) {
135             throw new DataSetException(e);
136         }
137 
138     }
139 
140     private void writeColumnNames() throws DataSetException, IOException {
141         logger.debug("writeColumnNames() - start");
142 
143         Column[] columns = _activeMetaData.getColumns();
144         for (int i = 0; i < columns.length; i++) {
145             String columnName = columns[i].getColumnName();
146             getWriter().write(columnName);
147             if (i < columns.length - 1) getWriter().write(FIELD_SEPARATOR);
148         }
149     }
150 
151     public void endTable() throws DataSetException {
152         logger.debug("endTable() - start");
153 
154         try {
155             getWriter().close();
156             tableList.add(_activeMetaData.getTableName());
157             _activeMetaData = null;
158         } catch (IOException e) {
159             throw new DataSetException(e);
160         }
161     }
162 
163     public void row(Object[] values) throws DataSetException {
164         logger.debug("row(values={}) - start", values);
165 
166         try {
167 
168             Column[] columns = _activeMetaData.getColumns();
169             for (int i = 0; i < columns.length; i++) {
170                 String columnName = columns[i].getColumnName();
171                 Object value = values[i];
172 
173                 // null
174                 if (value == null) {
175                     getWriter().write(NULL);
176                 }
177                 // none
178                 else if (value == ITable.NO_VALUE) {
179                     getWriter().write(NONE);
180                 }
181                 // values
182                 else {
183                     try {
184                         String stringValue = DataType.asString(value);
185                         final String quoted = quote(stringValue);
186                         getWriter().write(quoted);
187                     } catch (TypeCastException e) {
188                         throw new DataSetException("table=" +
189                                 _activeMetaData.getTableName() + ", row=" + i +
190                                 ", column=" + columnName +
191                                 ", value=" + value, e);
192                     }
193                 }
194                 if (i < columns.length - 1) getWriter().write(",");
195             }
196             getWriter().write(System.getProperty("line.separator"));
197         } catch (IOException e) {
198             throw new DataSetException(e);
199         }
200     }
201 
202     private String quote(String stringValue) {
203         logger.debug("quote(stringValue={}) - start", stringValue);
204 
205         return new StringBuffer(QUOTE).append(escape(stringValue)).append(QUOTE).toString();
206     }
207 
208     protected static String escape(String stringValue) {
209         logger.debug("escape(stringValue={}) - start", stringValue);
210 
211         char [] array = stringValue.toCharArray();
212         testExport = QUOTE.toCharArray()[0];
213         final char escape = ESCAPE.toCharArray()[0];
214         StringBuffer buffer = new StringBuffer();
215         for (int i = 0; i < array.length; i++) {
216             char c = array[i];
217             if (c == testExport || c == escape) {
218                 buffer.append('\\');
219             }
220             buffer.append(c);
221         }
222         return buffer.toString();
223     }
224 
225     public Writer getWriter() {
226         logger.debug("getWriter() - start");
227 
228         return writer;
229     }
230 
231     public void setWriter(Writer writer) {
232         logger.debug("setWriter(writer={}) - start", writer);
233 
234         this.writer = writer;
235     }
236 
237     public String getTheDirectory() {
238         logger.debug("getTheDirectory() - start");
239 
240         return theDirectory;
241     }
242 
243     public void setTheDirectory(String theDirectory) {
244         logger.debug("setTheDirectory(theDirectory={}) - start", theDirectory);
245 
246         this.theDirectory = theDirectory;
247     }
248 
249     public static void write(IDataSet dataset, File dest) throws DataSetException {
250         logger.debug("write(dataset={}, dest={}) - start", dataset, dest);
251 
252         CsvDataSetWriter writer = new CsvDataSetWriter(dest);
253         writer.write(dataset);
254     }
255 
256     protected void finalize() throws Throwable {
257         logger.debug("finalize() - start");
258 
259         if (getWriter() != null) {
260             getWriter().close();
261         }
262     }
263 }