diff --git a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerPreparedStatement.java b/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerPreparedStatement.java index 13e5c7ae5..78ed757fb 100644 --- a/src/main/java/com/microsoft/sqlserver/jdbc/SQLServerPreparedStatement.java +++ b/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; } /** diff --git a/src/test/java/com/microsoft/sqlserver/jdbc/unit/statement/StatementTest.java b/src/test/java/com/microsoft/sqlserver/jdbc/unit/statement/StatementTest.java index 5ac5116a1..78ded1edd 100644 --- a/src/test/java/com/microsoft/sqlserver/jdbc/unit/statement/StatementTest.java +++ b/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()) {