CsvDataSetWriter.java
/*
*
* The DbUnit Database Testing Framework
* Copyright (C)2002-2004, DbUnit.org
*
* This library is free software; you can redistribute it and/or
* modify it under the terms of the GNU Lesser General Public
* License as published by the Free Software Foundation; either
* version 2.1 of the License, or (at your option) any later version.
*
* This library is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
* Lesser General Public License for more details.
*
* You should have received a copy of the GNU Lesser General Public
* License along with this library; if not, write to the Free Software
* Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
*
*/
package org.dbunit.dataset.csv;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.File;
import java.io.IOException;
import java.io.PrintWriter;
import java.io.Writer;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import org.dbunit.dataset.Column;
import org.dbunit.dataset.DataSetException;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.ITable;
import org.dbunit.dataset.ITableMetaData;
import org.dbunit.dataset.datatype.DataType;
import org.dbunit.dataset.datatype.TypeCastException;
import org.dbunit.dataset.stream.DataSetProducerAdapter;
import org.dbunit.dataset.stream.IDataSetConsumer;
/**
* {@link IDataSetConsumer} that writes a dataset's tables and rows to CSV files.
*
* @author fede
* @since 24-set-2003 15.27.05
* @version $Revision$
*/
public class CsvDataSetWriter implements IDataSetConsumer {
/**
* Logger for this class
*/
private static final Logger logger = LoggerFactory.getLogger(CsvDataSetWriter.class);
/**
* todo: customizable separators (field, lines), manage the writers opened for each table
*/
public static final String NULL = "null";
private static final String NONE = "none";
private static final String FIELD_SEPARATOR = ", ";
private static final String QUOTE = "\"";
private static final String ESCAPE = "\\";
private Writer writer;
private ITableMetaData _activeMetaData;
private String theDirectory;
private static char testExport;
/** list of tables */
private List tableList;
/**
* Creates a writer that writes CSV files to the given directory.
*
* @param theDirectory the path of the directory to write CSV files to.
*/
public CsvDataSetWriter(String theDirectory) {
setTheDirectory(theDirectory);
}
/**
* Creates a writer that writes CSV files to the given directory.
*
* @param theDirectory the directory to write CSV files to.
*/
public CsvDataSetWriter(File theDirectory) {
setTheDirectory(theDirectory.getAbsolutePath());
}
/**
* Writes the given dataset's tables and rows to CSV files.
*
* @param dataSet the dataset to write.
* @throws DataSetException if writing the dataset fails.
*/
public void write(IDataSet dataSet) throws DataSetException {
logger.debug("write(dataSet={}) - start", dataSet);
DataSetProducerAdapter provider = new DataSetProducerAdapter(dataSet);
provider.setConsumer(this);
provider.produce();
}
/**
* {@inheritDoc}
*/
@Override
public void startDataSet() throws DataSetException {
logger.debug("startDataSet() - start");
try {
tableList = new LinkedList();
Paths.get(getTheDirectory()).toFile().mkdirs();
} catch (Exception e) {
throw new DataSetException("Error while creating the destination directory '" + getTheDirectory() + "'", e);
}
}
/**
* {@inheritDoc}
*/
@Override
public void endDataSet() throws DataSetException {
logger.debug("endDataSet() - start");
// write out table ordering file
File orderingFile = Paths.get(getTheDirectory(), CsvDataSet.TABLE_ORDERING_FILE).toFile();
PrintWriter pw = null;
try {
pw = new PrintWriter(Files.newBufferedWriter(orderingFile.toPath(), StandardCharsets.UTF_8));
for (Iterator fileNames = tableList.iterator(); fileNames.hasNext();) {
String file = (String) fileNames.next();
pw.println(file);
}
}
catch (IOException e) {
throw new DataSetException("problems writing the table ordering file", e);
}
finally {
if(pw != null){
pw.close();
}
}
}
/**
* {@inheritDoc}
*/
@Override
public void startTable(ITableMetaData metaData) throws DataSetException {
logger.debug("startTable(metaData={}) - start", metaData);
try {
_activeMetaData = metaData;
String tableName = _activeMetaData.getTableName();
setWriter(Files.newBufferedWriter(Paths.get(getTheDirectory(), tableName + ".csv"), StandardCharsets.UTF_8));
writeColumnNames();
getWriter().write(System.getProperty("line.separator"));
} catch (IOException e) {
throw new DataSetException(e);
}
}
private void writeColumnNames() throws DataSetException, IOException {
logger.debug("writeColumnNames() - start");
Column[] columns = _activeMetaData.getColumns();
for (int i = 0; i < columns.length; i++) {
String columnName = columns[i].getColumnName();
getWriter().write(columnName);
if (i < columns.length - 1) getWriter().write(FIELD_SEPARATOR);
}
}
public void endTable() throws DataSetException {
logger.debug("endTable() - start");
try {
getWriter().close();
tableList.add(_activeMetaData.getTableName());
_activeMetaData = null;
} catch (IOException e) {
throw new DataSetException(e);
}
}
public void row(Object[] values) throws DataSetException {
logger.debug("row(values={}) - start", values);
try {
Column[] columns = _activeMetaData.getColumns();
for (int i = 0; i < columns.length; i++) {
String columnName = columns[i].getColumnName();
Object value = values[i];
// null
if (value == null) {
getWriter().write(NULL);
}
// none
else if (value == ITable.NO_VALUE) {
getWriter().write(NONE);
}
// values
else {
try {
String stringValue = DataType.asString(value);
final String quoted = quote(stringValue);
getWriter().write(quoted);
} catch (TypeCastException e) {
throw new DataSetException("table=" +
_activeMetaData.getTableName() + ", row=" + i +
", column=" + columnName +
", value=" + value, e);
}
}
if (i < columns.length - 1) getWriter().write(",");
}
getWriter().write(System.getProperty("line.separator"));
} catch (IOException e) {
throw new DataSetException(e);
} catch (DataSetException e) {
// Rows already written for this table sit in the BufferedWriter until
// endTable() closes it; since a row failure here skips endTable()
// entirely, flush explicitly so they still reach disk.
flushWriterQuietly();
throw e;
}
}
private void flushWriterQuietly() {
logger.debug("flushWriterQuietly() - start");
try {
if (getWriter() != null) {
getWriter().flush();
}
} catch (IOException e) {
logger.warn("Failed to flush writer after a row error", e);
}
}
private String quote(String stringValue) {
logger.debug("quote(stringValue={}) - start", stringValue);
return new StringBuilder(QUOTE).append(escape(stringValue)).append(QUOTE).toString();
}
/**
* Escapes quote and backslash characters in the given string by prefixing them with a backslash.
*
* @param stringValue the string to escape.
* @return the escaped string.
*/
protected static String escape(String stringValue) {
logger.debug("escape(stringValue={}) - start", stringValue);
char [] array = stringValue.toCharArray();
testExport = QUOTE.toCharArray()[0];
final char escape = ESCAPE.toCharArray()[0];
StringBuilder buffer = new StringBuilder();
for (int i = 0; i < array.length; i++) {
char c = array[i];
if (c == testExport || c == escape) {
buffer.append('\\');
}
buffer.append(c);
}
return buffer.toString();
}
/**
* Returns the writer for the currently active table.
* @return the writer for the currently active table.
*/
public Writer getWriter() {
logger.debug("getWriter() - start");
return writer;
}
/**
* Sets the writer for the currently active table.
* @param writer the writer for the currently active table.
*/
public void setWriter(Writer writer) {
logger.debug("setWriter(writer={}) - start", writer);
this.writer = writer;
}
/**
* Returns the directory CSV files are written to.
* @return the directory CSV files are written to.
*/
public String getTheDirectory() {
logger.debug("getTheDirectory() - start");
return theDirectory;
}
/**
* Sets the directory CSV files are written to.
* @param theDirectory the directory CSV files are written to.
*/
public void setTheDirectory(String theDirectory) {
logger.debug("setTheDirectory(theDirectory={}) - start", theDirectory);
this.theDirectory = theDirectory;
}
/**
* Writes the given dataset's tables and rows as CSV files to the given directory.
*
* @param dataset the dataset to write.
* @param dest the directory to write CSV files to.
* @throws DataSetException if writing the dataset fails.
*/
public static void write(IDataSet dataset, File dest) throws DataSetException {
logger.debug("write(dataset={}, dest={}) - start", dataset, dest);
CsvDataSetWriter writer = new CsvDataSetWriter(dest);
writer.write(dataset);
}
protected void finalize() throws Throwable {
logger.debug("finalize() - start");
if (getWriter() != null) {
getWriter().close();
}
}
}