InsertIdentityOperation.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.ext.mssql;
import java.sql.Connection;
import java.sql.SQLException;
import java.sql.Statement;
import org.dbunit.DatabaseUnitException;
import org.dbunit.database.DatabaseConfig;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.dataset.Column;
import org.dbunit.dataset.DataSetException;
import org.dbunit.dataset.DefaultDataSet;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.ITable;
import org.dbunit.dataset.ITableIterator;
import org.dbunit.dataset.ITableMetaData;
import org.dbunit.dataset.Column.AutoIncrement;
import org.dbunit.dataset.filter.IColumnFilter;
import org.dbunit.operation.AbstractOperation;
import org.dbunit.operation.CompositeOperation;
import org.dbunit.operation.DatabaseOperation;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
/**
* This class disable the MS SQL Server automatic identifier generation for
* the execution of inserts.
* <p>
* If you are using the Microsoft driver (i.e.
* <code>com.microsoft.jdbc.sqlserver.SQLServerDriver</code>), you'll need to
* use the <code>SelectMethod=cursor</code> parameter in the JDBC connection
* string. Your databaseUrl would look something like the following:
* <p>
* <code>jdbc:microsoft:sqlserver://localhost:1433;DatabaseName=mydb;SelectMethod=cursor</code>
* <p>
* Thanks to Jeremy Stein who has submitted multiple patches.
*
* @author Manuel Laflamme
* @author Eric Pugh
* @author Last changed by: $Author$
* @version $Revision$ $Date$
* @since 1.4 (Apr 9, 2002)
*/
public class InsertIdentityOperation extends AbstractOperation
{
/**
* Logger for this class
*/
private static final Logger logger = LoggerFactory.getLogger(InsertIdentityOperation.class);
public static final DatabaseOperation INSERT =
new InsertIdentityOperation(DatabaseOperation.INSERT);
public static final DatabaseOperation CLEAN_INSERT =
new CompositeOperation(DatabaseOperation.DELETE_ALL,
new InsertIdentityOperation(DatabaseOperation.INSERT));
public static final DatabaseOperation REFRESH =
new InsertIdentityOperation(DatabaseOperation.REFRESH);
private static final IColumnFilter DEFAULT_IDENTITY_FILTER = new IColumnFilter()
{
public boolean accept(String tableName, Column column)
{
return column.getSqlTypeName().endsWith("identity");
}
};
/**
* Accepts columns that have one of the remarks
* <ul><li>GENERATED BY DEFAULT AS IDENTITY</li>
* <li>GENERATED ALWAYS AS IDENTITY</li></ul>
* set which is the SQL standard syntax to describe auto-generated key columns.
* Also accepts columns that have the auto-increment property set to <code>true</code> (note that
* it does not yet have the ability to check whether the column is a primary key col).
* @since 2.4.3
*/
public static final IColumnFilter IDENTITY_FILTER_EXTENDED = new IColumnFilter() {
public boolean accept(String tableName, Column column)
{
String remarks = column.getRemarks();
boolean isIdentityCol = (remarks != null) && (
remarks.indexOf("GENERATED BY DEFAULT AS IDENTITY") > -1 ||
remarks.indexOf("GENERATED ALWAYS AS IDENTITY") > -1
);
// If "remarks" did not give the appropriate hint, check the auto-increment property
if(!isIdentityCol)
{
//TODO Should we ensure that the column is a PrimaryKey column?
isIdentityCol = (AutoIncrement.YES == column.getAutoIncrement());
}
return isIdentityCol;
}
};
private final DatabaseOperation _operation;
/**
* Creates a new InsertIdentityOperation object that decorates the
* specified operation.
*/
public InsertIdentityOperation(DatabaseOperation operation)
{
_operation = operation;
}
boolean hasIdentityColumn(ITableMetaData metaData, IDatabaseConnection connection)
throws DataSetException
{
logger.debug("hasIdentityColumn(metaData={}, connection={}) - start", metaData, connection);
DatabaseConfig config = connection.getConfig();
IColumnFilter identityFilter = (IColumnFilter)config.getProperty(
DatabaseConfig.PROPERTY_IDENTITY_COLUMN_FILTER);
if (identityFilter == null)
{
identityFilter = DEFAULT_IDENTITY_FILTER;
}
// Verify if there is at least one identity column
Column[] columns = metaData.getColumns();
for (int i = 0; i < columns.length; i++)
{
if (identityFilter.accept(metaData.getTableName(), columns[i]))
{
return true;
}
}
return false;
}
////////////////////////////////////////////////////////////////////////////
// DatabaseOperation class
public void execute(IDatabaseConnection connection, IDataSet dataSet)
throws DatabaseUnitException, SQLException
{
logger.debug("execute(connection={}, dataSet={}) - start", connection, dataSet);
Connection jdbcConnection = connection.getConnection();
Statement statement = jdbcConnection.createStatement();
boolean wasAutoCommit = false;
try
{
IDataSet databaseDataSet = connection.createDataSet();
// Note that MSSQL has a different transaction strategy from oracle.
// By default the transaction is always in "autocommit=true" so
// that every statement is immediately committed. If a dbunit
// user does not want this behavior dbunit takes it into account
// here.
// INSERT_IDENTITY need to be enabled/disabled inside the
// same transaction
if (jdbcConnection.getAutoCommit() == true)
{
wasAutoCommit = true;
jdbcConnection.setAutoCommit(false);
}
// Execute decorated operation one table at a time
ITableIterator iterator = dataSet.iterator();
while(iterator.next())
{
ITable table = iterator.getTable();
String tableName = table.getTableMetaData().getTableName();
ITableMetaData metaData =
databaseDataSet.getTableMetaData(tableName);
// enable identity insert
boolean hasIdentityColumn = hasIdentityColumn(metaData, connection);
if (hasIdentityColumn)
{
final StringBuilder sqlBuffer = new StringBuilder(128);
sqlBuffer.append("SET IDENTITY_INSERT ");
sqlBuffer.append(getQualifiedName(connection.getSchema(),
metaData.getTableName(), connection));
sqlBuffer.append(" ON");
statement.execute(sqlBuffer.toString());
}
try
{
_operation.execute(connection, new DefaultDataSet(table));
}
finally
{
// disable identity insert
if (hasIdentityColumn)
{
final StringBuilder sqlBuffer = new StringBuilder(128);
sqlBuffer.append("SET IDENTITY_INSERT ");
sqlBuffer.append(getQualifiedName(connection.getSchema(),
metaData.getTableName(), connection));
sqlBuffer.append(" OFF");
statement.execute(sqlBuffer.toString());
}
if (wasAutoCommit)
{
jdbcConnection.commit();
}
}
}
}
finally
{
if(wasAutoCommit)
{
// Reset the autocommit property
jdbcConnection.setAutoCommit(true);
}
statement.close();
}
}
}