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;

agents/drivers/sqlserver-legacy/src/main/java/com/dbx/agent/sqlserverlegacy/SqlServerLegacyAgent.java

@@ -125,6 +125,11 @@ public final class SqlServerLegacyAgent extends ConfiguredJdbcAgent {
         return "SELECT 1";
     }
 
+    @Override
+    protected boolean advancePastUpdateCounts() {
+        return true;
+    }
+
     @Override
     protected void afterDisconnect() {
         sqlServer2000Mode = false;

agents/scripts/driver_release_packages_test.py

@@ -50,6 +50,7 @@ class DriverReleasePackagesTest(unittest.TestCase):
                 "zookeeper": "0.1.0",
                 "cassandra": "0.1.37",
                 "hive": "0.1.43",
+                "argo": "0.1.0",
                 "tdengine": "0.1.0",
                 "etcd": "0.1.40",
                 "etcd2": "0.1.0",

crates/dbx-core/src/query.rs

@@ -2929,7 +2929,13 @@ pub async fn execute_multi_core_with_options_for_client_and_progress_typed(
     state.touch_pool_activity(&pool_key).await;
     let _activity_touch = state.pool_activity_touch(pool_key.as_str());
 
-    let is_sqlserver = { matches!(state.pool_handle(&pool_key).await, Some(PoolKind::SqlServer(_))) };
+    let (is_sqlserver, is_sqlserver_agent) = {
+        match state.pool_handle(&pool_key).await {
+            Some(PoolKind::SqlServer(_)) => (true, false),
+            Some(PoolKind::Agent(_)) if db_type == Some(DatabaseType::SqlServer) => (false, true),
+            _ => (false, false),
+        }
+    };
 
     if is_sqlserver {
         return execute_multi_sqlserver(state, &pool_key, sql, cancel_token, options).await.map_err(Into::into);
@@ -2961,10 +2967,7 @@ pub async fn execute_multi_core_with_options_for_client_and_progress_typed(
         );
     }
 
-    let execution_plan = db_type.map_or_else(
+    let execution_plan = query_execution_plan(sql, db_type, is_sqlserver_agent);
-        || crate::sql::SqlExecutionPlan { statements: split_sql_statements(sql), stop_on_error: false },
-        |db_type| crate::sql::sql_execution_plan_for_database(sql, db_type),
-    );
     let continue_on_error = options.continue_on_error && !execution_plan.stop_on_error;
     let statements = execution_plan.statements;
     if statements.is_empty() {
@@ -3103,6 +3106,21 @@ pub async fn execute_multi_core_with_options_for_client_and_progress_typed(
     Ok(results)
 }
 
+fn query_execution_plan(
+    sql: &str,
+    db_type: Option<DatabaseType>,
+    preserve_sqlserver_batches: bool,
+) -> crate::sql::SqlExecutionPlan {
+    if preserve_sqlserver_batches && db_type == Some(DatabaseType::SqlServer) {
+        return crate::sql::SqlExecutionPlan { statements: split_sql_batches(sql), stop_on_error: false };
+    }
+
+    db_type.map_or_else(
+        || crate::sql::SqlExecutionPlan { statements: split_sql_statements(sql), stop_on_error: false },
+        |db_type| crate::sql::sql_execution_plan_for_database(sql, db_type),
+    )
+}
+
 fn single_statement_multi_result(
     result: Result<db::QueryResult, QueryExecutionError>,
     table_data_preview: bool,
@@ -6211,6 +6229,88 @@ for line in sys.stdin:
         let _ = std::fs::remove_dir_all(dir);
     }
 
+    async fn sqlserver_agent_echo_state(
+    ) -> (AppState, std::path::PathBuf, std::sync::Arc<crate::db::agent_driver::AgentRuntimeClient>) {
+        let dir = std::env::temp_dir().join(format!("dbx-query-sqlserver-agent-{}", uuid::Uuid::new_v4()));
+        std::fs::create_dir_all(&dir).unwrap();
+        let script_path = dir.join("agent.py");
+        std::fs::write(
+            &script_path,
+            r#"import json, sys
+print(json.dumps({'ready': True}), flush=True)
+for line in sys.stdin:
+    req = json.loads(line)
+    if req['method'] == 'handshake':
+        result = {'protocolVersion': 2, 'agentProtocolVersion': 2, 'capabilities': ['multi_session']}
+    elif req['method'] == 'execute_query':
+        result = {
+            'columns': ['sql'], 'column_types': ['nvarchar'], 'column_sortables': [],
+            'rows': [[req['params']['sql']]], 'affected_rows': 0, 'execution_time_ms': 1,
+            'truncated': False, 'session_id': None, 'has_more': False
+        }
+    else:
+        result = {}
+    print(json.dumps({'jsonrpc': '2.0', 'id': req['id'], 'result': result}), flush=True)
+"#,
+        )
+        .unwrap();
+
+        let python = if cfg!(windows) { "python" } else { "python3" };
+        let runtime = crate::db::agent_driver::AgentRuntimeClient::spawn(
+            crate::db::agent_driver::AgentLaunchSpec::new(python)
+                .with_args([script_path.to_string_lossy().to_string()]),
+            "test",
+        )
+        .await
+        .unwrap();
+        runtime.increment_session_count();
+
+        let storage = Storage::open(&dir.join("storage.db")).await.unwrap();
+        let state = AppState::new(storage);
+        state.configs.write().await.insert("conn-1".to_string(), test_connection_config(DatabaseType::SqlServer));
+        state
+            .update_connection_pools(|connections| {
+                connections.insert(
+                    "conn-1".to_string(),
+                    PoolKind::agent(crate::db::agent_driver::AgentDriverClient::shared_session(
+                        runtime.clone(),
+                        "session-1".to_string(),
+                    )),
+                );
+            })
+            .await;
+
+        (state, dir, runtime)
+    }
+
+    #[tokio::test]
+    async fn sqlserver_agent_multi_execution_sends_table_variable_script_as_one_batch() {
+        let (state, dir, runtime) = sqlserver_agent_echo_state().await;
+        let sql = "DECLARE @TargetTables TABLE (TableName NVARCHAR(128));\n\
+                   INSERT INTO @TargetTables VALUES ('Bill_Record');\n\
+                   SELECT TableName FROM @TargetTables;";
+
+        let results = execute_multi_core_with_options_for_client_and_progress_typed(
+            &state,
+            "conn-1",
+            "",
+            sql,
+            None,
+            None,
+            QueryExecutionOptions::default(),
+            None,
+        )
+        .await
+        .unwrap();
+
+        assert_eq!(results.len(), 1);
+        assert_eq!(results[0].result.rows, vec![vec![serde_json::Value::String(sql.to_string())]]);
+
+        runtime.kill();
+        drop(state);
+        let _ = std::fs::remove_dir_all(dir);
+    }
+
     async fn agent_error_state(
         disposition: &str,
     ) -> (AppState, std::path::PathBuf, std::sync::Arc<crate::db::agent_driver::AgentRuntimeClient>) {
@@ -7312,6 +7412,33 @@ for line in sys.stdin:
         assert_eq!(serialized.get("server_message"), Some(&serde_json::Value::Bool(true)));
     }
 
+    #[test]
+    fn sqlserver_agent_execution_plan_preserves_table_variable_batch() {
+        let sql = "DECLARE @TargetTables TABLE (TableName NVARCHAR(128));\n\
+                   INSERT INTO @TargetTables VALUES ('Bill_Record');\n\
+                   SELECT TableName FROM @TargetTables;";
+
+        let plan = query_execution_plan(sql, Some(DatabaseType::SqlServer), true);
+
+        assert_eq!(plan.statements, vec![sql]);
+    }
+
+    #[test]
+    fn sqlserver_agent_execution_plan_splits_only_on_go() {
+        let sql = "DECLARE @x TABLE (id INT);\nINSERT INTO @x VALUES (1);\nGO\nSELECT 2;";
+
+        let plan = query_execution_plan(sql, Some(DatabaseType::SqlServer), true);
+
+        assert_eq!(plan.statements, vec!["DECLARE @x TABLE (id INT);\nINSERT INTO @x VALUES (1);", "SELECT 2;"]);
+    }
+
+    #[test]
+    fn ordinary_agent_execution_plan_still_splits_semicolon_statements() {
+        let plan = query_execution_plan("SELECT 1; SELECT 2;", Some(DatabaseType::Dameng), false);
+
+        assert_eq!(plan.statements, vec!["SELECT 1", "SELECT 2"]);
+    }
+
     // Regression test for #6097: SQL Server queries share a single mutex-guarded
     // connection (see PoolKind::SqlServer), so a fast query can queue for seconds
     // behind another operation (e.g. autocomplete/schema metadata) holding that