Meta data caching batch query issue

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

@@ -927,11 +927,12 @@ private boolean reuseCachedHandle(boolean hasNewTypeDefinitions, boolean discard
 			cachedPreparedStatementHandle = null;
 		}
 
-		 // Check for new cache reference.
+        // Check for new cache reference.
         if (null == cachedPreparedStatementHandle) {
             PreparedStatementHandle cachedHandle = connection.getCachedPreparedStatementHandle(new Sha1HashKey(preparedSQL, preparedTypeDefinitions));
 
-            // If handle was found then re-use, only if AE is not on, or if it is on, make sure encryptionMetadataIsRetrieved is retrieved.
+            // If handle was found then re-use, only if AE is not on and is not a batch query with new type definitions (We shouldn't reuse handle
+            // if it is batch query and has new type definition, or if it is on, make sure encryptionMetadataIsRetrieved is retrieved.
             if (null != cachedHandle) {
                 if (!connection.isColumnEncryptionSettingEnabled()
                         || (connection.isColumnEncryptionSettingEnabled() && encryptionMetadataIsRetrieved)) {
@@ -944,7 +945,7 @@ private boolean reuseCachedHandle(boolean hasNewTypeDefinitions, boolean discard
             }
         }
         return false;
-	}
+    }
 
     private boolean doPrepExec(TDSWriter tdsWriter,
             Parameter[] params,
@@ -2585,10 +2586,9 @@ final void doExecutePreparedStatementBatch(PrepStmtBatchExecCmd batchCommand) th
             }
 
             // Retry execution if existing handle could not be re-used.
