DBUnitImporter.java

package org.jboss.seam.wiki.util;

import org.dbunit.database.DatabaseConfig;
import org.dbunit.database.DatabaseConnection;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.ReplacementDataSet;
import org.dbunit.dataset.datatype.DataType;
import org.dbunit.dataset.datatype.DataTypeException;
import org.dbunit.dataset.datatype.DefaultDataTypeFactory;
import org.dbunit.dataset.xml.FlatXmlDataSet;
import org.dbunit.operation.DatabaseOperation;

import javax.naming.InitialContext;
import javax.sql.DataSource;
import java.io.InputStream;
import java.net.URL;
import java.sql.Connection;
import java.sql.Types;
import java.util.ArrayList;
import java.util.List;

/**
 * Imports some data into the database with the help of DBUnit. This allows us to
 * use the same dataset files as in unit testing, but in the regular application startup during
 * development. Also helps to avoid maintaining the crude Hibernate import.sql file.
 *
 * @author Christian Bauer
 */
public class DBUnitImporter {

    public enum Database {
        hsql, mysql, postgresql
    }

    protected Database database;

    private String datasourceJndiName;

    String binaryDir;

    List<String> datasets = new ArrayList<String>();

    public void setDatabase(String database) {
        this.database = Database.valueOf(database);
    }

    public String getDatasourceJndiName() {
        return datasourceJndiName;
    }

    public void setDatasourceJndiName(String datasourceJndiName) {
        this.datasourceJndiName = datasourceJndiName;
    }

    public String getBinaryDir() {
        return binaryDir;
    }

    public void setBinaryDir(String binaryDir) {
        this.binaryDir = binaryDir;
    }

    public List<String> getDatasets() {
        return datasets;
    }

    public void setDatasets(List<String> datasets) {
        this.datasets = datasets;
    }

    public void importDatasets() throws Exception {

        List<DataSetOperation> dataSetOperations = new ArrayList<DataSetOperation>();

        for (String dataset : datasets) {
            dataSetOperations.add(new DataSetOperation(dataset));
        }

        IDatabaseConnection con = null;
        try {
            con = getConnection();
            disableReferentialIntegrity(con);
            for (DataSetOperation op : dataSetOperations) {
                prepareExecution(con, op);
                op.execute(con);
                afterExecution(con, op);
            }
            enableReferentialIntegrity(con);
        } finally {
            if (con != null) {
                try {
                    con.close();
                } catch (Exception ex) {
                    ex.printStackTrace(System.err);
                }
            }
        }

    }

    /**
     * Override this method if you want to provide your own DBUnit <tt>IDatabaseConnection</tt> instance.
     * <p/>
     * If you do not override this, default behavior is to use the * configured datasource name and
     * to obtain a connection with a JNDI lookup.
     *
     * @return a DBUnit database connection (wrapped)
     */
    protected IDatabaseConnection getConnection() {
        try {
            DataSource datasource = ((DataSource)new InitialContext().lookup(datasourceJndiName));

            // Get a JDBC connection from JNDI datasource
            Connection con = datasource.getConnection();
            IDatabaseConnection dbUnitCon = new DatabaseConnection(con);
            editConfig(dbUnitCon.getConfig());
            return dbUnitCon;
        } catch (Exception ex) {
            throw new RuntimeException(ex);
        }
    }

    /**
     * Override this method if you aren't using HSQL DB.
     * <p/>
     * Execute whatever statement is necessary to either defer or disable foreign
     * key constraint checking on the given database connection, which is used by
     * DBUnit to import datasets.
     *
     * @param con A DBUnit connection wrapper, which is used afterwards for dataset operations
     */
    protected void disableReferentialIntegrity(IDatabaseConnection con) {
        try {
            if (database.equals(Database.hsql)) {
                con.getConnection().prepareStatement("set referential_integrity FALSE").execute(); // HSQL DB
            } else if (database.equals(Database.mysql)) {
                con.getConnection().prepareStatement("set foreign_key_checks=0").execute(); // MySQL > 4.1.1
            } else if (database.equals(Database.postgresql)) {
                // See prepareExecution() and afterExecution()
            }
        } catch (Exception ex) {
            throw new RuntimeException(ex);
        }
    }

    /**
     * Override this method if you aren't using HSQL DB.
     * <p/>
     * Execute whatever statement is necessary to enable integrity constraint checks after
     * dataset operations.
     *
     * @param con A DBUnit connection wrapper, before it is used by the application again
     */
    protected void enableReferentialIntegrity(IDatabaseConnection con) {
        try {
            if (database.equals(Database.hsql)) {
                con.getConnection().prepareStatement("set referential_integrity TRUE").execute();  // HSQL DB
            } else if (database.equals(Database.mysql)) {
                con.getConnection().prepareStatement("set foreign_key_checks=1").execute(); // MySQL > 4.1.1
            }
        } catch (Exception ex) {
            throw new RuntimeException(ex);
        }
    }

