Fixed incorrect updateCount (#2013) (#2021)

* Fixed incorrect updateCount

* Formatting changes

* Additional testing and changes

* Removed debug print line

Author: tkyc

src/main/java/com/microsoft/sqlserver/jdbc/SQLServerPreparedStatement.java

@@ -2972,6 +2972,11 @@ final void doExecutePreparedStatementBatch(PrepStmtBatchExecCmd batchCommand) th
                                 if (null == batchCommand.batchException)
                                     batchCommand.batchException = e;
 
+                                if (batchCommand.batchException.getSQLState()
+                                        .equals(SQLState.STATEMENT_CANCELED.getSQLStateCode())) {
+                                    processBatch();
+                                    continue;
+                                }
                             }
 
                             // In batch execution, we have a special update count

src/test/java/com/microsoft/sqlserver/jdbc/unit/statement/BatchExecutionTest.java

@@ -4,17 +4,22 @@
  */
 package com.microsoft.sqlserver.jdbc.unit.statement;
 
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.fail;
 
 import java.lang.reflect.Field;
+import java.sql.BatchUpdateException;
+import java.sql.CallableStatement;
 import java.sql.Connection;
 import java.sql.PreparedStatement;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
+import java.util.Arrays;
 
+import com.microsoft.sqlserver.jdbc.SQLServerPreparedStatement;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Tag;
@@ -44,8 +49,86 @@
 @Tag(Constants.xAzureSQLDW)
 public class BatchExecutionTest extends AbstractTest {
 
-    static String ctstable1;
-    static String ctstable2;
+    private static String ctstable1;
+    private static String ctstable2;
+    private static String ctstable3;
+    private static String ctstable4;
+    private static String ctstable3Procedure1;
+
+    /**
+     * This tests the updateCount when the error query does cause a SQL state HY008.
+     *
+     * @throws Exception
+     */
+    @Test
+    public void testBatchUpdateCountFalseOnFirstPstmtPrepexec() throws Exception {
+        long[] expectedUpdateCount = {1, 1, 1, 1, -3, -3, -3, -3, -3, -3};
+        testBatchUpdateCountWith(10, 6, false, "prepexec", expectedUpdateCount);
+    }
+
+    /**
+     * This tests the updateCount when the error query does cause a SQL state HY008.
+     *
+     * @throws Exception
+     */
+    @Test
+    public void testBatchUpdateCountTrueOnFirstPstmtPrepexec() throws Exception {
+        long[] expectedUpdateCount = {1, 1, -3, -3, -3};
+        testBatchUpdateCountWith(5, 4, true, "prepexec", expectedUpdateCount);
+    }
+
+    /**
+     * This tests the updateCount when the error query does cause a SQL state HY008.
+     *
+     * @throws Exception
+     */
+    @Test
+    public void testBatchUpdateCountFalseOnFirstPstmtSpPrepare() throws Exception {
+        long[] expectedUpdateCount = {1, 1, 1, 1, -3, -3, -3, -3, -3, -3};
+        testBatchUpdateCountWith(10, 6, false, "prepare", expectedUpdateCount);
+    }
+
+    /**
+     * This tests the updateCount when the error query does cause a SQL state HY008.
+     *
+     * @throws Exception
+     */
+    @Test
+    public void testBatchUpdateCountTrueOnFirstPstmtSpPrepare() throws Exception {
+        long[] expectedUpdateCount = {1, 1, -3, -3, -3};
+        testBatchUpdateCountWith(5, 4, true, "prepare", expectedUpdateCount);
+    }
+
+    /**
+     * This tests the updateCount when the error query does not cause a SQL state HY008.
+     *
+     * @throws Exception
+     */
+    @Test
+    public void testBatchUpdateCount() throws Exception {
+        long[] expectedUpdateCount = {1, 1, 1, 1, -3, 1, 1, 1, 1, 1};
+
+        try (SQLServerConnection connection = PrepUtil.getConnection(connectionString)) {
+            try (SQLServerPreparedStatement pstmt = (SQLServerPreparedStatement) connection.prepareStatement(
+                    "insert into " + AbstractSQLGenerator.escapeIdentifier(ctstable4) + " values(?)")) {
+                for (int i = 1; i <= 10; i++) {
+                    if (i == 5) {
+                        pstmt.setInt(1, -1);
+                    } else {
+                        pstmt.setInt(1, i);
+                    }
+                    pstmt.addBatch();
+                }
+
+                try {
+                    pstmt.executeBatch();
+                } catch (BatchUpdateException e) {
+                    assertArrayEquals(expectedUpdateCount, e.getLargeUpdateCounts(),
+                            "Actual: " + Arrays.toString(e.getLargeUpdateCounts()));
+                }
+            }
+        }
+    }
 
     /**
      * testAddBatch1 and testExecutionBatch one looks similar except for the parameters being passed for select query.
@@ -139,6 +222,35 @@ public void testExecuteBatch1UseBulkCopyAPI() {
         testExecuteBatch1Internal("BulkCopy");
     }
 
+    private void testBatchUpdateCountWith(int numOfInserts, int errorQueryIndex,
+            boolean prepareOnFirstPreparedStatement, String prepareMethod,
+            long[] expectedUpdateCount) throws Exception {
+        try (SQLServerConnection connection = PrepUtil.getConnection(connectionString)) {
+            connection.setEnablePrepareOnFirstPreparedStatementCall(prepareOnFirstPreparedStatement);
+            connection.setPrepareMethod(prepareMethod);
+            try (CallableStatement cstmt = connection.prepareCall(
+                    AbstractSQLGenerator.escapeIdentifier(ctstable3Procedure1) + " @duration=?, @value=?")) {
+                cstmt.setQueryTimeout(7);
+                for (int i = 1; i <= numOfInserts; i++) {
+                    if (i == errorQueryIndex) {
+                        cstmt.setString(1, "00:00:14");
+                    } else {
+                        cstmt.setString(1, "00:00:00");
+                    }
+                    cstmt.setInt(2, i);
+                    cstmt.addBatch();
+                }
+
+                try {
+                    cstmt.executeBatch();
+                } catch (BatchUpdateException e) {
+                    assertArrayEquals(expectedUpdateCount, e.getLargeUpdateCounts(),
+                            "Actual: " + Arrays.toString(e.getLargeUpdateCounts()));
+                }
+            }
+        }
+    }
+
     private void testExecuteBatch1Internal(String mode) {
         int i = 0;
         int retValue[] = {0, 0, 0};
@@ -193,6 +305,17 @@ private void testExecuteBatch1Internal(String mode) {
         }
     }
 
+    private static void createProcedure() throws SQLException {
+        String sql1 = "CREATE PROCEDURE " + AbstractSQLGenerator.escapeIdentifier(ctstable3Procedure1) + "\n"
+                + "@value int,\n" + "@duration varchar(8)\n" + "AS\n" + "BEGIN\n" + "WAITFOR DELAY @duration;\n"
+                + "INSERT INTO " + AbstractSQLGenerator.escapeIdentifier(ctstable3) + " VALUES (@value);\n" + "END";
+
+        try (Connection connection = PrepUtil.getConnection(connectionString);
+                Statement stmt = (SQLServerStatement) connection.createStatement()) {
+            stmt.execute(sql1);
+        } ;
+    }
+
     private static void createTable() throws SQLException {
         try (Connection connection = PrepUtil.getConnection(connectionString + ";columnEncryptionSetting=Enabled;");
                 Statement stmt = (SQLServerStatement) connection.createStatement()) {
@@ -201,8 +324,13 @@ private static void createTable() throws SQLException {
             String sql2 = "create table " + AbstractSQLGenerator.escapeIdentifier(ctstable2)
                     + " (KEY_ID int,  COF_NAME varchar(32),  PRICE float, TYPE_ID int, primary key(KEY_ID), foreign key(TYPE_ID) references "
                     + AbstractSQLGenerator.escapeIdentifier(ctstable1) + ")";
+            String sql3 = "create table " + AbstractSQLGenerator.escapeIdentifier(ctstable3) + "(C1 int)";
+            String sql4 = "create table " + AbstractSQLGenerator.escapeIdentifier(ctstable4)
+                    + "(C1 int check (C1 > 0))";
             stmt.execute(sql1);
             stmt.execute(sql2);
+            stmt.execute(sql3);
+            stmt.execute(sql4);
 
             String sqlin2 = "insert into " + AbstractSQLGenerator.escapeIdentifier(ctstable1)
                     + " values (1,'COFFEE-Desc')";
@@ -315,15 +443,29 @@ public static void testSetup() throws TestAbortedException, Exception {
 
         ctstable1 = RandomUtil.getIdentifier("ctstable1");
         ctstable2 = RandomUtil.getIdentifier("ctstable2");
+        ctstable3 = RandomUtil.getIdentifier("ctstable3");
+        ctstable4 = RandomUtil.getIdentifier("ctstable4");
+        ctstable3Procedure1 = RandomUtil.getIdentifier("ctstable3Procedure1");
 
         dropTable();
         createTable();
+
+        dropProcedure();
+        createProcedure();
+    }
+
+    private static void dropProcedure() throws SQLException {
+        try (Statement stmt = connection.createStatement()) {
+            TestUtils.dropProcedureIfExists(AbstractSQLGenerator.escapeIdentifier(ctstable3Procedure1), stmt);
+        }
     }
 
     private static void dropTable() throws SQLException {
         try (Statement stmt = connection.createStatement()) {
             TestUtils.dropTableIfExists(AbstractSQLGenerator.escapeIdentifier(ctstable2), stmt);
             TestUtils.dropTableIfExists(AbstractSQLGenerator.escapeIdentifier(ctstable1), stmt);
+            TestUtils.dropTableIfExists(AbstractSQLGenerator.escapeIdentifier(ctstable3), stmt);
+            TestUtils.dropTableIfExists(AbstractSQLGenerator.escapeIdentifier(ctstable4), stmt);
         }
     }