1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22 package org.dbunit.dataset.csv;
23
24 import java.io.BufferedReader;
25 import java.io.File;
26 import java.io.IOException;
27 import java.io.InputStream;
28 import java.io.InputStreamReader;
29 import java.net.URL;
30 import java.util.ArrayList;
31 import java.util.Iterator;
32 import java.util.List;
33
34 import org.dbunit.dataset.Column;
35 import org.dbunit.dataset.DataSetException;
36 import org.dbunit.dataset.DefaultTableMetaData;
37 import org.dbunit.dataset.ITableMetaData;
38 import org.dbunit.dataset.common.handlers.IllegalInputCharacterException;
39 import org.dbunit.dataset.common.handlers.PipelineException;
40 import org.dbunit.dataset.datatype.DataType;
41 import org.dbunit.dataset.stream.DefaultConsumer;
42 import org.dbunit.dataset.stream.IDataSetConsumer;
43 import org.dbunit.dataset.stream.IDataSetProducer;
44 import org.slf4j.Logger;
45 import org.slf4j.LoggerFactory;
46
47
48
49
50
51
52
53 public class CsvProducer implements IDataSetProducer {
54
55
56
57
58 private static final Logger logger = LoggerFactory.getLogger(CsvProducer.class);
59
60 private static final IDataSetConsumer EMPTY_CONSUMER = new DefaultConsumer();
61 private IDataSetConsumer _consumer = EMPTY_CONSUMER;
62 private String _theDirectory;
63
64 public CsvProducer(String theDirectory) {
65 _theDirectory = theDirectory;
66 }
67
68 public CsvProducer(File theDirectory) {
69 _theDirectory = theDirectory.getAbsolutePath();
70 }
71
72 public void setConsumer(IDataSetConsumer consumer) throws DataSetException {
73 logger.debug("setConsumer(consumer) - start");
74
75 _consumer = consumer;
76 }
77
78 public void produce() throws DataSetException {
79 logger.debug("produce() - start");
80
81 File dir = new File(_theDirectory);
82
83 if (!dir.isDirectory()) {
84 throw new DataSetException("'" + _theDirectory + "' should be a directory");
85 }
86
87 _consumer.startDataSet();
88 try {
89 List tableSpecs = CsvProducer.getTables(dir.toURL(), CsvDataSet.TABLE_ORDERING_FILE);
90 for (Iterator tableIter = tableSpecs.iterator(); tableIter.hasNext();) {
91 String table = (String) tableIter.next();
92 try {
93 produceFromFile(new File(dir, table + ".csv"));
94 } catch (CsvParserException e) {
95 throw new DataSetException("error producing dataset for table '" + table + "'", e);
96 } catch (DataSetException e) {
97 throw new DataSetException("error producing dataset for table '" + table + "'", e);
98 }
99
100 }
101 _consumer.endDataSet();
102 } catch (IOException e) {
103 throw new DataSetException("error getting list of tables", e);
104 }
105 }
106
107 private void produceFromFile(File theDataFile) throws DataSetException, CsvParserException {
108 logger.debug("produceFromFile(theDataFile={}) - start", theDataFile);
109
110 try {
111 CsvParser parser = new CsvParserImpl();
112 List readData = parser.parse(theDataFile);
113 List readColumns = ((List) readData.get(0));
114 Column[] columns = new Column[readColumns.size()];
115
116 for (int i = 0; i < readColumns.size(); i++) {
117 String columnName = (String) readColumns.get(i);
118 columnName = columnName.trim();
119 columns[i] = new Column(columnName, DataType.UNKNOWN);
120 }
121
122 String tableName = theDataFile.getName().substring(0, theDataFile.getName().indexOf(".csv"));
123 ITableMetaData metaData = new DefaultTableMetaData(tableName, columns);
124 _consumer.startTable(metaData);
125 for (int i = 1 ; i < readData.size(); i++) {
126 List rowList = (List)readData.get(i);
127 Object[] row = rowList.toArray();
128 for(int col = 0; col < row.length; col++) {
129 row[col] = row[col].equals(CsvDataSetWriter.NULL) ? null : row[col];
130 }
131 _consumer.row(row);
132 }
133 _consumer.endTable();
134 } catch (PipelineException e) {
135 throw new DataSetException(e);
136 } catch (IllegalInputCharacterException e) {
137 throw new DataSetException(e);
138 } catch (IOException e) {
139 throw new DataSetException(e);
140 }
141 }
142
143
144
145
146
147
148 public static List getTables(URL base, String tableList) throws IOException {
149 logger.debug("getTables(base={}, tableList={}) - start", base, tableList);
150
151 List orderedNames = new ArrayList();
152 InputStream tableListStream = new URL(base, tableList).openStream();
153 BufferedReader reader = null;
154 try {
155 reader = new BufferedReader(new InputStreamReader(tableListStream));
156 String line = null;
157 while((line = reader.readLine()) != null) {
158 String table = line.trim();
159 if (table.length() > 0) {
160 orderedNames.add(table);
161 }
162 }
163 }
164 finally {
165 if(reader != null)
166 {
167 reader.close();
168 }
169 }
170 return orderedNames;
171 }
172
173 }