1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
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
47
48
49
50 public class CsvDataSetWriter implements IDataSetConsumer {
51
52
53
54
55 private static final Logger logger = LoggerFactory.getLogger(CsvDataSetWriter.class);
56
57
58
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
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
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
174 if (value == null) {
175 getWriter().write(NULL);
176 }
177
178 else if (value == ITable.NO_VALUE) {
179 getWriter().write(NONE);
180 }
181
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 }