001: /*
002: * InsertIdentityOperation.java
003: *
004: * The DbUnit Database Testing Framework
005: * Copyright (C)2002-2004, DbUnit.org
006: *
007: * This library is free software; you can redistribute it and/or
008: * modify it under the terms of the GNU Lesser General Public
009: * License as published by the Free Software Foundation; either
010: * version 2.1 of the License, or (at your option) any later version.
011: *
012: * This library is distributed in the hope that it will be useful,
013: * but WITHOUT ANY WARRANTY; without even the implied warranty of
014: * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
015: * Lesser General Public License for more details.
016: *
017: * You should have received a copy of the GNU Lesser General Public
018: * License along with this library; if not, write to the Free Software
019: * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
020: */
021:
022: package org.dbunit.ext.mssql;
023:
024: import org.slf4j.Logger;
025: import org.slf4j.LoggerFactory;
026:
027: import org.dbunit.DatabaseUnitException;
028: import org.dbunit.database.IDatabaseConnection;
029: import org.dbunit.database.DatabaseConfig;
030: import org.dbunit.dataset.*;
031: import org.dbunit.dataset.filter.IColumnFilter;
032: import org.dbunit.operation.AbstractOperation;
033: import org.dbunit.operation.CompositeOperation;
034: import org.dbunit.operation.DatabaseOperation;
035: import org.dbunit.operation.ExclusiveTransactionException;
036:
037: import java.sql.Connection;
038: import java.sql.SQLException;
039: import java.sql.Statement;
040:
041: /**
042: * This class disable the MS SQL Server automatic identifier generation for
043: * the execution of inserts.
044: * <p>
045: * If you are using the Microsoft driver (i.e.
046: * <code>com.microsoft.jdbc.sqlserver.SQLServerDriver</code>), you'll need to
047: * use the <code>SelectMethod=cursor</code> parameter in the JDBC connection
048: * string. Your databaseUrl would look something like the following:
049: * <p>
050: * <code>jdbc:microsoft:sqlserver://localhost:1433;DatabaseName=mydb;SelectMethod=cursor</code>
051: * <p>
052: * Thanks to Jeremy Stein who has submitted multiple patches.
053: *
054: * @author Manuel Laflamme
055: * @author Eric Pugh
056: * @version $Revision: 558 $
057: * @since Apr 9, 2002
058: */
059: public class InsertIdentityOperation extends AbstractOperation {
060:
061: /**
062: * Logger for this class
063: */
064: private static final Logger logger = LoggerFactory
065: .getLogger(InsertIdentityOperation.class);
066:
067: public static final String PROPERTY_IDENTITY_COLUMN_FILTER = "http://www.dbunit.org/properties/mssql/identityColumnFilter";
068:
069: public static final DatabaseOperation INSERT = new InsertIdentityOperation(
070: DatabaseOperation.INSERT);
071:
072: public static final DatabaseOperation CLEAN_INSERT = new CompositeOperation(
073: DatabaseOperation.DELETE_ALL, new InsertIdentityOperation(
074: DatabaseOperation.INSERT));
075:
076: public static final DatabaseOperation REFRESH = new InsertIdentityOperation(
077: DatabaseOperation.REFRESH);
078:
079: private static final IColumnFilter DEFAULT_IDENTITY_FILTER = new IColumnFilter() {
080: public boolean accept(String tableName, Column column) {
081: logger.debug("accept(tableName=" + tableName + ", column="
082: + column + ") - start");
083:
084: return column.getSqlTypeName().endsWith("identity");
085: }
086: };
087:
088: private final DatabaseOperation _operation;
089:
090: /**
091: * Creates a new InsertIdentityOperation object that decorates the
092: * specified operation.
093: */
094: public InsertIdentityOperation(DatabaseOperation operation) {
095: _operation = operation;
096: }
097:
098: private boolean hasIdentityColumn(ITableMetaData metaData,
099: IDatabaseConnection connection) throws DataSetException {
100: logger.debug("hasIdentityColumn(metaData=" + metaData
101: + ", connection=" + connection + ") - start");
102:
103: DatabaseConfig config = connection.getConfig();
104: IColumnFilter identityFilter = (IColumnFilter) config
105: .getProperty(PROPERTY_IDENTITY_COLUMN_FILTER);
106: if (identityFilter == null) {
107: identityFilter = DEFAULT_IDENTITY_FILTER;
108: }
109:
110: // Verify if there is at least one identity column
111: Column[] columns = metaData.getColumns();
112: for (int i = 0; i < columns.length; i++) {
113: if (identityFilter.accept(metaData.getTableName(),
114: columns[i])) {
115: return true;
116: }
117: }
118:
119: return false;
120: }
121:
122: ////////////////////////////////////////////////////////////////////////////
123: // DatabaseOperation class
124:
125: public void execute(IDatabaseConnection connection, IDataSet dataSet)
126: throws DatabaseUnitException, SQLException {
127: logger.debug("execute(connection=" + connection
128: + ", dataSet) - start");
129:
130: Connection jdbcConnection = connection.getConnection();
131: Statement statement = jdbcConnection.createStatement();
132:
133: try {
134: IDataSet databaseDataSet = connection.createDataSet();
135:
136: // INSERT_IDENTITY need to be enabled/disabled inside the
137: // same transaction
138: if (jdbcConnection.getAutoCommit() == false) {
139: throw new ExclusiveTransactionException();
140: }
141: jdbcConnection.setAutoCommit(false);
142:
143: // Execute decorated operation one table at a time
144: ITableIterator iterator = dataSet.iterator();
145: while (iterator.next()) {
146: ITable table = iterator.getTable();
147: String tableName = table.getTableMetaData()
148: .getTableName();
149:
150: ITableMetaData metaData = databaseDataSet
151: .getTableMetaData(tableName);
152:
153: // enable identity insert
154: boolean hasIdentityColumn = hasIdentityColumn(metaData,
155: connection);
156:
157: if (hasIdentityColumn) {
158: StringBuffer sqlBuffer = new StringBuffer(128);
159: sqlBuffer.append("SET IDENTITY_INSERT ");
160: sqlBuffer.append(getQualifiedName(connection
161: .getSchema(), metaData.getTableName(),
162: connection));
163: sqlBuffer.append(" ON");
164: statement.execute(sqlBuffer.toString());
165: }
166:
167: try {
168: _operation.execute(connection, new DefaultDataSet(
169: table));
170: } finally {
171: // disable identity insert
172: if (hasIdentityColumn) {
173: StringBuffer sqlBuffer = new StringBuffer(128);
174: sqlBuffer.append("SET IDENTITY_INSERT ");
175: sqlBuffer.append(getQualifiedName(connection
176: .getSchema(), metaData.getTableName(),
177: connection));
178: sqlBuffer.append(" OFF");
179: statement.execute(sqlBuffer.toString());
180: }
181: jdbcConnection.commit();
182: }
183: }
184: } finally {
185: jdbcConnection.setAutoCommit(true);
186: statement.close();
187: }
188: }
189: }
|