org.dbunit.ext.mssql.InsertIdentityOperation Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of dbunit Show documentation
dbUnit is a JUnit extension (also usable from Ant and Maven) targeted for database-driven projects that, among other things, puts your database into a known state between test runs. This is an excellent way to avoid the myriad of problems that can occur when one test case corrupts the database and causes subsequent tests to fail or exacerbate the damage.
There is a newer version: 2.8.0
Show newest version
/*
 *
 * 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.
 * 
 * If you are using the Microsoft driver (i.e.
 * com.microsoft.jdbc.sqlserver.SQLServerDriver), you'll need to
 * use the SelectMethod=cursor parameter in the JDBC connection
 * string. Your databaseUrl would look something like the following:
 * 

 * jdbc:microsoft:sqlserver://localhost:1433;DatabaseName=mydb;SelectMethod=cursor
 * 

 * 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
     * 
GENERATED BY DEFAULT AS IDENTITY
     * GENERATED ALWAYS AS IDENTITY
     * set which is the SQL standard syntax to describe auto-generated key columns.
     * Also accepts columns that have the auto-increment property set to true (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)
                {
                    StringBuffer sqlBuffer = new StringBuffer(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)
                    {
                        StringBuffer sqlBuffer = new StringBuffer(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();
        }
    }
}