/*
    *
    * 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.
   * <p>
   * If you are using the Microsoft driver (i.e.
   * <code>com.microsoft.jdbc.sqlserver.SQLServerDriver</code>), you'll need to
   * use the <code>SelectMethod=cursor</code> parameter in the JDBC connection
   * string. Your databaseUrl would look something like the following:
   * <p>
   * <code>jdbc:microsoft:sqlserver://localhost:1433;DatabaseName=mydb;SelectMethod=cursor</code>
   * <p>
   * 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
       * <ul><li>GENERATED BY DEFAULT AS IDENTITY</li>
       * <li>GENERATED ALWAYS AS IDENTITY</li></ul>
       * set which is the SQL standard syntax to describe auto-generated key columns.
       * Also accepts columns that have the auto-increment property set to <code>true</code> (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)
                 {
                     final StringBuilder sqlBuffer = new StringBuilder(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)
                     {
                         final StringBuilder sqlBuffer = new StringBuilder(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();
         }
     }
 }