1 21 22 package org.dbunit.ext.mssql; 23 24 import org.dbunit.DatabaseUnitException; 25 import org.dbunit.database.IDatabaseConnection; 26 import org.dbunit.database.DatabaseConfig; 27 import org.dbunit.dataset.*; 28 import org.dbunit.dataset.filter.IColumnFilter; 29 import org.dbunit.operation.AbstractOperation; 30 import org.dbunit.operation.CompositeOperation; 31 import org.dbunit.operation.DatabaseOperation; 32 import org.dbunit.operation.ExclusiveTransactionException; 33 34 import java.sql.Connection ; 35 import java.sql.SQLException ; 36 import java.sql.Statement ; 37 38 56 public class InsertIdentityOperation extends AbstractOperation 57 { 58 public static final String PROPERTY_IDENTITY_COLUMN_FILTER = 59 "http://www.dbunit.org/properties/mssql/identityColumnFilter"; 60 61 public static final DatabaseOperation INSERT = 62 new InsertIdentityOperation(DatabaseOperation.INSERT); 63 64 public static final DatabaseOperation CLEAN_INSERT = 65 new CompositeOperation(DatabaseOperation.DELETE_ALL, 66 new InsertIdentityOperation(DatabaseOperation.INSERT)); 67 68 public static final DatabaseOperation REFRESH = 69 new InsertIdentityOperation(DatabaseOperation.REFRESH); 70 71 private static final IColumnFilter DEFAULT_IDENTITY_FILTER = new IColumnFilter() 72 { 73 public boolean accept(String tableName, Column column) 74 { 75 return column.getSqlTypeName().endsWith("identity"); 76 } 77 }; 78 79 80 private final DatabaseOperation _operation; 81 82 86 public InsertIdentityOperation(DatabaseOperation operation) 87 { 88 _operation = operation; 89 } 90 91 private boolean hasIdentityColumn(ITableMetaData metaData, IDatabaseConnection connection) 92 throws DataSetException 93 { 94 DatabaseConfig config = connection.getConfig(); 95 IColumnFilter identityFilter = (IColumnFilter)config.getProperty( 96 PROPERTY_IDENTITY_COLUMN_FILTER); 97 if (identityFilter == null) 98 { 99 identityFilter = DEFAULT_IDENTITY_FILTER; 100 } 101 102 Column[] columns = metaData.getColumns(); 104 for (int i = 0; i < columns.length; i++) 105 { 106 if (identityFilter.accept(null, columns[i])) 107 { 108 return true; 109 } 110 } 111 112 return false; 113 } 114 115 118 public void execute(IDatabaseConnection connection, IDataSet dataSet) 119 throws DatabaseUnitException, SQLException 120 { 121 Connection jdbcConnection = connection.getConnection(); 122 Statement statement = jdbcConnection.createStatement(); 123 124 try 125 { 126 IDataSet databaseDataSet = connection.createDataSet(); 127 128 129 if (jdbcConnection.getAutoCommit() == false) 132 { 133 throw new ExclusiveTransactionException(); 134 } 135 jdbcConnection.setAutoCommit(false); 136 137 ITableIterator iterator = dataSet.iterator(); 139 while(iterator.next()) 140 { 141 ITable table = iterator.getTable(); 142 String tableName = table.getTableMetaData().getTableName(); 143 144 ITableMetaData metaData = 145 databaseDataSet.getTableMetaData(tableName); 146 147 boolean hasIdentityColumn = hasIdentityColumn(metaData, connection); 149 150 if (hasIdentityColumn) 151 { 152 StringBuffer sqlBuffer = new StringBuffer (128); 153 sqlBuffer.append("SET IDENTITY_INSERT "); 154 sqlBuffer.append(getQualifiedName(connection.getSchema(), 155 metaData.getTableName(), connection)); 156 sqlBuffer.append(" ON"); 157 statement.execute(sqlBuffer.toString()); 158 } 159 160 try 161 { 162 _operation.execute(connection, new DefaultDataSet(table)); 163 } 164 finally 165 { 166 if (hasIdentityColumn) 168 { 169 StringBuffer sqlBuffer = new StringBuffer (128); 170 sqlBuffer.append("SET IDENTITY_INSERT "); 171 sqlBuffer.append(getQualifiedName(connection.getSchema(), 172 metaData.getTableName(), connection)); 173 sqlBuffer.append(" OFF"); 174 statement.execute(sqlBuffer.toString()); 175 } 176 jdbcConnection.commit(); 177 } 178 } 179 } 180 finally 181 { 182 jdbcConnection.setAutoCommit(true); 183 statement.close(); 184 } 185 } 186 } 187 188 189 190 191 192 193 | Popular Tags |