diff --git a/agents/common/src/main/java/com/dbx/agent/AbstractJdbcAgent.java b/agents/common/src/main/java/com/dbx/agent/AbstractJdbcAgent.java index 87b4732428..0f65d08c6f 100644 --- a/agents/common/src/main/java/com/dbx/agent/AbstractJdbcAgent.java +++ b/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 { } diff --git a/agents/common/src/main/java/com/dbx/agent/JdbcExecutor.java b/agents/common/src/main/java/com/dbx/agent/JdbcExecutor.java index d135da8a90..09ccfebef4 100644 --- a/agents/common/src/main/java/com/dbx/agent/JdbcExecutor.java +++ b/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 setSchemaSql, + Supplier 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 setSchemaSql, + Supplier 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 targetSessions + ) { + return executePage( + conn, + sql, + schema, + setSchemaSql, + resetSchemaSql, + options, + valueReader, + statementMessageReader, + targetSessions, + false + ); + } + + private QueryPageResult executePage( + Connection conn, + String sql, + String schema, + Function setSchemaSql, + Supplier resetSchemaSql, + QueryPageOptions options, + ResultValueReader valueReader, + StatementMessageReader statementMessageReader, + ConcurrentHashMap 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 ); diff --git a/agents/common/src/test/java/com/dbx/agent/JdbcExecutorTest.java b/agents/common/src/test/java/com/dbx/agent/JdbcExecutorTest.java index 9244b6af4d..6f19c26393 100644 --- a/agents/common/src/test/java/com/dbx/agent/JdbcExecutorTest.java +++ b/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; diff --git a/agents/drivers/sqlserver-legacy/src/main/java/com/dbx/agent/sqlserverlegacy/SqlServerLegacyAgent.java b/agents/drivers/sqlserver-legacy/src/main/java/com/dbx/agent/sqlserverlegacy/SqlServerLegacyAgent.java index c31090d506..9d2cf591ab 100644 --- a/agents/drivers/sqlserver-legacy/src/main/java/com/dbx/agent/sqlserverlegacy/SqlServerLegacyAgent.java +++ b/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; diff --git a/agents/scripts/driver_release_packages_test.py b/agents/scripts/driver_release_packages_test.py index 3323e280a9..31b2162526 100644 --- a/agents/scripts/driver_release_packages_test.py +++ b/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", diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index 8e128216a7..4b679a6f1b 100644 --- a/crates/dbx-core/src/query.rs +++ b/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( - || 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 execution_plan = query_execution_plan(sql, db_type, is_sqlserver_agent); 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, + 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, 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) { + 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) { @@ -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