org.dbunit.operation.DatabaseOperationConnectionCallback Maven / Gradle / Ivy

package org.dbunit.operation;

import java.io.IOException;
import java.io.InputStream;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.Date;
import java.util.UUID;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.dbunit.DatabaseUnitException;
import org.dbunit.database.DatabaseConnectionCustom;
import org.dbunit.dataset.ITable;
import org.dbunit.dataset.ReplacementDataSet;
import org.dbunit.dataset.datatype.IDataTypeFactory;
import org.dbunit.dataset.xml.FlatXmlDataSetBuilder;
import org.dbunit.ext.mssql.InsertIdentityOperation;
import org.springframework.core.io.Resource;
import org.springframework.jdbc.core.ConnectionCallback;
import org.springframework.orm.jpa.vendor.Database;
import org.springframework.util.Assert;

public class DatabaseOperationConnectionCallback implements ConnectionCallback {
  private static final Log logger = LogFactory.getLog(DatabaseOperationConnectionCallback.class);
  private final Database database;
  private final Resource[] resources;

  public DatabaseOperationConnectionCallback(Database database, Resource... resources) {
    this.database = database;
    this.resources = resources;
  }

  /**
   * @see org.dbunit.dataset.xml.XmlDataSet#XmlDataSet(InputStream)
   */
  @Override
  public Void doInConnection(Connection connection) throws SQLException {
    String lob = UUID.randomUUID().toString();
    logger.info(String.format("%n%nUsing default lob: %s%n", lob));
    for (Resource resource : resources) {
      if (resource.exists()) {
        execute(connection, resource, lob);
      }
    }
    return null;
  }

  /**
   * @see java.nio.charset.StandardCharsets#UTF_8
   */
  private void execute(Connection connection, Resource resource, String lob) throws SQLException {
    try (InputStream inputStream = resource.getInputStream(); DatabaseConnectionCustom databaseConnection = new DatabaseConnectionCustom(connection)) {
      databaseConnection.getConfig().setProperty(org.dbunit.database.DatabaseConfig.PROPERTY_DATATYPE_FACTORY, getDatatypeFactory(this.database));
      databaseConnection.getConfig().setProperty(org.dbunit.database.DatabaseConfig.PROPERTY_ESCAPE_PATTERN, getEscapePattern(this.database));

      FlatXmlDataSetBuilder flatXmlDataSetBuilder = new FlatXmlDataSetBuilder();
      flatXmlDataSetBuilder.setColumnSensing(true);

      ReplacementDataSet replacementDataSet = new ReplacementDataSet(flatXmlDataSetBuilder.build(inputStream));
      replacementDataSet.addReplacementObject(ITable.NO_VALUE, null);
      replacementDataSet.addReplacementObject("[null]", null);
      replacementDataSet.addReplacementObject("[date]", new Date());
      replacementDataSet.addReplacementObject("[lob]", lob.getBytes());

      DatabaseOperation databaseOperation = DatabaseOperation.INSERT;
      if (this.database != null && this.database == Database.SQL_SERVER) {
        databaseOperation = new InsertIdentityOperation(databaseOperation);
      }
      databaseOperation.execute(databaseConnection, replacementDataSet);
    }
    catch (SQLException | IOException | DatabaseUnitException e) {
      throw new SQLException("DBunit execute fail...", e);
    }
  }

  // "'?'"
  private String getEscapePattern(Database database) {
    switch (database) {
    case H2:
      return "\"?\"";
    case MYSQL:
      return "`?`";
    case SQL_SERVER:
      return "[?]";
    default:
      return "?";
    }
  }

  /**
   * @see UnsupportedOperationException
   */
  private IDataTypeFactory getDatatypeFactory(Database database) {
    Assert.notNull(database, "'database' must not be null");
    switch (database) {
    case H2:
      return new org.dbunit.ext.h2.H2DataTypeFactory();
    case MYSQL:
      return new org.dbunit.ext.mysql.MySqlDataTypeFactory();
    case SQL_SERVER:
      return new org.dbunit.ext.mssql.MsSqlDataTypeFactory();
    default:
      return new org.dbunit.dataset.datatype.DefaultDataTypeFactory();
    }
  }
}