View Javadoc

1   /*
2    *
3    * The DbUnit Database Testing Framework
4    * Copyright (C)2002-2004, DbUnit.org
5    *
6    * This library is free software; you can redistribute it and/or
7    * modify it under the terms of the GNU Lesser General Public
8    * License as published by the Free Software Foundation; either
9    * version 2.1 of the License, or (at your option) any later version.
10   *
11   * This library is distributed in the hope that it will be useful,
12   * but WITHOUT ANY WARRANTY; without even the implied warranty of
13   * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
14   * Lesser General Public License for more details.
15   *
16   * You should have received a copy of the GNU Lesser General Public
17   * License along with this library; if not, write to the Free Software
18   * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
19   *
20   */
21  package org.dbunit.ext.mssql;
22  
23  import java.sql.Connection;
24  import java.sql.SQLException;
25  import java.sql.Statement;
26  
27  import org.dbunit.DatabaseUnitException;
28  import org.dbunit.database.DatabaseConfig;
29  import org.dbunit.database.IDatabaseConnection;
30  import org.dbunit.dataset.Column;
31  import org.dbunit.dataset.DataSetException;
32  import org.dbunit.dataset.DefaultDataSet;
33  import org.dbunit.dataset.IDataSet;
34  import org.dbunit.dataset.ITable;
35  import org.dbunit.dataset.ITableIterator;
36  import org.dbunit.dataset.ITableMetaData;
37  import org.dbunit.dataset.Column.AutoIncrement;
38  import org.dbunit.dataset.filter.IColumnFilter;
39  import org.dbunit.operation.AbstractOperation;
40  import org.dbunit.operation.CompositeOperation;
41  import org.dbunit.operation.DatabaseOperation;
42  import org.slf4j.Logger;
43  import org.slf4j.LoggerFactory;
44  
45  /**
46   * This class disable the MS SQL Server automatic identifier generation for
47   * the execution of inserts.
48   * <p>
49   * If you are using the Microsoft driver (i.e.
50   * <code>com.microsoft.jdbc.sqlserver.SQLServerDriver</code>), you'll need to
51   * use the <code>SelectMethod=cursor</code> parameter in the JDBC connection
52   * string. Your databaseUrl would look something like the following:
53   * <p>
54   * <code>jdbc:microsoft:sqlserver://localhost:1433;DatabaseName=mydb;SelectMethod=cursor</code>
55   * <p>
56   * Thanks to Jeremy Stein who has submitted multiple patches.
57   *
58   * @author Manuel Laflamme
59   * @author Eric Pugh
60   * @author Last changed by: $Author: gommma $
61   * @version $Revision: 934 $ $Date: 2008-12-31 23:29:25 +0100 (mer, 31 dic 2008) $
62   * @since 1.4 (Apr 9, 2002)
63   */
64  public class InsertIdentityOperation extends AbstractOperation
65  {
66  
67      /**
68       * Logger for this class
69       */
70      private static final Logger logger = LoggerFactory.getLogger(InsertIdentityOperation.class);
71  
72      public static final String PROPERTY_IDENTITY_COLUMN_FILTER =
73              "http://www.dbunit.org/properties/mssql/identityColumnFilter";
74  
75      public static final DatabaseOperation INSERT =
76              new InsertIdentityOperation(DatabaseOperation.INSERT);
77  
78      public static final DatabaseOperation CLEAN_INSERT =
79              new CompositeOperation(DatabaseOperation.DELETE_ALL,
80                      new InsertIdentityOperation(DatabaseOperation.INSERT));
81  
82      public static final DatabaseOperation REFRESH =
83              new InsertIdentityOperation(DatabaseOperation.REFRESH);
84  
85      private static final IColumnFilter DEFAULT_IDENTITY_FILTER = new IColumnFilter()
86      {
87          public boolean accept(String tableName, Column column)
88          {
89              return column.getSqlTypeName().endsWith("identity");
90          }
91      };
92  
93      
94      /**
95       * Accepts columns that have one of the remarks
96       * <ul><li>GENERATED BY DEFAULT AS IDENTITY</li>
97       * <li>GENERATED ALWAYS AS IDENTITY</li></ul>
98       * set which is the SQL standard syntax to describe auto-generated key columns.
99       * Also accepts columns that have the auto-increment property set to <code>true</code> (note that
100      * it does not yet have the ability to check whether the column is a primary key col).
101      * @since 2.4.3
102      */
103     public static final IColumnFilter IDENTITY_FILTER_EXTENDED = new IColumnFilter() {
104     
105         public boolean accept(String tableName, Column column) 
106         {
107             String remarks = column.getRemarks();
108             boolean isIdentityCol = (remarks != null) && ( 
109                     remarks.indexOf("GENERATED BY DEFAULT AS IDENTITY") > -1 || 
110                     remarks.indexOf("GENERATED ALWAYS AS IDENTITY") > -1
111                     );
112             
113             // If "remarks" did not give the appropriate hint, check the auto-increment property
114             if(!isIdentityCol)
115             {
116                 //TODO Should we ensure that the column is a PrimaryKey column?
117                 isIdentityCol = (AutoIncrement.YES == column.getAutoIncrement());
118             }
119             
120             return isIdentityCol;
121         }
122     };
123 
124     
125     private final DatabaseOperation _operation;
126 
127     /**
128      * Creates a new InsertIdentityOperation object that decorates the
129      * specified operation.
130      */
131     public InsertIdentityOperation(DatabaseOperation operation)
132     {
133         _operation = operation;
134     }
135 
136     boolean hasIdentityColumn(ITableMetaData metaData, IDatabaseConnection connection)
137             throws DataSetException
138     {
139         logger.debug("hasIdentityColumn(metaData={}, connection={}) - start", metaData, connection);
140 
141         DatabaseConfig config = connection.getConfig();
142         IColumnFilter identityFilter = (IColumnFilter)config.getProperty(
143                 PROPERTY_IDENTITY_COLUMN_FILTER);
144         if (identityFilter == null)
145         {
146             identityFilter = DEFAULT_IDENTITY_FILTER;
147         }
148 
149         // Verify if there is at least one identity column
150         Column[] columns = metaData.getColumns();
151         for (int i = 0; i < columns.length; i++)
152         {
153             if (identityFilter.accept(metaData.getTableName(), columns[i]))
154             {
155                 return true;
156             }
157         }
158 
159         return false;
160     }
161 
162     ////////////////////////////////////////////////////////////////////////////
163     // DatabaseOperation class
164 
165     public void execute(IDatabaseConnection connection, IDataSet dataSet)
166             throws DatabaseUnitException, SQLException
167     {
168         logger.debug("execute(connection={}, dataSet={}) - start", connection, dataSet);
169 
170         Connection jdbcConnection = connection.getConnection();
171         Statement statement = jdbcConnection.createStatement();
172 
173         boolean wasAutoCommit = false;
174         try
175         {
176             IDataSet databaseDataSet = connection.createDataSet();
177             
178             // Note that MSSQL has a different transaction strategy from oracle.
179             // By default the transaction is always in "autocommit=true" so
180             // that every statement is immediately committed. If a dbunit
181             // user does not want this behavior dbunit takes it into account
182             // here.
183             
184             // INSERT_IDENTITY need to be enabled/disabled inside the
185             // same transaction
186             if (jdbcConnection.getAutoCommit() == true)
187             {
188                 wasAutoCommit = true;
189                 jdbcConnection.setAutoCommit(false);
190             }
191 
192             // Execute decorated operation one table at a time
193             ITableIterator iterator = dataSet.iterator();
194             while(iterator.next())
195             {
196                 ITable table = iterator.getTable();
197                 String tableName = table.getTableMetaData().getTableName();
198 
199                 ITableMetaData metaData =
200                         databaseDataSet.getTableMetaData(tableName);
201 
202                 // enable identity insert
203                 boolean hasIdentityColumn = hasIdentityColumn(metaData, connection);
204 
205                 if (hasIdentityColumn)
206                 {
207                     StringBuffer sqlBuffer = new StringBuffer(128);
208                     sqlBuffer.append("SET IDENTITY_INSERT ");
209                     sqlBuffer.append(getQualifiedName(connection.getSchema(),
210                             metaData.getTableName(), connection));
211                     sqlBuffer.append(" ON");
212                     statement.execute(sqlBuffer.toString());
213                 }
214 
215                 try
216                 {
217                     _operation.execute(connection, new DefaultDataSet(table));
218                 }
219                 finally
220                 {
221                     // disable identity insert
222                     if (hasIdentityColumn)
223                     {
224                         StringBuffer sqlBuffer = new StringBuffer(128);
225                         sqlBuffer.append("SET IDENTITY_INSERT ");
226                         sqlBuffer.append(getQualifiedName(connection.getSchema(),
227                                 metaData.getTableName(), connection));
228                         sqlBuffer.append(" OFF");
229                         statement.execute(sqlBuffer.toString());
230                     }
231                     jdbcConnection.commit();
232                 }
233             }
234         }
235         finally
236         {
237             if(wasAutoCommit)
238             {
239                 // Reset the autocommit property
240                 jdbcConnection.setAutoCommit(true);
241             }
242             statement.close();
243         }
244     }
245 }