Fixed incorrect updateCount (#2013)

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);
  }
  }