package de.flower.rmt.test;
import de.flower.common.util.IO;
import org.apache.commons.lang3.Validate;
import org.dbunit.database.DatabaseDataSourceConnection;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.dataset.DataSetException;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.xml.FlatXmlDataSet;
import org.dbunit.dataset.xml.FlatXmlDataSetBuilder;
import org.dbunit.operation.DatabaseOperation;
import org.dbunit.operation.TransactionOperation;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import javax.sql.DataSource;
import java.io.FileOutputStream;
import java.sql.SQLException;
import java.util.Map;
/**
* @author flowerrrr
*/
public class Database {
private final static Logger log = LoggerFactory.getLogger(Database.class);
private DataSource dataSource;
private Map<String, Object> properties;
public Database(Map<String, Object> properties) {
this.properties = properties;
}
public static IDataSet createDataSet(String classpathResource) {
try {
return new FlatXmlDataSetBuilder().build(IO.loadClasspathResourceAsStream(classpathResource));
} catch (DataSetException e) {
throw new RuntimeException(e);
}
}
public void deleteAll(String[] tables) {
IDatabaseConnection connection;
try {
connection = getDatabaseConnection();
IDataSet dataset = connection.createDataSet(tables);
deleteAll(dataset);
} catch (Exception e) {
throw new RuntimeException(e);
}
}
public void deleteAll(IDataSet dataSet) {
log.info("Delete all. Dataset: " + dataSet);
execute(DatabaseOperation.DELETE_ALL, dataSet);
}
public void cleanInsert(IDataSet dataSet) {
log.info("Clean inserting dataset: " + dataSet);
execute(DatabaseOperation.CLEAN_INSERT, dataSet);
}
public void export(String filename, String[] tables) {
try {
IDatabaseConnection connection = getDatabaseConnection();
IDataSet dataset = connection.createDataSet(tables);
FlatXmlDataSet.write(dataset, new FileOutputStream(filename));
} catch (Exception e) {
throw new RuntimeException(e);
}
}
private void execute(DatabaseOperation operation, IDataSet dataSet) {
IDatabaseConnection connection;
try {
connection = getDatabaseConnection();
TransactionOperation transOp = new TransactionOperation(operation);
transOp.execute(connection, dataSet);
connection.close();
} catch (Exception e) {
throw new RuntimeException(e);
}
}
private IDatabaseConnection getDatabaseConnection() {
IDatabaseConnection connection;
try {
connection = new DatabaseDataSourceConnection(Validate.notNull(dataSource));
} catch (SQLException e) {
throw new RuntimeException(e);
}
return configureConnection(connection);
}
private IDatabaseConnection configureConnection(IDatabaseConnection connection) {
for (String key : properties.keySet()) {
connection.getConfig().setProperty(key, properties.get(key));
}
return connection;
}
public DataSource getDataSource() {
return dataSource;
}
public void setDataSource(final DataSource dataSource) {
this.dataSource = dataSource;
}
}