    /**
     * Override this method if you require DBUnit configuration features or additional properties.
     * <p>
     * Called after a connection has been obtaind and before the connection is used. Can be a
     * NOOP method if no additional settings are necessary for your DBUnit/DBMS setup.
     *
     * @param config A DBUnit <tt>DatabaseConfig</tt> object for setting properties and features
     */
    protected void editConfig(DatabaseConfig config) {

        if (database.equals(Database.hsql)) {
            // DBUnit/HSQL bugfix
            // http://www.carbonfive.com/community/archives/2005/07/dbunit_hsql_and.html
            config.setProperty(DatabaseConfig.PROPERTY_DATATYPE_FACTORY, new DefaultDataTypeFactory() {
                public DataType createDataType(int sqlType, String sqlTypeName)
                  throws DataTypeException {
                   if (sqlType == Types.BOOLEAN) {
                      return DataType.BOOLEAN;
                    }
                   return super.createDataType(sqlType, sqlTypeName);
                 }
            });
        }

    }

    /**
     * Resolves the binary dir location with the help of the classloader, we need the
     * absolute full path of that directory.
     *
     * @return String full absolute path of the binary directory
     */
    protected String getBinaryDirFullpath() {
        if (binaryDir == null) {
            throw new RuntimeException("Please set binaryDir property to location of binary test files");
        }
        URL url = Thread.currentThread().getContextClassLoader().getResource(getBinaryDir());
        if (url == null) {
            throw new RuntimeException("Could not find full path with classloader of binaryDir: " + getBinaryDir());
        }
        return url.toString();
    }

    /**
     * Callback for each operation, useful if extra preparation of data/tables is necessary.
     *
     * @param con A DBUnit connection wrapper
     * @param operation The operation to be executed, call <tt>getDataSet()</tt> to access the data.
     */
    protected void prepareExecution(IDatabaseConnection con, DataSetOperation operation) {
        try {
            // PostgreSQL has no global switch to turn off foreign key checks, needs to be toggled on each table
            if (database.equals(Database.postgresql)) {
                for (String tableName : operation.getDataSet().getTableNames()) {
                    con.getConnection().prepareStatement("alter table " + tableName + " disable trigger all").execute();
                }
            }
        } catch (Exception ex) {
            throw new RuntimeException(ex);
        }
    }

    /**
     * Callback for each operation, useful if extra preparation of data/tables is necessary.
     *
     * @param con A DBUnit connection wrapper
     * @param operation The operation that was executed, call <tt>getDataSet()</tt> to access the data.
     */
    protected void afterExecution(IDatabaseConnection con, DataSetOperation operation) {
        try {
            // PostgreSQL has no global switch to turn off foreign key checks, needs to be toggled on each table
            if (database.equals(Database.postgresql)) {
                for (String tableName : operation.getDataSet().getTableNames()) {
                    con.getConnection().prepareStatement("alter table " + tableName + " enable trigger all").execute();
                }
            }
        } catch (Exception ex) {
            throw new RuntimeException(ex);
        }
    }

    protected class DataSetOperation {
        String dataSetLocation;
        ReplacementDataSet dataSet;
        DatabaseOperation operation;

        /**
         * Defaults to <tt>DatabaseOperation.CLEAN_INSERT</tt>
         */
        public DataSetOperation(String dataSetLocation){
            this(dataSetLocation, DatabaseOperation.INSERT);
        }

        public DataSetOperation(String dataSetLocation, DatabaseOperation operation) {
            // Load the base dataset file
            InputStream input = Thread.currentThread().getContextClassLoader().getResourceAsStream(dataSetLocation);
            try {
                this.dataSet = new ReplacementDataSet( new FlatXmlDataSet(input) );
            } catch (Exception ex) {
                throw new RuntimeException("Could not load dataset for import: " + dataSetLocation, ex);
            }
            this.dataSet.addReplacementObject("[NULL]", null);
            this.dataSet.addReplacementSubstring("[BINARY_DIR]", getBinaryDirFullpath());
            this.operation = operation;
            this.dataSetLocation = dataSetLocation;
        }

        public IDataSet getDataSet() {
            return dataSet;
        }

        public DatabaseOperation getOperation() {
            return operation;
        }

        public void execute(IDatabaseConnection connection) {
            try {
                this.operation.execute(connection, dataSet);
            } catch (Exception ex) {
                throw new RuntimeException(ex);
            }
        }

        public String toString() {
            // TODO: This is not pretty because DBUnit's DatabaseOperation doesn't implement toString() properly
            return operation.getClass() + " with dataset: " + dataSetLocation;
        }
    }

}