fix(sqlserver): keep table variable scope in agent batches

agents/common/src/main/java/com/dbx/agent/AbstractJdbcAgent.java

@@ -196,7 +196,9 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
             options.getMaxRows(),
             options.getFetchSize(),
             options.getTimeoutSecs(),
-            resultValueReader()
+            resultValueReader(),
+            JdbcExecutor.StatementMessageReader.NONE,
+            advancePastUpdateCounts()
         );
     }
 
@@ -211,7 +213,8 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
             this::setSchemaSQL,
             this::resetSchemaSQL,
             options,
-            resultValueReader()
+            resultValueReader(),
+            advancePastUpdateCounts()
         );
     }
 
@@ -529,6 +532,10 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
     protected void beforeQueryExecution(Connection connection, int timeoutSecs) throws Exception {
     }
 
+    protected boolean advancePastUpdateCounts() {
+        return false;
+    }
+
     protected void beforePooledConnectionReturn(Connection connection) throws Exception {
     }
 

agents/common/src/main/java/com/dbx/agent/JdbcExecutor.java

@@ -102,6 +102,34 @@ public final class JdbcExecutor {
         int timeoutSecs,
         ResultValueReader valueReader,
         StatementMessageReader statementMessageReader
+    ) {
+        return execute(
+            conn,
+            sql,
+            schema,
+            setSchemaSql,
+            resetSchemaSql,
+            maxRows,
+            fetchSize,
+            timeoutSecs,
+            valueReader,
+            statementMessageReader,
+            false
+        );
+    }
+
+    public QueryResult execute(
+        Connection conn,
+        String sql,
+        String schema,
+        Function<String, String> setSchemaSql,
+        Supplier<String> resetSchemaSql,
+        int maxRows,
+        Integer fetchSize,
+        int timeoutSecs,
+        ResultValueReader valueReader,
+        StatementMessageReader statementMessageReader,
+        boolean advancePastUpdateCounts
     ) {
         return unchecked(() -> {
             String trimmedSql = trimSql(sql);
@@ -122,17 +150,31 @@ public final class JdbcExecutor {
                 // Do not translate them to Connection.commit(), which requires autoCommit=false.
                 boolean hasResultSet = stmt.execute(trimmedSql);
                 long elapsed = System.currentTimeMillis() - start;
+                long affectedRows;
+                if (advancePastUpdateCounts) {
+                    affectedRows = 0L;
+                    while (!hasResultSet) {
+                        int updateCount = stmt.getUpdateCount();
+                        if (updateCount < 0) {
+                            break;
+                        }
+                        affectedRows += updateCount;
+                        hasResultSet = stmt.getMoreResults();
+                    }
+                } else {
+                    int updateCount = hasResultSet ? -1 : stmt.getUpdateCount();
+                    affectedRows = updateCount >= 0 ? updateCount : 0L;
+                }
                 QueryResult result;
                 if (hasResultSet) {
                     try (ResultSet rs = stmt.getResultSet()) {
                         result = readResultSet(rs, elapsed, effectiveMaxRows, valueReader);
                     }
                 } else {
-                    int updateCount = stmt.getUpdateCount();
                     result = new QueryResult(
                         Collections.emptyList(),
                         Collections.emptyList(),
-                        updateCount >= 0 ? updateCount : 0,
+                        affectedRows,
                         elapsed,
                         false
                     );
@@ -240,6 +282,30 @@ public final class JdbcExecutor {
         return executePage(conn, sql, schema, setSchemaSql, resetSchemaSql, options, valueReader, StatementMessageReader.NONE, sessions);
     }
 
+    public QueryPageResult executePage(
+        Connection conn,
+        String sql,
+        String schema,
+        Function<String, String> setSchemaSql,
+        Supplier<String> resetSchemaSql,
+        QueryPageOptions options,
+        ResultValueReader valueReader,
+        boolean advancePastUpdateCounts
+    ) {
+        return executePage(
+            conn,
+            sql,
+            schema,
+            setSchemaSql,
+            resetSchemaSql,
+            options,
+            valueReader,
+            StatementMessageReader.NONE,
+            sessions,
+            advancePastUpdateCounts
+        );
+    }
+
     public QueryPageResult startTableRead(
         Connection conn,
         String sql,
@@ -273,6 +339,32 @@ public final class JdbcExecutor {
         ResultValueReader valueReader,
         StatementMessageReader statementMessageReader,
         ConcurrentHashMap<String, QuerySession> targetSessions
+    ) {
+        return executePage(
+            conn,
+            sql,
+            schema,
+            setSchemaSql,
+            resetSchemaSql,
+            options,
+            valueReader,
+            statementMessageReader,
+            targetSessions,
+            false
+        );
+    }
+
+    private QueryPageResult executePage(
+        Connection conn,
+        String sql,
+        String schema,
+        Function<String, String> setSchemaSql,
+        Supplier<String> resetSchemaSql,
+        QueryPageOptions options,
+        ResultValueReader valueReader,
+        StatementMessageReader statementMessageReader,
+        ConcurrentHashMap<String, QuerySession> targetSessions,
+        boolean advancePastUpdateCounts
     ) {
         return unchecked(() -> {
             expireIdleSessions(targetSessions, System.currentTimeMillis(), QUERY_SESSION_IDLE_TIMEOUT_MILLIS);
@@ -293,12 +385,26 @@ public final class JdbcExecutor {
                 // JDBC transaction APIs are reserved for executeTransaction.
                 boolean hasResultSet = stmt.execute(trimmedSql);
                 long elapsed = System.currentTimeMillis() - start;
+                long affectedRows;
+                if (advancePastUpdateCounts) {
+                    affectedRows = 0L;
+                    while (!hasResultSet) {
+                        int updateCount = stmt.getUpdateCount();
+                        if (updateCount < 0) {
+                            break;
+                        }
+                        affectedRows += updateCount;
+                        hasResultSet = stmt.getMoreResults();
+                    }
+                } else {
+                    int updateCount = hasResultSet ? -1 : stmt.getUpdateCount();
+                    affectedRows = updateCount >= 0 ? updateCount : 0L;
+                }
                 if (!hasResultSet) {
-                    int updateCount = stmt.getUpdateCount();
                     QueryResult result = new QueryResult(
                         Collections.emptyList(),
                         Collections.emptyList(),
-                        updateCount >= 0 ? updateCount : 0,
+                        affectedRows,
                         elapsed,
                         false
                     );

agents/common/src/test/java/com/dbx/agent/JdbcExecutorTest.java

@@ -242,6 +242,47 @@ class JdbcExecutorTest {
         assertEquals(Arrays.asList(Arrays.asList(1, "Ada")), result.getRows());
     }
 
+    @Test
+    void executeAdvancesPastUpdateCountsToTheFollowingResultSet() {
+        CountingResultSetFixture fixture = countingResultSet(new Object[][]{{42, "Bill_Record"}});
+
+        QueryResult result = JdbcExecutor.INSTANCE.execute(
+            updateThenResultSetConnection(1, fixture.resultSet()),
+            "DECLARE @tables TABLE (id INT, name NVARCHAR(128)); INSERT INTO @tables VALUES (42, 'Bill_Record'); SELECT * FROM @tables;",
+            "",
+            schema -> "",
+            () -> "",
+            100,
+            null,
+            0,
+            JdbcExecutor.INSTANCE::defaultResultValue,
+            JdbcExecutor.StatementMessageReader.NONE,
+            true
+        );
+
+        assertEquals(Arrays.asList("id", "name"), result.getColumns());
+        assertEquals(Arrays.asList(Arrays.asList(42, "Bill_Record")), result.getRows());
+    }
+
+    @Test
+    void executePageAdvancesPastUpdateCountsToTheFollowingResultSet() {
+        CountingResultSetFixture fixture = countingResultSet(new Object[][]{{42, "Bill_Record"}});
+
+        QueryPageResult result = JdbcExecutor.INSTANCE.executePage(
+            updateThenResultSetConnection(1, fixture.resultSet()),
+            "DECLARE @tables TABLE (id INT, name NVARCHAR(128)); INSERT INTO @tables VALUES (42, 'Bill_Record'); SELECT * FROM @tables;",
+            "",
+            schema -> "",
+            () -> "",
+            new QueryPageOptions(100, null, 100),
+            JdbcExecutor.INSTANCE::defaultResultValue,
+            true
+        );
+
+        assertEquals(Arrays.asList("id", "name"), result.getColumns());
+        assertEquals(Arrays.asList(Arrays.asList(42, "Bill_Record")), result.getRows());
+    }
+
     @Test
     void executeStillPropagatesStatementErrors() {
         SQLException failure = new SQLException("permission denied", "42000", 229);
@@ -467,6 +508,42 @@ class JdbcExecutorTest {
         );
     }
 
+    private static Connection updateThenResultSetConnection(int updateCount, ResultSet resultSet) {
+        AtomicInteger resultIndex = new AtomicInteger();
+        InvocationHandler statementHandler = (Object unused, Method method, Object[] args) -> {
+            switch (method.getName()) {
+                case "execute":
+                    return false;
+                case "getUpdateCount":
+                    return resultIndex.get() == 0 ? updateCount : -1;
+                case "getMoreResults":
+                    return resultIndex.incrementAndGet() == 1;
+                case "getResultSet":
+                    return resultSet;
+                case "close":
+                    return null;
+                default:
+                    return defaultValue(method.getReturnType());
+            }
+        };
+        Statement statement = (Statement) Proxy.newProxyInstance(
+            Statement.class.getClassLoader(),
+            new Class<?>[]{Statement.class},
+            statementHandler
+        );
+        InvocationHandler connectionHandler = (Object unused, Method method, Object[] args) -> {
+            if (method.getName().equals("createStatement")) {
+                return statement;
+            }
+            return defaultValue(method.getReturnType());
+        };
+        return (Connection) Proxy.newProxyInstance(
+            Connection.class.getClassLoader(),
+            new Class<?>[]{Connection.class},
+            connectionHandler
+        );
+    }
+
     private static Object defaultValue(Class<?> type) {
         if (type == Boolean.TYPE) {
             return false;