Fix swallowed update count after error in mixed SQL batch execution

Author: muskan124947

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

@@ -1579,7 +1579,9 @@ boolean onDone(TDSReader tdsReader) throws SQLServerException {
                     // the DML/DDL update count is not valid, and this result should be skipped
                     // unless it's for a batch where it's ok to report a "done without count"
                     // status (Statement.SUCCESS_NO_INFO)
-                    if (-1 == doneToken.getUpdateCount() && EXECUTE_BATCH != executeMethod)
+
+                    // Prevent driver from skipping a failed DONE token and losing the next statement’s update count.
+                    if (-1 == doneToken.getUpdateCount() && EXECUTE_BATCH != executeMethod && !doneToken.isError())
                         return true;
 
                     if (-1 != doneToken.getUpdateCount() && EXECUTE_QUERY == executeMethod)

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

@@ -3455,6 +3455,129 @@ public void testPreparedStatementInsertMultipleValuesWithTrigger() throws SQLExc
             }
         }
 
+        /**
+         * Tests execute a mixed SQL batch (INSERT → INSERT error → INSERT → SELECT) using
+         * {@link Statement#execute(String)} and verifies correct result traversal after
+         * a primary key violation.
+         * 
+         * After catching the expected error, the test calls {@link Statement#getMoreResults()}
+         * to continue processing the batch and validates that the update count of the
+         * subsequent successful INSERT is not swallowed before reaching the SELECT.
+         * 
+         * Validates that update counts are reported correctly even after an error occurs in the batch.
+         * Github issue #2850 : Statement.execute() skips valid update count after catching SQLException in mixed batch execution
+         *
+         * @throws SQLException
+         */
+        @Test
+        public void testBatchUpdateCount() throws SQLException {
+            // Create separate test tables to avoid conflicts with existing setup
+            String testTable = AbstractSQLGenerator.escapeIdentifier(RandomUtil.getIdentifier("UpdateCountTestTable"));
+            
+            try (Connection conn = getConnection();
+                    Statement stmt = conn.createStatement()) {
+                        
+                TestUtils.dropTableIfExists(testTable, stmt);
+
+                // Create table
+                String createTableSql = String.format("""
+                        CREATE TABLE %s (
+                            id int primary key, column_name varchar(100))
+                        """, testTable);
+
+                stmt.execute(createTableSql);
+
+                // SQL Batch breakdown:
+                // 1. INSERT (Success)
+                // 2. INSERT (Failure - Primary Key Conflict Error 2627)
+                // 3. INSERT (Success - BUT this update count gets swallowed by the driver bug)
+                // 4. SELECT (Result Set)
+                
+                String sqlBatch =
+                    "insert into " + testTable + " values (1, 'test'); " +
+                    "insert into " + testTable + " values (1, 'test'); " +
+                    "insert into " + testTable + " values (2, 'test'); " +
+                    "select * from " + testTable + ";";
+                    
+                boolean hasResult = stmt.execute(sqlBatch);
+                int resultCount = 0;
+                List<Integer> updateCounts = new ArrayList<>();
+                int resultSetCount = 0;
+                boolean exceptionCaught = false;
+
+                while (true) {
+                    resultCount++;
+                    try {
+                        // Standard JDBC processing logic
+                        if (hasResult) {
+                            try (ResultSet rs = stmt.getResultSet()) {
+                                resultSetCount++;;
+                                int rowCount = 0;
+                                while (rs.next()) {
+                                    rowCount++;
+                                }
+                                // Verify we get the expected 2 rows in the final SELECT
+                                if (resultSetCount == 1) {
+                                    assertEquals(2, rowCount, "Final SELECT should return 2 rows");
+                                }
+                            }
+                        } else {
+                            int updateCount = stmt.getUpdateCount();
+                            if (updateCount == -1) {
+                                break; // Exit loop - no more results
+                            }
+                            updateCounts.add(updateCount);
+                        }
+
+                        // Attempt to get the next result
+                        hasResult = stmt.getMoreResults();
+
+                    } catch (SQLException e) {
+                        exceptionCaught = true;
+                        assertEquals(2627, e.getErrorCode(), "Expected primary key violation error");
+                        assertTrue(e.getMessage().contains("PRIMARY KEY constraint"), 
+                                  "Expected primary key constraint violation message");
+
+                        // ================= Core Recovery Logic =================
+                        // The driver throws an exception for the 2nd SQL (Duplicate Key).
+                        // We catch it and try to move to the next result (3rd SQL).
+                        try {
+                            // Force move pointer to continue processing the batch
+                            hasResult = stmt.getMoreResults();
+                        } catch (Exception ex) {
+                            fail("Failed to recover from batch exception: " + ex.getMessage());
+                        }
+                        // =======================================================
+                    }
+                }
+
+                // Verify test results
+                assertTrue(exceptionCaught, "Expected exception for duplicate key was not caught");
+                assertEquals(1, resultSetCount, "Should have processed exactly 1 ResultSet");
+
+                // This is the key assertion that will fail with the current bug:
+                // We should get 2 update counts (first INSERT and third INSERT)
+                // But due to the bug, we only get 1 update count (the first INSERT)
+                // The third INSERT's update count gets swallowed by the driver
+                assertEquals(2, updateCounts.size(), 
+                            "Should have 2 update counts (1st INSERT success + 3rd INSERT success), " +
+                            "but driver bug causes 3rd INSERT update count to be swallowed");
+                
+                // Verify the update counts are correct
+                assertEquals(Integer.valueOf(1), updateCounts.get(0), "First INSERT should affect 1 row");
+                assertEquals(Integer.valueOf(1), updateCounts.get(1), "Third INSERT should affect 1 row");
+
+                // Verify final table state
+                try (ResultSet rs = stmt.executeQuery("SELECT COUNT(*) FROM " + testTable)) {
+                    assertTrue(rs.next());
+                    assertEquals(2, rs.getInt(1), "Table should contain exactly 2 rows");
+                }
+
+                TestUtils.dropTableIfExists(testTable, stmt);
+
+            }
+        }
+
         @AfterEach
         public void terminate() {
             try (Connection con = getConnection(); Statement stmt = con.createStatement()) {