-            for(int attempt = 1; attempt <= 2; ++attempt) {
-
+            for(int attempt = 1; attempt <= 2; ++attempt) {               
                 try {
-
+                    
                     // Re-use handle if available, requires parameter definitions which are not available until here.
                     if (reuseCachedHandle(hasNewTypeDefinitions, 1 < attempt)) {
                         hasNewTypeDefinitions = false;
@@ -2648,9 +2648,13 @@ final void doExecutePreparedStatementBatch(PrepStmtBatchExecCmd batchCommand) th
 
                                 // Retry if invalid handle exception.
                                 if (retryBasedOnFailedReuseOfCachedHandle(e, attempt)) {
-                                    // Reset number of batches prepared.
+                                    // Reset number of batches prepare and reset the prepared type definitions
                                     numBatchesPrepared = numBatchesExecuted;
-                                    retry = true;
+                                    paramValues = batchParamValues.get(numBatchesPrepared);
+                                    for (int i = 0; i < paramValues.length; i++)
+                                        batchParam[i] = paramValues[i];
+                                    buildPreparedStrings(batchParam, false);
+                                    retry = true;                                    
                                     break;
                                 }
 

src/test/java/com/microsoft/sqlserver/jdbc/preparedStatement/BatchExecutionWithNullTest.java

@@ -0,0 +1,140 @@
+/*
+ * Microsoft JDBC Driver for SQL Server
+ * 
+ * Copyright(c) Microsoft Corporation All rights reserved.
+ * 
+ * This program is made available under the terms of the MIT License. See the LICENSE file in the project root for more information.
+ */
+package com.microsoft.sqlserver.jdbc.preparedStatement;
+
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assumptions.assumeTrue;
+
+import java.sql.Connection;
+import java.sql.DriverManager;
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.sql.Types;
+
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.platform.runner.JUnitPlatform;
+import org.junit.runner.RunWith;
+import org.opentest4j.TestAbortedException;
+
+import com.microsoft.sqlserver.jdbc.SQLServerStatement;
+import com.microsoft.sqlserver.testframework.AbstractTest;
+import com.microsoft.sqlserver.testframework.DBConnection;
+import com.microsoft.sqlserver.testframework.Utils;
+
+@RunWith(JUnitPlatform.class)
+public class BatchExecutionWithNullTest extends AbstractTest {
+
+    static Statement stmt = null;
+    static Connection connection = null;
+    static PreparedStatement pstmt = null;
+    static PreparedStatement pstmt1 = null;
+    static ResultSet rs = null;
+
+    /**
+     * Test with combination of setString and setNull which cause the "Violation of PRIMARY KEY constraint and internally 
+     * "Could not find prepared statement with handle X" error.
+     * @throws SQLException
+     */
+    @Test
+    public void testAddBatch2() throws SQLException {
+        // try {
+        String sPrepStmt = "insert into esimple (id, name) values (?, ?)";
+        int updateCountlen = 0;
+        int key = 42;
+
+        // this is the minimum sequence, I've found to trigger the error
+        pstmt = connection.prepareStatement(sPrepStmt);
+        pstmt.setInt(1, key++);
+        pstmt.setNull(2, Types.VARCHAR);
+        pstmt.addBatch();
+
+        pstmt.setInt(1, key++);
+        pstmt.setString(2, "FOO");
+        pstmt.addBatch();
+
+        pstmt.setInt(1, key++);
+        pstmt.setNull(2, Types.VARCHAR);
+        pstmt.addBatch();
+
+        int[] updateCount = pstmt.executeBatch();
+        updateCountlen += updateCount.length;
+
+        pstmt.setInt(1, key++);
+        pstmt.setString(2, "BAR");
+        pstmt.addBatch();
+
+        pstmt.setInt(1, key++);
+        pstmt.setNull(2, Types.VARCHAR);
+        pstmt.addBatch();
+
+        updateCount = pstmt.executeBatch();
+        updateCountlen += updateCount.length;
+
+        assertTrue(updateCountlen == 5, "addBatch does not add the SQL Statements to Batch ,call to addBatch failed");
+
+        String sPrepStmt1 = "select count(*) from esimple";
+
+        pstmt1 = connection.prepareStatement(sPrepStmt1);
+        rs = pstmt1.executeQuery();
+        rs.next();
+        assertTrue(rs.getInt(1) == 5, "affected rows does not match with batch size. Insert failed");
+        pstmt1.close();
+
+    }
+
+    /**
+     * Tests the same as addBatch2, only with AE on the connection string
+     * 
+     * @throws SQLException
+     */
+    @Test
+    public void testAddbatch2AEOnConnection() throws SQLException {
+        connection = DriverManager.getConnection(connectionString + ";columnEncryptionSetting=Enabled;");
+        testAddBatch2();
+    }
+
+    @BeforeEach
+    public void testSetup() throws TestAbortedException, Exception {
+        assumeTrue(13 <= new DBConnection(connectionString).getServerVersion(),
+                "Aborting test case as SQL Server version is not compatible with Always encrypted ");
+
+        connection = DriverManager.getConnection(connectionString);
+        SQLServerStatement stmt = (SQLServerStatement) connection.createStatement();
+        Utils.dropTableIfExists("esimple", stmt);
+        String sql1 = "create table esimple (id integer not null, name varchar(255), constraint pk_esimple primary key (id))";
+        stmt.execute(sql1);
+        stmt.close();
+    }
+
+    @AfterAll
+    public static void terminateVariation() throws SQLException {
+
+        SQLServerStatement stmt = (SQLServerStatement) connection.createStatement();
+        Utils.dropTableIfExists("esimple", stmt);
+
+        if (null != connection) {
+            connection.close();
+        }
+        if (null != pstmt) {
+            pstmt.close();
+        }
+        if (null != pstmt1) {
+            pstmt1.close();
+        }
+        if (null != stmt) {
+            stmt.close();
+        }
+        if (null != rs) {
+            rs.close();
+        }
+    }
+}

src/test/java/com/microsoft/sqlserver/jdbc/preparedStatement/RegressionTest.java

@@ -7,11 +7,16 @@
  */
 package com.microsoft.sqlserver.jdbc.preparedStatement;
 
+import static org.junit.jupiter.api.Assertions.fail;
+
 import java.sql.Connection;
 import java.sql.DriverManager;
 import java.sql.PreparedStatement;
+import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
+import java.util.LinkedHashMap;
+import java.util.Map;
 
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
@@ -21,8 +26,6 @@
 import com.microsoft.sqlserver.testframework.AbstractTest;
 import com.microsoft.sqlserver.testframework.Utils;
 
-import static org.junit.jupiter.api.Assertions.fail;
-
 /**
  * Tests with sql queries using preparedStatement without parameters
  * 
@@ -201,6 +204,115 @@ public void grantTest() throws SQLException {
         }
     }
 
+    /**
+     * Test with large string and batch
+     * 
+     * @throws SQLException
+     */
+    @Test
+    public void batchWithLargeStringTest() throws SQLException {
+        Statement stmt = con.createStatement();
+        PreparedStatement pstmt = null;
+        ResultSet rs = null;
+        Utils.dropTableIfExists("TEST_TABLE", stmt);
+
+        con.setAutoCommit(false);
+
+        // create a table with two columns
+        boolean createPrimaryKey = false;
+        try {
+            stmt.execute("if object_id('TEST_TABLE', 'U') is not null\ndrop table TEST_TABLE;");
+            if (createPrimaryKey) {
+                stmt.execute("create table TEST_TABLE ( ID int, DATA nvarchar(max), primary key (ID) );");
+            }
+            else {
+                stmt.execute("create table TEST_TABLE ( ID int, DATA nvarchar(max) );");
+            }
+        }
+        catch (Exception e) {
+            fail(e.toString());
+        }
+
+        con.commit();
+
+        // build a String with 4001 characters
+        StringBuilder stringBuilder = new StringBuilder();
+        for (int i = 0; i < 4001; i++) {
+            stringBuilder.append('c');
+        }
+        String largeString = stringBuilder.toString();
+
+        String[] values = {"a", "b", largeString, "d", "e"};
+        // insert five rows into the table; use a batch for each row
+        try {
+            pstmt = con.prepareStatement("insert into TEST_TABLE values (?,?)");
+            // 0,a
+            pstmt.setInt(1, 0);
+            pstmt.setNString(2, values[0]);
+            pstmt.addBatch();
+
+            // 1,b
+            pstmt.setInt(1, 1);
+            pstmt.setNString(2, values[1]);
+            pstmt.addBatch();
+
+            // 2,ccc...
+            pstmt.setInt(1, 2);
+            pstmt.setNString(2, values[2]);
+            pstmt.addBatch();
+
+            // 3,d
+            pstmt.setInt(1, 3);
+            pstmt.setNString(2, values[3]);
+            pstmt.addBatch();
+
+            // 4,e
+            pstmt.setInt(1, 4);
+            pstmt.setNString(2, values[4]);
+            pstmt.addBatch();
+
+            pstmt.executeBatch();
+        }
+        catch (Exception e) {
+            fail(e.toString());
+        }
+        connection.commit();
+
+        // check the data in the table
+        Map<Integer, String> selectedValues = new LinkedHashMap<>();
+        int id = 0;
+        try {
+            pstmt = con.prepareStatement("select * from TEST_TABLE;");
+            try {
+                rs = pstmt.executeQuery();
+                int i = 0;
+                while (rs.next()) {
+                    id = rs.getInt(1);
+                    String data = rs.getNString(2);
+                    if (selectedValues.containsKey(id)) {
+                        fail("Found duplicate id: " + id + " ,actual values is : " + values[i++] + " data is: " + data);
+                    }
+                    selectedValues.put(id, data);
+                }
+            }
+            finally {
+                if (null != rs) {
+                    rs.close();
+                }
+            }
+        }
+        finally {
+            Utils.dropTableIfExists("TEST_TABLE", stmt);
+            if (null != pstmt) {
+                pstmt.close();
+            }
+            if (null != stmt) {
+                stmt.close();
+            }
+        }
+
+    }
+
     /**
      * Cleanup after test
      * 
@@ -210,6 +322,7 @@ public void grantTest() throws SQLException {
     public static void cleanup() throws SQLException {
         Statement stmt = con.createStatement();
         Utils.dropTableIfExists("x", stmt);
+        Utils.dropTableIfExists("TEST_TABLE", stmt);
         if (null != stmt) {
             stmt.close();
         }