InsertIdentityOperation.java

  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;

  23. import java.sql.Connection;
  24. import java.sql.SQLException;
  25. import java.sql.Statement;

  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;

  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$
  61.  * @version $Revision$ $Date$
  62.  * @since 1.4 (Apr 9, 2002)
  63.  */
  64. public class InsertIdentityOperation extends AbstractOperation
  65. {

  67.     /**
  68.      * Logger for this class
  69.      */
  70.     private static final Logger logger = LoggerFactory.getLogger(InsertIdentityOperation.class);

  72.     public static final DatabaseOperation INSERT =
  73.             new InsertIdentityOperation(DatabaseOperation.INSERT);

  75.     public static final DatabaseOperation CLEAN_INSERT =
  76.             new CompositeOperation(DatabaseOperation.DELETE_ALL,
  77.                     new InsertIdentityOperation(DatabaseOperation.INSERT));

  79.     public static final DatabaseOperation REFRESH =
  80.             new InsertIdentityOperation(DatabaseOperation.REFRESH);

  82.     private static final IColumnFilter DEFAULT_IDENTITY_FILTER = new IColumnFilter()
  83.     {
  84.         public boolean accept(String tableName, Column column)
  85.         {
  86.             return column.getSqlTypeName().endsWith("identity");
  87.         }
  88.     };

  90.     
  91.     /**
  92.      * Accepts columns that have one of the remarks
  93.      * <ul><li>GENERATED BY DEFAULT AS IDENTITY</li>
  94.      * <li>GENERATED ALWAYS AS IDENTITY</li></ul>
  95.      * set which is the SQL standard syntax to describe auto-generated key columns.
  96.      * Also accepts columns that have the auto-increment property set to <code>true</code> (note that
  97.      * it does not yet have the ability to check whether the column is a primary key col).
  98.      * @since 2.4.3
  99.      */
  100.     public static final IColumnFilter IDENTITY_FILTER_EXTENDED = new IColumnFilter() {
  101.     
  102.         public boolean accept(String tableName, Column column) 
  103.         {
  104.             String remarks = column.getRemarks();
  105.             boolean isIdentityCol = (remarks != null) && ( 
  106.                     remarks.indexOf("GENERATED BY DEFAULT AS IDENTITY") > -1 || 
  107.                     remarks.indexOf("GENERATED ALWAYS AS IDENTITY") > -1
  108.                     );
  109.             
  110.             // If "remarks" did not give the appropriate hint, check the auto-increment property
  111.             if(!isIdentityCol)
  112.             {
  113.                 //TODO Should we ensure that the column is a PrimaryKey column?
  114.                 isIdentityCol = (AutoIncrement.YES == column.getAutoIncrement());
  115.             }
  116.             
  117.             return isIdentityCol;
  118.         }
  119.     };

  121.     
  122.     private final DatabaseOperation _operation;

  124.     /**
  125.      * Creates a new InsertIdentityOperation object that decorates the
  126.      * specified operation.
  127.      */
  128.     public InsertIdentityOperation(DatabaseOperation operation)
  129.     {
  130.         _operation = operation;
  131.     }

  133.     boolean hasIdentityColumn(ITableMetaData metaData, IDatabaseConnection connection)
  134.             throws DataSetException
  135.     {
  136.         logger.debug("hasIdentityColumn(metaData={}, connection={}) - start", metaData, connection);

  138.         DatabaseConfig config = connection.getConfig();
  139.         IColumnFilter identityFilter = (IColumnFilter)config.getProperty(
  140.                 DatabaseConfig.PROPERTY_IDENTITY_COLUMN_FILTER);
  141.         if (identityFilter == null)
  142.         {
  143.             identityFilter = DEFAULT_IDENTITY_FILTER;
  144.         }

  146.         // Verify if there is at least one identity column
  147.         Column[] columns = metaData.getColumns();
  148.         for (int i = 0; i < columns.length; i++)
  149.         {
  150.             if (identityFilter.accept(metaData.getTableName(), columns[i]))
  151.             {
  152.                 return true;
  153.             }
  154.         }

  156.         return false;
  157.     }

  159.     ////////////////////////////////////////////////////////////////////////////
  160.     // DatabaseOperation class

  162.     public void execute(IDatabaseConnection connection, IDataSet dataSet)
  163.             throws DatabaseUnitException, SQLException
  164.     {
  165.         logger.debug("execute(connection={}, dataSet={}) - start", connection, dataSet);

  167.         Connection jdbcConnection = connection.getConnection();
  168.         Statement statement = jdbcConnection.createStatement();

  170.         boolean wasAutoCommit = false;
  171.         try
  172.         {
  173.             IDataSet databaseDataSet = connection.createDataSet();
  174.             
  175.             // Note that MSSQL has a different transaction strategy from oracle.
  176.             // By default the transaction is always in "autocommit=true" so
  177.             // that every statement is immediately committed. If a dbunit
  178.             // user does not want this behavior dbunit takes it into account
  179.             // here.
  180.             
  181.             // INSERT_IDENTITY need to be enabled/disabled inside the
  182.             // same transaction
  183.             if (jdbcConnection.getAutoCommit() == true)
  184.             {
  185.                 wasAutoCommit = true;
  186.                 jdbcConnection.setAutoCommit(false);
  187.             }

  189.             // Execute decorated operation one table at a time
  190.             ITableIterator iterator = dataSet.iterator();
  191.             while(iterator.next())
  192.             {
  193.                 ITable table = iterator.getTable();
  194.                 String tableName = table.getTableMetaData().getTableName();

  196.                 ITableMetaData metaData =
  197.                         databaseDataSet.getTableMetaData(tableName);

  199.                 // enable identity insert
  200.                 boolean hasIdentityColumn = hasIdentityColumn(metaData, connection);

  202.                 if (hasIdentityColumn)
  203.                 {
  204.                     final StringBuilder sqlBuffer = new StringBuilder(128);
  205.                     sqlBuffer.append("SET IDENTITY_INSERT ");
  206.                     sqlBuffer.append(getQualifiedName(connection.getSchema(),
  207.                             metaData.getTableName(), connection));
  208.                     sqlBuffer.append(" ON");
  209.                     statement.execute(sqlBuffer.toString());
  210.                 }

  212.                 try
  213.                 {
  214.                     _operation.execute(connection, new DefaultDataSet(table));
  215.                 }
  216.                 finally
  217.                 {
  218.                     // disable identity insert
  219.                     if (hasIdentityColumn)
  220.                     {
  221.                         final StringBuilder sqlBuffer = new StringBuilder(128);
  222.                         sqlBuffer.append("SET IDENTITY_INSERT ");
  223.                         sqlBuffer.append(getQualifiedName(connection.getSchema(),
  224.                                 metaData.getTableName(), connection));
  225.                         sqlBuffer.append(" OFF");
  226.                         statement.execute(sqlBuffer.toString());
  227.                     }
  228.                     if (wasAutoCommit)
  229.                     {
  230.                         jdbcConnection.commit();
  231.                     }
  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. }
Created with JaCoCo 0.8.3.201901230119