Fix PreparedStatement getGeneratedKeys() failure with insert triggers (#2740) (#2742)

* Fix PreparedStatement getGeneratedKeys() failure with triggers (#2740)

* added test case

Author: Ananya2

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

@@ -792,11 +792,14 @@ boolean onRetValue(TDSReader tdsReader) throws SQLServerException {
      * Override TDS token processing behavior for PreparedStatement.
      * For regular Statement, the execute API for INSERT requires reading an additional explicit 
      * TDS_DONE token that contains the actual update count returned by the server.
-     * PreparedStatement does not require this additional token processing.
+     * PreparedStatement does not require this additional token processing, unless
+     * generated keys were requested (which requires processing additional TDS tokens).
      */
     @Override
     protected boolean hasUpdateCountTDSTokenForInsertCmd() {
-        return false;
+        // When generated keys are requested, we need to process additional TDS tokens
+        // to properly locate the ResultSet containing the generated keys
+        return bRequestedGeneratedKeys;
     }
 
     /**

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

@@ -2797,6 +2797,31 @@ public void testPrepStmtExecuteUpdateInsertAndGenKeys() {
             }
         }
 
+        /**
+         * Tests execute using PreparedStatement for Insert followed by getGenerateKeys
+         *
+         * @throws Exception
+         */
+        @Test
+        public void testPrepStmtExecuteInsertAndGenKeys() {
+            try (Connection con = getConnection()) {
+                String sql = "INSERT INTO " + tableName + " (NAME) VALUES('test');";
+                try(PreparedStatement stmt = con.prepareStatement(sql,PreparedStatement.RETURN_GENERATED_KEYS)) {
+                    stmt.execute();
+                    int updateCount = stmt.getUpdateCount();
+                    assertEquals(updateCount, 1, "updateCount should have been 1, but received : " + updateCount);
+                    try (ResultSet generatedKeys = stmt.getGeneratedKeys()) {
+                        if (generatedKeys.next()) {
+                            int id = generatedKeys.getInt(1);
+                            assertEquals(id, 4, "id should have been 4, but received : " + id);
+                        }
+                    }
+                }
+            } catch (SQLException e) {
+                fail(TestResource.getResource("R_unexpectedException") + e.getMessage());
+            }
+        }
+
         /**
          * Tests executeUpdate using PreparedStatement for Insert followed by getGenerateKeys
          *
@@ -3324,6 +3349,62 @@ public void testPreparedStatementExecuteDelAndSelect() throws SQLException {
             }
         }
 
+        /**
+         * Tests PreparedStatement with triggers and generated keys to validate PR #2742 fix.
+         * This test validates that both update counts work correctly AND getGeneratedKeys() 
+         * works when triggers are involved.
+         *
+         * @throws SQLException
+         */
+        @Test
+        public void testPreparedStatementWithTriggersAndGeneratedKeys() throws SQLException {
+            // Create separate test tables to avoid conflicts with existing setup
+            String testTableA = AbstractSQLGenerator.escapeIdentifier(RandomUtil.getIdentifier("TriggerTestTableA"));
+            String testTableB = AbstractSQLGenerator.escapeIdentifier(RandomUtil.getIdentifier("TriggerTestTableB"));
+            String testTrigger = AbstractSQLGenerator.escapeIdentifier(RandomUtil.getIdentifier("TriggerTestTrigger"));
+
+            try (Connection conn = getConnection();
+                 Statement stmt = conn.createStatement()) {
+
+                // Cleanup any existing objects
+                TestUtils.dropTriggerIfExists(testTrigger, stmt);
+                TestUtils.dropTableIfExists(testTableB, stmt);
+                TestUtils.dropTableIfExists(testTableA, stmt);
+
+                // Create schema
+                stmt.executeUpdate("CREATE TABLE " + testTableA + " (ID int NOT NULL IDENTITY(1,1) PRIMARY KEY, NAME varchar(32))");
+                stmt.executeUpdate("CREATE TABLE " + testTableB + " (ID int NOT NULL IDENTITY(1,1) PRIMARY KEY)");
+                stmt.executeUpdate("CREATE TRIGGER " + testTrigger + " ON " + testTableA + " FOR INSERT AS "
+                        + "INSERT INTO " + testTableB + " DEFAULT VALUES");
+
+                // Insert row into TABLE_A requesting generated keys
+                String sql = "INSERT INTO " + testTableA + " (NAME) VALUES (?)";
+                try (PreparedStatement ps = conn.prepareStatement(sql, new String[]{"ID"})) {
+                    ps.setString(1, "test");
+
+                    // Execute the insert + trigger
+                    ps.execute();
+
+                    // Validate update count is correct (should be 1 for the INSERT)
+                    int updateCount = ps.getUpdateCount();
+                    assertEquals(1, updateCount, "Update count should be 1 for single INSERT");
+
+                    // Validate generated keys can be retrieved (this was broken before the fix)
+                    try (ResultSet rs = ps.getGeneratedKeys()) {
+                        assertTrue(rs.next(), "Generated keys ResultSet should have at least one row");
+                        int generatedKey = rs.getInt(1);
+                        assertTrue(generatedKey > 0, "Generated key should be a positive integer, got: " + generatedKey);
+                    }
+                }
+
+                // Cleanup
+                TestUtils.dropTriggerIfExists(testTrigger, stmt);
+                TestUtils.dropTableIfExists(testTableB, stmt);
+                TestUtils.dropTableIfExists(testTableA, stmt);
+
+            }
+        }
+
         @AfterEach
         public void terminate() {
             try (Connection con = getConnection(); Statement stmt = con.createStatement()) {