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  
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   * @author Federico Spinazzi
49   * @author Last changed by: $Author$
50   * @version $Revision$ $Date$
51   * @since 1.5 (Sep 17, 2003)
52   */
53  public class CsvProducer implements IDataSetProducer {
54  
55      /**
56       * Logger for this class
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 	 * Get a list of tables that this producer will create
145 	 * @return a list of Strings, where each item is a CSV file relative to the base URL
146 	 * @throws IOException when IO on the base URL has issues.
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 }