use futures::StreamExt; use mysql_async::prelude::Queryable; use serde::{Deserialize, Serialize}; use sqlparser::ast::{ visit_relations_mut, Ident, ObjectName, ObjectNamePart, ObjectType, Statement, TableFactor, VisitMut, VisitorMut, }; use sqlparser::dialect::{GenericDialect, PostgreSqlDialect}; use sqlparser::parser::Parser; use std::collections::{HashMap, HashSet}; use std::future::Future; use std::ops::ControlFlow; use std::sync::{ atomic::{AtomicU64, Ordering}, Arc, }; use std::time::Duration; use tokio::time::timeout; use tokio_util::sync::CancellationToken; use crate::agent_recovery::{RecoveryDecision, RecoveryPolicy, RecoveryScope}; use crate::connection::{AppState, PoolKind, TransactionSession, TxnConnection}; use crate::database_capabilities; use crate::db; use crate::db::agent_driver::{AgentCallError, AgentErrorStage, AgentOperationOutcome}; use crate::models::connection::{ConnectionConfig, DatabaseType}; use crate::query_execution_sql::{is_write_sql, strip_sql_comments_and_literals}; use crate::sql::{split_sql_batches, split_sql_statements, starts_with_executable_sql_keyword_for_database}; use crate::sql_dialect::{resolve_for_db, CAP_TRANSACTIONAL_DDL}; use crate::sql_risk::{classify_sql_risk_for_database, SqlRisk}; pub const QUERY_TIMEOUT: Duration = Duration::from_secs(30); pub const MAX_ROWS: usize = 10000; pub const AGENT_PROTOCOL_MAX_ROWS: usize = i32::MAX as usize; pub const QUERY_CANCELED: &str = "Query canceled"; pub const METADATA_POOL_BUSY_ERROR: &str = "DBX metadata pool is busy; please retry"; /// Returns true when a metadata request failed because all pool/client slots /// were temporarily occupied. This is deliberately separate from /// `is_connection_error`: a saturated pool is healthy and must not trigger a /// shared-pool reconnect. pub fn is_pool_saturation_error(err: &str) -> bool { let lower = err.to_lowercase(); lower.contains("connection pool checkout timed out [stage=wait") || lower.contains("dbx metadata pool is busy") } /// Fallback when a Mongo connection hits the generic SQL executor instead of the shell path. /// Wording must match packages/mongo-shell `MONGO_SHELL_COMMAND_HINT` /// (desktop/CLI diagnose first; this is only the Rust SQL-executor backstop). const MONGO_SHELL_COMMAND_HINT: &str = "Use MongoDB shell-style commands, for example: db.collection.find({}).limit(100), db.collection.aggregate([]), db.collection.aggregate([], { explain: true }), db.version(), db.collection.countDocuments({}), db.collection.distinct(\"field\"), db.collection.getIndexes(), db.collection.createIndex({...}), db.createUser({...}), or db.collection.insertOne({...})."; const SQL_OMITTED_ERROR_CONTEXT: &str = "SQL text omitted from user-facing error; enable debug SQL diagnostics to inspect the original statement."; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum PoolErrorAction { Keep, Discard, ReconnectAndRetry, } #[derive(Debug, Clone)] pub enum QueryExecutionError { Agent(AgentCallError), DuckDb { code: String, message: String }, Canceled { stage: AgentErrorStage, operation_outcome: AgentOperationOutcome }, Timeout(String), Sql(String), Legacy(String), } impl QueryExecutionError { pub fn into_legacy_string(self) -> String { match self { Self::Agent(error) => error.into_legacy_string(), Self::DuckDb { message, .. } => message, Self::Canceled { .. } => canceled_error(), Self::Timeout(error) => error, Self::Sql(error) => error, Self::Legacy(error) => error, } } pub fn into_backend_error(self) -> crate::backend_error::BackendError { match self { Self::Agent(error) => crate::backend_error::BackendError::from_agent_call_error(&error), Self::DuckDb { code, message } => { crate::backend_error::BackendError::from_duckdb_worker_error(&code, &message) } Self::Canceled { stage, operation_outcome } => { crate::backend_error::BackendError::from_canceled(stage, operation_outcome) } Self::Timeout(error) => crate::backend_error::BackendError::from_timeout_detail(&error), Self::Sql(error) => crate::backend_error::BackendError::from_sql_detail(&error), Self::Legacy(error) => crate::backend_error::BackendError::from_legacy_string(&error), } } fn with_omitted_sql_context(self, sql: &str) -> Self { match self { Self::Agent(error) => Self::Agent(error), Self::DuckDb { code, message } => { Self::DuckDb { code, message: query_error_with_omitted_sql_context(&message, sql) } } canceled @ Self::Canceled { .. } => canceled, Self::Timeout(error) => Self::Timeout(query_error_with_omitted_sql_context(&error, sql)), Self::Sql(error) => Self::Sql(append_typed_sql_error_context(&error, sql)), Self::Legacy(error) => Self::Legacy(query_error_with_omitted_sql_context(&error, sql)), } } fn with_context(self, context: &str) -> Self { match self { Self::Agent(error) => Self::Agent(error), Self::DuckDb { code, message } => Self::DuckDb { code, message: format!("{message}; {context}") }, canceled @ Self::Canceled { .. } => canceled, Self::Timeout(error) => Self::Timeout(format!("{error}; {context}")), Self::Sql(error) => Self::Sql(format!("{error}; {context}")), Self::Legacy(error) => Self::Legacy(format!("{error}; {context}")), } } fn as_agent_error(&self) -> Option<&AgentCallError> { match self { Self::Agent(error) => Some(error), Self::DuckDb { .. } | Self::Canceled { .. } | Self::Timeout(_) | Self::Sql(_) | Self::Legacy(_) => None, } } } impl std::fmt::Display for QueryExecutionError { fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { Self::Agent(error) => error.fmt(formatter), Self::DuckDb { message, .. } => formatter.write_str(message), Self::Canceled { .. } => formatter.write_str(QUERY_CANCELED), Self::Timeout(error) => formatter.write_str(error), Self::Sql(error) => formatter.write_str(error), Self::Legacy(error) => formatter.write_str(error), } } } impl From for QueryExecutionError { fn from(error: AgentCallError) -> Self { Self::Agent(error) } } impl From for QueryExecutionError { fn from(error: String) -> Self { Self::Legacy(error) } } impl From<&str> for QueryExecutionError { fn from(error: &str) -> Self { Self::Legacy(error.to_string()) } } fn query_error_with_omitted_sql_context(error: &str, _sql: &str) -> String { crate::db::agent_driver::append_legacy_error_context(error, SQL_OMITTED_ERROR_CONTEXT) } fn append_typed_sql_error_context(error: &str, _sql: &str) -> String { if error.contains(SQL_OMITTED_ERROR_CONTEXT) { return error.to_string(); } let separator = if error.trim_start().starts_with("Server error:") { " " } else { "\n" }; format!("{error}{separator}{SQL_OMITTED_ERROR_CONTEXT}") } /// A multi-statement result with metadata intended for query clients. /// /// `execution_error` is emitted for synthesized per-statement errors so clients /// can distinguish them from a successful result column named `Error`. /// `statement_index` is emitted only after a concrete statement starts running. /// `server_message` is emitted only for SQL Server TDS informational messages. #[derive(Debug, Clone, Serialize)] pub struct ExecuteMultiResult { #[serde(flatten)] pub result: db::QueryResult, #[serde(default, skip_serializing_if = "Vec::is_empty")] pub large_value_cells: Vec, #[serde(skip_serializing_if = "is_false")] pub execution_error: bool, #[serde(skip_serializing_if = "Option::is_none")] pub statement_index: Option, #[serde(skip_serializing_if = "Option::is_none")] pub error: Option, #[serde(skip_serializing_if = "is_false")] pub server_message: bool, } #[derive(Debug, Clone, PartialEq)] pub struct ExecuteMultiProgress { pub statement_index: usize, pub completed: usize, pub total: usize, pub success: bool, pub execution_time_ms: u128, pub affected_rows: u64, pub error: Option, } pub type ExecuteMultiProgressCallback = Arc; fn report_execute_multi_progress( progress: Option<&ExecuteMultiProgressCallback>, statement_index: usize, total: usize, result: &db::QueryResult, success: bool, error: Option, ) { if let Some(progress) = progress { progress(ExecuteMultiProgress { statement_index, completed: statement_index + 1, total, success, execution_time_ms: result.execution_time_ms, affected_rows: result.affected_rows, error, }); } } impl ExecuteMultiResult { fn execution_error(result: db::QueryResult) -> Self { let error = error_from_query_result(&result); Self { result, large_value_cells: Vec::new(), execution_error: true, statement_index: None, error, server_message: false, } } fn execution_error_with_index(result: db::QueryResult, statement_index: usize) -> Self { let error = error_from_query_result(&result); Self { result, large_value_cells: Vec::new(), execution_error: true, statement_index: Some(statement_index), error, server_message: false, } } fn execution_error_with_backend( result: db::QueryResult, statement_index: Option, error: crate::backend_error::BackendError, ) -> Self { Self { result, large_value_cells: Vec::new(), execution_error: true, statement_index, error: Some(error), server_message: false, } } fn success_with_index(result: db::QueryResult, statement_index: usize) -> Self { Self { result, large_value_cells: Vec::new(), execution_error: false, statement_index: Some(statement_index), error: None, server_message: false, } } fn success_with_index_and_large_values( mut result: db::QueryResult, statement_index: usize, mut large_value_cells: Vec, table_data_preview: bool, ) -> Self { let server_cells = if table_data_preview { remap_large_value_cells_around_server_markers(&result, &mut large_value_cells); extract_server_large_value_markers(&mut result) } else { Vec::new() }; Self { result, large_value_cells: merge_large_value_cells(large_value_cells, server_cells), execution_error: false, statement_index: Some(statement_index), error: None, server_message: false, } } fn success_with_index_and_optional_server_large_values( result: db::QueryResult, statement_index: usize, table_data_preview: bool, ) -> Self { Self::success_with_index_and_large_values(result, statement_index, Vec::new(), table_data_preview) } fn success_with_optional_server_large_values(mut result: db::QueryResult, table_data_preview: bool) -> Self { let large_value_cells = if table_data_preview { extract_server_large_value_markers(&mut result) } else { Vec::new() }; Self { result, large_value_cells, execution_error: false, statement_index: None, error: None, server_message: false, } } pub fn without_error_detail(mut self) -> Self { self.error = self.error.map(crate::backend_error::BackendError::without_detail); self } fn into_query_result(self) -> db::QueryResult { self.result } } const SERVER_LARGE_VALUE_UNKNOWN_BYTES: usize = 64 * 1024 * 1024; #[derive(Clone, Copy, PartialEq, Eq)] enum ServerLargeValuePreviewKind { Text, Binary, Vector, Deferred, } #[derive(Clone, Copy)] struct ServerLargeValueMarker { result_index: usize, source_index: usize, preview_kind: Option, source_type: Option<&'static str>, } #[derive(Clone, Copy)] struct ServerLargeValueMarkerValue { kind: ServerLargeValuePreviewKind, preview_size: usize, original_bytes: Option, } fn server_large_value_alias( suffix: &str, ) -> Option<(usize, Option, Option<&'static str>)> { if let Ok(source_index) = suffix.parse::() { return Some((source_index, None, None)); } let (kind, source_index) = suffix.split_once('_')?; let source_index = source_index.parse::().ok()?; let (preview_kind, source_type) = match kind { "T" => (ServerLargeValuePreviewKind::Text, None), "B" => (ServerLargeValuePreviewKind::Binary, None), "V" => (ServerLargeValuePreviewKind::Vector, Some("vector")), "J" => (ServerLargeValuePreviewKind::Text, Some("json")), "K" => (ServerLargeValuePreviewKind::Text, Some("jsonb")), "S" => (ServerLargeValuePreviewKind::Text, Some("tsvector")), "C" => (ServerLargeValuePreviewKind::Deferred, Some("clob")), "N" => (ServerLargeValuePreviewKind::Deferred, Some("nclob")), "L" => (ServerLargeValuePreviewKind::Deferred, Some("blob")), "F" => (ServerLargeValuePreviewKind::Deferred, Some("bfile")), _ => return None, }; Some((source_index, Some(preview_kind), source_type)) } fn server_large_value_marker(value: &serde_json::Value) -> Option { let mut parts = value.as_str()?.split(':'); let kind = parts.next()?; let preview_size = parts.next()?; let kind = match kind { "T" => ServerLargeValuePreviewKind::Text, "B" => ServerLargeValuePreviewKind::Binary, "V" => ServerLargeValuePreviewKind::Vector, "D" => ServerLargeValuePreviewKind::Deferred, _ => return None, }; let original_bytes = parts.next().and_then(|value| value.parse::().ok()); Some(ServerLargeValueMarkerValue { kind, preview_size: preview_size.parse::().ok()?.max(1), original_bytes }) } fn truncate_server_large_value_preview( value: &mut serde_json::Value, kind: ServerLargeValuePreviewKind, preview_size: usize, ) -> bool { let serde_json::Value::String(text) = value else { return false; }; if matches!(kind, ServerLargeValuePreviewKind::Deferred) { return true; } if matches!(kind, ServerLargeValuePreviewKind::Vector) { let truncated = text.chars().count() > preview_size; let vector_text = if truncated { let prefix: String = text.chars().take(preview_size).collect(); let Some(last_separator) = prefix.rfind(',') else { return false; }; format!("{}]", &prefix[..last_separator]) } else { text.clone() }; let Ok(vector) = serde_json::from_str::>(&vector_text) else { return false; }; *value = serde_json::Value::Array(vector); return truncated; } let truncate_at = match kind { ServerLargeValuePreviewKind::Text => text.char_indices().nth(preview_size).map(|(index, _)| index), ServerLargeValuePreviewKind::Binary => text .strip_prefix("0x") .filter(|hex| hex.len() > preview_size.saturating_mul(2)) .map(|_| 2usize.saturating_add(preview_size.saturating_mul(2))), ServerLargeValuePreviewKind::Vector | ServerLargeValuePreviewKind::Deferred => unreachable!(), }; let Some(truncate_at) = truncate_at else { return false; }; text.truncate(truncate_at); text.push_str("..."); true } fn server_large_value_markers(result: &db::QueryResult) -> Vec { let mut markers = Vec::new(); for (result_index, column) in result.columns.iter().enumerate() { let Some((source_index, preview_kind, source_type)) = column .strip_prefix(crate::sql_dialect::DBX_LARGE_VALUE_BYTES_COLUMN_PREFIX) .and_then(server_large_value_alias) else { continue; }; let expected_source_index = result_index.checked_sub(markers.len() + 1); if expected_source_index == Some(source_index) { markers.push(ServerLargeValueMarker { result_index, source_index, preview_kind, source_type }); } } markers } fn extract_server_large_value_markers(result: &mut db::QueryResult) -> Vec { let markers = server_large_value_markers(result); if markers.is_empty() { return Vec::new(); } let mut large_value_cells = Vec::new(); for marker in &markers { let source_result_index = marker.result_index.saturating_sub(1); let row_kind = result .rows .iter() .find_map(|row| row.get(marker.result_index).and_then(server_large_value_marker).map(|value| value.kind)); let source_type = marker.source_type.or_else(|| { (marker.preview_kind.or(row_kind) == Some(ServerLargeValuePreviewKind::Vector)).then_some("vector") }); if let (Some(source_type), Some(column_type)) = (source_type, result.column_types.get_mut(source_result_index)) { *column_type = source_type.to_string(); } } for (row_index, row) in result.rows.iter_mut().enumerate() { for marker in &markers { let marker_value = row.get(marker.result_index).and_then(server_large_value_marker); let source_result_index = marker.result_index.saturating_sub(1); if marker_value.is_some_and(|value| { row.get_mut(source_result_index) .is_some_and(|source| truncate_server_large_value_preview(source, value.kind, value.preview_size)) }) { large_value_cells.push(db::LargeValueCell { row_index, column_index: marker.source_index, original_bytes: marker_value .and_then(|value| value.original_bytes) .unwrap_or(SERVER_LARGE_VALUE_UNKNOWN_BYTES), }); } } } let removed: std::collections::HashSet = markers.iter().map(|marker| marker.result_index).collect(); let retained_index = |index: usize| index - removed.iter().filter(|removed_index| **removed_index < index).count(); result.spatial_columns.retain_mut(|column| { if removed.contains(&column.column_index) { return false; } column.column_index = retained_index(column.column_index); true }); for row in &mut result.rows { let mut index = 0; row.retain(|_| { let keep = !removed.contains(&index); index += 1; keep }); } for row in &mut result.spatial_values { let mut index = 0; row.retain(|_| { let keep = !removed.contains(&index); index += 1; keep }); } let mut index = 0; result.columns.retain(|_| { let keep = !removed.contains(&index); index += 1; keep }); let mut index = 0; result.column_types.retain(|_| { let keep = !removed.contains(&index); index += 1; keep }); let mut index = 0; result.column_sortables.retain(|_| { let keep = !removed.contains(&index); index += 1; keep }); large_value_cells } fn remap_large_value_cells_around_server_markers(result: &db::QueryResult, cells: &mut Vec) { let removed: std::collections::HashSet = server_large_value_markers(result).into_iter().map(|marker| marker.result_index).collect(); if removed.is_empty() { return; } cells.retain_mut(|cell| { if removed.contains(&cell.column_index) { return false; } cell.column_index -= removed.iter().filter(|index| **index < cell.column_index).count(); true }); } fn merge_large_value_cells( mut driver_cells: Vec, server_cells: Vec, ) -> Vec { if driver_cells.is_empty() { return server_cells; } if server_cells.is_empty() { return driver_cells; } let mut driver_indexes = driver_cells .iter() .enumerate() .map(|(index, cell)| ((cell.row_index, cell.column_index), index)) .collect::>(); for server_cell in server_cells { let key = (server_cell.row_index, server_cell.column_index); if let Some(index) = driver_indexes.get(&key).copied() { driver_cells[index] = server_cell; } else { driver_indexes.insert(key, driver_cells.len()); driver_cells.push(server_cell); } } driver_cells } impl From for ExecuteMultiResult { fn from(result: db::QueryResult) -> Self { Self { result, large_value_cells: Vec::new(), execution_error: false, statement_index: None, error: None, server_message: false, } } } impl From for ExecuteMultiResult { fn from(result: db::sqlserver::SqlServerBatchResult) -> Self { Self { result: result.result, large_value_cells: Vec::new(), execution_error: false, statement_index: None, error: None, server_message: result.server_message, } } } fn is_false(value: &bool) -> bool { !*value } fn error_from_query_result(result: &db::QueryResult) -> Option { result.rows.first()?.first()?.as_str().map(crate::backend_error::BackendError::from_legacy_string) } /// Unified database operation execution budget. /// query_timeout = None only means SQL execution has no upper limit; /// checkout/connect/recycle/cancel/cleanup always have hard upper limits and cannot be disabled. #[derive(Debug, Clone)] pub struct DbOperationBudget { pub checkout_timeout: Duration, pub connect_timeout: Duration, pub recycle_timeout: Duration, pub query_timeout: Option, pub cancel_timeout: Duration, pub cleanup_timeout: Duration, } impl DbOperationBudget { /// Build an execution budget from connection config. /// checkout/connect/recycle use connect_timeout_secs (clamped to 1s minimum, 300s maximum). /// query_timeout follows resolve_query_timeout semantics (Some(0) -> None). /// cancel/cleanup are fixed values and cannot be disabled. pub fn from_config(connect_timeout_secs: u64, query_timeout_secs: Option) -> Self { let infra_timeout = Duration::from_secs(connect_timeout_secs.clamp(1, 300)); Self { checkout_timeout: infra_timeout, connect_timeout: infra_timeout, recycle_timeout: infra_timeout, query_timeout: resolve_query_timeout(query_timeout_secs), cancel_timeout: Duration::from_secs(5), cleanup_timeout: Duration::from_secs(3), } } pub fn from_connection_config(config: &ConnectionConfig) -> Self { Self::from_config(config.effective_connect_timeout_secs(), Some(config.query_timeout_secs)) } /// Use global default values (when no connection config is available). pub fn with_defaults() -> Self { let default_infra = db::connection_timeout(); Self { checkout_timeout: default_infra, connect_timeout: default_infra, recycle_timeout: default_infra, query_timeout: Some(QUERY_TIMEOUT), cancel_timeout: Duration::from_secs(5), cleanup_timeout: Duration::from_secs(3), } } } /// Check read-only protection for a connection, blocking write SQL statements. /// Only clones the connection name when read-only mode is active, avoiding /// unnecessary allocations otherwise. /// Uses config_for_pool_key to correctly resolve configs when pool_key includes /// a database suffix (e.g., "prod:app" → config stored under "prod"). pub async fn check_read_only_for_connection(state: &AppState, pool_key: &str, sql: &str) -> Result<(), String> { let connection = { let configs = state.configs.read().await; crate::connection::config_for_pool_key(pool_key, &configs) .filter(|config| config.read_only) .map(|config| (config.name.clone(), config.db_type)) }; if let Some((name, database_type)) = connection { crate::query_execution_sql::check_read_only(sql, &name, database_type)?; } Ok(()) } /// Check read-only protection for a connection across multiple SQL statements. pub async fn check_read_only_for_connection_multi( state: &AppState, pool_key: &str, statements: &[impl AsRef], ) -> Result<(), String> { let connection = { let configs = state.configs.read().await; crate::connection::config_for_pool_key(pool_key, &configs) .filter(|config| config.read_only) .map(|config| (config.name.clone(), config.db_type)) }; if let Some((name, database_type)) = connection { for sql in statements { crate::query_execution_sql::check_read_only(sql.as_ref(), &name, database_type)?; } } Ok(()) } /// Check whether a connection has read-only mode enabled, returning the connection name if so. /// This uses connection_id directly (not pool_key), so it is safe to call at command entry points /// before any pool key is constructed. pub async fn connection_readonly_name(state: &AppState, connection_id: &str) -> Option { state.configs.read().await.get(connection_id).filter(|c| c.read_only).map(|c| c.name.clone()) } async fn connection_is_mongodb(state: &AppState, connection_id: &str) -> bool { let configs = state.configs.read().await; configs.get(connection_id).is_some_and(|config| config.db_type == DatabaseType::MongoDb) } async fn connection_database_type(state: &AppState, connection_id: &str) -> Option { let configs = state.configs.read().await; configs.get(connection_id).map(|config| config.db_type) } async fn connection_mysql_query_dialect(state: &AppState, connection_id: &str) -> db::mysql::MySqlQueryDialect { let configs = state.configs.read().await; configs .get(connection_id) .map(|config| db::mysql::MySqlQueryDialect::for_connection(config.db_type, config.driver_profile.as_deref())) .unwrap_or_default() } async fn connection_mysql_catalog_dialect( state: &AppState, connection_id: &str, ) -> Option { let configs = state.configs.read().await; configs .get(connection_id) .and_then(|config| db::mysql::mysql_catalog_dialect(config.db_type, config.driver_profile.as_deref())) } async fn connection_mysql_catalog_dialect_for_pool_key( state: &AppState, pool_key: &str, ) -> Option { let configs = state.configs.read().await; crate::connection::config_for_pool_key(pool_key, &configs) .and_then(|config| db::mysql::mysql_catalog_dialect(config.db_type, config.driver_profile.as_deref())) } async fn connection_database_type_for_pool_key(state: &AppState, pool_key: &str) -> Option { let configs = state.configs.read().await; configs .iter() .filter(|(connection_id, _)| { pool_key.strip_prefix(connection_id.as_str()).is_some_and(|rest| rest.is_empty() || rest.starts_with(':')) }) .max_by_key(|(connection_id, _)| connection_id.len()) .map(|(_, config)| config.db_type) } fn schema_for_execution_context(db_type: Option, schema: Option<&str>) -> Option<&str> { if matches!(db_type, Some(DatabaseType::Iris)) { None } else { schema } } fn sql_for_execution_context(db_type: Option, sql: &str, schema: Option<&str>) -> String { let Some(schema) = schema.map(str::trim).filter(|schema| !schema.is_empty()) else { return sql.to_string(); }; match db_type { Some(DatabaseType::Iris) => qualify_iris_unqualified_dml(sql, schema).unwrap_or_else(|| sql.to_string()), Some(DatabaseType::Kingbase) => { qualify_kingbase_unqualified_relations(sql, schema).unwrap_or_else(|| sql.to_string()) } _ => sql.to_string(), } } fn qualify_iris_unqualified_dml(sql: &str, schema: &str) -> Option { let dialect = GenericDialect {}; let mut statements = Parser::parse_sql(&dialect, sql).ok()?; if statements.is_empty() { return None; } let mut changed = false; for statement in &mut statements { if !statement_uses_schema_context(statement) { continue; } let cte_names = statement_cte_names(statement); let _ = visit_relations_mut(statement, |name| { if qualify_unqualified_relation_name(name, schema, &cte_names) { changed = true; } ControlFlow::<()>::Continue(()) }); } changed.then(|| statements.iter().map(ToString::to_string).collect::>().join("; ")) } fn qualify_kingbase_unqualified_relations(sql: &str, schema: &str) -> Option { let dialect = PostgreSqlDialect {}; let mut statements = Parser::parse_sql(&dialect, sql).ok()?; if statements.is_empty() { return None; } let mut changed = false; for statement in &mut statements { if !statement_uses_schema_context(statement) { continue; } let cte_names = statement_cte_names(statement); let mut qualifier = KingbaseRelationQualifier { schema, cte_names: &cte_names, parameterized_table_depth: 0, changed: false }; let _ = statement.visit(&mut qualifier); changed |= qualifier.changed; } changed.then(|| statements.iter().map(ToString::to_string).collect::>().join("; ")) } struct KingbaseRelationQualifier<'a> { schema: &'a str, cte_names: &'a HashSet, parameterized_table_depth: usize, changed: bool, } impl VisitorMut for KingbaseRelationQualifier<'_> { type Break = (); fn pre_visit_table_factor(&mut self, table_factor: &mut TableFactor) -> ControlFlow { if matches!(table_factor, TableFactor::Table { args: Some(_), .. }) { self.parameterized_table_depth += 1; } ControlFlow::Continue(()) } fn post_visit_table_factor(&mut self, table_factor: &mut TableFactor) -> ControlFlow { if matches!(table_factor, TableFactor::Table { args: Some(_), .. }) { self.parameterized_table_depth = self.parameterized_table_depth.saturating_sub(1); } ControlFlow::Continue(()) } fn post_visit_relation(&mut self, relation: &mut ObjectName) -> ControlFlow { if self.parameterized_table_depth == 0 && qualify_unqualified_relation_name(relation, self.schema, self.cte_names) { self.changed = true; } ControlFlow::Continue(()) } } fn statement_uses_schema_context(statement: &Statement) -> bool { matches!( statement, Statement::Query(_) | Statement::Insert(_) | Statement::Update(_) | Statement::Delete(_) | Statement::Truncate(_) ) } fn qualify_unqualified_relation_name(name: &mut ObjectName, schema: &str, cte_names: &HashSet) -> bool { let [ObjectNamePart::Identifier(table)] = name.0.as_slice() else { return false; }; if cte_names.contains(&table.value.to_ascii_uppercase()) { return false; } let table = table.clone(); name.0 = vec![ObjectNamePart::Identifier(Ident::with_quote('"', schema)), ObjectNamePart::Identifier(table)]; true } fn statement_cte_names(statement: &Statement) -> HashSet { let mut names = HashSet::new(); collect_statement_cte_names(statement, &mut names); names } fn collect_statement_cte_names(statement: &Statement, names: &mut HashSet) { match statement { Statement::Query(query) => collect_query_cte_names(query, names), Statement::Insert(insert) => { if let Some(source) = &insert.source { collect_query_cte_names(source, names); } } _ => {} } } fn collect_query_cte_names(query: &sqlparser::ast::Query, names: &mut HashSet) { if let Some(with) = &query.with { for cte in &with.cte_tables { names.insert(cte.alias.name.value.to_ascii_uppercase()); collect_query_cte_names(&cte.query, names); } } } fn qualifies_unqualified_agent_relations(db_type: Option) -> bool { matches!(db_type, Some(DatabaseType::Iris | DatabaseType::Kingbase)) } #[derive(Clone, Copy, Debug, Default, Deserialize, PartialEq, Eq)] #[serde(rename_all = "snake_case")] pub enum QueryExecutionMode { #[default] Standard, Simple, PostgresReadOnlyTransaction, } #[derive(Clone, Debug, Default)] pub struct QueryExecutionOptions { pub max_rows: Option, pub fetch_size: Option, pub page_size: Option, pub row_offset: Option, pub max_result_bytes: Option, /// Result columns that must stay exact because clients use them as stable /// row identifiers when fetching full large-cell values on demand. pub result_key_columns: Vec, /// Enables extraction of hidden server-side preview metadata. This must be /// set only for generated table-data preview SQL, never arbitrary queries. pub table_data_preview: bool, /// Doris / StarRocks catalog selected for this query tab. pub catalog: Option, pub result_session_id: Option, pub client_session_id: Option, /// Query timeout in seconds. `None` uses the default (30s). /// `Some(0)` disables the timeout entirely. pub timeout_secs: Option, pub execution_id: Option, /// When `Some(true)`, multiple statements are executed within a single transaction /// (BEGIN … COMMIT) instead of auto-commit mode. `None` and `Some(false)` behave /// identically — auto-commit for each statement. pub use_transaction: Option, /// When `true`, multi-statement execution continues after a statement error instead /// of stopping at the first failure. Connection-level failures always stop the batch. pub continue_on_error: bool, /// Explicit low-level execution path. `Simple` is currently used by SQL Server /// SHOWPLAN so the source SQL bypasses result-set probing and query rewriting. /// `PostgresReadOnlyTransaction` executes on an isolated client session and /// always rolls the transaction back after the result is collected. pub execution_mode: QueryExecutionMode, } fn validate_query_execution_mode( db_type: Option, sql: &str, options: &QueryExecutionOptions, ) -> Result<(), String> { if options.execution_mode != QueryExecutionMode::PostgresReadOnlyTransaction { return Ok(()); } if db_type != Some(DatabaseType::Postgres) { return Err("PostgreSQL read-only transaction mode requires a PostgreSQL connection".to_string()); } if options.client_session_id.as_deref().is_none_or(|session_id| session_id.trim().is_empty()) { return Err("PostgreSQL read-only transaction mode requires an isolated client session".to_string()); } if crate::sql::split_sql_statements_for_database(sql, DatabaseType::Postgres).len() != 1 { return Err("PostgreSQL read-only transaction mode requires exactly one statement".to_string()); } Ok(()) } fn query_result_row_limit(max_rows: Option) -> usize { max_rows.unwrap_or(MAX_ROWS).max(1) } pub fn truncate_result(result: db::QueryResult) -> db::QueryResult { truncate_result_with_max_rows(result, None) } pub fn truncate_result_with_max_rows(mut result: db::QueryResult, max_rows: Option) -> db::QueryResult { let row_limit = query_result_row_limit(max_rows); if result.rows.len() > row_limit { result.rows.truncate(row_limit); result.truncated = true; } result } fn normalize_query_result_for_js(mut result: db::QueryResult) -> db::QueryResult { result.rows = result.rows.into_iter().map(|row| row.into_iter().map(db::json_value_for_js).collect()).collect(); result } pub fn agent_execute_query_params( sql: &str, database: Option<&str>, schema: Option<&str>, options: QueryExecutionOptions, ) -> serde_json::Value { let mut params = serde_json::json!({ "sql": sql, "maxRows": agent_protocol_row_count(options.max_rows.unwrap_or(MAX_ROWS)), "deferLobs": options.table_data_preview, }); if let Some(database) = database.map(str::trim).filter(|database| !database.is_empty()) { params["database"] = serde_json::json!(database); } if let Some(schema) = schema { params["schema"] = serde_json::json!(schema); } if let Some(fetch_size) = options.fetch_size { params["fetchSize"] = serde_json::json!(agent_protocol_row_count(fetch_size)); } if let Some(row_offset) = options.row_offset { params["rowOffset"] = serde_json::json!(agent_protocol_row_offset(row_offset)); } if let Some(timeout_secs) = options.timeout_secs { params["timeoutSecs"] = serde_json::json!(timeout_secs); } params } pub fn agent_execute_query_page_params( sql: &str, database: Option<&str>, schema: Option<&str>, options: QueryExecutionOptions, ) -> serde_json::Value { let mut params = serde_json::json!({ "sql": sql, "pageSize": agent_protocol_row_count(options.page_size.unwrap_or(MAX_ROWS)), "maxRows": agent_protocol_row_count(options.max_rows.unwrap_or(MAX_ROWS)), "deferLobs": options.table_data_preview, }); if let Some(database) = database.map(str::trim).filter(|database| !database.is_empty()) { params["database"] = serde_json::json!(database); } if let Some(schema) = schema { params["schema"] = serde_json::json!(schema); } if let Some(fetch_size) = options.fetch_size { params["fetchSize"] = serde_json::json!(agent_protocol_row_count(fetch_size)); } if let Some(row_offset) = options.row_offset { params["rowOffset"] = serde_json::json!(agent_protocol_row_offset(row_offset)); } if let Some(timeout_secs) = options.timeout_secs { params["timeoutSecs"] = serde_json::json!(timeout_secs); } params } pub fn agent_fetch_query_page_params(session_id: &str, page_size: usize) -> serde_json::Value { serde_json::json!({ "sessionId": session_id, "pageSize": agent_protocol_row_count(page_size), }) } fn agent_protocol_row_count(value: usize) -> usize { value.clamp(1, AGENT_PROTOCOL_MAX_ROWS) } fn agent_protocol_row_offset(value: usize) -> usize { value.min(AGENT_PROTOCOL_MAX_ROWS) } pub fn agent_close_query_session_params(session_id: &str) -> serde_json::Value { serde_json::json!({ "sessionId": session_id, }) } pub fn is_connection_error(err: &str) -> bool { let lower = err.to_lowercase(); if is_dbx_query_timeout_error(&lower) || is_agent_rpc_timeout_error(&lower) || is_pool_saturation_error(err) { return false; } lower.contains("connection") || lower.contains("broken pipe") || lower.contains("reset by peer") || lower.contains("timed out") || (lower.contains("pool") && lower.contains("timeout")) || lower.contains("closed") || lower.contains("关闭的连接") || lower.contains("连接已关闭") || lower.contains("网络通信异常") || lower.contains("通信异常") || lower.contains("communications link failure") || lower.contains("sqlrecoverableexception") || lower.contains("sqlnontransientconnectionexception") || lower.contains("sqltransientconnectionexception") || lower.contains("eof") || lower.contains("i/o error") || lower.contains("input/output error") || lower.contains("not connected") || lower.contains("end-of-file") || lower.contains("idle") || lower.contains("agent stdin not available") || lower.contains("agent stdout not available") || lower.contains("agent runtime terminated") || lower.contains("agent runtime is unavailable") || lower.contains("agent runtime unavailable") || lower.contains("failed to write to agent stdin") || lower.contains("failed to flush agent stdin") || lower.contains("communicating with the server") || is_os_connection_error(&lower) } fn is_dbx_query_timeout_error(lower: &str) -> bool { lower.starts_with("query timed out after ") } fn is_agent_rpc_timeout_error(lower: &str) -> bool { lower.starts_with("agent rpc call timed out ") } fn is_schema_reset_cleanup_error(lower: &str) -> bool { lower.contains("schema.reset cleanup failed") } fn is_postgres_transaction_cleanup_error(lower: &str) -> bool { lower.contains("postgresql read-only transaction cleanup failed") } fn should_discard_agent_pool_after_error(err: &str) -> bool { crate::db::agent_driver::agent_recovery_decision(err, RecoveryScope::UserOperation).discards_session() } pub fn pool_error_action(db_type: Option, err: &str) -> PoolErrorAction { if db_type.is_some_and(|db_type| database_capabilities::is_agent_type(&db_type)) { return if should_discard_agent_pool_after_error(err) { // Agent user operations are never replayed from an error hint because // the database-side outcome may be unknown. PoolErrorAction::Discard } else { PoolErrorAction::Keep }; } let lower = err.to_lowercase(); if db::sqlserver::is_driver_panic_error(err) || (is_dbx_query_timeout_error(&lower) && should_discard_pool_after_query_timeout(db_type)) || is_schema_reset_cleanup_error(&lower) || is_postgres_transaction_cleanup_error(&lower) { return PoolErrorAction::Discard; } if is_connection_error(err) { PoolErrorAction::ReconnectAndRetry } else { PoolErrorAction::Keep } } fn should_continue_batch_after_error(continue_on_error: bool, action: PoolErrorAction) -> bool { // A broken connection cannot safely execute the remaining statements even when // the user explicitly enabled continue-on-error. continue_on_error && action == PoolErrorAction::Keep } fn options_for_sequential_statements( options: &QueryExecutionOptions, statement_count: usize, db_type: Option, ) -> QueryExecutionOptions { let mut statement_options = options.clone(); if statement_count <= 1 || db_type != Some(DatabaseType::Kingbase) || statement_options.result_session_id.is_some() { return statement_options; } if let Some(page_size) = statement_options.page_size.take() { let page_size = page_size.max(1); statement_options.max_rows = Some(statement_options.max_rows.map_or(page_size, |max_rows| max_rows.min(page_size))); } statement_options } fn should_discard_pool_after_query_timeout(db_type: Option) -> bool { let Some(db_type) = db_type else { return false; }; database_capabilities::is_agent_type(&db_type) || matches!( db_type, DatabaseType::Mysql | DatabaseType::Postgres | DatabaseType::Redshift | DatabaseType::Gaussdb | DatabaseType::Kwdb | DatabaseType::OpenGauss | DatabaseType::Questdb | DatabaseType::Doris | DatabaseType::StarRocks | DatabaseType::ManticoreSearch | DatabaseType::ClickHouse | DatabaseType::SqlServer | DatabaseType::Rqlite | DatabaseType::Turso | DatabaseType::CloudflareD1 | DatabaseType::Elasticsearch | DatabaseType::Easysearch | DatabaseType::Qdrant | DatabaseType::Milvus | DatabaseType::Weaviate | DatabaseType::ChromaDb | DatabaseType::InfluxDb | DatabaseType::VictoriaMetrics ) } pub fn should_discard_pool_after_error(db_type: Option, err: &str) -> bool { matches!(pool_error_action(db_type, err), PoolErrorAction::Discard | PoolErrorAction::ReconnectAndRetry) } async fn discard_pool_after_error(state: &AppState, pool_key: &str, db_type: Option, error: &str) { let action = pool_error_action(db_type, error); if !matches!(action, PoolErrorAction::Discard | PoolErrorAction::ReconnectAndRetry) { return; } let replace_agent_runtime = db_type.is_some_and(|db_type| database_capabilities::is_agent_type(&db_type)) && crate::db::agent_driver::agent_recovery_decision(error, RecoveryScope::UserOperation).replaces_runtime(); if replace_agent_runtime { state.detach_pool_by_key(pool_key, true).await; } else { state.remove_pool_by_key(pool_key).await; } } async fn discard_agent_pool_after_typed_error( state: &AppState, pool_key: &str, client: &Arc, error: &AgentCallError, scope: RecoveryScope, ) { discard_agent_pool_after_decision(state, pool_key, client, RecoveryPolicy::decide(error, scope)).await; } async fn discard_agent_pool_after_decision( state: &AppState, pool_key: &str, client: &Arc, decision: RecoveryDecision, ) { if decision.discards_session() { state.detach_agent_pool_if_current(pool_key, client, decision.replaces_runtime()).await; } } fn query_pool_error_action(db_type: Option, sql: &str, err: &str) -> PoolErrorAction { match pool_error_action(db_type, err) { // A connection error does not prove that the database did not receive // a write. Only replay statements already accepted by the read-only // protection classifier; writes discard the stale pool without retry. PoolErrorAction::ReconnectAndRetry if is_write_sql(sql) => PoolErrorAction::Discard, action => action, } } fn query_execution_error_action( db_type: Option, sql: &str, error: &QueryExecutionError, ) -> PoolErrorAction { if let Some(agent_error) = error.as_agent_error() { return if RecoveryPolicy::decide(agent_error, RecoveryScope::UserOperation).discards_session() { PoolErrorAction::Discard } else { PoolErrorAction::Keep }; } match error { QueryExecutionError::Canceled { .. } => PoolErrorAction::Keep, QueryExecutionError::DuckDb { message, .. } => query_pool_error_action(db_type, sql, message), QueryExecutionError::Timeout(message) | QueryExecutionError::Sql(message) | QueryExecutionError::Legacy(message) => query_pool_error_action(db_type, sql, message), QueryExecutionError::Agent(_) => unreachable!("Agent errors return above"), } } fn is_os_connection_error(lower: &str) -> bool { let os_error_codes = ["10053", "10054", "10057", "10058", "10060", "10061"]; if let Some(pos) = lower.find("os error ") { let after = &lower[pos + 9..]; let code: String = after.chars().take_while(|c| c.is_ascii_digit()).collect(); return os_error_codes.contains(&code.as_str()); } false } pub fn timeout_error() -> String { timeout_error_for(QUERY_TIMEOUT) } fn timeout_error_for(timeout_duration: Duration) -> String { let seconds = timeout_duration.as_secs().max(1); format!("Query timed out after {seconds} seconds") } pub fn canceled_error() -> String { QUERY_CANCELED.to_string() } fn canceled_query_execution_error() -> QueryExecutionError { QueryExecutionError::Canceled { stage: AgentErrorStage::Cancel, operation_outcome: AgentOperationOutcome::Unknown } } fn pre_dispatch_canceled_query_execution_error() -> QueryExecutionError { QueryExecutionError::Canceled { stage: AgentErrorStage::Request, operation_outcome: AgentOperationOutcome::NotStarted, } } fn postgres_transaction_statement_error( statement_index: usize, message: &str, sql: &str, is_server_error: bool, ) -> QueryExecutionError { let detail = query_error_with_omitted_sql_context(&format!("Statement {statement_index} failed: {message}"), sql); let lower = message.to_ascii_lowercase(); if is_dbx_query_timeout_error(&lower) { QueryExecutionError::Timeout(detail) } else if message == QUERY_CANCELED { canceled_query_execution_error() } else if is_server_error { QueryExecutionError::Sql(detail) } else { QueryExecutionError::Legacy(detail) } } pub(crate) struct StreamProgressClock { started_at: tokio::time::Instant, last_progress_ms: AtomicU64, } impl StreamProgressClock { pub(crate) fn new() -> Self { Self { started_at: tokio::time::Instant::now(), last_progress_ms: AtomicU64::new(0) } } pub(crate) fn mark(&self) { self.last_progress_ms.store(self.started_at.elapsed().as_millis() as u64, Ordering::Relaxed); } fn elapsed_since_progress(&self) -> Duration { let last_progress_ms = self.last_progress_ms.load(Ordering::Relaxed); let elapsed_ms = self.started_at.elapsed().as_millis() as u64; Duration::from_millis(elapsed_ms.saturating_sub(last_progress_ms)) } } pub(crate) async fn await_stream_with_progress_timeout( stream_future: F, timeout: Option, progress_clock: Arc, cancel_token: Option<&CancellationToken>, timeout_message: String, ) -> Result where F: Future>, { let Some(timeout) = timeout else { return match cancel_token { Some(token) => { tokio::select! { biased; _ = token.cancelled() => Err(canceled_error()), result = stream_future => result, } } None => stream_future.await, }; }; tokio::pin!(stream_future); loop { // Query timeout is an inactivity budget, not a cap on total stream duration. let remaining = timeout.saturating_sub(progress_clock.elapsed_since_progress()); if remaining.is_zero() { return Err(timeout_message.clone()); } let sleep = tokio::time::sleep(remaining); tokio::pin!(sleep); match cancel_token { Some(token) => { tokio::select! { biased; _ = token.cancelled() => return Err(canceled_error()), result = &mut stream_future => return result, _ = &mut sleep => {}, } } None => { tokio::select! { biased; result = &mut stream_future => return result, _ = &mut sleep => {}, } } } if progress_clock.elapsed_since_progress() >= timeout { return Err(timeout_message); } } } pub fn is_canceled(cancel_token: &Option) -> bool { cancel_token.as_ref().map(|token| token.is_cancelled()).unwrap_or(false) } pub async fn wait_for_query(cancel_token: Option, future: F) -> Result where F: Future>, { wait_for_query_with_timeout(cancel_token, QUERY_TIMEOUT, future).await } pub async fn wait_for_query_with_timeout( cancel_token: Option, timeout_duration: Duration, future: F, ) -> Result where F: Future>, { wait_for_result_with_timeout(cancel_token, timeout_duration, future).await } async fn wait_for_result_with_timeout( cancel_token: Option, timeout_duration: Duration, future: F, ) -> Result where F: Future>, { if let Some(token) = cancel_token { tokio::select! { biased; _ = token.cancelled() => Err(canceled_error()), result = timeout(timeout_duration, future) => result.map_err(|_| timeout_error_for(timeout_duration))?, } } else { timeout(timeout_duration, future).await.map_err(|_| timeout_error_for(timeout_duration))? } } /// Like `wait_for_query_with_timeout` but with an optional timeout. /// `None` means no timeout (only cancellation can stop the query). pub async fn wait_for_query_opt( cancel_token: Option, timeout_duration: Option, future: F, ) -> Result where F: Future>, { wait_for_result_opt(cancel_token, timeout_duration, future).await } async fn wait_for_result_opt( cancel_token: Option, timeout_duration: Option, future: F, ) -> Result where F: Future>, { match timeout_duration { Some(d) => wait_for_result_with_timeout(cancel_token, d, future).await, None => match cancel_token { Some(token) => { tokio::select! { biased; _ = token.cancelled() => Err(canceled_error()), result = future => result, } } None => future.await, }, } } async fn wait_for_value_opt( cancel_token: Option, timeout_duration: Option, future: F, ) -> Result where F: Future, { match timeout_duration { Some(timeout_duration) => { if let Some(token) = cancel_token { tokio::select! { biased; _ = token.cancelled() => Err(canceled_error()), result = timeout(timeout_duration, future) => result.map_err(|_| timeout_error_for(timeout_duration)), } } else { timeout(timeout_duration, future).await.map_err(|_| timeout_error_for(timeout_duration)) } } None => match cancel_token { Some(token) => { tokio::select! { biased; _ = token.cancelled() => Err(canceled_error()), result = future => Ok(result), } } None => Ok(future.await), }, } } /// Locks a mutex-guarded shared connection (e.g. SQL Server's single connection /// per pool key) and reports how long the caller waited for the lock alongside /// the guard. Callers should fold the returned wait time into any execution-time /// metric they report, since a driver-level timer that only starts once the lock /// is held cannot see time spent queued behind another operation on the same /// connection. async fn lock_shared_client_with_wait<'a, T>( client: &'a Arc>, cancel_token: Option, timeout_duration: Option, ) -> Result<(tokio::sync::MutexGuard<'a, T>, u128), String> { let started_at = std::time::Instant::now(); let guard = wait_for_value_opt(cancel_token, timeout_duration, client.lock()).await?; Ok((guard, started_at.elapsed().as_millis())) } async fn sqlserver_pool_is_current( state: &AppState, pool_key: &str, client: &Arc>, ) -> bool { let connections = state.connections.read().await; matches!(connections.get(pool_key), Some(PoolKind::SqlServer(current)) if Arc::ptr_eq(current, client)) } fn resolve_query_timeout(timeout_secs: Option) -> Option { match timeout_secs { Some(0) => None, Some(n) => Some(Duration::from_secs(n)), None => Some(QUERY_TIMEOUT), } } fn query_pool_database<'a>(database: &'a str, catalog: Option<&str>) -> Option<&'a str> { if database.is_empty() || catalog.is_some() { None } else { Some(database) } } fn postgres_prefers_text_protocol(db_type: Option) -> bool { db_type == Some(DatabaseType::Redshift) } pub async fn operation_budget_for_pool_key( state: &AppState, pool_key: &str, query_timeout: Option, ) -> DbOperationBudget { let mut budget = configured_operation_budget_for_pool_key(state, pool_key).await; budget.query_timeout = query_timeout; budget } async fn configured_operation_budget_for_pool_key(state: &AppState, pool_key: &str) -> DbOperationBudget { let configs = state.configs.read().await; crate::connection::config_for_pool_key(pool_key, &configs) .map(DbOperationBudget::from_connection_config) .unwrap_or_else(DbOperationBudget::with_defaults) } fn oceanbase_mysql_session_timeout_sql(config: Option<&ConnectionConfig>, timeout_secs: Option) -> Option { let config = config?; let timeout_secs = timeout_secs.unwrap_or(config.query_timeout_secs); crate::db::oceanbase_mysql::query_timeout_sql(config, timeout_secs) } async fn apply_oceanbase_mysql_session_timeout( state: &AppState, pool_key: &str, conn: &mut mysql_async::Conn, timeout_secs: Option, ) -> Result<(), String> { let sql = { let configs = state.configs.read().await; oceanbase_mysql_session_timeout_sql(crate::connection::config_for_pool_key(pool_key, &configs), timeout_secs) }; if let Some(sql) = sql { // OceanBase enforces query timeouts through a session variable; set it // on the checked-out connection in case the pooled session was reset. conn.query_drop(&sql).await.map_err(|err| format!("Failed to apply OceanBase query timeout: {err}"))?; } Ok(()) } #[allow(clippy::too_many_arguments)] async fn do_execute_typed( state: &AppState, pool_key: &str, mysql_dialect: db::mysql::MySqlQueryDialect, database: Option<&str>, sql: &str, schema: Option<&str>, cancel_token: Option, options: QueryExecutionOptions, ) -> Result { crate::sql_diagnostics::debug_sql("do_execute", sql); if let Some(execution_id) = options.execution_id.as_deref() { state.running_queries.set_pool_key(execution_id, pool_key.to_string()); } state.touch_pool_activity(pool_key).await; let _activity_touch = state.pool_activity_touch(pool_key); let query_timeout = resolve_query_timeout(options.timeout_secs); let read_only_connection = { let configs = state.configs.read().await; let config = crate::connection::config_for_pool_key(pool_key, &configs); config.filter(|config| config.read_only).map(|config| (config.name.clone(), config.db_type)) }; let operation_budget = operation_budget_for_pool_key(state, pool_key, query_timeout).await; if let Some((name, database_type)) = read_only_connection { crate::query_execution_sql::check_read_only(sql, &name, database_type)?; } let pool_db_type = connection_database_type_for_pool_key(state, pool_key).await; let mysql_catalog_dialect = connection_mysql_catalog_dialect_for_pool_key(state, pool_key).await; let connections = state.connections.read().await; let pool = connections.get(pool_key).ok_or("Connection not found")?; let mut typed_agent_error = None; #[cfg(feature = "duckdb-sidecar")] let mut typed_duckdb_error = None; let result: Result = match pool { #[cfg(feature = "duckdb-sidecar")] PoolKind::DuckDbWorker(client) => { let client = client.clone(); if let Some(ref execution_id) = options.execution_id { let cancel_client = client.clone(); state.running_queries.register_interrupt(execution_id, move || { let cancel_client = cancel_client.clone(); tokio::spawn(async move { if let Err(error) = cancel_client.cancel().await { log::warn!("Failed to cancel DuckDB worker query: {error}"); } }); }); } let sql = sql.to_string(); let database = database.map(str::to_string); let max_rows = options.max_rows; drop(connections); match client.execute_typed(database, sql, max_rows, cancel_token, query_timeout).await { Ok(result) => Ok(result), Err(error) => { let is_control_error = error.message == QUERY_CANCELED || is_dbx_query_timeout_error(&error.message.to_ascii_lowercase()); if !is_control_error { typed_duckdb_error = Some(error.clone()); } Err(error.message) } } } #[cfg(not(feature = "duckdb-sidecar"))] PoolKind::DuckDbWorker(_) => { return Err("DuckDB worker support is not compiled in this build".into()); } PoolKind::Mysql(p, mode) => { let p = p.clone(); let bare = *mode == crate::connection::MysqlMode::Bare; let max_rows = options.max_rows; let max_result_bytes = options.max_result_bytes.filter(|value| *value > 0); drop(connections); let mut conn = match db::mysql::get_conn_with_health_check_with_cancel( &p, operation_budget.checkout_timeout, operation_budget.cleanup_timeout, cancel_token.as_ref(), ) .await { Ok(conn) => conn, Err(err) if err == QUERY_CANCELED => { state.remove_pool_by_key(pool_key).await; return Err(err.into()); } Err(err) => return Err(err.into()), }; let connection_id = conn.id(); if let Some(ref execution_id) = options.execution_id { let kill_opts = conn.opts().clone(); state.running_queries.register_interrupt(execution_id, move || { let kill_opts = kill_opts.clone(); tokio::spawn(async move { if let Err(error) = db::mysql::kill_query_with_opts(kill_opts, connection_id).await { log::warn!("Failed to cancel MySQL query {connection_id}: {error}"); } }); }); } apply_oceanbase_mysql_session_timeout(state, pool_key, &mut conn, options.timeout_secs).await?; wait_for_result_opt( cancel_token.clone(), query_timeout, db::mysql::apply_catalog_database_context( &mut conn, mysql_catalog_dialect, options.catalog.as_deref(), database.unwrap_or_default(), ), ) .await?; wait_for_result_opt( cancel_token, query_timeout, db::mysql::execute_query_on_conn_with_limits( &mut conn, sql, bare, max_rows, max_result_bytes, &options.result_key_columns, mysql_dialect, options.execution_id.as_deref(), ), ) .await .map(|result| result.result) } PoolKind::Postgres(p) => { let p = p.clone(); let schema = schema.map(|s| s.to_string()); let max_rows = options.max_rows; let prefer_text_protocol = postgres_prefers_text_protocol(pool_db_type); let execution_mode = options.execution_mode; let cancel_context = state.get_postgres_cancel_context(pool_key).await; drop(connections); if execution_mode == QueryExecutionMode::PostgresReadOnlyTransaction { db::postgres::execute_query_in_read_only_transaction_with_rollback( &p, schema.as_deref(), sql, max_rows, cancel_token, operation_budget.clone(), cancel_context, ) .await } else if let Some(schema) = schema { db::postgres::execute_query_with_schema_and_max_rows_and_cancel( &p, &schema, sql, max_rows, cancel_token, operation_budget.clone(), cancel_context, prefer_text_protocol, ) .await } else { db::postgres::execute_query_with_max_rows_and_cancel( &p, sql, max_rows, cancel_token, operation_budget.clone(), cancel_context, prefer_text_protocol, ) .await } } PoolKind::Sqlite(p) => { let p = p.clone(); let max_rows = options.max_rows; drop(connections); wait_for_query_opt(cancel_token, query_timeout, db::sqlite::execute_query_with_max_rows(&p, sql, max_rows)) .await } PoolKind::Rqlite(client) => { let client = client.clone(); let max_rows = options.max_rows; drop(connections); wait_for_query_opt( cancel_token, query_timeout, db::rqlite_driver::execute_query_with_max_rows(&client, sql, max_rows), ) .await } PoolKind::Turso(client) => { let client = client.clone(); let max_rows = options.max_rows; drop(connections); wait_for_query_opt( cancel_token, query_timeout, db::turso_driver::execute_query_with_max_rows(&client, sql, max_rows), ) .await } PoolKind::CloudflareD1(client) => { let client = client.clone(); let max_rows = options.max_rows; drop(connections); wait_for_query_opt( cancel_token, query_timeout, db::cloudflare_d1_driver::execute_query_with_max_rows(&client, sql, max_rows), ) .await } PoolKind::ClickHouse(client) => { let client = client.clone(); let database = pool_key.split(':').nth(1).unwrap_or("default").to_string(); let max_rows = options.max_rows; drop(connections); let result = wait_for_query_opt( cancel_token, query_timeout, db::clickhouse_driver::execute_query_with_max_rows(&client, &database, sql, max_rows), ) .await .map(|result| truncate_result_with_max_rows(result, max_rows)); if matches!(result.as_ref(), Err(err) if should_discard_pool_after_error(pool_db_type, err)) { state.remove_pool_by_key(pool_key).await; } result } PoolKind::SqlServer(client) => { let client = client.clone(); let max_rows = options.max_rows; let execution_mode = options.execution_mode; drop(connections); let (mut client, lock_wait_ms) = match lock_shared_client_with_wait(&client, cancel_token.clone(), None).await { Ok(value) => value, Err(err) => return Err(err.into()), }; let execution = async { if execution_mode == QueryExecutionMode::Simple { let mut results = db::sqlserver::execute_simple_batch_with_max_rows(&mut client, sql, max_rows).await?; Ok(results.remove(0)) } else { db::sqlserver::execute_query_with_max_rows(&mut client, sql, max_rows).await } }; let result = wait_for_query_opt(cancel_token, query_timeout, execution) .await .map(|result| truncate_result_with_max_rows(result, max_rows)) .map(|mut result| { result.execution_time_ms += lock_wait_ms; result }); drop(client); if matches!(result.as_ref(), Err(err) if should_discard_pool_after_error(pool_db_type, err)) { state.remove_pool_by_key(pool_key).await; } result } PoolKind::Elasticsearch(client) => { let client = client.clone(); let sql = sql.to_string(); let max_rows = options.max_rows; drop(connections); let result = wait_for_query_opt( cancel_token, query_timeout, db::elasticsearch_driver::execute_rest_query(&client, &sql), ) .await .map(|result| truncate_result_with_max_rows(result, max_rows)); if matches!(result.as_ref(), Err(err) if should_discard_pool_after_error(pool_db_type, err)) { state.remove_pool_by_key(pool_key).await; } result } PoolKind::Easysearch(client) => { let client = client.clone(); let sql = sql.to_string(); let max_rows = options.max_rows; drop(connections); let result = wait_for_query_opt( cancel_token, query_timeout, db::easysearch_driver::execute_rest_query(&client, &sql), ) .await .map(|result| truncate_result_with_max_rows(result, max_rows)); if matches!(result.as_ref(), Err(err) if should_discard_pool_after_error(pool_db_type, err)) { state.remove_pool_by_key(pool_key).await; } result } PoolKind::Meilisearch(client) => { let client = client.clone(); let sql = sql.to_string(); let max_rows = options.max_rows; drop(connections); let result = wait_for_query_opt( cancel_token, query_timeout, db::meilisearch_driver::execute_rest_query(&client, &sql), ) .await .map(|result| truncate_result_with_max_rows(result, max_rows)); if matches!(result.as_ref(), Err(err) if should_discard_pool_after_error(pool_db_type, err)) { state.remove_pool_by_key(pool_key).await; } result } PoolKind::VectorDb(client) => { let client = client.clone(); let sql = sql.to_string(); let max_rows = options.max_rows; drop(connections); let result = wait_for_query_opt(cancel_token, query_timeout, db::vector_driver::execute_rest_query(&client, &sql)) .await .map(|result| truncate_result_with_max_rows(result, max_rows)); if matches!(result.as_ref(), Err(err) if should_discard_pool_after_error(pool_db_type, err)) { state.remove_pool_by_key(pool_key).await; } result } PoolKind::Redis(_) => Err("Use Redis-specific commands".to_string()), PoolKind::MongoDb(_) => Err(MONGO_SHELL_COMMAND_HINT.to_string()), PoolKind::MessageQueue => Err("Use Message Queue-specific commands".to_string()), #[cfg(feature = "mq-admin")] PoolKind::Mqtt(_) => Err("Use MQTT-specific commands".to_string()), PoolKind::Nacos => Err("Use Nacos-specific commands".to_string()), PoolKind::InfluxDb(client) => { let client = client.clone(); let database = pool_key.split(':').nth(1).unwrap_or("default").to_string(); let max_rows = options.max_rows; drop(connections); let result = wait_for_query_opt( cancel_token, query_timeout, db::influxdb_driver::execute_query(&client, &database, sql), ) .await .map(|result| truncate_result_with_max_rows(result, max_rows)); if matches!(result.as_ref(), Err(err) if should_discard_pool_after_error(pool_db_type, err)) { state.remove_pool_by_key(pool_key).await; } result } PoolKind::VictoriaMetrics(client) => { let client = client.clone(); let max_rows = options.max_rows; drop(connections); let result = wait_for_query_opt( cancel_token, query_timeout, db::victoriametrics_driver::execute_query(&client, sql), ) .await .map(|result| truncate_result_with_max_rows(result, max_rows)); if matches!(result.as_ref(), Err(err) if should_discard_pool_after_error(pool_db_type, err)) { state.remove_pool_by_key(pool_key).await; } result } PoolKind::Agent(client) => { let client = client.clone(); let source_client = client.clone(); let sql = sql_for_execution_context(pool_db_type, sql, schema); let database = database.map(|s| s.to_string()); let schema = schema_for_execution_context(pool_db_type, schema).map(|s| s.to_string()); let max_rows = options.max_rows; let rpc_timeout = query_timeout; drop(connections); if is_canceled(&cancel_token) { return Err(canceled_error().into()); } let cancel_for_agent = cancel_token.clone(); let result = async move { let mut client = match cancel_for_agent.as_ref() { Some(token) => { tokio::select! { biased; _ = token.cancelled() => return Err(AgentCallError::Canceled { stage: AgentErrorStage::Cancel, operation_outcome: AgentOperationOutcome::Unknown, }), guard = client.lock() => guard, } } None => client.lock().await, }; if let Some(session_id) = options.result_session_id.as_deref() { let params = agent_fetch_query_page_params(session_id, options.page_size.unwrap_or(MAX_ROWS)); client .fetch_query_page_typed_with_timeout_and_cancel(params, rpc_timeout, cancel_for_agent.clone()) .await } else if options.page_size.is_some() { let params = agent_execute_query_page_params(&sql, database.as_deref(), schema.as_deref(), options); client .execute_query_page_typed_with_timeout_and_cancel(params, rpc_timeout, cancel_for_agent.clone()) .await } else { let params = agent_execute_query_params(&sql, database.as_deref(), schema.as_deref(), options); client .execute_query_typed_with_timeout_and_cancel(params, rpc_timeout, cancel_for_agent.clone()) .await } } .await .map(|result| truncate_result_with_max_rows(result, max_rows)); if let Err(err) = result.as_ref() { discard_agent_pool_after_typed_error( state, pool_key, &source_client, err, RecoveryScope::UserOperation, ) .await; } typed_agent_error = result.as_ref().err().cloned(); result.map_err(AgentCallError::into_legacy_string) } PoolKind::ExternalDriver { config, session, .. } => { let config = config.clone(); let session = session.clone(); let sql = sql.to_string(); let schema = schema.map(str::to_string); let database = database.unwrap_or_else(|| config.effective_database().unwrap_or("")).to_string(); let max_rows = options.max_rows; let plugin_timeout = query_timeout; drop(connections); wait_for_query_opt(cancel_token, query_timeout, async move { if let Some(session_id) = options.result_session_id.as_deref() { let params = external_driver_fetch_query_page_params( config.as_ref(), session_id, options.page_size.unwrap_or(MAX_ROWS), ); session.invoke_with_timeout::("fetchQueryPage", params, plugin_timeout).await } else if options.page_size.is_some() { let params = external_driver_query_params(config.as_ref(), &sql, &database, schema.as_deref(), &options); invoke_external_driver_query_page(session.as_ref(), params, plugin_timeout).await } else { let params = external_driver_query_params(config.as_ref(), &sql, &database, schema.as_deref(), &options); session.invoke_with_timeout::("executeQuery", params, plugin_timeout).await } }) .await .map(|result| truncate_result_with_max_rows(result, max_rows)) } PoolKind::HBase(_) => Err("SQL execution is not supported for HBase connections".to_string()), PoolKind::DynamoDb(client) => { let client = client.clone(); let sql = sql.to_string(); let max_rows = options.max_rows.unwrap_or(MAX_ROWS); drop(connections); // Keep the AWS SDK cold-path future off this already-large query dispatcher stack. let execution = Box::pin(db::dynamodb_driver::execute_statement(&client, &sql, max_rows)); wait_for_query_opt(cancel_token, query_timeout, execution).await } PoolKind::Consul(_) => Err("SQL execution is not supported for Consul connections".to_string()), }; result .map(normalize_query_result_for_js) .map_err(|error| { #[cfg(feature = "duckdb-sidecar")] if let Some(duckdb_error) = typed_duckdb_error { return QueryExecutionError::DuckDb { code: duckdb_error.code, message: duckdb_error.message }; } typed_agent_error.map_or_else(|| QueryExecutionError::Legacy(error), QueryExecutionError::Agent) }) .map_err(|error| classify_query_error(pool_db_type, error)) } fn classify_query_error(db_type: Option, error: QueryExecutionError) -> QueryExecutionError { match error { QueryExecutionError::Legacy(message) if message == QUERY_CANCELED => canceled_query_execution_error(), QueryExecutionError::Legacy(message) if is_dbx_query_timeout_error(&message.to_ascii_lowercase()) => { QueryExecutionError::Timeout(message) } QueryExecutionError::Legacy(message) if is_native_sql_server_error(db_type, &message) => { QueryExecutionError::Sql(message) } other => other, } } fn is_native_sql_server_error(db_type: Option, message: &str) -> bool { let message = message.trim_start(); match db_type { Some(DatabaseType::Postgres) => message.starts_with("ERROR:"), Some(DatabaseType::Mysql) => message.starts_with("Server error: `ERROR "), _ => false, } } pub async fn do_execute( state: &AppState, pool_key: &str, mysql_dialect: db::mysql::MySqlQueryDialect, database: Option<&str>, sql: &str, schema: Option<&str>, cancel_token: Option, options: QueryExecutionOptions, ) -> Result { do_execute_typed(state, pool_key, mysql_dialect, database, sql, schema, cancel_token, options) .await .map_err(QueryExecutionError::into_legacy_string) } async fn invoke_external_driver_query_page( session: &crate::plugins::PluginDriverSession, params: serde_json::Value, plugin_timeout: Option, ) -> Result { match session.invoke_with_timeout::("executeQueryPage", params.clone(), plugin_timeout).await { Ok(result) => Ok(result), Err(error) if is_external_driver_method_unsupported(&error, "executeQueryPage") => { // Plugins installed by older DBX releases predate cursor pagination. Keep // basic queries usable until the user updates the plugin, without retrying // actual JDBC/SQL failures that may have side effects. log::warn!("[query][external-driver] executeQueryPage unsupported; falling back to executeQuery"); session.invoke_with_timeout::("executeQuery", params, plugin_timeout).await } Err(error) => Err(error), } } fn is_external_driver_method_unsupported(error: &str, method: &str) -> bool { let normalized = error.to_ascii_lowercase(); let method = method.to_ascii_lowercase(); normalized.contains(&method) && (normalized.contains("unsupported jdbc plugin method") || normalized.contains("unknown method") || normalized.contains("method not found")) } fn external_driver_query_params( config: &crate::models::connection::ConnectionConfig, sql: &str, database: &str, schema: Option<&str>, options: &QueryExecutionOptions, ) -> serde_json::Value { let mut params = serde_json::json!({ "connection": config, "sql": sql, "database": database, "schema": schema, "maxRows": options.max_rows.unwrap_or(MAX_ROWS), }); if let Some(fetch_size) = options.fetch_size { params["fetchSize"] = serde_json::json!(fetch_size); } if let Some(row_offset) = options.row_offset { params["rowOffset"] = serde_json::json!(agent_protocol_row_offset(row_offset)); } if let Some(timeout_secs) = options.timeout_secs { params["timeoutSecs"] = serde_json::json!(timeout_secs); } if let Some(page_size) = options.page_size { params["pageSize"] = serde_json::json!(page_size); } params } fn external_driver_fetch_query_page_params( config: &crate::models::connection::ConnectionConfig, session_id: &str, page_size: usize, ) -> serde_json::Value { serde_json::json!({ "connection": config, "sessionId": session_id, "pageSize": page_size, }) } pub async fn execute_sql_statement( state: &AppState, connection_id: &str, database: &str, sql: &str, schema: Option<&str>, cancel_token: Option, ) -> Result { execute_sql_statement_with_options( state, connection_id, database, sql, schema, cancel_token, QueryExecutionOptions::default(), ) .await } pub async fn execute_sql_statement_with_options_typed( state: &AppState, connection_id: &str, database: &str, sql: &str, schema: Option<&str>, cancel_token: Option, options: QueryExecutionOptions, ) -> Result { // MongoDB connections use shell-style commands dispatched through the // frontend parser. Queries that fall through to the generic SQL executor // (e.g. typos) must be rejected before any pool/key creation so that // session-scoped pools do not leak MongoDB Clients and SSH tunnels. if connection_is_mongodb(state, connection_id).await { return Err(MONGO_SHELL_COMMAND_HINT.into()); } let db_type = connection_database_type(state, connection_id).await; validate_query_execution_mode(db_type, sql, &options)?; let has_executable_sql = db_type.map_or_else( || crate::sql::has_executable_sql(sql), |db_type| crate::sql::has_executable_sql_for_database(sql, db_type), ); if !has_executable_sql { return Ok(empty_query_result(0)); } if let Some(target_database) = postgres_drop_database_target(db_type, sql) { return execute_postgres_drop_database(state, connection_id, &target_database, sql, cancel_token, options) .await .map_err(Into::into); } // When a query tab has a client session, keep even database-less execution // on that tab-scoped pool so connection-level state (for example MySQL @vars) // survives across runs. let pool_database = query_pool_database(database, options.catalog.as_deref()); let pool_key = state .get_or_create_pool_for_session(connection_id, pool_database, options.client_session_id.as_deref()) .await .map_err(|e| query_error_with_omitted_sql_context(&e, sql))?; if is_canceled(&cancel_token) { return Err(pre_dispatch_canceled_query_execution_error()); } let mysql_dialect = connection_mysql_query_dialect(state, connection_id).await; let result = do_execute_typed( state, &pool_key, mysql_dialect, Some(database), sql, schema, cancel_token.clone(), options.clone(), ) .await; let with_sql_context = |result: Result| { result.map_err(|error| error.with_omitted_sql_context(sql)) }; let action = result.as_ref().err().map(|error| query_execution_error_action(db_type, sql, error)); match action { Some(PoolErrorAction::ReconnectAndRetry) if !is_canceled(&cancel_token) => { let pool_database = query_pool_database(database, options.catalog.as_deref()); let new_key = state .reconnect_pool_for_session(connection_id, pool_database, options.client_session_id.as_deref()) .await .map_err(|e| query_error_with_omitted_sql_context(&e, sql))?; with_sql_context( do_execute_typed(state, &new_key, mysql_dialect, Some(database), sql, schema, cancel_token, options) .await, ) } Some(PoolErrorAction::Discard) => { // Agent execution owns structured quarantine/runtime replacement before // returning. Native drivers retain the existing caller-side cleanup. if !db_type.is_some_and(|db_type| database_capabilities::is_agent_type(&db_type)) { state.remove_pool_by_key(&pool_key).await; } with_sql_context(result) } _ => with_sql_context(result), } } pub async fn execute_sql_statement_with_options( state: &AppState, connection_id: &str, database: &str, sql: &str, schema: Option<&str>, cancel_token: Option, options: QueryExecutionOptions, ) -> Result { execute_sql_statement_with_options_typed(state, connection_id, database, sql, schema, cancel_token, options) .await .map_err(QueryExecutionError::into_legacy_string) } async fn execute_postgres_drop_database( state: &AppState, connection_id: &str, target_database: &str, sql: &str, cancel_token: Option, options: QueryExecutionOptions, ) -> Result { state.close_database_pool(connection_id, Some(target_database)).await?; let admin_database = postgres_drop_database_admin_database(target_database); let pool_key = state .get_or_create_pool_for_session(connection_id, Some(admin_database), options.client_session_id.as_deref()) .await .map_err(|e| query_error_with_omitted_sql_context(&e, sql))?; if let Some(execution_id) = options.execution_id.as_deref() { state.running_queries.set_pool_key(execution_id, pool_key.clone()); } state.touch_pool_activity(&pool_key).await; let _activity_touch = state.pool_activity_touch(pool_key.as_str()); if is_canceled(&cancel_token) { return Err(canceled_error()); } check_read_only_for_connection(state, &pool_key, sql).await?; let pool = { let connections = state.connections.read().await; match connections.get(&pool_key) { Some(PoolKind::Postgres(pool)) => pool.clone(), Some(_) => return Err("DROP DATABASE reconnect did not create a PostgreSQL connection".to_string()), None => return Err("Connection not found".to_string()), } }; let query_timeout = resolve_query_timeout(options.timeout_secs); let max_rows = options.max_rows; wait_for_query_opt(cancel_token, query_timeout, async { db::postgres::terminate_current_user_database_backends(&pool, target_database).await?; db::postgres::execute_query_with_max_rows(&pool, sql, max_rows).await }) .await } fn postgres_drop_database_target(db_type: Option, sql: &str) -> Option { if db_type != Some(DatabaseType::Postgres) { return None; } parse_drop_database_target(sql) } fn postgres_drop_database_admin_database(target_database: &str) -> &'static str { if target_database.eq_ignore_ascii_case("postgres") { "template1" } else { "postgres" } } fn parse_drop_database_target(sql: &str) -> Option { let dialect = PostgreSqlDialect {}; let statements = Parser::parse_sql(&dialect, sql).ok()?; let [Statement::Drop { object_type, names, .. }] = statements.as_slice() else { return None; }; if *object_type != ObjectType::Database || names.len() != 1 { return None; } let parts = &names[0].0; if parts.len() != 1 { return None; } parts[0].as_ident().map(|ident| ident.value.clone()) } pub async fn close_query_session( state: &AppState, connection_id: &str, database: &str, session_id: &str, client_session_id: Option<&str>, catalog: Option<&str>, ) -> Result { let pool_database = query_pool_database(database, catalog); let pool_key = state.get_or_create_pool_for_session(connection_id, pool_database, client_session_id).await?; let connections = state.connections.read().await; let pool = connections.get(&pool_key).ok_or("Connection not found")?; match pool { PoolKind::Agent(client) => { let client = client.clone(); drop(connections); let mut client = client.lock().await; client.close_query_session(session_id).await } PoolKind::ExternalDriver { config, session, .. } => { let config = config.clone(); let session = session.clone(); drop(connections); let params = external_driver_fetch_query_page_params(config.as_ref(), session_id, 1); session .invoke::("closeQuerySession", params) .await .map(|value| value.get("ok").and_then(|ok| ok.as_bool()).unwrap_or(false)) } _ => Ok(false), } } pub async fn execute_multi_core( state: &AppState, connection_id: &str, database: &str, sql: &str, schema: Option<&str>, cancel_token: Option, ) -> Result, String> { execute_multi_core_with_options( state, connection_id, database, sql, schema, cancel_token, QueryExecutionOptions::default(), ) .await } pub async fn execute_multi_core_with_options( state: &AppState, connection_id: &str, database: &str, sql: &str, schema: Option<&str>, cancel_token: Option, options: QueryExecutionOptions, ) -> Result, String> { execute_multi_core_with_options_for_client(state, connection_id, database, sql, schema, cancel_token, options) .await .map(|results| results.into_iter().map(ExecuteMultiResult::into_query_result).collect()) } /// Execute a SQL batch and retain client-facing metadata for synthesized errors. pub async fn execute_multi_core_with_options_for_client( state: &AppState, connection_id: &str, database: &str, sql: &str, schema: Option<&str>, cancel_token: Option, options: QueryExecutionOptions, ) -> Result, String> { execute_multi_core_with_options_for_client_typed(state, connection_id, database, sql, schema, cancel_token, options) .await .map_err(QueryExecutionError::into_legacy_string) } /// Execute a SQL batch for a client while preserving typed failures. pub async fn execute_multi_core_with_options_for_client_typed( state: &AppState, connection_id: &str, database: &str, sql: &str, schema: Option<&str>, cancel_token: Option, options: QueryExecutionOptions, ) -> Result, QueryExecutionError> { execute_multi_core_with_options_for_client_and_progress_typed( state, connection_id, database, sql, schema, cancel_token, options, None, ) .await } /// Executes a SQL batch and reports each completed statement to the optional callback. pub async fn execute_multi_core_with_options_for_client_and_progress( state: &AppState, connection_id: &str, database: &str, sql: &str, schema: Option<&str>, cancel_token: Option, options: QueryExecutionOptions, progress: Option, ) -> Result, String> { execute_multi_core_with_options_for_client_and_progress_typed( state, connection_id, database, sql, schema, cancel_token, options, progress, ) .await .map_err(QueryExecutionError::into_legacy_string) } /// Executes a SQL batch without erasing structured query errors at transport boundaries. pub async fn execute_multi_core_with_options_for_client_and_progress_typed( state: &AppState, connection_id: &str, database: &str, sql: &str, schema: Option<&str>, cancel_token: Option, options: QueryExecutionOptions, progress: Option, ) -> Result, QueryExecutionError> { let pool_database = query_pool_database(database, options.catalog.as_deref()); // Reject MongoDB queries that fall through to the generic executor. if connection_is_mongodb(state, connection_id).await { return Err(MONGO_SHELL_COMMAND_HINT.into()); } let db_type = connection_database_type(state, connection_id).await; validate_query_execution_mode(db_type, sql, &options)?; if options.execution_mode == QueryExecutionMode::PostgresReadOnlyTransaction { let result = execute_sql_statement_with_options(state, connection_id, database, sql, schema, cancel_token, options) .await?; return Ok(vec![result.into()]); } let pool_key = state .get_or_create_pool_for_session(connection_id, pool_database, options.client_session_id.as_deref()) .await .map_err(|e| query_error_with_omitted_sql_context(&e, sql))?; if let Some(execution_id) = options.execution_id.as_deref() { state.running_queries.set_pool_key(execution_id, pool_key.clone()); } state.touch_pool_activity(&pool_key).await; let _activity_touch = state.pool_activity_touch(pool_key.as_str()); let is_sqlserver = { let connections = state.connections.read().await; matches!(connections.get(&pool_key), Some(PoolKind::SqlServer(_))) }; if is_sqlserver { return execute_multi_sqlserver(state, &pool_key, sql, cancel_token, options).await.map_err(Into::into); } let is_http_sqlite = { let configs = state.configs.read().await; configs .get(connection_id) .is_some_and(|c| matches!(c.db_type, DatabaseType::Turso | DatabaseType::CloudflareD1)) }; // HTTP SQLite providers send all statements in one request so the provider // can preserve batch ordering and atomicity. if is_http_sqlite { let table_data_preview = options.table_data_preview; return single_statement_multi_result( execute_sql_statement_with_options_typed( state, connection_id, database, sql, schema, cancel_token, options, ) .await, table_data_preview, ); } let statements = db_type.map_or_else( || split_sql_statements(sql), |db_type| crate::sql::split_sql_statements_for_database(sql, db_type), ); if statements.is_empty() { return Ok(vec![empty_query_result(0).into()]); } // When use_transaction is explicitly true and we have multiple statements, // route through the transaction wrapper instead of the sequential auto-commit loop. if options.use_transaction == Some(true) && statements.len() > 1 { let result = execute_statements_in_transaction_typed( state, connection_id, database, &statements, schema, options.catalog.as_deref(), ) .await?; return Ok(vec![result.into()]); } let mysql_pool = { let connections = state.connections.read().await; match connections.get(&pool_key) { Some(PoolKind::Mysql(pool, mode)) => Some((pool.clone(), *mode)), _ => None, } }; if statements.len() == 1 && !mysql_single_statement_uses_batch_route( db_type, mysql_pool.is_some(), &statements[0], options.max_result_bytes, ) { let single_sql = statements.into_iter().next().unwrap_or_default(); let table_data_preview = options.table_data_preview; return single_statement_multi_result( execute_sql_statement_with_options_typed( state, connection_id, database, &single_sql, schema, cancel_token, options, ) .await, table_data_preview, ); } if let Some((pool, mode)) = mysql_pool { // Read-only check for MySQL batch path check_read_only_for_connection_multi(state, &pool_key, &statements).await?; let mysql_dialect = connection_mysql_query_dialect(state, connection_id).await; let mysql_catalog_dialect = connection_mysql_catalog_dialect(state, connection_id).await; return execute_multi_mysql( state, &pool_key, db_type, &pool, mode, mysql_dialect, mysql_catalog_dialect, database, &statements, cancel_token, options, progress.as_ref(), ) .await .map_err(Into::into); } // Kingbase Go keeps one physical connection per Agent session, so an open // result cursor prevents the next statement from acquiring that connection. // Multi-result execution therefore reads a bounded first page for each // Kingbase statement without retaining cursors. let statement_options = options_for_sequential_statements(&options, statements.len(), db_type); let mut results = Vec::with_capacity(statements.len()); for (statement_index, stmt) in statements.iter().enumerate() { if is_canceled(&cancel_token) { let error = canceled_query_execution_error(); results.push(ExecuteMultiResult::execution_error_with_backend( error_query_result(error.clone().into_legacy_string()), Some(statement_index), error.into_backend_error(), )); break; } match execute_sql_statement_with_options_typed( state, connection_id, database, stmt, schema, cancel_token.clone(), statement_options.clone(), ) .await { Ok(r) => { report_execute_multi_progress(progress.as_ref(), statement_index, statements.len(), &r, true, None); results.push(ExecuteMultiResult::success_with_index_and_optional_server_large_values( r, statement_index, options.table_data_preview, )); } Err(error) => { let action = query_execution_error_action(db_type, stmt, &error); let result = error_query_result(error.clone().into_legacy_string()); let backend_error = error.into_backend_error(); report_execute_multi_progress( progress.as_ref(), statement_index, statements.len(), &result, false, Some(backend_error.clone()), ); results.push(ExecuteMultiResult::execution_error_with_backend( result, Some(statement_index), backend_error, )); if !should_continue_batch_after_error(options.continue_on_error, action) { break; } } } } Ok(results) } fn single_statement_multi_result( result: Result, table_data_preview: bool, ) -> Result, QueryExecutionError> { result.map(|result| vec![ExecuteMultiResult::success_with_optional_server_large_values(result, table_data_preview)]) } fn mysql_single_statement_uses_batch_route( db_type: Option, has_mysql_pool: bool, sql: &str, max_result_bytes: Option, ) -> bool { has_mysql_pool && (max_result_bytes.is_some_and(|value| value > 0) || (db_type == Some(DatabaseType::Mysql) && starts_with_executable_sql_keyword_for_database(sql, &["CALL"], DatabaseType::Mysql))) } trait MysqlBatchStatementExecutor { fn table_data_preview(&self) -> bool { false } async fn execute_statement(&mut self, statement: &str) -> Result, String>; async fn execute_non_result_batch(&mut self, statements: &[String]) -> db::mysql::MySqlNonResultBatchOutcome { let mut results = Vec::with_capacity(statements.len()); for statement in statements { match self.execute_statement(statement).await { Ok(statement_results) if statement_results.len() == 1 => { results.push(statement_results.into_iter().next().expect("single MySQL batch result").result); } Ok(_) => { return db::mysql::MySqlNonResultBatchOutcome { results, error: Some("A non-result MySQL batch statement returned multiple results.".to_string()), }; } Err(error) => return db::mysql::MySqlNonResultBatchOutcome { results, error: Some(error) }, } } db::mysql::MySqlNonResultBatchOutcome { results, error: None } } } struct MysqlBatchConnection<'a> { conn: &'a mut mysql_async::Conn, cancel_token: Option, query_timeout: Option, bare: bool, max_rows: Option, max_result_bytes: Option, result_key_columns: &'a [String], table_data_preview: bool, dialect: db::mysql::MySqlQueryDialect, diagnostic_trace_id: Option<&'a str>, } impl MysqlBatchStatementExecutor for MysqlBatchConnection<'_> { fn table_data_preview(&self) -> bool { self.table_data_preview } async fn execute_statement(&mut self, statement: &str) -> Result, String> { wait_for_result_opt( self.cancel_token.clone(), self.query_timeout, db::mysql::execute_query_results_on_conn_with_limits( &mut *self.conn, statement, self.bare, self.max_rows, self.max_result_bytes, self.result_key_columns, self.dialect, self.diagnostic_trace_id, ), ) .await } async fn execute_non_result_batch(&mut self, statements: &[String]) -> db::mysql::MySqlNonResultBatchOutcome { let sql = statements.join(";\n"); match wait_for_result_opt( self.cancel_token.clone(), self.query_timeout, db::mysql::execute_non_result_batch_on_conn(&mut *self.conn, &sql, statements.len()), ) .await { Ok(outcome) => outcome, Err(error) => db::mysql::MySqlNonResultBatchOutcome { results: Vec::new(), error: Some(error) }, } } } const MYSQL_MULTI_STATEMENT_BATCH_MAX_STATEMENTS: usize = 50; const MYSQL_MULTI_STATEMENT_BATCH_MAX_BYTES: usize = 4 * 1024 * 1024; fn mysql_batch_pool_error_action(db_type: Option, error: &str) -> PoolErrorAction { if error == QUERY_CANCELED { // Dropping an in-flight COM_QUERY future can leave unread packets on the // connection. Do not return that connection to the pool. PoolErrorAction::Discard } else { pool_error_action(db_type, error) } } fn mysql_non_result_batch_end( statements: &[String], start: usize, dialect: db::mysql::MySqlQueryDialect, max_bytes: usize, ) -> usize { let Some(first) = statements.get(start) else { return start; }; if !db::mysql::is_batchable_non_result_query(first, dialect) { return start + 1; } let mut end = start; let mut byte_len = 0usize; while let Some(statement) = statements.get(end) { if end - start >= MYSQL_MULTI_STATEMENT_BATCH_MAX_STATEMENTS || !db::mysql::is_batchable_non_result_query(statement, dialect) { break; } let next_len = byte_len.saturating_add(statement.len()).saturating_add(2); if end > start && next_len > max_bytes { break; } byte_len = next_len; end += 1; } end.max(start + 1) } fn mysql_non_result_pipeline_enabled( statement_count: usize, continue_on_error: bool, mode: crate::connection::MysqlMode, ) -> bool { statement_count > 1 && !continue_on_error && mode == crate::connection::MysqlMode::Normal } async fn execute_mysql_batch_statements( executor: &mut E, statements: &[String], db_type: Option, mysql_dialect: db::mysql::MySqlQueryDialect, cancel_token: Option, continue_on_error: bool, pipeline_non_result_max_bytes: Option, progress: Option<&ExecuteMultiProgressCallback>, ) -> (Vec, Option) where E: MysqlBatchStatementExecutor, { let mut results = Vec::with_capacity(statements.len()); let mut statement_index = 0usize; let table_data_preview = executor.table_data_preview(); while statement_index < statements.len() { if is_canceled(&cancel_token) { results.push(ExecuteMultiResult::execution_error(error_query_result(canceled_error()))); return (results, None); } let batch_end = if let Some(max_bytes) = pipeline_non_result_max_bytes { mysql_non_result_batch_end(statements, statement_index, mysql_dialect, max_bytes) } else { statement_index + 1 }; if batch_end > statement_index + 1 { let outcome = executor.execute_non_result_batch(&statements[statement_index..batch_end]).await; let completed = outcome.results.len(); for (offset, result) in outcome.results.into_iter().enumerate() { results.push(ExecuteMultiResult::success_with_index(result, statement_index + offset)); } if completed > 0 { let last_index = statement_index + completed - 1; let last_result = &results.last().expect("completed MySQL batch result").result; report_execute_multi_progress(progress, last_index, statements.len(), last_result, true, None); } if let Some(error) = outcome.error { let failed_index = statement_index + completed; let action = mysql_batch_pool_error_action(db_type, &error); let result = error_query_result(error.clone()); let backend_error = crate::backend_error::BackendError::from_legacy_backend(&error); report_execute_multi_progress( progress, failed_index, statements.len(), &result, false, Some(backend_error.clone()), ); results.push(ExecuteMultiResult::execution_error_with_backend( result, Some(failed_index), backend_error, )); return (results, Some(action)); } statement_index = batch_end; continue; } let statement = &statements[statement_index]; match executor.execute_statement(statement).await { Ok(statement_results) => { if let Some(result) = statement_results.last() { report_execute_multi_progress( progress, statement_index, statements.len(), &result.result, true, None, ); } results.extend(statement_results.into_iter().map(|result| { ExecuteMultiResult::success_with_index_and_large_values( result.result, statement_index, result.large_value_cells, table_data_preview, ) })); } Err(err) => { let action = mysql_batch_pool_error_action(db_type, &err); let result = error_query_result(err.clone()); report_execute_multi_progress( progress, statement_index, statements.len(), &result, false, Some(crate::backend_error::BackendError::from_legacy_backend(&err)), ); results.push(ExecuteMultiResult::execution_error_with_index(result, statement_index)); // Statement errors are safe to collect, but connection-level failures leave // the protocol state unusable and must still trigger pool cleanup. if !should_continue_batch_after_error(continue_on_error, action) { return (results, Some(action)); } } } statement_index += 1; } (results, None) } #[allow(clippy::too_many_arguments)] async fn execute_multi_mysql( state: &AppState, pool_key: &str, db_type: Option, pool: &db::mysql::MySqlPool, mode: crate::connection::MysqlMode, dialect: db::mysql::MySqlQueryDialect, catalog_dialect: Option, database: &str, statements: &[String], cancel_token: Option, options: QueryExecutionOptions, progress: Option<&ExecuteMultiProgressCallback>, ) -> Result, String> { let trace_id = options.execution_id.as_deref().unwrap_or("none"); let total_started_at = std::time::Instant::now(); let query_timeout = resolve_query_timeout(options.timeout_secs); let operation_budget = operation_budget_for_pool_key(state, pool_key, query_timeout).await; let bare = mode == crate::connection::MysqlMode::Bare; let max_rows = options.max_rows; let max_result_bytes = options.max_result_bytes.filter(|value| *value > 0); let pipeline_non_result_statements = mysql_non_result_pipeline_enabled(statements.len(), options.continue_on_error, mode); let checkout_started_at = std::time::Instant::now(); let mut conn = match db::mysql::get_conn_with_health_check_with_cancel( pool, operation_budget.checkout_timeout, operation_budget.cleanup_timeout, cancel_token.as_ref(), ) .await { Ok(conn) => conn, Err(err) => { if matches!(pool_error_action(db_type, &err), PoolErrorAction::Discard | PoolErrorAction::ReconnectAndRetry) || err == QUERY_CANCELED { state.remove_pool_by_key(pool_key).await; } return Ok(vec![ExecuteMultiResult::execution_error(error_query_result(err))]); } }; let checkout_ms = checkout_started_at.elapsed().as_millis(); apply_oceanbase_mysql_session_timeout(state, pool_key, &mut conn, options.timeout_secs).await?; let catalog_started_at = std::time::Instant::now(); wait_for_result_opt( cancel_token.clone(), query_timeout, db::mysql::apply_catalog_database_context(&mut conn, catalog_dialect, options.catalog.as_deref(), database), ) .await?; let catalog_ms = catalog_started_at.elapsed().as_millis(); let pipeline_non_result_max_bytes = if pipeline_non_result_statements { db::mysql::max_allowed_packet_on_conn(&mut conn) .await .ok() .and_then(db::mysql::mysql_sql_statement_hard_limit) .map(|limit| limit.min(MYSQL_MULTI_STATEMENT_BATCH_MAX_BYTES)) } else { None }; let mut executor = MysqlBatchConnection { conn: &mut conn, cancel_token: cancel_token.clone(), query_timeout, bare, max_rows, max_result_bytes, result_key_columns: &options.result_key_columns, table_data_preview: options.table_data_preview, dialect, diagnostic_trace_id: options.execution_id.as_deref(), }; let statements_started_at = std::time::Instant::now(); let (results, error_action) = execute_mysql_batch_statements( &mut executor, statements, db_type, dialect, cancel_token, options.continue_on_error, pipeline_non_result_max_bytes, progress, ) .await; let statements_ms = statements_started_at.elapsed().as_millis(); drop(executor); log::info!( "[query][mysql-batch] trace_id={} checkout_ms={} catalog_ms={} statements_ms={} total_ms={} result_count={} row_counts={:?}", trace_id, checkout_ms, catalog_ms, statements_ms, total_started_at.elapsed().as_millis(), results.len(), results.iter().map(|result| result.result.rows.len()).collect::>() ); if matches!(error_action, Some(PoolErrorAction::Discard | PoolErrorAction::ReconnectAndRetry)) { state.remove_pool_by_key(pool_key).await; } Ok(results) } fn error_query_result(message: String) -> db::QueryResult { db::QueryResult { columns: vec!["Error".to_string()], column_types: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], rows: vec![vec![serde_json::Value::String(message)]], affected_rows: 0, execution_time_ms: 0, truncated: false, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: Vec::new(), } } fn empty_query_result(execution_time_ms: u128) -> db::QueryResult { db::QueryResult { columns: vec![], column_types: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], rows: vec![], affected_rows: 0, execution_time_ms, truncated: false, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: Vec::new(), } } fn sqlserver_batch_results(results: Vec) -> Vec { results.into_iter().map(ExecuteMultiResult::from).collect() } async fn execute_multi_sqlserver( state: &AppState, pool_key: &str, sql: &str, cancel_token: Option, options: QueryExecutionOptions, ) -> Result, String> { let batches = split_sql_batches(sql); // Read-only check for SQL Server batch path check_read_only_for_connection_multi(state, pool_key, &batches).await?; let mut all_results = Vec::new(); let max_rows = options.max_rows; let query_timeout = resolve_query_timeout(options.timeout_secs); let execution_mode = options.execution_mode; for batch in &batches { if is_canceled(&cancel_token) { let error = canceled_error(); all_results.push(ExecuteMultiResult::execution_error_with_backend( error_query_result(error.clone()), None, crate::backend_error::BackendError::from_legacy_backend(&error), )); break; } let connections = state.connections.read().await; let pool = connections.get(pool_key).ok_or("Connection not found")?; let client = match pool { PoolKind::SqlServer(c) => c.clone(), _ => return Err("Expected SQL Server connection".to_string()), }; drop(connections); let (mut client_guard, lock_wait_ms) = match lock_shared_client_with_wait(&client, cancel_token.clone(), query_timeout).await { Ok(value) => value, Err(err) => { all_results.push(ExecuteMultiResult::execution_error_with_backend( error_query_result(err.clone()), None, crate::backend_error::BackendError::from_legacy_backend(&err), )); break; } }; if !sqlserver_pool_is_current(state, pool_key, &client).await { let error = "SQL Server connection was reset while waiting for the query lock; please retry.".to_string(); all_results.push(ExecuteMultiResult::execution_error_with_backend( error_query_result(error.clone()), None, crate::backend_error::BackendError::from_legacy_backend(&error), )); break; } let execution = async { if execution_mode == QueryExecutionMode::Simple { db::sqlserver::execute_simple_batch_with_max_rows_metadata(&mut client_guard, batch, max_rows).await } else { db::sqlserver::execute_batch_with_max_rows_metadata(&mut client_guard, batch, max_rows).await } }; let result = wait_for_result_opt(cancel_token.clone(), query_timeout, execution).await; drop(client_guard); match result { Ok(results) => all_results.extend(sqlserver_batch_results(results).into_iter().map(|mut item| { item.result.execution_time_ms += lock_wait_ms; item })), Err(e) => { let action = pool_error_action(Some(DatabaseType::SqlServer), &e); all_results.push(ExecuteMultiResult::execution_error_with_backend( error_query_result(e.clone()), None, crate::backend_error::BackendError::from_legacy_backend(&e), )); if matches!(action, PoolErrorAction::Discard | PoolErrorAction::ReconnectAndRetry) { state.remove_pool_by_key(pool_key).await; } if !should_continue_batch_after_error(options.continue_on_error, action) { break; } } } } if all_results.is_empty() { all_results.push(empty_query_result(0).into()); } Ok(all_results) } async fn execute_multi_agent( client: &mut db::agent_driver::AgentDriverClient, database: Option<&str>, statements: &[String], schema: Option<&str>, timeout_secs: Option, ) -> Result { client.execute_batch_typed(database, statements, schema, resolve_query_timeout(timeout_secs)).await } pub async fn execute_statements( state: &AppState, connection_id: &str, database: &str, statements: &[String], schema: Option<&str>, timeout_secs: Option, ) -> Result { let sql_ctx = statements.first().map(|s| s.as_str()).unwrap_or(""); let pool_key = if database.is_empty() { connection_id.to_string() } else { state .get_or_create_pool(connection_id, Some(database)) .await .map_err(|e| query_error_with_omitted_sql_context(&e, sql_ctx))? }; let mut total_affected: u64 = 0; let start = std::time::Instant::now(); let mysql_dialect = connection_mysql_query_dialect(state, connection_id).await; let agent_client = { let conns = state.connections.read().await; match conns.get(&pool_key) { Some(PoolKind::Agent(client)) => Some(client.clone()), _ => None, } }; if let Some(client) = agent_client { let source_client = client.clone(); check_read_only_for_connection_multi(state, &pool_key, statements).await?; let db_type = connection_database_type_for_pool_key(state, &pool_key).await; let execution_schema = schema_for_execution_context(db_type, schema); let rewritten_statements; let statements = if qualifies_unqualified_agent_relations(db_type) { rewritten_statements = statements.iter().map(|sql| sql_for_execution_context(db_type, sql, schema)).collect::>(); rewritten_statements.as_slice() } else { statements }; let mut client = client.lock().await; let database = if database.trim().is_empty() { None } else { Some(database) }; let result = execute_multi_agent(&mut client, database, statements, execution_schema, timeout_secs).await; drop(client); match result { Ok(result) => return Ok(db::QueryResult { execution_time_ms: start.elapsed().as_millis(), ..result }), Err(err) => { if is_agent_execute_batch_unsupported(&err.to_string()) { log::warn!( "Agent does not support execute_batch; falling back to statement-by-statement execution" ); } else { discard_agent_pool_after_typed_error( state, &pool_key, &source_client, &err, RecoveryScope::UserOperation, ) .await; return Err(query_error_with_omitted_sql_context(&err.into_legacy_string(), sql_ctx)); } } } } for (i, sql) in statements.iter().enumerate() { match do_execute( state, &pool_key, mysql_dialect, Some(database), sql, schema, None, QueryExecutionOptions { timeout_secs, ..Default::default() }, ) .await { Ok(result) => { total_affected += result.affected_rows; } Err(e) => { let db_type = connection_database_type(state, connection_id).await; match pool_error_action(db_type, &e) { PoolErrorAction::ReconnectAndRetry => { let db_opt = if database.is_empty() { None } else { Some(database) }; let _ = state.reconnect_pool(connection_id, db_opt).await; } PoolErrorAction::Discard if !db_type.is_some_and(|db_type| database_capabilities::is_agent_type(&db_type)) => { let _ = state.remove_pool_by_key(&pool_key).await; } PoolErrorAction::Discard | PoolErrorAction::Keep => {} } let error = crate::db::agent_driver::append_legacy_error_context( &e, &format!("Statement {} failed; previous {} statement(s) may have been committed.", i + 1, i), ); return Err(query_error_with_omitted_sql_context(&error, sql)); } } } Ok(db::QueryResult { columns: vec![], column_types: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], rows: vec![], affected_rows: total_affected, execution_time_ms: start.elapsed().as_millis(), truncated: false, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: Vec::new(), }) } fn is_agent_execute_batch_unsupported(error: &str) -> bool { let lower = error.to_lowercase(); lower.contains("execute_batch") && (lower.contains("unknown method") || lower.contains("method not found")) } /// Deploy result for Schema Diff: single-connection transactional execution. /// On statement failure the transaction is rolled back by the underlying path. #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct SchemaDiffDeployResult { pub transaction_id: String, pub status: String, pub participants: Vec, pub created_at: String, pub updated_at: String, pub executed_count: usize, pub statement_count: usize, pub error: Option, pub metadata: serde_json::Value, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum SchemaDiffAtomicity { GuaranteedRollback, PartialEffectsPossible, } impl SchemaDiffAtomicity { fn ddl_atomic(self) -> bool { matches!(self, Self::GuaranteedRollback) } } fn database_supports_transactional_ddl(db_type: DatabaseType) -> bool { if resolve_for_db(db_type).has_capability(CAP_TRANSACTIONAL_DDL) { return true; } // SQLite-family DDL is transactional even when dialect registry omits the flag. matches!(db_type, DatabaseType::Sqlite | DatabaseType::Rqlite | DatabaseType::Turso | DatabaseType::CloudflareD1) } fn classify_schema_diff_atomicity( db_type: Option, statements: &[String], has_transactional_path: bool, ) -> SchemaDiffAtomicity { if !has_transactional_path { return SchemaDiffAtomicity::PartialEffectsPossible; } let mut contains_ddl = false; for statement in statements { let risk = match db_type { Some(db_type) => classify_sql_risk_for_database(statement, db_type).ok(), None => None, }; match risk { Some(SqlRisk::Ddl) => { contains_ddl = true; } Some(SqlRisk::Write | SqlRisk::ReadOnly) => {} Some(SqlRisk::Transaction) | None => { return SchemaDiffAtomicity::PartialEffectsPossible; } } } if !contains_ddl { return SchemaDiffAtomicity::GuaranteedRollback; } match db_type { Some(db_type) if database_supports_transactional_ddl(db_type) => SchemaDiffAtomicity::GuaranteedRollback, _ => SchemaDiffAtomicity::PartialEffectsPossible, } } fn executed_count_before_error(error: &str, statement_count: usize) -> usize { let Some(rest) = error.strip_prefix("Statement ") else { return statement_count; }; let digits: String = rest.chars().take_while(|ch| ch.is_ascii_digit()).collect(); let Ok(statement_number) = digits.parse::() else { return statement_count; }; statement_number.saturating_sub(1).min(statement_count) } fn is_destructive_schema_diff_statement(statement: &str) -> bool { let normalized = strip_sql_comments_and_literals(statement).to_ascii_uppercase(); let tokens = normalized .split(|ch: char| !(ch.is_ascii_alphanumeric() || ch == '_')) .filter(|token| !token.is_empty()) .collect::>(); match tokens.first().copied() { Some("DROP" | "TRUNCATE") => true, Some("ALTER") => tokens.iter().skip(1).any(|token| *token == "DROP"), _ => false, } } /// Pure failure mapping used by deploy and unit tests. fn schema_diff_failure_outcome( atomicity: SchemaDiffAtomicity, error: &str, statement_count: usize, ) -> (crate::two_phase_commit::TransactionStatus, usize) { if atomicity.ddl_atomic() { (crate::two_phase_commit::TransactionStatus::RolledBack, 0) } else { (crate::two_phase_commit::TransactionStatus::Mixed, executed_count_before_error(error, statement_count)) } } fn pool_kind_has_transactional_path(pool: &PoolKind) -> bool { match pool { PoolKind::Postgres(_) | PoolKind::Mysql(_, _) | PoolKind::Sqlite(_) | PoolKind::CloudflareD1(_) | PoolKind::ClickHouse(_) | PoolKind::Rqlite(_) | PoolKind::Turso(_) | PoolKind::SqlServer(_) | PoolKind::Agent(_) => true, PoolKind::MessageQueue | PoolKind::Nacos | PoolKind::Consul(_) | PoolKind::HBase(_) | PoolKind::DuckDbWorker(_) | PoolKind::Redis(_) | PoolKind::MongoDb(_) | PoolKind::DynamoDb(_) | PoolKind::Elasticsearch(_) | PoolKind::Easysearch(_) | PoolKind::Meilisearch(_) | PoolKind::VectorDb(_) | PoolKind::InfluxDb(_) | PoolKind::VictoriaMetrics(_) | PoolKind::ExternalDriver { .. } => false, #[cfg(feature = "mq-admin")] PoolKind::Mqtt(_) => false, } } /// Execute Schema Diff deploy SQL as one real single-connection transaction. /// /// - Uses [`execute_statements_in_transaction`] so partial success rolls back. /// - Returns a structured result (never re-executes statements to probe status). /// - Comment-only / empty scripts succeed as `committed` with zero statements. /// - When the target path cannot guarantee DDL atomicity (MySQL/Oracle DDL /// auto-commit, `TxPath::None`, etc.), a failure reports `mixed` and /// `executed_count` reflects the statements that were issued before the /// error, so the caller can warn the user that partial effects may persist. pub async fn execute_schema_diff_deploy( state: &AppState, connection_id: &str, database: &str, statements: &[String], schema: Option<&str>, destructive_confirmed: bool, ) -> SchemaDiffDeployResult { let tx_id = format!("deploy_{}", uuid::Uuid::new_v4()); let now = chrono::Utc::now().to_rfc3339(); let db_type = connection_database_type(state, connection_id).await; let parsed: Vec = statements .iter() .flat_map(|s| { db_type.map_or_else(|| split_sql_statements(s), |dt| crate::sql::split_sql_statements_for_database(s, dt)) }) .map(|s| s.trim().to_string()) .filter(|s| { !s.is_empty() && !s.lines().all(|line| { let t = line.trim(); t.is_empty() || t.starts_with("--") }) }) .collect(); let participant = crate::two_phase_commit::ParticipantInfo { id: "connection".to_string(), name: format!("{connection_id}/{database}"), role: "database".to_string(), }; let destructive_statement_count = parsed.iter().filter(|statement| is_destructive_schema_diff_statement(statement)).count(); if destructive_statement_count > 0 && !destructive_confirmed { return SchemaDiffDeployResult { transaction_id: tx_id, status: crate::two_phase_commit::TransactionStatus::RolledBack.as_str().to_string(), participants: vec![participant], created_at: now.clone(), updated_at: now, executed_count: 0, statement_count: parsed.len(), error: Some("Destructive schema diff SQL requires explicit confirmation".to_string()), metadata: serde_json::json!({ "source": "schema_diff_deploy", "mode": "single_connection_tx", "blocked": "destructive_confirmation_required", "destructive_statement_count": destructive_statement_count, }), }; } if parsed.is_empty() { return SchemaDiffDeployResult { transaction_id: tx_id, status: crate::two_phase_commit::TransactionStatus::Committed.as_str().to_string(), participants: vec![participant], created_at: now.clone(), updated_at: now, executed_count: 0, statement_count: 0, error: None, metadata: serde_json::json!({"source": "schema_diff_deploy", "mode": "single_connection_tx"}), }; } let pool_key = if database.is_empty() { connection_id.to_string() } else { match state.get_or_create_pool(connection_id, Some(database)).await { Ok(key) => key, Err(_) => { return SchemaDiffDeployResult { transaction_id: tx_id.clone(), status: crate::two_phase_commit::TransactionStatus::RolledBack.as_str().to_string(), participants: vec![participant], created_at: now.clone(), updated_at: chrono::Utc::now().to_rfc3339(), executed_count: 0, statement_count: parsed.len(), error: Some("Connection not available for deploy".to_string()), metadata: serde_json::json!({"source": "schema_diff_deploy", "mode": "single_connection_tx"}), }; } } }; let has_transactional_path = { let conns = state.connections.read().await; conns.get(&pool_key).is_some_and(pool_kind_has_transactional_path) }; let atomicity = classify_schema_diff_atomicity(db_type, &parsed, has_transactional_path); match execute_statements_in_transaction_on_pool(state, &pool_key, connection_id, database, &parsed, schema, None) .await { Ok(result) => SchemaDiffDeployResult { transaction_id: tx_id, status: crate::two_phase_commit::TransactionStatus::Committed.as_str().to_string(), participants: vec![participant], created_at: now.clone(), updated_at: chrono::Utc::now().to_rfc3339(), executed_count: parsed.len(), statement_count: parsed.len(), error: None, metadata: serde_json::json!({ "source": "schema_diff_deploy", "mode": "single_connection_tx", "affected_rows": result.affected_rows, "execution_time_ms": result.execution_time_ms, }), }, Err(e) => { let (status, executed_count) = schema_diff_failure_outcome(atomicity, &e, parsed.len()); SchemaDiffDeployResult { transaction_id: tx_id, status: status.as_str().to_string(), participants: vec![participant], created_at: now.clone(), updated_at: chrono::Utc::now().to_rfc3339(), executed_count, statement_count: parsed.len(), error: Some(e.clone()), metadata: serde_json::json!({ "source": "schema_diff_deploy", "mode": "single_connection_tx", "ddl_atomic": atomicity.ddl_atomic(), "atomicity": match atomicity { SchemaDiffAtomicity::GuaranteedRollback => "guaranteed_rollback", SchemaDiffAtomicity::PartialEffectsPossible => "partial_effects_possible", }, "error": e, }), } } } } /// Execute multiple SQL statements within a single transaction. /// For pooled drivers (Postgres/MySQL), uses the driver transaction API. /// For SQLite and already-single-connection drivers (ClickHouse/SqlServer/Agent), /// uses explicit BEGIN/COMMIT/ROLLBACK on the shared connection. /// For databases that don't support explicit transactions (Redis, MongoDB, Oracle), /// executes statements sequentially without transaction. /// If BEGIN fails, returns an error instead of silently falling back to auto-commit. pub async fn execute_statements_in_transaction( state: &AppState, connection_id: &str, database: &str, statements: &[String], schema: Option<&str>, catalog: Option<&str>, ) -> Result { execute_statements_in_transaction_typed(state, connection_id, database, statements, schema, catalog) .await .map_err(QueryExecutionError::into_legacy_string) } /// Execute multiple SQL statements transactionally while retaining typed failures. pub async fn execute_statements_in_transaction_typed( state: &AppState, connection_id: &str, database: &str, statements: &[String], schema: Option<&str>, catalog: Option<&str>, ) -> Result { let sql_ctx = statements.first().map(|s| s.as_str()).unwrap_or(""); let pool_database = query_pool_database(database, catalog); let pool_key = state .get_or_create_pool(connection_id, pool_database) .await .map_err(|e| query_error_with_omitted_sql_context(&e, sql_ctx))?; execute_statements_in_transaction_on_pool_typed( state, &pool_key, connection_id, database, statements, schema, catalog, ) .await } /// Execute multiple SQL statements transactionally on an already-resolved pool. /// This preserves session-scoped pools used by long-running imports. pub async fn execute_statements_in_transaction_on_pool( state: &AppState, pool_key: &str, connection_id: &str, database: &str, statements: &[String], schema: Option<&str>, catalog: Option<&str>, ) -> Result { execute_statements_in_transaction_on_pool_typed( state, pool_key, connection_id, database, statements, schema, catalog, ) .await .map_err(QueryExecutionError::into_legacy_string) } /// Execute a transaction on an already-resolved pool without erasing typed failures. pub async fn execute_statements_in_transaction_on_pool_typed( state: &AppState, pool_key: &str, connection_id: &str, database: &str, statements: &[String], schema: Option<&str>, catalog: Option<&str>, ) -> Result { // Read-only check: intercept all transaction paths before dispatching check_read_only_for_connection_multi(state, pool_key, statements).await?; let start = std::time::Instant::now(); let db_type = connection_database_type(state, connection_id).await; let mysql_catalog_dialect = connection_mysql_catalog_dialect(state, connection_id).await; let operation_budget = configured_operation_budget_for_pool_key(state, pool_key).await; // Clone the pool handle within the lock, then drop it before any async work. let path = { let conns = state.connections.read().await; conns.get(pool_key).map(|p| match p { PoolKind::Postgres(pg) => TxPath::Pg(pg.clone()), PoolKind::Mysql(mp, _mode) => TxPath::Mysql(mp.clone(), false), PoolKind::Sqlite(sq) => TxPath::Sqlite(sq.clone()), PoolKind::CloudflareD1(client) => TxPath::CloudflareD1(client.clone()), PoolKind::ClickHouse(_) | PoolKind::Rqlite(_) | PoolKind::Turso(_) | PoolKind::SqlServer(_) => { TxPath::Explicit } PoolKind::Agent(client) => TxPath::Agent(client.clone()), PoolKind::MessageQueue | PoolKind::Nacos | PoolKind::Consul(_) | PoolKind::HBase(_) => TxPath::None, #[cfg(feature = "mq-admin")] PoolKind::Mqtt(_) => TxPath::None, PoolKind::DuckDbWorker(_) | PoolKind::Redis(_) | PoolKind::MongoDb(_) | PoolKind::DynamoDb(_) | PoolKind::Elasticsearch(_) | PoolKind::Easysearch(_) | PoolKind::Meilisearch(_) | PoolKind::VectorDb(_) | PoolKind::InfluxDb(_) | PoolKind::VictoriaMetrics(_) | PoolKind::ExternalDriver { .. } => TxPath::None, }) }; let result = match path { Some(TxPath::Pg(pool)) => { let cancel_context = state.get_postgres_cancel_context(pool_key).await; exec_tx_pg_inner(pool, statements, schema, start, operation_budget.clone(), cancel_context).await } Some(TxPath::Mysql(pool, _bare)) => exec_tx_mysql_inner( state, pool_key, pool, statements, start, operation_budget.clone(), mysql_catalog_dialect, catalog, database, ) .await .map_err(Into::into), Some(TxPath::Sqlite(pool)) => exec_tx_sqlite_inner(pool, statements, start).await.map_err(Into::into), Some(TxPath::CloudflareD1(client)) => { let sql = statements.join(";\n"); wait_for_query_opt( None, operation_budget.query_timeout, db::cloudflare_d1_driver::execute_query_with_max_rows(&client, &sql, None), ) .await .map_err(Into::into) } Some(TxPath::Agent(client)) => { let result = exec_tx_agent_inner(client.clone(), db_type, Some(database), statements, schema, start).await; if let Err(error) = result.as_ref() { discard_agent_pool_after_typed_error(state, pool_key, &client, error, RecoveryScope::UserOperation) .await; } return result.map_err(QueryExecutionError::Agent); } Some(TxPath::Explicit) => { let mysql_dialect = connection_mysql_query_dialect(state, connection_id).await; exec_tx_explicit_inner(state, pool_key, mysql_dialect, Some(database), statements, schema, start) .await .map_err(Into::into) } Some(TxPath::None) => { let mysql_dialect = connection_mysql_query_dialect(state, connection_id).await; exec_tx_none_inner(state, pool_key, mysql_dialect, Some(database), statements, schema, start) .await .map_err(Into::into) } None => Err("Connection not found for transaction".to_string().into()), }; if let Err(err) = result.as_ref() { discard_pool_after_error(state, pool_key, db_type, &err.to_string()).await; } result } /// Owned pool variants for safe dispatch across async boundaries. enum TxPath { Pg(deadpool_postgres::Pool), Mysql(db::mysql::MySqlPool, bool), Sqlite(db::sqlite::SqliteHandle), CloudflareD1(db::cloudflare_d1_driver::CloudflareD1Client), Agent(Arc), Explicit, None, } // Each of these acquires a dedicated connection and runs all statements within // BEGIN ... COMMIT/ROLLBACK, guaranteeing a single physical connection. async fn exec_tx_pg_inner( pool: deadpool_postgres::Pool, statements: &[String], schema: Option<&str>, start: std::time::Instant, budget: DbOperationBudget, cancel_context: Option, ) -> Result { let mut client = db::postgres::checkout_postgres_client(&pool, None, budget.checkout_timeout) .await .map_err(|e| format!("Failed to acquire connection: {}", e))?; let had_schema = schema.is_some(); if let Some(s) = schema { db::postgres::set_postgres_search_path( &client, s, db::postgres::PostgresSearchPathContext::Transaction, budget.recycle_timeout, ) .await .map_err(|e| format!("SET search_path failed: {}", e))?; } let tx_result = exec_tx_pg_statements(&mut client, statements, &budget, cancel_context).await; // Always reset search_path so the connection is clean when returned to the pool let reset_result = if had_schema { db::postgres::execute_postgres_infra_statement( &client, "RESET search_path", budget.cleanup_timeout, "schema.reset", ) .await .map_err(|err| QueryExecutionError::Legacy(format!("PostgreSQL schema.reset cleanup failed: {err}"))) } else { Ok(0) }; match (tx_result, reset_result) { (Ok(total_affected), Ok(_)) => Ok(db::QueryResult { columns: vec![], column_types: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], rows: vec![], affected_rows: total_affected, execution_time_ms: start.elapsed().as_millis(), truncated: false, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: Vec::new(), }), (Err(e), Ok(_)) => Err(e), (Ok(_), Err(reset_err)) => Err(reset_err), (Err(e), Err(reset_err)) => Err(e.with_context(&reset_err.to_string())), } } async fn exec_tx_pg_statements( client: &mut deadpool_postgres::Client, statements: &[String], budget: &DbOperationBudget, cancel_context: Option, ) -> Result { let tx = tokio::time::timeout(budget.recycle_timeout, client.transaction()) .await .map_err(|_| { format!("Failed to begin transaction: timed out after {} seconds", budget.recycle_timeout.as_secs()) })? .map_err(|e| format!("Failed to begin transaction: {}", e))?; let mut total_affected: u64 = 0; for (i, sql) in statements.iter().enumerate() { let pg_cancel_token = tx.client().cancel_token(); let mut is_server_error = false; let affected = db::postgres::wait_postgres_operation( pg_cancel_token, cancel_context.clone(), budget.query_timeout, budget.cancel_timeout, async { tx.execute(sql, &[]).await.map_err(|error| { is_server_error = error.as_db_error().is_some(); error.to_string() }) }, ) .await .map_err(|e| postgres_transaction_statement_error(i + 1, &e, sql, is_server_error))?; total_affected += affected; } tokio::time::timeout(budget.cleanup_timeout, tx.commit()) .await .map_err(|_| format!("COMMIT timed out after {} seconds", budget.cleanup_timeout.as_secs()))? .map_err(|e| format!("COMMIT failed: {}", e))?; Ok(total_affected) } async fn exec_tx_mysql_inner( state: &AppState, pool_key: &str, pool: db::mysql::MySqlPool, statements: &[String], start: std::time::Instant, budget: DbOperationBudget, catalog_dialect: Option, catalog: Option<&str>, database: &str, ) -> Result { let mut conn = db::mysql::get_conn_with_health_check_with_timeout(&pool, budget.checkout_timeout).await?; apply_oceanbase_mysql_session_timeout(state, pool_key, &mut conn, None).await?; db::mysql::apply_catalog_database_context(&mut conn, catalog_dialect, catalog, database).await?; mysql_query_drop_with_timeout( &mut conn, "START TRANSACTION", budget.recycle_timeout, "Failed to begin transaction", ) .await?; let mut total_affected: u64 = 0; for (i, sql) in statements.iter().enumerate() { match mysql_query_iter_with_timeout(&mut conn, sql, budget.query_timeout).await { Ok(affected) => total_affected += affected, Err(e) => { let _ = mysql_query_drop_with_timeout(&mut conn, "ROLLBACK", budget.cleanup_timeout, "ROLLBACK failed") .await; return Err(query_error_with_omitted_sql_context(&format!("Statement {} failed: {}", i + 1, e), sql)); } } } mysql_query_drop_with_timeout(&mut conn, "COMMIT", budget.cleanup_timeout, "COMMIT failed").await?; Ok(db::QueryResult { columns: vec![], column_types: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], rows: vec![], affected_rows: total_affected, execution_time_ms: start.elapsed().as_millis(), truncated: false, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: Vec::new(), }) } async fn mysql_query_drop_with_timeout( conn: &mut mysql_async::Conn, sql: &str, timeout_duration: Duration, context: &str, ) -> Result<(), String> { tokio::time::timeout(timeout_duration, conn.query_drop(sql)) .await .map_err(|_| format!("{context}: timed out after {} seconds", timeout_duration.as_secs()))? .map_err(|e| format!("{context}: {e}")) } async fn mysql_query_iter_with_timeout( conn: &mut mysql_async::Conn, sql: &str, timeout_duration: Option, ) -> Result { match timeout_duration { Some(timeout_duration) => tokio::time::timeout(timeout_duration, conn.query_iter(sql)) .await .map_err(|_| format!("Query timed out after {} seconds", timeout_duration.as_secs()))? .map(|result| result.affected_rows()) .map_err(|e| e.to_string()), None => conn.query_iter(sql).await.map(|result| result.affected_rows()).map_err(|e| e.to_string()), } } async fn exec_tx_sqlite_inner( pool: db::sqlite::SqliteHandle, statements: &[String], start: std::time::Instant, ) -> Result { let statements = statements.to_vec(); tokio::task::spawn_blocking(move || { pool.with_connection(|conn| { conn.execute_batch("BEGIN").map_err(|e| format!("Failed to begin transaction: {}", e))?; let mut total_affected: u64 = 0; for (i, sql) in statements.iter().enumerate() { match conn.execute_batch(sql) { Ok(_) => total_affected += conn.changes(), Err(e) => { let _ = conn.execute_batch("ROLLBACK"); return Err(query_error_with_omitted_sql_context( &format!("Statement {} failed: {}", i + 1, e), sql, )); } } } conn.execute_batch("COMMIT").map_err(|e| format!("COMMIT failed: {}", e))?; Ok(db::QueryResult { columns: vec![], column_types: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], rows: vec![], affected_rows: total_affected, execution_time_ms: start.elapsed().as_millis(), truncated: false, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: Vec::new(), }) }) }) .await .map_err(|e| e.to_string())? } async fn exec_tx_explicit_inner( state: &AppState, pool_key: &str, mysql_dialect: db::mysql::MySqlQueryDialect, database: Option<&str>, statements: &[String], schema: Option<&str>, start: std::time::Instant, ) -> Result { do_execute( state, pool_key, mysql_dialect, database, "BEGIN TRANSACTION", schema, None, QueryExecutionOptions::default(), ) .await .map_err(|e| format!("Failed to begin transaction: {}", e))?; let mut total_affected: u64 = 0; for (i, sql) in statements.iter().enumerate() { match do_execute(state, pool_key, mysql_dialect, database, sql, schema, None, QueryExecutionOptions::default()) .await { Ok(result) => { total_affected += result.affected_rows; } Err(e) => { if let Err(rb_err) = do_execute( state, pool_key, mysql_dialect, database, "ROLLBACK", schema, None, QueryExecutionOptions::default(), ) .await { log::error!("ROLLBACK failed after statement {} error: {}", i + 1, rb_err); } return Err(query_error_with_omitted_sql_context(&format!("Statement {} failed: {}", i + 1, e), sql)); } } } do_execute(state, pool_key, mysql_dialect, database, "COMMIT", schema, None, QueryExecutionOptions::default()) .await .map_err(|e| format!("COMMIT failed: {}", e))?; Ok(db::QueryResult { columns: vec![], column_types: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], rows: vec![], affected_rows: total_affected, execution_time_ms: start.elapsed().as_millis(), truncated: false, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: Vec::new(), }) } async fn exec_tx_agent_inner( client: Arc, db_type: Option, database: Option<&str>, statements: &[String], schema: Option<&str>, start: std::time::Instant, ) -> Result { let execution_schema = schema_for_execution_context(db_type, schema); let rewritten_statements; let statements = if qualifies_unqualified_agent_relations(db_type) { rewritten_statements = statements.iter().map(|sql| sql_for_execution_context(db_type, sql, schema)).collect::>(); rewritten_statements.as_slice() } else { statements }; let mut client = client.lock().await; let result: db::QueryResult = client.execute_transaction_typed(database, statements, execution_schema).await?; Ok(db::QueryResult { execution_time_ms: start.elapsed().as_millis(), ..result }) } async fn exec_tx_none_inner( state: &AppState, pool_key: &str, mysql_dialect: db::mysql::MySqlQueryDialect, database: Option<&str>, statements: &[String], schema: Option<&str>, start: std::time::Instant, ) -> Result { let mut total_affected: u64 = 0; for (i, sql) in statements.iter().enumerate() { log::info!("[query][tx-none:statement:start] index={}", i + 1); match do_execute(state, pool_key, mysql_dialect, database, sql, schema, None, QueryExecutionOptions::default()) .await { Ok(result) => { total_affected += result.affected_rows; log::info!("[query][tx-none:statement:done] index={} affected_rows={}", i + 1, result.affected_rows); } Err(e) => { log::warn!("Statement {} failed (no transaction support): {}", i + 1, e); return Err(query_error_with_omitted_sql_context( &format!("Statement {} failed: {}. No transaction support for this database type.", i + 1, e), sql, )); } } } Ok(db::QueryResult { columns: vec![], column_types: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], rows: vec![], affected_rows: total_affected, execution_time_ms: start.elapsed().as_millis(), truncated: false, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: Vec::new(), }) } /// Start a manual transaction session, holding a connection from the pool. /// Returns a transaction session ID that must be passed to subsequent calls. pub async fn begin_manual_transaction( state: &AppState, connection_id: &str, database: &str, schema: Option<&str>, catalog: Option<&str>, ) -> Result { begin_transaction_session(state, connection_id, database, schema, catalog, false).await } /// Start a read-only, repeatable snapshot for a database backup. pub async fn begin_database_backup_snapshot( state: &AppState, connection_id: &str, database: &str, ) -> Result { begin_transaction_session(state, connection_id, database, None, None, true).await } fn postgres_transaction_begin_sql(consistent_snapshot: bool) -> &'static str { if consistent_snapshot { "BEGIN TRANSACTION ISOLATION LEVEL REPEATABLE READ READ ONLY" } else { "BEGIN" } } fn mysql_transaction_begin_sql_candidates(consistent_snapshot: bool) -> &'static [&'static str] { if consistent_snapshot { &[ "START TRANSACTION WITH CONSISTENT SNAPSHOT, READ ONLY", "START TRANSACTION WITH CONSISTENT SNAPSHOT", "START TRANSACTION", ] } else { &["START TRANSACTION"] } } fn mysql_transaction_isolation_sql(consistent_snapshot: bool) -> Option<&'static str> { consistent_snapshot.then_some("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ") } fn mysql_error_is_syntax_error(error: &mysql_async::Error) -> bool { matches!(error, mysql_async::Error::Server(server_error) if server_error.code == 1064) } async fn begin_transaction_session( state: &AppState, connection_id: &str, database: &str, schema: Option<&str>, catalog: Option<&str>, consistent_snapshot: bool, ) -> Result { let mysql_catalog_dialect = connection_mysql_catalog_dialect(state, connection_id).await; let pool_database = query_pool_database(database, catalog); let pool_key = state.get_or_create_pool(connection_id, pool_database).await?; // Clone the pool handle under a brief read lock, then drop the lock before // any async I/O — same pattern as do_execute throughout this file. enum TxnPoolHandle { Postgres(deadpool_postgres::Pool), Mysql(db::mysql::MySqlPool), } let pool_handle = { let connections = state.connections.read().await; match connections.get(&pool_key).ok_or("Connection not found")? { PoolKind::Postgres(pg) => TxnPoolHandle::Postgres(pg.clone()), PoolKind::Mysql(mp, _) => TxnPoolHandle::Mysql(mp.clone()), _ => return Err("Manual transaction is not supported for this database type".to_string()), } }; // connections lock released here let txn_conn = match pool_handle { TxnPoolHandle::Postgres(pg_pool) => { let conn = pg_pool.get().await.map_err(|e| format!("Failed to get Postgres connection: {e}"))?; let begin_sql = postgres_transaction_begin_sql(consistent_snapshot); conn.execute_typed(begin_sql, &[]).await.map_err(|e| format!("BEGIN failed: {e}"))?; if let Some(schema) = schema { db::postgres::set_postgres_search_path( &conn, schema, db::postgres::PostgresSearchPathContext::LocalTransaction, db::connection_timeout(), ) .await .map_err(|e| format!("SET search_path failed: {e}"))?; } TxnConnection::Postgres(Box::new(conn)) } TxnPoolHandle::Mysql(mysql_pool) => { let mut conn = mysql_pool.get_conn().await.map_err(|e| format!("Failed to get MySQL connection: {e}"))?; db::mysql::apply_catalog_database_context(&mut conn, mysql_catalog_dialect, catalog, database).await?; if let Some(isolation_sql) = mysql_transaction_isolation_sql(consistent_snapshot) { conn.query_drop(isolation_sql).await.map_err(|e| format!("SET TRANSACTION failed: {e}"))?; } let mut syntax_errors = Vec::new(); for begin_sql in mysql_transaction_begin_sql_candidates(consistent_snapshot) { match conn.query_drop(*begin_sql).await { Ok(()) => { syntax_errors.clear(); break; } Err(error) if mysql_error_is_syntax_error(&error) => { syntax_errors.push(format!("{begin_sql}: {error}")); } Err(error) => return Err(format!("START TRANSACTION failed: {error}")), } } if !syntax_errors.is_empty() { return Err(format!("START TRANSACTION failed for all compatible forms: {}", syntax_errors.join("; "))); } TxnConnection::Mysql(conn) } }; let txn_session_id = uuid::Uuid::new_v4().to_string(); let session = TransactionSession { connection: Arc::new(tokio::sync::Mutex::new(txn_conn)), pool_key: pool_key.clone(), last_activity: std::time::Instant::now(), busy: false, connection_id: connection_id.to_string(), database: database.to_string(), schema: schema.map(|s| s.to_string()), }; { let mut sessions = state.transaction_sessions.write().await; sessions.insert(txn_session_id.clone(), session); } // Schedule idle timeout watcher spawn_txn_idle_watcher(state, txn_session_id.clone()); log::info!("[query][manual_txn:begin] session_id={}", txn_session_id); Ok(txn_session_id) } pub struct ManualTransactionKeepAlive { task: tokio::task::JoinHandle<()>, sessions: Arc>>, txn_session_id: String, } impl Drop for ManualTransactionKeepAlive { fn drop(&mut self) { self.task.abort(); spawn_txn_idle_watcher_for_sessions(Arc::clone(&self.sessions), self.txn_session_id.clone()); } } /// Keep an existing transaction session alive while a caller prepares work for /// that session. The caller must retain the returned guard for the full period; /// dropping it restores the normal five-minute idle rollback behavior. pub async fn keep_manual_transaction_alive( state: &AppState, txn_session_id: &str, ) -> Result { { let mut sessions = state.transaction_sessions.write().await; let session = sessions.get_mut(txn_session_id).ok_or_else(|| { "Transaction session not found or expired; it may have been auto-rolled back due to inactivity".to_string() })?; if !session.busy { session.last_activity = std::time::Instant::now(); } } let sessions = Arc::clone(&state.transaction_sessions); let keep_alive_sessions = Arc::clone(&sessions); let txn_session_id = txn_session_id.to_string(); let keep_alive_txn_session_id = txn_session_id.clone(); let task = tokio::spawn(async move { const KEEPALIVE_INTERVAL: Duration = Duration::from_secs(30); loop { tokio::time::sleep(KEEPALIVE_INTERVAL).await; let should_continue = { let mut guard = sessions.write().await; if let Some(session) = guard.get_mut(&keep_alive_txn_session_id) { if !session.busy { session.last_activity = std::time::Instant::now(); } true } else { false } }; if !should_continue { break; } } }); Ok(ManualTransactionKeepAlive { task, sessions: keep_alive_sessions, txn_session_id }) } /// Execute SQL within an existing manual transaction session. pub async fn execute_in_manual_transaction( state: &AppState, txn_session_id: &str, sql: &str, _database: &str, _schema: Option<&str>, max_rows: Option, ) -> Result, String> { const TXN_IDLE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(300); // Resolve statements and validate before taking the per-session connection // lock. The session stays visible in the map so close/disconnect cleanup can // remove it and roll back once the current DB operation releases the lock. let (pool_key, connection_id) = { let sessions = state.transaction_sessions.read().await; let session = sessions .get(txn_session_id) .ok_or("Transaction session not found or expired; it may have been auto-rolled back due to inactivity")?; (session.pool_key.clone(), session.connection_id.clone()) }; let db_type = connection_database_type(state, &connection_id).await; let statements = db_type.map_or_else( || split_sql_statements(sql), |db_type| crate::sql::split_sql_statements_for_database(sql, db_type), ); if statements.is_empty() { return Ok(vec![empty_query_result(0)]); } // Read-only check while the session is still in the map. If this fails the // session remains intact. check_read_only_for_connection_multi(state, &pool_key, &statements).await?; let connection = { let mut sessions = state.transaction_sessions.write().await; let Some(session) = sessions.get_mut(txn_session_id) else { return Err( "Transaction session not found or expired; it may have been auto-rolled back due to inactivity" .to_string(), ); }; if session.busy { return Err("Transaction session is already executing".to_string()); } if session.last_activity.elapsed() > TXN_IDLE_TIMEOUT { let session = sessions.remove(txn_session_id).expect("session exists"); Some(session.connection) } else { session.busy = true; session.last_activity = std::time::Instant::now(); None } }; if let Some(connection) = connection { let mut conn = connection.lock().await; let _ = rollback_manual_txn_connection(&mut conn).await; return Err("Transaction was auto-rolled back due to 5 minutes of inactivity".to_string()); } let connection = { let sessions = state.transaction_sessions.read().await; sessions .get(txn_session_id) .map(|session| Arc::clone(&session.connection)) .ok_or("Transaction session not found or expired; it may have been auto-rolled back due to inactivity")? }; let row_limit = max_rows.unwrap_or(MAX_ROWS).max(1); let mut results = Vec::with_capacity(statements.len()); let mut conn = connection.lock().await; for (i, statement) in statements.iter().enumerate() { let result = match &mut *conn { TxnConnection::Postgres(conn) => { execute_manual_txn_postgres_statement(conn.as_ref(), statement, row_limit).await } TxnConnection::Mysql(conn) => execute_manual_txn_mysql_statement(conn, statement, row_limit).await, }; match result { Ok(query_result) => results.push(query_result), Err(e) => { // Statement failure ends the transaction. If another cleanup path // already removed the session, it owns the final rollback. let should_rollback = { let mut sessions = state.transaction_sessions.write().await; sessions.remove(txn_session_id).is_some() }; if should_rollback { let _ = rollback_manual_txn_connection(&mut conn).await; } return Err(format!("Statement {} failed: {}. Transaction was auto-rolled back.", i + 1, e)); } } } drop(conn); let should_watch = { let mut sessions = state.transaction_sessions.write().await; if let Some(session) = sessions.get_mut(txn_session_id) { session.busy = false; session.last_activity = std::time::Instant::now(); true } else { false } }; if should_watch { spawn_txn_idle_watcher(state, txn_session_id.to_string()); } Ok(results) } /// Stream a read query through an existing transaction without materializing /// the whole result set. The callback runs once per batch while the same held /// connection and transaction snapshot remain active. pub async fn stream_rows_in_manual_transaction( state: &AppState, txn_session_id: &str, sql: &str, batch_size: usize, mut on_batch: F, ) -> Result where F: FnMut(Vec>) -> Result<(), String> + Send, { const TXN_IDLE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(300); let expired_connection = { let mut sessions = state.transaction_sessions.write().await; let Some(session) = sessions.get_mut(txn_session_id) else { return Err( "Transaction session not found or expired; it may have been auto-rolled back due to inactivity" .to_string(), ); }; if session.busy { return Err("Transaction session is already executing".to_string()); } if session.last_activity.elapsed() > TXN_IDLE_TIMEOUT { Some(sessions.remove(txn_session_id).expect("session exists").connection) } else { session.busy = true; session.last_activity = std::time::Instant::now(); None } }; if let Some(connection) = expired_connection { let mut conn = connection.lock().await; let _ = rollback_manual_txn_connection(&mut conn).await; return Err("Transaction was auto-rolled back due to 5 minutes of inactivity".to_string()); } let connection = { let sessions = state.transaction_sessions.read().await; sessions .get(txn_session_id) .map(|session| Arc::clone(&session.connection)) .ok_or("Transaction session not found or expired; it may have been auto-rolled back due to inactivity")? }; let batch_size = batch_size.max(1); let mut conn = connection.lock().await; let stream_result = match &mut *conn { TxnConnection::Postgres(conn) => { let mut batch = Vec::with_capacity(batch_size); let mut total_rows = 0_u64; let result = db::postgres::stream_select_query_inner_unnamed(conn, sql, None, &mut |item| { if let db::postgres::PostgresQueryStreamItem::Row(row) = item { batch.push(row); total_rows += 1; if batch.len() >= batch_size { on_batch(std::mem::take(&mut batch))?; batch = Vec::with_capacity(batch_size); } } Ok(()) }) .await; match result { Ok(_) if !batch.is_empty() => on_batch(batch).map(|_| total_rows), Ok(_) => Ok(total_rows), Err(error) => Err(error), } } TxnConnection::Mysql(conn) => match conn.query_iter(sql).await { Ok(mut result) => match result.stream::().await { Ok(Some(mut stream)) => { let mut batch = Vec::with_capacity(batch_size); let mut total_rows = 0_u64; let mut error = None; while let Some(row_result) = stream.next().await { match row_result { Ok(row) => { batch.push( (0..row.len()).map(|index| db::mysql::mysql_value_to_json(&row, index)).collect(), ); total_rows += 1; if batch.len() >= batch_size { if let Err(err) = on_batch(std::mem::take(&mut batch)) { error = Some(err); break; } batch = Vec::with_capacity(batch_size); } } Err(err) => { error = Some(format!("Query failed: {err}")); break; } } } if error.is_none() && !batch.is_empty() { if let Err(err) = on_batch(batch) { error = Some(err); } } error.map_or(Ok(total_rows), Err) } Ok(None) => Err("Empty result set stream".to_string()), Err(err) => Err(format!("Query failed: {err}")), }, Err(err) => Err(format!("Query failed: {err}")), }, }; if let Err(err) = &stream_result { let should_rollback = state.transaction_sessions.write().await.remove(txn_session_id).is_some(); if should_rollback { let _ = rollback_manual_txn_connection(&mut conn).await; } return Err(format!("{err}. Transaction was auto-rolled back.")); } drop(conn); let should_watch = { let mut sessions = state.transaction_sessions.write().await; if let Some(session) = sessions.get_mut(txn_session_id) { session.busy = false; session.last_activity = std::time::Instant::now(); true } else { false } }; if should_watch { spawn_txn_idle_watcher(state, txn_session_id.to_string()); } stream_result } async fn rollback_manual_txn_connection(conn: &mut TxnConnection) -> Result<(), String> { match conn { TxnConnection::Postgres(conn) => { conn.execute_typed("ROLLBACK", &[]).await.map_err(|e| format!("ROLLBACK failed: {e}"))?; } TxnConnection::Mysql(conn) => { conn.query_drop("ROLLBACK").await.map_err(|e| format!("ROLLBACK failed: {e}"))?; } } Ok(()) } /// Spawn a background task that removes and rolls back a transaction session /// after 5 minutes of inactivity. The task does not hold the global lock across /// I/O: it briefly checks the map, and if the session exists and is expired, /// removes it, drops the lock, then rolls back the held connection. /// /// Safety: if multiple watchers exist for the same session ID (e.g. due to /// a race), only the one that actually finds the session in the map and /// observes an elapsed time >= timeout will remove and roll back. Others /// will see a missing session or a non-expired one and exit harmlessly. fn spawn_txn_idle_watcher(state: &AppState, txn_session_id: String) { let sessions = Arc::clone(&state.transaction_sessions); spawn_txn_idle_watcher_for_sessions(sessions, txn_session_id); } fn spawn_txn_idle_watcher_for_sessions( sessions: Arc>>, txn_session_id: String, ) { tokio::spawn(async move { const TXN_IDLE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(300); tokio::time::sleep(TXN_IDLE_TIMEOUT).await; let removed: Option = { let mut guard = sessions.write().await; match guard.get(&txn_session_id) { Some(session) if !session.busy && session.last_activity.elapsed() >= TXN_IDLE_TIMEOUT => { guard.remove(&txn_session_id) } _ => None, } }; if let Some(session) = removed { let mut conn = session.connection.lock().await; let _ = rollback_manual_txn_connection(&mut conn).await; log::info!( "[query][manual_txn:idle_timeout] session_id={} auto-rolled back after 5 minutes of inactivity", txn_session_id ); } }); } async fn execute_manual_txn_postgres_statement( conn: &deadpool_postgres::Object, sql: &str, row_limit: usize, ) -> Result { if db::postgres::postgres_statement_returns_rows(sql) { db::postgres::execute_select_query_unnamed(conn, sql, std::time::Instant::now(), row_limit).await } else { let affected = conn.execute_typed(sql, &[]).await.map_err(|e| format!("Query failed: {e}"))?; Ok(db::QueryResult { columns: vec![], column_types: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], rows: vec![], affected_rows: affected, execution_time_ms: 0, truncated: false, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: Vec::new(), }) } } async fn execute_manual_txn_mysql_statement( conn: &mut mysql_async::Conn, sql: &str, row_limit: usize, ) -> Result { if db::mysql::is_result_set_query(sql, db::mysql::MySqlQueryDialect::default()) { let start = std::time::Instant::now(); let mut result = conn.query_iter(sql).await.map_err(|e| format!("Query failed: {e}"))?; let columns: Vec = result.columns_ref().iter().map(|c| c.name_str().to_string()).collect(); let column_types: Vec = result.columns_ref().iter().map(db::mysql::mysql_column_type_name).collect(); let mut data: Vec> = Vec::with_capacity(row_limit.min(1024)); let mut stream = result .stream::() .await .map_err(|e| format!("Query failed: {e}"))? .ok_or_else(|| "Empty result set stream".to_string())?; let mut truncated = false; while let Some(row) = stream.next().await { if data.len() >= row_limit { truncated = true; break; } let row = row.map_err(|e| format!("Query failed: {e}"))?; data.push((0..row.len()).map(|i| db::mysql::mysql_value_to_json(&row, i)).collect()); } Ok(db::QueryResult { columns, column_types, column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], rows: data, affected_rows: 0, execution_time_ms: start.elapsed().as_millis(), truncated, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: Vec::new(), }) } else { let result = conn.query_iter(sql).await.map_err(|e| format!("Query failed: {e}"))?; let affected_rows = result.affected_rows(); result.drop_result().await.map_err(|e| format!("Query failed: {e}"))?; Ok(db::QueryResult { columns: vec![], column_types: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], rows: vec![], affected_rows, execution_time_ms: 0, truncated: false, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: Vec::new(), }) } } /// Commit an existing manual transaction session. pub async fn commit_manual_transaction(state: &AppState, txn_session_id: &str) -> Result { let session = { let mut sessions = state.transaction_sessions.write().await; sessions.remove(txn_session_id).ok_or("Transaction session not found")? }; let mut conn = session.connection.lock().await; match &mut *conn { TxnConnection::Postgres(conn) => { conn.execute_typed("COMMIT", &[]).await.map_err(|e| format!("COMMIT failed: {e}"))?; } TxnConnection::Mysql(conn) => { conn.query_drop("COMMIT").await.map_err(|e| format!("COMMIT failed: {e}"))?; } } log::info!("[query][manual_txn:commit] session_id={}", txn_session_id); Ok(db::QueryResult { columns: vec![], column_types: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], rows: vec![], affected_rows: 0, execution_time_ms: 0, truncated: false, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: Vec::new(), }) } /// Rollback an existing manual transaction session. pub async fn rollback_manual_transaction(state: &AppState, txn_session_id: &str) -> Result { let session = { let mut sessions = state.transaction_sessions.write().await; sessions.remove(txn_session_id).ok_or("Transaction session not found")? }; let mut conn = session.connection.lock().await; rollback_manual_txn_connection(&mut conn).await?; log::info!("[query][manual_txn:rollback] session_id={}", txn_session_id); Ok(db::QueryResult { columns: vec![], column_types: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], rows: vec![], affected_rows: 0, execution_time_ms: 0, truncated: false, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: Vec::new(), }) } #[cfg(test)] mod tests { use super::*; #[test] fn redshift_queries_prefer_text_protocol() { assert!(postgres_prefers_text_protocol(Some(DatabaseType::Redshift))); assert!(!postgres_prefers_text_protocol(Some(DatabaseType::Postgres))); assert!(!postgres_prefers_text_protocol(None)); } #[test] fn schema_diff_destructive_detection_covers_drop_and_alter_drop() { assert!(is_destructive_schema_diff_statement("DROP INDEX idx_users_email ON users")); assert!(is_destructive_schema_diff_statement("TRUNCATE TABLE audit_log")); assert!(is_destructive_schema_diff_statement( "ALTER TABLE users DROP COLUMN legacy_code, DROP INDEX idx_legacy" )); } #[test] fn schema_diff_destructive_detection_ignores_comments_literals_and_identifiers() { assert!(!is_destructive_schema_diff_statement("-- DROP INDEX idx_fake\nSELECT 1")); assert!(!is_destructive_schema_diff_statement("SELECT 'DROP TABLE users'")); assert!(!is_destructive_schema_diff_statement("ALTER TABLE \"DROP INDEX audit\" ADD COLUMN note TEXT")); } #[cfg(unix)] use crate::db::agent_driver::{AgentDriverClient, AgentLaunchSpec}; use crate::models::connection::{default_redis_key_separator, ConnectionConfig, DatabaseType}; #[cfg(unix)] use crate::plugins::{ InstalledPlugin, PluginDriverManifest, PluginDriverSession, PluginManifest, PluginRuntimeEnv, }; use crate::storage::Storage; #[cfg(unix)] async fn spawn_agent_batch_timeout_test_client() -> (AgentDriverClient, tempfile::NamedTempFile) { use std::io::Write; let mut script = tempfile::NamedTempFile::new().unwrap(); write!( script, r#"import json import sys import time print(json.dumps({{"ready": True}}), flush=True) for line in sys.stdin: request = json.loads(line) statements = request.get("params", {{}}).get("statements", []) time.sleep(1.2 if statements == ["slow"] else 0.05) result = {{ "columns": [], "column_types": [], "column_sortables": [], "rows": [], "affected_rows": 1, "execution_time_ms": 50, "truncated": False, "session_id": None, "has_more": False }} print(json.dumps({{"jsonrpc": "2.0", "id": request["id"], "result": result}}), flush=True) "# ) .unwrap(); script.flush().unwrap(); let client = AgentDriverClient::spawn( AgentLaunchSpec::new("python3").with_args([script.path().to_string_lossy().to_string()]), ) .await .unwrap(); (client, script) } #[cfg(unix)] #[tokio::test] async fn execute_multi_agent_zero_timeout_waits_for_response() { let (mut client, _script) = spawn_agent_batch_timeout_test_client().await; let result = execute_multi_agent(&mut client, None, &["fast".to_string()], None, Some(0)).await.unwrap(); assert_eq!(result.affected_rows, 1); } #[cfg(unix)] #[tokio::test] async fn execute_multi_agent_positive_timeout_still_expires() { let (mut client, _script) = spawn_agent_batch_timeout_test_client().await; let error = execute_multi_agent(&mut client, None, &["slow".to_string()], None, Some(1)).await.unwrap_err(); assert!(matches!( error, AgentCallError::Timeout { stage: AgentErrorStage::Execute, operation_outcome: AgentOperationOutcome::Unknown, } )); } #[cfg(unix)] #[tokio::test] async fn execute_multi_agent_default_timeout_keeps_normal_execution() { let (mut client, _script) = spawn_agent_batch_timeout_test_client().await; let result = execute_multi_agent(&mut client, None, &["fast".to_string()], None, None).await.unwrap(); assert_eq!(result.affected_rows, 1); assert_eq!(resolve_query_timeout(None), Some(QUERY_TIMEOUT)); } #[test] fn external_catalog_queries_do_not_bind_database_during_pool_creation() { assert_eq!(query_pool_database("bi", Some("paimon_catalog")), None); assert_eq!(query_pool_database("bi", None), Some("bi")); assert_eq!(query_pool_database("", None), None); } #[tokio::test] async fn query_and_transaction_paths_resolve_catalog_dialect_from_connection() { let dir = std::env::temp_dir().join(format!("dbx-catalog-dialect-{}", uuid::Uuid::new_v4())); std::fs::create_dir_all(&dir).unwrap(); let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); let state = AppState::new(storage); let mut doris = test_connection_config(DatabaseType::Doris); doris.id = "doris".to_string(); let mut starrocks = test_connection_config(DatabaseType::StarRocks); starrocks.id = "starrocks".to_string(); { let mut configs = state.configs.write().await; configs.insert(doris.id.clone(), doris); configs.insert(starrocks.id.clone(), starrocks); } assert_eq!( connection_mysql_catalog_dialect_for_pool_key(&state, "doris:bi").await, Some(db::mysql::MySqlCatalogDialect::Doris) ); assert_eq!( connection_mysql_catalog_dialect(&state, "starrocks").await, Some(db::mysql::MySqlCatalogDialect::StarRocks) ); drop(state); let _ = std::fs::remove_dir_all(dir); } #[test] fn schema_diff_atomicity_marks_mysql_ddl_as_partial() { let atomicity = classify_schema_diff_atomicity( Some(DatabaseType::Mysql), &["CREATE TABLE users (id INT)".to_string(), "ALTER TABLE users ADD COLUMN name VARCHAR(32)".to_string()], true, ); assert_eq!(atomicity, SchemaDiffAtomicity::PartialEffectsPossible); } #[test] fn schema_diff_atomicity_marks_oracle_ddl_as_partial() { let atomicity = classify_schema_diff_atomicity( Some(DatabaseType::Oracle), &["CREATE TABLE users (id INT)".to_string(), "ALTER TABLE users ADD name VARCHAR2(32)".to_string()], true, ); assert_eq!(atomicity, SchemaDiffAtomicity::PartialEffectsPossible); } #[test] fn schema_diff_atomicity_keeps_postgres_ddl_atomic() { let atomicity = classify_schema_diff_atomicity( Some(DatabaseType::Postgres), &["CREATE TABLE users (id INT)".to_string(), "ALTER TABLE users ADD COLUMN name VARCHAR(32)".to_string()], true, ); assert_eq!(atomicity, SchemaDiffAtomicity::GuaranteedRollback); } #[test] fn schema_diff_atomicity_keeps_dml_atomic_when_tx_path_exists() { let atomicity = classify_schema_diff_atomicity( Some(DatabaseType::Mysql), &["INSERT INTO users VALUES (1)".to_string(), "UPDATE users SET active = 1 WHERE id = 1".to_string()], true, ); assert_eq!(atomicity, SchemaDiffAtomicity::GuaranteedRollback); } #[test] fn schema_diff_atomicity_marks_missing_transaction_path_as_partial() { let atomicity = classify_schema_diff_atomicity( Some(DatabaseType::Postgres), &["CREATE TABLE users (id INT)".to_string()], false, ); assert_eq!(atomicity, SchemaDiffAtomicity::PartialEffectsPossible); } #[test] fn executed_count_before_error_uses_failing_statement_index() { assert_eq!(executed_count_before_error("Statement 1 failed: syntax error", 3), 0); assert_eq!(executed_count_before_error("Statement 2 failed: syntax error", 3), 1); assert_eq!(executed_count_before_error("Statement 3 failed: syntax error", 3), 2); } #[test] fn schema_diff_failure_outcome_rolls_back_when_atomic() { let (status, executed) = schema_diff_failure_outcome(SchemaDiffAtomicity::GuaranteedRollback, "Statement 2 failed: syntax error", 3); assert_eq!(status, crate::two_phase_commit::TransactionStatus::RolledBack); assert_eq!(executed, 0); } #[test] fn schema_diff_failure_outcome_reports_mixed_with_partial_count() { let (status, executed) = schema_diff_failure_outcome( SchemaDiffAtomicity::PartialEffectsPossible, "Statement 2 failed: syntax error", 3, ); assert_eq!(status, crate::two_phase_commit::TransactionStatus::Mixed); assert_eq!(executed, 1); } #[test] fn schema_diff_atomicity_keeps_sqlite_ddl_atomic() { let atomicity = classify_schema_diff_atomicity( Some(DatabaseType::Sqlite), &["CREATE TABLE users (id INTEGER)".to_string()], true, ); assert_eq!(atomicity, SchemaDiffAtomicity::GuaranteedRollback); } /// MySQL: first DDL may already commit; second fails → mixed + executed_count = 1. #[test] fn mysql_second_ddl_failure_maps_to_mixed_with_partial_executed_count() { let stmts = ["CREATE TABLE t1 (id INT)".to_string(), "CREATE TABLE t2 (id INT)".to_string()]; let atomicity = classify_schema_diff_atomicity(Some(DatabaseType::Mysql), &stmts, true); assert_eq!(atomicity, SchemaDiffAtomicity::PartialEffectsPossible); let (status, executed) = schema_diff_failure_outcome(atomicity, "Statement 2 failed: table already exists", stmts.len()); assert_eq!(status, crate::two_phase_commit::TransactionStatus::Mixed); assert_eq!(executed, 1); } /// Oracle: same non-transactional DDL semantics as MySQL for deploy status. #[test] fn oracle_second_ddl_failure_maps_to_mixed_with_partial_executed_count() { let stmts = ["CREATE TABLE t1 (id NUMBER)".to_string(), "ALTER TABLE t1 ADD name VARCHAR2(32)".to_string()]; let atomicity = classify_schema_diff_atomicity(Some(DatabaseType::Oracle), &stmts, true); assert_eq!(atomicity, SchemaDiffAtomicity::PartialEffectsPossible); let (status, executed) = schema_diff_failure_outcome(atomicity, "Statement 2 failed: ORA-00942", stmts.len()); assert_eq!(status, crate::two_phase_commit::TransactionStatus::Mixed); assert_eq!(executed, 1); } /// Postgres transactional DDL: second fails → rolled_back + executed_count = 0. #[test] fn postgres_second_ddl_failure_maps_to_rolled_back_zero_executed() { let stmts = ["CREATE TABLE t1 (id INT)".to_string(), "CREATE TABLE t2 (id INT)".to_string()]; let atomicity = classify_schema_diff_atomicity(Some(DatabaseType::Postgres), &stmts, true); assert_eq!(atomicity, SchemaDiffAtomicity::GuaranteedRollback); let (status, executed) = schema_diff_failure_outcome(atomicity, "Statement 2 failed: relation already exists", stmts.len()); assert_eq!(status, crate::two_phase_commit::TransactionStatus::RolledBack); assert_eq!(executed, 0); } #[test] fn query_execution_mode_deserializes_simple_client_value() { let mode: QueryExecutionMode = serde_json::from_str("\"simple\"").unwrap(); assert_eq!(mode, QueryExecutionMode::Simple); assert_eq!(QueryExecutionMode::default(), QueryExecutionMode::Standard); } #[test] fn query_execution_mode_deserializes_postgres_read_only_transaction() { let mode: QueryExecutionMode = serde_json::from_str("\"postgres_read_only_transaction\"").unwrap(); assert_eq!(mode, QueryExecutionMode::PostgresReadOnlyTransaction); } #[test] fn postgres_read_only_transaction_requires_postgres_and_isolated_session() { let mut options = QueryExecutionOptions { execution_mode: QueryExecutionMode::PostgresReadOnlyTransaction, ..Default::default() }; assert!(validate_query_execution_mode(Some(DatabaseType::Mysql), "SELECT 1", &options).is_err()); assert!(validate_query_execution_mode(Some(DatabaseType::Postgres), "SELECT 1", &options).is_err()); options.client_session_id = Some("tab:explain:execution".to_string()); assert_eq!(validate_query_execution_mode(Some(DatabaseType::Postgres), "SELECT 1", &options), Ok(())); assert!(validate_query_execution_mode(Some(DatabaseType::Postgres), "SELECT 1; SELECT 2", &options).is_err()); } fn test_connection_config(db_type: DatabaseType) -> ConnectionConfig { ConnectionConfig { docs_notes_path: None, id: "conn-1".to_string(), name: "Connection".to_string(), note: String::new(), db_type, driver_profile: None, driver_label: None, url_params: None, agent_java_options: Vec::new(), host: "localhost".to_string(), port: 0, username: String::new(), password: String::new(), database: None, default_schema: None, visible_databases: None, visible_schemas: None, show_system_schemas: false, attached_databases: Vec::new(), init_script: None, color: None, transport_layers: Vec::new(), connect_timeout_secs: 10, query_timeout_secs: 30, idle_timeout_secs: 60, keepalive_interval_secs: 30, ssl: false, ca_cert_path: String::new(), client_cert_path: String::new(), client_key_path: String::new(), sysdba: false, oracle_connection_type: None, connection_string: None, redis_connection_mode: None, redis_sentinel_master: String::new(), redis_sentinel_nodes: String::new(), redis_sentinel_username: String::new(), redis_sentinel_password: String::new(), redis_sentinel_tls: false, redis_cluster_nodes: String::new(), redis_key_separator: default_redis_key_separator(), redis_scan_page_size: None, redis_database_aliases: Default::default(), etcd_endpoints: String::new(), gbase_server: String::new(), informix_server: String::new(), external_config: None, jdbc_driver_class: None, jdbc_driver_paths: Vec::new(), one_time: false, save_password: true, read_only: false, is_production: false, production_databases: vec![], database_info: None, } } #[cfg(feature = "dynamodb")] #[tokio::test] #[ignore = "requires DBX_DYNAMODB_ENDPOINT and an orders table"] async fn live_dynamodb_editor_scan_serializes_one_thousand_rows() { let endpoint = std::env::var("DBX_DYNAMODB_ENDPOINT").expect("DBX_DYNAMODB_ENDPOINT is required"); let (ssl, address) = endpoint .strip_prefix("https://") .map(|address| (true, address)) .or_else(|| endpoint.strip_prefix("http://").map(|address| (false, address))) .expect("DynamoDB endpoint must start with http:// or https://"); let (host, port) = address.rsplit_once(':').expect("DynamoDB endpoint must include a port"); let dir = std::env::temp_dir().join(format!("dbx-query-dynamodb-{}", uuid::Uuid::new_v4())); std::fs::create_dir_all(&dir).unwrap(); let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); let state = AppState::new(storage); let mut config = test_connection_config(DatabaseType::DynamoDb); config.host = host.to_string(); config.port = port.parse().expect("valid DynamoDB port"); config.username = "dummy".to_string(); config.password = "dummy".to_string(); config.database = Some("us-east-1".to_string()); config.ssl = ssl; let client = db::dynamodb_driver::connect(&config, host, config.port).unwrap(); db::dynamodb_driver::test_connection(&client, Duration::from_secs(5)).await.unwrap(); state.configs.write().await.insert(config.id.clone(), config.clone()); state.connections.write().await.insert(config.id.clone(), PoolKind::DynamoDb(client)); let results = execute_multi_core_with_options_for_client_and_progress_typed( &state, &config.id, "us-east-1", "DBX DYNAMODB SCAN\ntable: \"orders\"\nlimit: 1000", None, None, QueryExecutionOptions { max_rows: Some(1000), ..Default::default() }, None, ) .await .unwrap(); assert_eq!(results.len(), 1); assert_eq!(results[0].result.rows.len(), 1000); let serialized = serde_json::to_string(&results).unwrap(); assert!(!serialized.is_empty()); drop(state); let _ = std::fs::remove_dir_all(dir); } async fn agent_error_state( disposition: &str, ) -> (AppState, std::path::PathBuf, std::sync::Arc) { let dir = std::env::temp_dir().join(format!("dbx-query-agent-error-{}", uuid::Uuid::new_v4())); std::fs::create_dir_all(&dir).unwrap(); let script_path = dir.join("agent.py"); std::fs::write( &script_path, format!( r#"import json, sys print(json.dumps({{'ready': True}}), flush=True) for line in sys.stdin: req = json.loads(line) if req['method'] == 'handshake': response = {{ 'jsonrpc': '2.0', 'id': req['id'], 'result': {{'protocolVersion': 2, 'agentProtocolVersion': 2, 'capabilities': ['multi_session']}} }} elif req['method'] in ('execute_query', 'execute_batch', 'execute_transaction'): response = {{ 'jsonrpc': '2.0', 'id': req['id'], 'error': {{ 'code': -1, 'message': 'injected Agent failure', 'data': {{ 'category': 'resource', 'retryable': False, 'sessionDisposition': '{disposition}', 'stage': 'execute' }} }} }} else: response = {{'jsonrpc': '2.0', 'id': req['id'], 'result': {{}}}} print(json.dumps(response), 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::Dameng)); state.connections.write().await.insert( "conn-1".to_string(), PoolKind::agent(crate::db::agent_driver::AgentDriverClient::shared_session( runtime.clone(), "session-1".to_string(), )), ); (state, dir, runtime) } #[tokio::test] async fn agent_query_replace_runtime_error_detaches_pool_and_stops_runtime() { let (state, dir, runtime) = agent_error_state("replace_runtime").await; let error = execute_sql_statement(&state, "conn-1", "", "SELECT 1", None, None).await.unwrap_err(); assert!(error.contains("injected Agent failure")); assert!(!state.connections.read().await.contains_key("conn-1")); assert!(runtime.is_failed()); runtime.kill(); let _ = std::fs::remove_dir_all(dir); } #[tokio::test] async fn single_statement_multi_core_keeps_query_error_typed() { let (state, dir, runtime) = agent_error_state("keep").await; let error = execute_multi_core_with_options_for_client_and_progress_typed( &state, "conn-1", "", "SELECT 1", None, None, QueryExecutionOptions::default(), None, ) .await .unwrap_err(); assert!(matches!(error, QueryExecutionError::Agent(_))); runtime.kill(); let _ = std::fs::remove_dir_all(dir); } #[tokio::test] async fn native_pre_dispatch_cancellation_stays_typed() { let dir = std::env::temp_dir().join(format!("dbx-query-native-cancel-{}", uuid::Uuid::new_v4())); std::fs::create_dir_all(&dir).unwrap(); let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); let state = AppState::new(storage); let connection_id = "sqlite-cancel"; let sqlite = db::sqlite::connect_path_create_if_missing(dir.join("query.db").to_str().unwrap()).await.unwrap(); state.connections.write().await.insert(connection_id.to_string(), PoolKind::Sqlite(sqlite)); state.configs.write().await.insert(connection_id.to_string(), test_connection_config(DatabaseType::Sqlite)); let cancel_token = CancellationToken::new(); cancel_token.cancel(); let error = execute_sql_statement_with_options_typed( &state, connection_id, "", "SELECT 1", None, Some(cancel_token), QueryExecutionOptions::default(), ) .await .unwrap_err(); assert!(matches!( &error, QueryExecutionError::Canceled { stage: AgentErrorStage::Request, operation_outcome: AgentOperationOutcome::NotStarted, } )); let backend_error = error.into_backend_error(); assert_eq!(backend_error.code(), "DBX-JDBC-2003"); assert_eq!(backend_error.source(), crate::backend_error::BackendErrorSource::LegacyBackend); assert_eq!(backend_error.operation_outcome(), crate::backend_error::BackendOperationOutcome::NotStarted); let _ = std::fs::remove_dir_all(dir); } #[tokio::test] async fn agent_pre_dispatch_cancellation_stays_local_and_typed() { let (state, dir, runtime) = agent_error_state("keep").await; let cancel_token = CancellationToken::new(); cancel_token.cancel(); let error = execute_sql_statement_with_options_typed( &state, "conn-1", "", "SELECT 1", None, Some(cancel_token), QueryExecutionOptions::default(), ) .await .unwrap_err(); assert!(matches!( &error, QueryExecutionError::Canceled { stage: AgentErrorStage::Request, operation_outcome: AgentOperationOutcome::NotStarted, } )); let backend_error = error.into_backend_error(); assert_eq!(backend_error.code(), "DBX-JDBC-2003"); assert_eq!(backend_error.source(), crate::backend_error::BackendErrorSource::LegacyBackend); assert_eq!(backend_error.operation_outcome(), crate::backend_error::BackendOperationOutcome::NotStarted); assert!(!runtime.is_failed()); runtime.kill(); let _ = std::fs::remove_dir_all(dir); } #[tokio::test] async fn agent_transaction_replace_runtime_error_detaches_pool_and_stops_runtime() { let (state, dir, runtime) = agent_error_state("replace_runtime").await; let error = execute_statements_in_transaction_on_pool( &state, "conn-1", "conn-1", "", &["UPDATE test_table SET value = 1".to_string()], None, None, ) .await .unwrap_err(); assert!(error.contains("injected Agent failure")); assert!(!state.connections.read().await.contains_key("conn-1")); assert!(runtime.is_failed()); runtime.kill(); let _ = std::fs::remove_dir_all(dir); } #[tokio::test] async fn agent_batch_replace_runtime_error_detaches_pool_and_stops_runtime() { let (state, dir, runtime) = agent_error_state("replace_runtime").await; let error = execute_statements(&state, "conn-1", "", &["UPDATE test_table SET value = 1".to_string()], None, None) .await .unwrap_err(); assert!(error.contains("injected Agent failure")); assert!(!state.connections.read().await.contains_key("conn-1")); assert!(runtime.is_failed()); runtime.kill(); let _ = std::fs::remove_dir_all(dir); } #[tokio::test] async fn agent_quarantine_error_removes_only_target_pool() { let (state, dir, runtime) = agent_error_state("quarantine").await; let error = execute_sql_statement(&state, "conn-1", "", "SELECT 1", None, None).await.unwrap_err(); assert!(error.contains("injected Agent failure")); assert!(!state.connections.read().await.contains_key("conn-1")); assert!(!runtime.is_failed()); runtime.kill(); let _ = std::fs::remove_dir_all(dir); } struct FakeMysqlBatchExecutor { outcomes: std::collections::VecDeque, String>>, executed: Vec, } impl MysqlBatchStatementExecutor for FakeMysqlBatchExecutor { async fn execute_statement(&mut self, statement: &str) -> Result, String> { self.executed.push(statement.to_string()); self.outcomes.pop_front().expect("test outcome for statement") } } fn mysql_query_result(result: db::QueryResult) -> db::mysql::MySqlQueryResult { db::mysql::MySqlQueryResult { result, large_value_cells: Vec::new() } } fn mysql_batch_result(result: db::QueryResult) -> Result, String> { Ok(vec![mysql_query_result(result)]) } struct FakePipelinedMysqlBatchExecutor { batch_outcomes: std::collections::VecDeque, statement_outcomes: std::collections::VecDeque, String>>, batches: Vec>, statements: Vec, } impl MysqlBatchStatementExecutor for FakePipelinedMysqlBatchExecutor { async fn execute_statement(&mut self, statement: &str) -> Result, String> { self.statements.push(statement.to_string()); self.statement_outcomes.pop_front().expect("test outcome for single statement") } async fn execute_non_result_batch(&mut self, statements: &[String]) -> db::mysql::MySqlNonResultBatchOutcome { self.batches.push(statements.to_vec()); self.batch_outcomes.pop_front().expect("test outcome for pipelined statements") } } async fn assert_sqlite_batch_error_behavior(failure_first: bool, continue_on_error: bool) { let dir = std::env::temp_dir().join(format!("dbx-query-batch-error-{}", uuid::Uuid::new_v4())); std::fs::create_dir_all(&dir).unwrap(); let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); let state = AppState::new(storage); let connection_id = "sqlite-batch"; let sqlite = db::sqlite::connect_path_create_if_missing(dir.join("query.db").to_str().unwrap()).await.unwrap(); state.connections.write().await.insert(connection_id.to_string(), PoolKind::Sqlite(sqlite)); state.configs.write().await.insert(connection_id.to_string(), test_connection_config(DatabaseType::Sqlite)); let sql = if failure_first { "INSERT INTO missing_table VALUES (1); CREATE TABLE executed_after_error (id INTEGER);" } else { "CREATE TABLE before_error (id INTEGER); INSERT INTO missing_table VALUES (1); CREATE TABLE executed_after_error (id INTEGER);" }; let results = execute_multi_core_with_options( &state, connection_id, "", sql, None, None, QueryExecutionOptions { continue_on_error, ..Default::default() }, ) .await .unwrap(); let error_index = usize::from(!failure_first); assert_eq!(results[error_index].columns, vec!["Error"]); assert_eq!( results.len(), if failure_first { 1 + usize::from(continue_on_error) } else { 2 + usize::from(continue_on_error) } ); let table_check = execute_sql_statement( &state, connection_id, "", "SELECT name FROM sqlite_master WHERE type = 'table' AND name = 'executed_after_error'", None, None, ) .await .unwrap(); assert_eq!(!table_check.rows.is_empty(), continue_on_error); } #[test] fn agent_execute_batch_unsupported_detects_case_insensitive_method_errors() { assert!(is_agent_execute_batch_unsupported("Agent RPC error (-1): unknown method: execute_batch")); assert!(is_agent_execute_batch_unsupported("Agent RPC error (-1): Unknown method: execute_batch")); assert!(is_agent_execute_batch_unsupported("Agent RPC error (-32601): Method not found: execute_batch")); } #[test] fn agent_execute_batch_unsupported_ignores_unrelated_errors() { assert!(!is_agent_execute_batch_unsupported("ORA-00955: name is already used by an existing object")); assert!(!is_agent_execute_batch_unsupported("Agent RPC error (-1): unknown method: execute_query")); } #[test] fn query_pool_error_policy_retries_reads_but_not_writes() { assert_eq!( query_pool_error_action(Some(DatabaseType::Postgres), "SELECT * FROM users", "connection reset by peer"), PoolErrorAction::ReconnectAndRetry ); assert_eq!( query_pool_error_action( Some(DatabaseType::Postgres), "UPDATE users SET active = true", "connection reset by peer" ), PoolErrorAction::Discard ); } #[tokio::test] async fn sqlite_batch_stops_when_the_first_statement_fails() { assert_sqlite_batch_error_behavior(true, false).await; } #[tokio::test] async fn sqlite_batch_continues_when_the_first_statement_fails_and_enabled() { assert_sqlite_batch_error_behavior(true, true).await; } #[tokio::test] async fn sqlite_batch_stops_when_a_middle_statement_fails() { assert_sqlite_batch_error_behavior(false, false).await; } #[tokio::test] async fn sqlite_batch_continues_when_a_middle_statement_fails_and_enabled() { assert_sqlite_batch_error_behavior(false, true).await; } #[tokio::test] async fn mysql_batch_stops_after_the_first_statement_error() { let statements = vec!["first".to_string(), "fails".to_string(), "must-not-run".to_string()]; let mut executor = FakeMysqlBatchExecutor { outcomes: std::collections::VecDeque::from([ mysql_batch_result(empty_query_result(0)), Err("Duplicate entry".to_string()), mysql_batch_result(empty_query_result(0)), ]), executed: Vec::new(), }; let (results, error_action) = execute_mysql_batch_statements( &mut executor, &statements, Some(DatabaseType::Mysql), db::mysql::MySqlQueryDialect::default(), None, false, None, None, ) .await; assert_eq!(executor.executed, vec!["first", "fails"]); assert_eq!(results.len(), 2); assert_eq!(results[0].statement_index, Some(0)); assert_eq!(results[1].statement_index, Some(1)); assert!(results[1].execution_error); assert_eq!(error_action, Some(PoolErrorAction::Keep)); } #[tokio::test] async fn mysql_batch_reports_progress_for_each_completed_statement() { let statements = vec!["first".to_string(), "fails".to_string(), "must-not-run".to_string()]; let mut executor = FakeMysqlBatchExecutor { outcomes: std::collections::VecDeque::from([ mysql_batch_result(empty_query_result(0)), Err("Duplicate entry".to_string()), ]), executed: Vec::new(), }; let progress_events = Arc::new(std::sync::Mutex::new(Vec::new())); let progress: ExecuteMultiProgressCallback = { let progress_events = Arc::clone(&progress_events); Arc::new(move |event| progress_events.lock().unwrap().push(event)) }; let (results, error_action) = execute_mysql_batch_statements( &mut executor, &statements, Some(DatabaseType::Mysql), db::mysql::MySqlQueryDialect::default(), None, false, None, Some(&progress), ) .await; assert_eq!(executor.executed, vec!["first", "fails"]); assert_eq!(results.len(), 2); assert_eq!( *progress_events.lock().unwrap(), vec![ ExecuteMultiProgress { statement_index: 0, completed: 1, total: 3, success: true, execution_time_ms: 0, affected_rows: 0, error: None, }, ExecuteMultiProgress { statement_index: 1, completed: 2, total: 3, success: false, execution_time_ms: 0, affected_rows: 0, error: Some(crate::backend_error::BackendError::from_legacy_backend("Duplicate entry")), }, ] ); assert_eq!(error_action, Some(PoolErrorAction::Keep)); } #[tokio::test] async fn mysql_batch_pipelines_adjacent_non_result_statements() { let statements = vec![ "SET @batch = 1".to_string(), "INSERT INTO users(id) VALUES (1)".to_string(), "INSERT INTO users(id) VALUES (2)".to_string(), "SELECT COUNT(*) FROM users".to_string(), ]; let mut executor = FakePipelinedMysqlBatchExecutor { batch_outcomes: std::collections::VecDeque::from([db::mysql::MySqlNonResultBatchOutcome { results: vec![empty_query_result(2), empty_query_result(3)], error: None, }]), statement_outcomes: std::collections::VecDeque::from([ mysql_batch_result(empty_query_result(1)), mysql_batch_result(db::QueryResult { columns: vec!["COUNT(*)".to_string()], column_types: vec!["BIGINT".to_string()], column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], rows: vec![vec![serde_json::json!(2)]], affected_rows: 0, execution_time_ms: 4, truncated: false, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: Vec::new(), }), ]), batches: Vec::new(), statements: Vec::new(), }; let progress_events = Arc::new(std::sync::Mutex::new(Vec::new())); let progress: ExecuteMultiProgressCallback = { let progress_events = Arc::clone(&progress_events); Arc::new(move |event| progress_events.lock().unwrap().push(event)) }; let (results, error_action) = execute_mysql_batch_statements( &mut executor, &statements, Some(DatabaseType::Mysql), db::mysql::MySqlQueryDialect::default(), None, false, Some(MYSQL_MULTI_STATEMENT_BATCH_MAX_BYTES), Some(&progress), ) .await; assert_eq!(executor.batches, vec![statements[1..3].to_vec()]); assert_eq!(executor.statements, vec![statements[0].clone(), statements[3].clone()]); assert_eq!(results.len(), 4); assert_eq!( results.iter().map(|result| result.statement_index).collect::>(), vec![Some(0), Some(1), Some(2), Some(3)] ); assert_eq!( progress_events.lock().unwrap().iter().map(|event| event.completed).collect::>(), vec![1, 3, 4] ); assert_eq!(error_action, None); } #[tokio::test] async fn mysql_pipelined_batch_reports_the_first_failed_statement() { let statements = vec![ "INSERT INTO users(id) VALUES (1)".to_string(), "INSERT INTO users(id) VALUES (1)".to_string(), "INSERT INTO users(id) VALUES (2)".to_string(), ]; let mut executor = FakePipelinedMysqlBatchExecutor { batch_outcomes: std::collections::VecDeque::from([db::mysql::MySqlNonResultBatchOutcome { results: vec![empty_query_result(1)], error: Some("Duplicate entry".to_string()), }]), statement_outcomes: std::collections::VecDeque::new(), batches: Vec::new(), statements: Vec::new(), }; let (results, error_action) = execute_mysql_batch_statements( &mut executor, &statements, Some(DatabaseType::Mysql), db::mysql::MySqlQueryDialect::default(), None, false, Some(MYSQL_MULTI_STATEMENT_BATCH_MAX_BYTES), None, ) .await; assert_eq!(executor.batches, vec![statements]); assert!(executor.statements.is_empty()); assert_eq!(results.len(), 2); assert_eq!(results[0].statement_index, Some(0)); assert_eq!(results[1].statement_index, Some(1)); assert!(results[1].execution_error); assert_eq!(error_action, Some(PoolErrorAction::Keep)); } #[tokio::test] async fn mysql_pipelined_batch_discards_a_cancelled_connection() { let statements = vec!["INSERT INTO users(id) VALUES (1)".to_string(), "INSERT INTO users(id) VALUES (2)".to_string()]; let mut executor = FakePipelinedMysqlBatchExecutor { batch_outcomes: std::collections::VecDeque::from([db::mysql::MySqlNonResultBatchOutcome { results: Vec::new(), error: Some(QUERY_CANCELED.to_string()), }]), statement_outcomes: std::collections::VecDeque::new(), batches: Vec::new(), statements: Vec::new(), }; let (results, error_action) = execute_mysql_batch_statements( &mut executor, &statements, Some(DatabaseType::Mysql), db::mysql::MySqlQueryDialect::default(), None, false, Some(MYSQL_MULTI_STATEMENT_BATCH_MAX_BYTES), None, ) .await; assert_eq!(results.len(), 1); assert_eq!(results[0].statement_index, Some(0)); assert!(results[0].execution_error); assert_eq!(error_action, Some(PoolErrorAction::Discard)); } #[test] fn mysql_non_result_batches_respect_the_statement_limit() { let statements = (0..51).map(|index| format!("INSERT INTO users(id) VALUES ({index})")).collect::>(); assert_eq!( mysql_non_result_batch_end( &statements, 0, db::mysql::MySqlQueryDialect::default(), MYSQL_MULTI_STATEMENT_BATCH_MAX_BYTES, ), MYSQL_MULTI_STATEMENT_BATCH_MAX_STATEMENTS ); assert_eq!( mysql_non_result_batch_end( &statements, MYSQL_MULTI_STATEMENT_BATCH_MAX_STATEMENTS, db::mysql::MySqlQueryDialect::default(), MYSQL_MULTI_STATEMENT_BATCH_MAX_BYTES, ), 51 ); } #[test] fn mysql_non_result_batches_respect_the_byte_limit() { let payload = "x".repeat(1_500_000); let statements = (0..3) .map(|index| format!("INSERT INTO users(id, payload) VALUES ({index}, '{payload}')")) .collect::>(); assert_eq!( mysql_non_result_batch_end( &statements, 0, db::mysql::MySqlQueryDialect::default(), MYSQL_MULTI_STATEMENT_BATCH_MAX_BYTES, ), 2 ); } #[test] fn mysql_single_call_uses_multi_result_route() { assert!(mysql_single_statement_uses_batch_route( Some(DatabaseType::Mysql), true, "# generated call\nCALL testA()", None, )); } #[test] fn ordinary_mysql_single_statements_keep_singular_route() { for sql in ["SELECT 1", "SHOW TABLES", "UPDATE users SET active = 1"] { assert!(!mysql_single_statement_uses_batch_route(Some(DatabaseType::Mysql), true, sql, None)); } } #[test] fn mysql_call_route_requires_native_mysql_type_and_pool() { assert!(!mysql_single_statement_uses_batch_route(Some(DatabaseType::Doris), true, "CALL testA()", None,)); assert!(!mysql_single_statement_uses_batch_route(Some(DatabaseType::Mysql), false, "CALL testA()", None,)); } #[test] fn mysql_result_byte_limit_keeps_existing_batch_route() { assert!(mysql_single_statement_uses_batch_route( Some(DatabaseType::Mysql), true, "SELECT * FROM users", Some(1024), )); assert!(!mysql_single_statement_uses_batch_route( Some(DatabaseType::Mysql), true, "SELECT * FROM users", Some(0), )); } #[test] fn single_mysql_batch_route_never_probes_non_result_pipeline_limits() { assert!(!mysql_non_result_pipeline_enabled(1, false, crate::connection::MysqlMode::Normal)); assert!(mysql_non_result_pipeline_enabled(2, false, crate::connection::MysqlMode::Normal)); assert!(!mysql_non_result_pipeline_enabled(2, true, crate::connection::MysqlMode::Normal)); assert!(!mysql_non_result_pipeline_enabled(2, false, crate::connection::MysqlMode::Bare)); } #[tokio::test] async fn mysql_batch_preserves_multiple_result_sets_from_one_statement() { let statements = vec!["CALL testA()".to_string(), "UPDATE users SET active = 1".to_string()]; let result_set = |value| db::QueryResult { columns: vec!["value".to_string()], column_types: vec!["INT".to_string()], column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], rows: vec![vec![serde_json::json!(value)]], affected_rows: 0, execution_time_ms: 1, truncated: false, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: Vec::new(), }; let mut executor = FakeMysqlBatchExecutor { outcomes: std::collections::VecDeque::from([ Ok(vec![ mysql_query_result(result_set(1)), mysql_query_result(result_set(2)), mysql_query_result(result_set(3)), ]), mysql_batch_result(empty_query_result(1)), ]), executed: Vec::new(), }; let (results, error_action) = execute_mysql_batch_statements( &mut executor, &statements, Some(DatabaseType::Mysql), db::mysql::MySqlQueryDialect::default(), None, false, None, None, ) .await; assert_eq!(executor.executed, statements); assert_eq!(results.len(), 4); assert_eq!( results.iter().map(|result| result.statement_index).collect::>(), vec![Some(0), Some(0), Some(0), Some(1)] ); assert_eq!( results[..3].iter().map(|result| result.result.rows[0][0].clone()).collect::>(), vec![serde_json::json!(1), serde_json::json!(2), serde_json::json!(3)] ); assert_eq!(error_action, None); } #[tokio::test] async fn mysql_batch_stops_when_the_first_statement_fails() { let statements = vec!["fails".to_string(), "must-not-run".to_string()]; let mut executor = FakeMysqlBatchExecutor { outcomes: std::collections::VecDeque::from([ Err("Duplicate entry".to_string()), mysql_batch_result(empty_query_result(0)), ]), executed: Vec::new(), }; let (results, error_action) = execute_mysql_batch_statements( &mut executor, &statements, Some(DatabaseType::Mysql), db::mysql::MySqlQueryDialect::default(), None, false, None, None, ) .await; assert_eq!(executor.executed, vec!["fails"]); assert_eq!(results.len(), 1); assert!(results[0].execution_error); assert_eq!(error_action, Some(PoolErrorAction::Keep)); } #[tokio::test] async fn mysql_batch_continues_after_statement_errors_when_enabled() { let statements = vec!["first".to_string(), "fails".to_string(), "third".to_string()]; let mut executor = FakeMysqlBatchExecutor { outcomes: std::collections::VecDeque::from([ mysql_batch_result(empty_query_result(0)), Err("Duplicate entry".to_string()), mysql_batch_result(empty_query_result(0)), ]), executed: Vec::new(), }; let (results, error_action) = execute_mysql_batch_statements( &mut executor, &statements, Some(DatabaseType::Mysql), db::mysql::MySqlQueryDialect::default(), None, true, None, None, ) .await; assert_eq!(executor.executed, statements); assert_eq!(results.len(), 3); assert_eq!( results.iter().map(|result| result.statement_index).collect::>(), vec![Some(0), Some(1), Some(2)] ); assert!(results[1].execution_error); assert_eq!(error_action, None); } #[tokio::test] async fn mysql_batch_continues_when_the_first_statement_fails_and_enabled() { let statements = vec!["fails".to_string(), "second".to_string()]; let mut executor = FakeMysqlBatchExecutor { outcomes: std::collections::VecDeque::from([ Err("Duplicate entry".to_string()), mysql_batch_result(empty_query_result(0)), ]), executed: Vec::new(), }; let (results, error_action) = execute_mysql_batch_statements( &mut executor, &statements, Some(DatabaseType::Mysql), db::mysql::MySqlQueryDialect::default(), None, true, None, None, ) .await; assert_eq!(executor.executed, statements); assert_eq!(results.len(), 2); assert!(results[0].execution_error); assert_eq!(error_action, None); } #[tokio::test] async fn mysql_batch_stops_on_connection_errors_when_continue_is_enabled() { let statements = vec!["first".to_string(), "disconnects".to_string(), "must-not-run".to_string()]; let mut executor = FakeMysqlBatchExecutor { outcomes: std::collections::VecDeque::from([ mysql_batch_result(empty_query_result(0)), Err("connection reset by peer".to_string()), mysql_batch_result(empty_query_result(0)), ]), executed: Vec::new(), }; let (results, error_action) = execute_mysql_batch_statements( &mut executor, &statements, Some(DatabaseType::Mysql), db::mysql::MySqlQueryDialect::default(), None, true, None, None, ) .await; assert_eq!(executor.executed, vec!["first", "disconnects"]); assert_eq!(results.len(), 2); assert_eq!(results.iter().map(|result| result.statement_index).collect::>(), vec![Some(0), Some(1)]); assert!(results[1].execution_error); assert_eq!(error_action, Some(PoolErrorAction::ReconnectAndRetry)); } #[test] fn execute_multi_result_serializes_client_metadata_only_when_present() { let success = serde_json::to_value(ExecuteMultiResult::from(empty_query_result(0))).unwrap(); assert!(success.get("execution_error").is_none()); assert!(success.get("statement_index").is_none()); assert!(success.get("server_message").is_none()); assert!(success.get("large_value_cells").is_none()); let mut error_column = empty_query_result(0); error_column.columns = vec!["Error".to_string()]; error_column.rows = vec![vec![serde_json::json!("valid query value")]]; let error_column = serde_json::to_value(ExecuteMultiResult::success_with_index(error_column, 0)).unwrap(); assert!(error_column.get("execution_error").is_none()); assert!(error_column.get("error").is_none()); let failure = serde_json::to_value(ExecuteMultiResult::execution_error_with_index( error_query_result("failed".to_string()), 2, )) .unwrap(); assert_eq!(failure.get("execution_error"), Some(&serde_json::Value::Bool(true))); assert_eq!(failure.get("statement_index"), Some(&serde_json::json!(2))); assert_eq!(failure.get("columns"), Some(&serde_json::json!(["Error"]))); assert_eq!( failure.get("error").and_then(|value| value.get("code")), Some(&serde_json::json!("DBX-LEGACY-0001")) ); let redacted = serde_json::to_value( ExecuteMultiResult::execution_error_with_index(error_query_result("safe failure".to_string()), 2) .without_error_detail(), ) .unwrap(); assert!(redacted.get("error").and_then(|value| value.get("detail")).is_none()); } #[test] fn execute_multi_result_serializes_large_value_metadata_only_when_present() { let result = ExecuteMultiResult::success_with_index_and_large_values( empty_query_result(1), 0, vec![db::LargeValueCell { row_index: 2, column_index: 3, original_bytes: 65_536 }], false, ); let serialized = serde_json::to_value(result).unwrap(); assert_eq!( serialized.get("large_value_cells"), Some(&serde_json::json!([{"row_index": 2, "column_index": 3, "original_bytes": 65_536}])) ); } #[test] fn large_value_cell_merge_replaces_driver_entries_and_appends_server_entries() { let driver_cells = vec![ db::LargeValueCell { row_index: 1, column_index: 2, original_bytes: 10 }, db::LargeValueCell { row_index: 3, column_index: 4, original_bytes: 20 }, ]; let server_cells = vec![ db::LargeValueCell { row_index: 1, column_index: 2, original_bytes: 100 }, db::LargeValueCell { row_index: 5, column_index: 6, original_bytes: 200 }, ]; assert_eq!( merge_large_value_cells(driver_cells, server_cells), vec![ db::LargeValueCell { row_index: 1, column_index: 2, original_bytes: 100 }, db::LargeValueCell { row_index: 3, column_index: 4, original_bytes: 20 }, db::LargeValueCell { row_index: 5, column_index: 6, original_bytes: 200 }, ] ); } #[test] fn large_value_cell_merge_preserves_single_source_inputs() { let driver_cell = db::LargeValueCell { row_index: 1, column_index: 2, original_bytes: 10 }; let server_cell = db::LargeValueCell { row_index: 3, column_index: 4, original_bytes: 20 }; assert_eq!(merge_large_value_cells(vec![driver_cell.clone()], Vec::new()), vec![driver_cell]); assert_eq!(merge_large_value_cells(Vec::new(), vec![server_cell.clone()]), vec![server_cell]); } #[test] fn large_value_cell_merge_handles_large_disjoint_inputs() { const CELL_COUNT: usize = 100_000; let driver_cells = (0..CELL_COUNT) .map(|row_index| db::LargeValueCell { row_index, column_index: 0, original_bytes: 10 }) .collect(); let server_cells = (0..CELL_COUNT) .map(|row_index| db::LargeValueCell { row_index, column_index: 1, original_bytes: 20 }) .collect(); let merged = merge_large_value_cells(driver_cells, server_cells); assert_eq!(merged.len(), CELL_COUNT * 2); assert_eq!(merged[CELL_COUNT].row_index, 0); assert_eq!(merged[CELL_COUNT * 2 - 1].row_index, CELL_COUNT - 1); } #[test] fn sqlserver_batch_results_do_not_claim_statement_indexes() { assert_eq!(split_sql_batches("SELECT 1; SELECT 2;").len(), 1); let results = sqlserver_batch_results(vec![ db::sqlserver::SqlServerBatchResult { result: empty_query_result(1), server_message: false }, db::sqlserver::SqlServerBatchResult { result: empty_query_result(2), server_message: true }, ]); assert_eq!(results.iter().map(|result| result.statement_index).collect::>(), vec![None, None]); assert!(!results[0].server_message); assert!(results[1].server_message); let serialized = serde_json::to_value(&results[1]).unwrap(); assert_eq!(serialized.get("server_message"), Some(&serde_json::Value::Bool(true))); } // 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 // same connection. Before this fix, `execution_time_ms` was measured only // from inside db::sqlserver's own timers, which start *after* the lock is // acquired — so that queueing time was invisible to the user, producing a // reported duration (e.g. "6-8ms") wildly smaller than what they actually // waited (e.g. "10s"). `lock_shared_client_with_wait` is the exact helper // both PoolKind::SqlServer call sites now use to fold that wait back in. #[tokio::test] async fn lock_shared_client_with_wait_reports_time_queued_behind_another_holder() { let client = Arc::new(tokio::sync::Mutex::new(0u8)); let holder_guard = client.lock().await; let waiter_client = client.clone(); let waiter = tokio::spawn(async move { let (_guard, wait_ms) = lock_shared_client_with_wait(&waiter_client, None, None).await.unwrap(); wait_ms }); // Give the spawned task a chance to actually start waiting on the lock // before the holder releases it, so the measured wait is meaningful. tokio::time::sleep(Duration::from_millis(150)).await; drop(holder_guard); let wait_ms = waiter.await.unwrap(); assert!(wait_ms >= 150, "expected queued wait time to be captured, got {wait_ms}ms"); } #[tokio::test] async fn lock_shared_client_with_wait_is_near_zero_when_uncontended() { let client = Arc::new(tokio::sync::Mutex::new(0u8)); let (_guard, wait_ms) = lock_shared_client_with_wait(&client, None, None).await.unwrap(); assert!(wait_ms < 50, "expected an uncontended lock to report negligible wait, got {wait_ms}ms"); } #[test] fn query_execution_error_preserves_structured_catalog_identity() { let error = QueryExecutionError::Agent(AgentCallError::Structured { rpc_code: -1, message: "connection lost".to_string(), context: crate::db::agent_driver::AgentErrorContext { contract_version: 1, category: crate::db::agent_driver::AgentErrorCategory::Connection, retryable: false, session_disposition: crate::db::agent_driver::AgentSessionDisposition::Quarantine, stage: AgentErrorStage::Execute, operation_outcome: AgentOperationOutcome::Unknown, agent_session_id: Some("session-1".to_string()), sql_state: None, vendor_code: None, exception_class: None, }, }); assert_eq!(error.into_backend_error().code(), "DBX-JDBC-1002"); } #[test] fn duckdb_worker_error_preserves_catalog_identity_and_detail() { let error = QueryExecutionError::DuckDb { code: "duckdb_execute_failed".to_string(), message: "Catalog Error: Table missing_table does not exist".to_string(), }; let backend_error = error.into_backend_error(); assert_eq!(backend_error.code(), "DBX-JDBC-4001"); assert_eq!(backend_error.detail(), Some("Catalog Error: Table missing_table does not exist")); } #[test] fn query_timeout_preserves_timeout_catalog_identity_and_detail() { let error = classify_query_error( Some(DatabaseType::Postgres), QueryExecutionError::Legacy("Query timed out after 1 seconds".to_string()), ) .with_omitted_sql_context("SELECT pg_sleep(10)"); let backend_error = error.into_backend_error(); assert_eq!(backend_error.code(), "DBX-JDBC-2002"); assert_eq!( backend_error.message_params().get("stage"), Some(&crate::backend_error::BackendMessageParam::String("execute".to_string())) ); assert_eq!( backend_error.detail(), Some("Query timed out after 1 seconds\nSQL text omitted from user-facing error; enable debug SQL diagnostics to inspect the original statement.") ); } #[test] fn postgres_server_error_preserves_sql_catalog_identity_and_detail() { let error = classify_query_error( Some(DatabaseType::Postgres), QueryExecutionError::Legacy("ERROR: relation \"dbx_table_that_does_not_exist\" does not exist".to_string()), ) .with_omitted_sql_context("SELECT * FROM dbx_table_that_does_not_exist"); let backend_error = error.into_backend_error(); assert_eq!(backend_error.code(), "DBX-JDBC-4001"); assert_eq!( backend_error.message_params().get("stage"), Some(&crate::backend_error::BackendMessageParam::String("execute".to_string())) ); assert_eq!( backend_error.detail(), Some( "ERROR: relation \"dbx_table_that_does_not_exist\" does not exist\nSQL text omitted from user-facing error; enable debug SQL diagnostics to inspect the original statement." ) ); } #[test] fn mysql_server_error_preserves_sql_catalog_identity_and_detail() { let error = classify_query_error( Some(DatabaseType::Mysql), QueryExecutionError::Legacy( "Server error: `ERROR 1064 (42000): You have an error in your SQL syntax`".to_string(), ), ) .with_omitted_sql_context("SELECT 111 AS first_value FROM DUAL"); let backend_error = error.into_backend_error(); assert_eq!(backend_error.code(), "DBX-JDBC-4001"); assert_eq!(backend_error.message_key(), "backendErrors.jdbc.sqlFailed"); assert_eq!( backend_error.message_params().get("stage"), Some(&crate::backend_error::BackendMessageParam::String("execute".to_string())) ); assert_eq!( backend_error.detail(), Some( "Server error: `ERROR 1064 (42000): You have an error in your SQL syntax` SQL text omitted from user-facing error; enable debug SQL diagnostics to inspect the original statement." ) ); } #[test] fn single_statement_multi_result_preserves_sql_error_type() { let error = classify_query_error( Some(DatabaseType::Postgres), QueryExecutionError::Legacy("ERROR: relation \"dbx_table_that_does_not_exist\" does not exist".to_string()), ) .with_omitted_sql_context("SELECT * FROM dbx_table_that_does_not_exist"); let error = single_statement_multi_result(Err(error), false).unwrap_err(); let backend_error = error.into_backend_error(); assert_eq!(backend_error.code(), "DBX-JDBC-4001"); assert_eq!(backend_error.message_key(), "backendErrors.jdbc.sqlFailed"); assert_eq!( backend_error.detail(), Some( "ERROR: relation \"dbx_table_that_does_not_exist\" does not exist\nSQL text omitted from user-facing error; enable debug SQL diagnostics to inspect the original statement." ) ); } #[test] fn postgres_transaction_statement_error_preserves_sql_catalog_identity() { let error = postgres_transaction_statement_error( 1, "ERROR: relation \"dbx_table_that_does_not_exist\" does not exist", "SELECT * FROM dbx_table_that_does_not_exist", true, ); let backend_error = error.into_backend_error(); assert_eq!(backend_error.code(), "DBX-JDBC-4001"); assert_eq!(backend_error.message_key(), "backendErrors.jdbc.sqlFailed"); assert_eq!( backend_error.detail(), Some( "Statement 1 failed: ERROR: relation \"dbx_table_that_does_not_exist\" does not exist\nSQL text omitted from user-facing error; enable debug SQL diagnostics to inspect the original statement." ) ); } #[test] fn postgres_transaction_uses_driver_fact_not_error_wording() { let error = postgres_transaction_statement_error( 1, "ERROR: relation \"connection_closed\" does not exist", "SELECT * FROM connection_closed", true, ); assert_eq!(error.into_backend_error().code(), "DBX-JDBC-4001"); } #[test] fn sequential_multi_cancellation_uses_canceled_catalog() { let backend_error = canceled_query_execution_error().into_backend_error(); assert_eq!(backend_error.code(), "DBX-JDBC-2003"); assert_eq!(backend_error.message_key(), "backendErrors.jdbc.operationCanceled"); } #[test] fn external_driver_method_unsupported_detects_legacy_plugin_errors() { assert!(is_external_driver_method_unsupported( "Unsupported JDBC plugin method: executeQueryPage", "executeQueryPage" )); assert!(is_external_driver_method_unsupported( "Plugin RPC error (-32601): Method not found: executeQueryPage", "executeQueryPage" )); assert!(is_external_driver_method_unsupported("Unknown method executeQueryPage", "executeQueryPage")); } #[test] fn external_driver_method_unsupported_ignores_query_and_other_method_errors() { assert!(!is_external_driver_method_unsupported( "The JDBC driver does not support this SQL operation", "executeQueryPage" )); assert!(!is_external_driver_method_unsupported( "Unsupported JDBC plugin method: listTables", "executeQueryPage" )); assert!(!is_external_driver_method_unsupported( "Unknown column executeQueryPage in field list", "executeQueryPage" )); } #[cfg(unix)] #[tokio::test] async fn external_driver_query_page_falls_back_to_legacy_execute_query() { use std::os::unix::fs::PermissionsExt; let dir = std::env::temp_dir().join(format!("dbx-legacy-jdbc-plugin-test-{}", uuid::Uuid::new_v4())); std::fs::create_dir_all(&dir).unwrap(); let executable = dir.join("plugin.sh"); let calls = dir.join("calls.log"); std::fs::write( &executable, format!( "#!/bin/sh\nwhile IFS= read -r line; do\n id=$(printf '%s' \"$line\" | sed -E 's/.*\"id\":([0-9]+).*/\\1/')\n case \"$line\" in\n *'\"method\":\"executeQueryPage\"'*)\n echo executeQueryPage >> '{}'\n printf '{{\"id\":%s,\"error\":{{\"message\":\"Unsupported JDBC plugin method: executeQueryPage\"}}}}\\n' \"$id\"\n ;;\n *'\"method\":\"executeQuery\"'*)\n echo executeQuery >> '{}'\n printf '{{\"id\":%s,\"result\":{{\"columns\":[\"value\"],\"rows\":[[42]],\"affected_rows\":0,\"execution_time_ms\":1,\"truncated\":false}}}}\\n' \"$id\"\n ;;\n esac\ndone\n", calls.display(), calls.display() ), ) .unwrap(); let mut permissions = std::fs::metadata(&executable).unwrap().permissions(); permissions.set_mode(0o755); std::fs::set_permissions(&executable, permissions).unwrap(); let plugin = InstalledPlugin { manifest: PluginManifest { id: "jdbc".to_string(), name: "JDBC".to_string(), version: "legacy".to_string(), protocol_version: 1, description: String::new(), executable: Some("plugin.sh".to_string()), drivers: vec![PluginDriverManifest { id: "jdbc".to_string(), label: "JDBC".to_string(), kind: "external".to_string(), database_type: Some("jdbc".to_string()), }], }, path: dir.clone(), }; let session = PluginDriverSession::start_for_test(plugin, "jdbc".to_string(), PluginRuntimeEnv::default()) .await .expect("legacy plugin should start"); let result = invoke_external_driver_query_page( &session, serde_json::json!({ "sql": "SELECT 42", "pageSize": 100 }), Some(Duration::from_secs(5)), ) .await .expect("legacy executeQuery fallback should succeed"); assert_eq!(result.columns, vec!["value"]); assert_eq!(result.rows, vec![vec![serde_json::json!(42)]]); assert_eq!(std::fs::read_to_string(&calls).unwrap(), "executeQueryPage\nexecuteQuery\n"); session.shutdown().await; let _ = std::fs::remove_dir_all(dir); } #[cfg(unix)] #[tokio::test] async fn external_driver_query_page_does_not_retry_jdbc_errors() { use std::os::unix::fs::PermissionsExt; let dir = std::env::temp_dir().join(format!("dbx-jdbc-query-error-test-{}", uuid::Uuid::new_v4())); std::fs::create_dir_all(&dir).unwrap(); let executable = dir.join("plugin.sh"); let calls = dir.join("calls.log"); std::fs::write( &executable, format!( "#!/bin/sh\nwhile IFS= read -r line; do\n id=$(printf '%s' \"$line\" | sed -E 's/.*\"id\":([0-9]+).*/\\1/')\n echo request >> '{}'\n printf '{{\"id\":%s,\"error\":{{\"message\":\"Incorrect syntax near SELECT\"}}}}\\n' \"$id\"\ndone\n", calls.display() ), ) .unwrap(); let mut permissions = std::fs::metadata(&executable).unwrap().permissions(); permissions.set_mode(0o755); std::fs::set_permissions(&executable, permissions).unwrap(); let plugin = InstalledPlugin { manifest: PluginManifest { id: "jdbc".to_string(), name: "JDBC".to_string(), version: "current".to_string(), protocol_version: 1, description: String::new(), executable: Some("plugin.sh".to_string()), drivers: vec![PluginDriverManifest { id: "jdbc".to_string(), label: "JDBC".to_string(), kind: "external".to_string(), database_type: Some("jdbc".to_string()), }], }, path: dir.clone(), }; let session = PluginDriverSession::start_for_test(plugin, "jdbc".to_string(), PluginRuntimeEnv::default()) .await .expect("plugin should start"); let error = invoke_external_driver_query_page( &session, serde_json::json!({ "sql": "SELECT broken", "pageSize": 100 }), Some(Duration::from_secs(5)), ) .await .expect_err("JDBC query errors must be returned without retrying"); assert_eq!(error, "Incorrect syntax near SELECT"); assert_eq!(std::fs::read_to_string(&calls).unwrap(), "request\n"); session.shutdown().await; let _ = std::fs::remove_dir_all(dir); } #[test] fn oceanbase_mysql_session_timeout_sql_uses_connection_timeout_by_default() { let mut config = test_connection_config(DatabaseType::Mysql); config.driver_profile = Some("oceanbase".to_string()); config.query_timeout_secs = 300_000; assert_eq!( oceanbase_mysql_session_timeout_sql(Some(&config), None), Some("SET ob_query_timeout = 300000000000".to_string()) ); } #[test] fn oceanbase_mysql_session_timeout_sql_prefers_execution_timeout_override() { let mut config = test_connection_config(DatabaseType::Mysql); config.driver_profile = Some("oceanbase".to_string()); config.query_timeout_secs = 30; assert_eq!( oceanbase_mysql_session_timeout_sql(Some(&config), Some(600)), Some("SET ob_query_timeout = 600000000".to_string()) ); } #[test] fn oceanbase_mysql_session_timeout_sql_skips_plain_mysql() { let config = test_connection_config(DatabaseType::Mysql); assert_eq!(oceanbase_mysql_session_timeout_sql(Some(&config), Some(600)), None); } #[tokio::test] async fn wait_for_query_returns_cancelled_when_token_is_cancelled() { let token = CancellationToken::new(); token.cancel(); let result = wait_for_query(Some(token), async { tokio::time::sleep(Duration::from_secs(30)).await; Ok(db::QueryResult { columns: vec![], column_types: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], rows: vec![], affected_rows: 0, execution_time_ms: 0, truncated: false, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: Vec::new(), }) }) .await; assert_eq!(result.unwrap_err(), QUERY_CANCELED); } #[tokio::test] async fn wait_for_query_without_token_still_times_out() { let result = wait_for_query_with_timeout(None, Duration::from_millis(10), async { tokio::time::sleep(Duration::from_secs(1)).await; Ok(db::QueryResult { columns: vec![], column_types: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], rows: vec![], affected_rows: 0, execution_time_ms: 0, truncated: false, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: Vec::new(), }) }) .await; assert_eq!(result.unwrap_err(), timeout_error_for(Duration::from_millis(10))); } #[tokio::test] async fn wait_for_value_opt_times_out_while_waiting_for_lock() { let lock = tokio::sync::Mutex::new(()); let _guard = lock.lock().await; let result = wait_for_value_opt(None, Some(Duration::from_millis(10)), lock.lock()).await; assert_eq!(result.unwrap_err(), timeout_error_for(Duration::from_millis(10))); } #[tokio::test] async fn wait_for_value_opt_can_cancel_while_waiting_for_lock() { let lock = tokio::sync::Mutex::new(()); let _guard = lock.lock().await; let token = CancellationToken::new(); token.cancel(); let result = wait_for_value_opt(Some(token), Some(Duration::from_secs(30)), lock.lock()).await; assert_eq!(result.unwrap_err(), QUERY_CANCELED); } #[test] fn db_operation_budget_from_config() { let budget = DbOperationBudget::from_config(10, Some(30)); assert_eq!(budget.checkout_timeout, Duration::from_secs(10)); assert_eq!(budget.connect_timeout, Duration::from_secs(10)); assert_eq!(budget.recycle_timeout, Duration::from_secs(10)); assert_eq!(budget.query_timeout, Some(Duration::from_secs(30))); assert_eq!(budget.cancel_timeout, Duration::from_secs(5)); assert_eq!(budget.cleanup_timeout, Duration::from_secs(3)); } #[test] fn db_operation_budget_from_connection_config_uses_connection_settings() { let mut config = test_connection_config(DatabaseType::Postgres); config.connect_timeout_secs = 12; config.query_timeout_secs = 0; let budget = DbOperationBudget::from_connection_config(&config); assert_eq!(budget.checkout_timeout, Duration::from_secs(12)); assert_eq!(budget.connect_timeout, Duration::from_secs(12)); assert_eq!(budget.recycle_timeout, Duration::from_secs(12)); assert_eq!(budget.query_timeout, None); assert_eq!(budget.cancel_timeout, Duration::from_secs(5)); assert_eq!(budget.cleanup_timeout, Duration::from_secs(3)); } #[test] fn db_operation_budget_query_timeout_zero_means_no_limit() { let budget = DbOperationBudget::from_config(10, Some(0)); assert_eq!(budget.query_timeout, None); // Infrastructure timeouts still have hard limits assert_eq!(budget.checkout_timeout, Duration::from_secs(10)); assert_eq!(budget.cancel_timeout, Duration::from_secs(5)); } #[test] fn db_operation_budget_query_timeout_zero_keeps_transaction_infra_limits() { let mut config = test_connection_config(DatabaseType::Mysql); config.connect_timeout_secs = 7; config.query_timeout_secs = 0; let budget = DbOperationBudget::from_connection_config(&config); assert_eq!(budget.query_timeout, None); assert_eq!(budget.checkout_timeout, Duration::from_secs(7)); assert_eq!(budget.recycle_timeout, Duration::from_secs(7)); assert_eq!(budget.cleanup_timeout, Duration::from_secs(3)); } #[test] fn db_operation_budget_clamps_infra_timeout() { let budget = DbOperationBudget::from_config(0, Some(30)); assert_eq!(budget.checkout_timeout, Duration::from_secs(1)); // clamped to min 1s let budget = DbOperationBudget::from_config(600, Some(30)); assert_eq!(budget.checkout_timeout, Duration::from_secs(300)); // clamped to max 300s } #[test] fn db_operation_budget_with_defaults() { let budget = DbOperationBudget::with_defaults(); assert_eq!(budget.checkout_timeout, db::connection_timeout()); assert_eq!(budget.query_timeout, Some(QUERY_TIMEOUT)); } #[test] fn is_connection_error_detects_english_messages() { assert!(is_connection_error("connection reset")); assert!(is_connection_error("broken pipe")); assert!(is_connection_error("reset by peer")); assert!(is_connection_error("Connection timed out")); assert!(is_connection_error("socket closed")); assert!(is_connection_error("unexpected eof")); assert!(is_connection_error("Error occurred while creating a new object: error communicating with the server")); } #[test] fn query_error_context_omits_raw_sql_and_is_not_duplicated() { let sql = "select 'secret-123' as token"; let error = query_error_with_omitted_sql_context("driver rejected statement", sql); assert!(error.contains("driver rejected statement")); assert!(error.contains(SQL_OMITTED_ERROR_CONTEXT)); assert!(!error.contains("secret-123")); assert!(!error.contains("SQL:")); let repeated = query_error_with_omitted_sql_context(&error, sql); assert_eq!(repeated.matches(SQL_OMITTED_ERROR_CONTEXT).count(), 1); } #[test] fn typed_sql_error_context_keeps_driver_text_on_one_line() { let error = QueryExecutionError::Sql("Server error: `ERROR 1064 (42000): syntax error`".to_string()) .with_omitted_sql_context("SELECT * FROM users") .into_backend_error(); assert_eq!( error.detail(), Some("Server error: `ERROR 1064 (42000): syntax error` SQL text omitted from user-facing error; enable debug SQL diagnostics to inspect the original statement.") ); } #[test] fn reconnect_retry_error_context_omits_raw_sql() { let sql = "select 'secret-123' as token"; let reconnect_error = query_error_with_omitted_sql_context("connection reset after reconnect", sql); assert!(reconnect_error.contains("connection reset after reconnect")); assert!(reconnect_error.contains(SQL_OMITTED_ERROR_CONTEXT)); assert!(!reconnect_error.contains("secret-123")); } #[test] fn execute_statements_error_omits_raw_sql() { let sql = "select 'secret-token' as t"; let err = query_error_with_omitted_sql_context( &format!( "Statement {} failed: {}. Previous {} statement(s) may have been committed.", 2, "driver error", 1 ), sql, ); assert!(err.contains("driver error")); assert!(err.contains(SQL_OMITTED_ERROR_CONTEXT)); assert!(!err.contains("secret-token")); assert!(!err.contains("SQL:")); assert!(err.contains("Statement 2 failed:")); } #[test] fn batch_transaction_error_omits_raw_sql() { let sql = "delete from users where id = 'secret-id'"; let err = query_error_with_omitted_sql_context( &format!("Statement {} failed: {}. No transaction support for this database type.", 3, "batch error"), sql, ); assert!(err.contains("batch error")); assert!(err.contains(SQL_OMITTED_ERROR_CONTEXT)); assert!(!err.contains("secret-id")); assert!(err.contains("Statement 3 failed:")); } #[test] fn is_connection_error_detects_oracle_idle_timeout() { assert!(is_connection_error("ORA-02396: exceeded maximum idle time, please connect again")); assert!(is_connection_error( "Agent RPC error (-32603): ORA-02396: exceeded maximum idle time, please connect again" )); assert!(is_connection_error("ORA-03113: end-of-file on communication channel")); assert!(is_connection_error("ORA-03114: not connected to Oracle")); assert!(is_connection_error("ORA-03135: connection lost contact")); assert!(is_connection_error("Agent RPC error (-1): java.sql.SQLRecoverableException: 关闭的连接")); assert!(is_connection_error("java.sql.SQLRecoverableException: 连接已关闭")); } #[test] fn is_connection_error_detects_localized_io_errors() { assert!(is_connection_error("I/O error: 远程主机强迫关闭了一个现有的连接。 (os error 10054)")); assert!(is_connection_error( "I/O error: 由于连接方在一段时间后没有正确答复或连接的主机没有反应，连接尝试失败。 (os error 10060)" )); assert!(is_connection_error("Agent RPC error (-1): dm.jdbc.driver.DMException: 网络通信异常")); assert!(is_connection_error( "Agent RPC error (-1): java.sql.SQLRecoverableException: IO 错误: Got minus one from a read call" )); assert!(is_connection_error( "Agent RPC error (-1): com.mysql.cj.jdbc.exceptions.CommunicationsException: Communications link failure" )); } #[test] fn is_connection_error_detects_os_error_codes() { assert!(is_connection_error("os error 10053")); assert!(is_connection_error("os error 10054")); assert!(is_connection_error("os error 10060")); assert!(is_connection_error("os error 10061")); } #[test] fn is_connection_error_rejects_non_connection_errors() { assert!(!is_connection_error("Query timed out after 30 seconds")); assert!(!is_connection_error("ORA-00942: table or view does not exist")); assert!(!is_connection_error("syntax error at position 5")); assert!(!is_connection_error("os error 13")); } #[test] fn is_connection_error_distinguishes_checkout_saturation_from_connection_timeouts() { // deadpool-postgres PoolError::Timeout messages (contain "pool" + "timeout" but not "timed out") assert!(is_connection_error("pool wait timeout")); assert!(is_connection_error("pool create timeout")); assert!(is_connection_error("pool recycle timeout")); // checkout helper timeout messages assert!(!is_connection_error("PostgreSQL connection pool checkout timed out [stage=wait, timeout_ms=5000]")); assert!(is_connection_error("PostgreSQL connection pool checkout timed out [stage=create, timeout_ms=5000]")); assert!(is_connection_error("PostgreSQL connection pool checkout timed out [stage=recycle, timeout_ms=5000]")); assert!(!is_connection_error("MySQL connection pool checkout timed out [stage=wait, timeout_ms=5000]")); assert!(!is_connection_error(METADATA_POOL_BUSY_ERROR)); assert!(is_connection_error("MySQL ping timed out")); assert!(is_connection_error("MySQL kill connection checkout timed out")); assert!(is_connection_error("MySQL KILL QUERY timed out")); assert!(is_pool_saturation_error( "PostgreSQL connection pool checkout timed out [stage=wait, timeout_ms=5000]" )); assert!(is_pool_saturation_error("MySQL connection pool checkout timed out [stage=wait, timeout_ms=5000]")); assert!(!is_pool_saturation_error( "PostgreSQL connection pool checkout timed out [stage=create, timeout_ms=5000]" )); assert!(is_pool_saturation_error(METADATA_POOL_BUSY_ERROR)); assert_eq!( pool_error_action( Some(DatabaseType::Mysql), "MySQL connection pool checkout timed out [stage=wait, timeout_ms=5000]" ), PoolErrorAction::Keep ); assert_eq!( pool_error_action( Some(DatabaseType::Postgres), "PostgreSQL connection pool checkout timed out [stage=wait, timeout_ms=5000]" ), PoolErrorAction::Keep ); } #[test] fn pool_error_action_discards_sqlserver_driver_panic_without_retry() { let err = format!("{} the current client will be rebuilt.", db::sqlserver::SQLSERVER_DRIVER_PANIC_ERROR_PREFIX); assert_eq!(pool_error_action(Some(DatabaseType::SqlServer), &err), PoolErrorAction::Discard); assert!(should_discard_pool_after_error(Some(DatabaseType::SqlServer), &err)); assert!(!is_connection_error(&err)); } #[test] fn pool_error_action_discards_sqlserver_timeout_without_retry() { let err = "Query timed out after 30 seconds"; assert_eq!(pool_error_action(Some(DatabaseType::SqlServer), err), PoolErrorAction::Discard); assert_eq!(pool_error_action(Some(DatabaseType::Mysql), err), PoolErrorAction::Discard); assert_eq!(pool_error_action(Some(DatabaseType::Postgres), err), PoolErrorAction::Discard); assert_eq!(pool_error_action(Some(DatabaseType::ClickHouse), err), PoolErrorAction::Discard); assert_eq!(pool_error_action(Some(DatabaseType::Oracle), err), PoolErrorAction::Discard); assert_eq!(pool_error_action(Some(DatabaseType::Sqlite), err), PoolErrorAction::Keep); assert_eq!(pool_error_action(Some(DatabaseType::DuckDb), err), PoolErrorAction::Keep); } #[test] fn pool_error_action_discards_schema_reset_cleanup_without_retry() { let err = "PostgreSQL schema.reset cleanup failed: PostgreSQL schema.reset timed out after 3 seconds"; assert_eq!(pool_error_action(Some(DatabaseType::Postgres), err), PoolErrorAction::Discard); assert_eq!(pool_error_action(Some(DatabaseType::OpenGauss), err), PoolErrorAction::Discard); assert!(should_discard_pool_after_error(Some(DatabaseType::Postgres), err)); } #[test] fn pool_error_action_discards_postgres_read_only_transaction_cleanup_without_retry() { let err = "PostgreSQL read-only transaction cleanup failed: PostgreSQL explain_analyze.rollback timed out after 3 seconds"; assert_eq!(pool_error_action(Some(DatabaseType::Postgres), err), PoolErrorAction::Discard); assert!(should_discard_pool_after_error(Some(DatabaseType::Postgres), err)); } #[test] fn pool_error_action_reconnects_non_agent_errors_but_never_replays_agent_operations() { let err = "connection reset by peer"; assert_eq!(pool_error_action(Some(DatabaseType::SqlServer), err), PoolErrorAction::ReconnectAndRetry); assert_eq!(pool_error_action(Some(DatabaseType::Postgres), err), PoolErrorAction::ReconnectAndRetry); let dameng_err = "Agent RPC error (-1): dm.jdbc.driver.DMException: 网络通信异常"; assert_eq!(pool_error_action(Some(DatabaseType::Dameng), dameng_err), PoolErrorAction::Discard); } #[test] fn external_driver_query_params_include_database_and_schema_context() { let config = ConnectionConfig { docs_notes_path: None, id: "jdbc-1".to_string(), name: "JDBC".to_string(), note: String::new(), db_type: DatabaseType::Jdbc, driver_profile: None, driver_label: None, url_params: None, agent_java_options: Vec::new(), host: "localhost".to_string(), port: 0, username: String::new(), password: String::new(), database: None, default_schema: None, visible_databases: None, visible_schemas: None, show_system_schemas: false, attached_databases: Vec::new(), init_script: None, color: None, transport_layers: Vec::new(), connect_timeout_secs: 5, query_timeout_secs: 30, idle_timeout_secs: 60, keepalive_interval_secs: 0, ssl: false, ca_cert_path: String::new(), client_cert_path: String::new(), client_key_path: String::new(), sysdba: false, oracle_connection_type: None, connection_string: Some("jdbc:h2:mem:test".to_string()), redis_connection_mode: None, redis_sentinel_master: String::new(), redis_sentinel_nodes: String::new(), redis_sentinel_username: String::new(), redis_sentinel_password: String::new(), redis_sentinel_tls: false, redis_cluster_nodes: String::new(), redis_key_separator: default_redis_key_separator(), redis_scan_page_size: None, redis_database_aliases: Default::default(), etcd_endpoints: String::new(), gbase_server: String::new(), informix_server: String::new(), external_config: None, jdbc_driver_class: None, jdbc_driver_paths: Vec::new(), one_time: false, save_password: true, read_only: false, is_production: false, production_databases: vec![], database_info: None, }; let params = external_driver_query_params( &config, "SELECT * FROM events", "analytics", Some("app"), &QueryExecutionOptions { max_rows: Some(500), fetch_size: Some(250), row_offset: Some(100), timeout_secs: Some(600), ..Default::default() }, ); assert_eq!(params["connection"]["id"], "jdbc-1"); assert_eq!(params["sql"], "SELECT * FROM events"); assert_eq!(params["database"], "analytics"); assert_eq!(params["schema"], "app"); assert_eq!(params["maxRows"], 500); assert_eq!(params["fetchSize"], 250); assert_eq!(params["rowOffset"], 100); assert_eq!(params["timeoutSecs"], 600); } #[test] fn agent_execute_query_params_include_row_and_fetch_limits() { let params = agent_execute_query_params( "SELECT * FROM events", Some("analytics"), Some("app"), QueryExecutionOptions { max_rows: Some(500), fetch_size: Some(250), row_offset: Some(100), timeout_secs: Some(600), table_data_preview: true, ..Default::default() }, ); assert_eq!(params["sql"], "SELECT * FROM events"); assert_eq!(params["database"], "analytics"); assert_eq!(params["schema"], "app"); assert_eq!(params["maxRows"], 500); assert_eq!(params["fetchSize"], 250); assert_eq!(params["rowOffset"], 100); assert_eq!(params["timeoutSecs"], 600); assert_eq!(params["deferLobs"], true); } #[test] fn iris_execution_context_omits_schema() { assert_eq!(schema_for_execution_context(Some(DatabaseType::Iris), Some("SQLUser")), None); assert_eq!(schema_for_execution_context(Some(DatabaseType::Oracle), Some("APP")), Some("APP")); assert_eq!(schema_for_execution_context(None, Some("APP")), Some("APP")); } #[test] fn iris_execution_context_qualifies_unqualified_dml_tables() { assert_eq!( sql_for_execution_context(Some(DatabaseType::Iris), "SELECT * FROM TABLES", Some("INFORMATION_SCHEMA")), "SELECT * FROM \"INFORMATION_SCHEMA\".TABLES" ); let qualified_join = sql_for_execution_context( Some(DatabaseType::Iris), "SELECT * FROM orders o JOIN customers c ON c.id = o.customer_id", Some("Sales"), ); assert!(qualified_join.contains("FROM \"Sales\".orders")); assert!(qualified_join.contains("JOIN \"Sales\".customers")); assert!(qualified_join.contains("c.id = o.customer_id")); assert_eq!( sql_for_execution_context(Some(DatabaseType::Iris), "SELECT * FROM INFORMATION_SCHEMA.TABLES", Some("APP")), "SELECT * FROM INFORMATION_SCHEMA.TABLES" ); } #[test] fn iris_execution_context_qualifies_nested_dml_tables_but_not_ctes() { assert_eq!( sql_for_execution_context( Some(DatabaseType::Iris), "WITH recent AS (SELECT * FROM events) SELECT * FROM recent WHERE EXISTS (SELECT 1 FROM audits)", Some("APP") ), "WITH recent AS (SELECT * FROM \"APP\".events) SELECT * FROM recent WHERE EXISTS (SELECT 1 FROM \"APP\".audits)" ); assert_eq!( sql_for_execution_context( Some(DatabaseType::Iris), "INSERT INTO events SELECT * FROM staging_events", Some("APP") ), "INSERT INTO \"APP\".events SELECT * FROM \"APP\".staging_events" ); assert_eq!( sql_for_execution_context( Some(DatabaseType::Iris), "UPDATE events SET status = 'done' WHERE id IN (SELECT event_id FROM audit_events)", Some("APP") ), "UPDATE \"APP\".events SET status = 'done' WHERE id IN (SELECT event_id FROM \"APP\".audit_events)" ); } #[test] fn iris_execution_context_leaves_ddl_and_unparseable_sql_unchanged() { assert_eq!( sql_for_execution_context(Some(DatabaseType::Iris), "CREATE TABLE events (id INT)", Some("APP")), "CREATE TABLE events (id INT)" ); assert_eq!( sql_for_execution_context(Some(DatabaseType::Iris), "SELECT %ID FROM", Some("APP")), "SELECT %ID FROM" ); assert_eq!( sql_for_execution_context(Some(DatabaseType::Postgres), "SELECT * FROM events", Some("APP")), "SELECT * FROM events" ); } #[test] fn kingbase_execution_context_qualifies_only_unqualified_relations() { assert_eq!( sql_for_execution_context(Some(DatabaseType::Kingbase), "SELECT * FROM sys_user", Some("app")), "SELECT * FROM \"app\".sys_user" ); assert_eq!( sql_for_execution_context(Some(DatabaseType::Kingbase), "SELECT * FROM other_schema.sys_user", Some("app")), "SELECT * FROM other_schema.sys_user" ); let mixed = sql_for_execution_context( Some(DatabaseType::Kingbase), "SELECT pg_typeof(u.id) FROM generate_series(1, 2) AS n JOIN sys_user u ON true", Some("app"), ); assert!(mixed.contains("FROM generate_series(1, 2) AS n"), "{mixed}"); assert!(mixed.contains("JOIN \"app\".sys_user u"), "{mixed}"); assert!(mixed.contains("pg_typeof(u.id)"), "{mixed}"); } #[test] fn kingbase_execution_context_preserves_ctes_functions_types_and_unsupported_sql() { assert_eq!( sql_for_execution_context( Some(DatabaseType::Kingbase), "WITH current_user AS (SELECT * FROM sys_user) SELECT * FROM current_user", Some("APP") ), "WITH current_user AS (SELECT * FROM \"APP\".sys_user) SELECT * FROM current_user" ); assert_eq!( sql_for_execution_context( Some(DatabaseType::Kingbase), "SELECT pg_typeof(1::int), current_user", Some("APP") ), "SELECT pg_typeof(1::int), current_user" ); assert_eq!( sql_for_execution_context(Some(DatabaseType::Kingbase), "CREATE TABLE sys_user (id INT)", Some("APP")), "CREATE TABLE sys_user (id INT)" ); assert_eq!( sql_for_execution_context(Some(DatabaseType::Kingbase), "SELECT * FROM", Some("APP")), "SELECT * FROM" ); } #[test] fn parses_postgres_drop_database_target() { assert_eq!(parse_drop_database_target("DROP DATABASE vaultwarden;"), Some("vaultwarden".to_string())); assert_eq!(parse_drop_database_target("drop database if exists \"app db\";"), Some("app db".to_string())); assert_eq!( parse_drop_database_target("/*x*/ DROP DATABASE \"app\"\"db\" -- trailing\n;"), Some("app\"db".to_string()) ); } #[test] fn ignores_non_single_drop_database_statements() { assert_eq!(parse_drop_database_target("DROP TABLE vaultwarden;"), None); assert_eq!(parse_drop_database_target("DROP DATABASE vaultwarden; SELECT 1;"), None); assert_eq!(parse_drop_database_target("DROP DATABASE 123bad;"), None); } #[test] fn chooses_safe_postgres_drop_database_admin_database() { assert_eq!(postgres_drop_database_admin_database("vaultwarden"), "postgres"); assert_eq!(postgres_drop_database_admin_database("postgres"), "template1"); } #[test] fn agent_execute_query_params_default_to_safety_row_limit() { let params = agent_execute_query_params("SELECT * FROM events", None, None, QueryExecutionOptions::default()); assert_eq!(params["sql"], "SELECT * FROM events"); assert!(params.get("database").is_none()); assert!(params.get("schema").is_none()); assert_eq!(params["maxRows"], MAX_ROWS); assert!(params.get("fetchSize").is_none()); assert!(params.get("rowOffset").is_none()); assert!(params.get("timeoutSecs").is_none()); } #[test] fn agent_execute_query_page_params_include_page_fetch_and_safety_limits() { let params = agent_execute_query_page_params( "SELECT * FROM events", Some("analytics"), Some("app"), QueryExecutionOptions { page_size: Some(500), fetch_size: Some(250), row_offset: Some(100), timeout_secs: Some(600), table_data_preview: true, ..Default::default() }, ); assert_eq!(params["sql"], "SELECT * FROM events"); assert_eq!(params["database"], "analytics"); assert_eq!(params["schema"], "app"); assert_eq!(params["pageSize"], 500); assert_eq!(params["fetchSize"], 250); assert_eq!(params["rowOffset"], 100); assert_eq!(params["timeoutSecs"], 600); assert_eq!(params["maxRows"], MAX_ROWS); assert_eq!(params["deferLobs"], true); } #[test] fn agent_query_row_counts_are_clamped_to_java_signed_int_range() { let oversized = AGENT_PROTOCOL_MAX_ROWS.saturating_add(1); let params = agent_execute_query_page_params( "SELECT * FROM events", None, None, QueryExecutionOptions { page_size: Some(oversized), fetch_size: Some(oversized), max_rows: Some(oversized), row_offset: Some(oversized), ..Default::default() }, ); assert_eq!(params["pageSize"], AGENT_PROTOCOL_MAX_ROWS); assert_eq!(params["fetchSize"], AGENT_PROTOCOL_MAX_ROWS); assert_eq!(params["maxRows"], AGENT_PROTOCOL_MAX_ROWS); assert_eq!(params["rowOffset"], AGENT_PROTOCOL_MAX_ROWS); assert_eq!(agent_fetch_query_page_params("session-1", oversized)["pageSize"], AGENT_PROTOCOL_MAX_ROWS); } #[test] fn agent_fetch_query_page_params_include_session_and_page_size() { let params = agent_fetch_query_page_params("session-1", 500); assert_eq!(params["sessionId"], "session-1"); assert_eq!(params["pageSize"], 500); } #[test] fn agent_close_query_session_params_include_session() { let params = agent_close_query_session_params("session-1"); assert_eq!(params["sessionId"], "session-1"); } #[test] fn multi_statement_execution_does_not_retain_query_cursors() { let options = QueryExecutionOptions { max_rows: Some(100_000), fetch_size: Some(100), page_size: Some(100), timeout_secs: Some(30), ..Default::default() }; let adjusted = options_for_sequential_statements(&options, 2, Some(DatabaseType::Kingbase)); assert_eq!(adjusted.page_size, None); assert_eq!(adjusted.max_rows, Some(100)); assert_eq!(adjusted.fetch_size, Some(100)); assert_eq!(adjusted.timeout_secs, Some(30)); } #[test] fn multi_statement_execution_preserves_smaller_row_limit() { let options = QueryExecutionOptions { max_rows: Some(25), page_size: Some(100), ..Default::default() }; let adjusted = options_for_sequential_statements(&options, 2, Some(DatabaseType::Kingbase)); assert_eq!(adjusted.page_size, None); assert_eq!(adjusted.max_rows, Some(25)); } #[test] fn single_statement_and_existing_cursor_execution_keep_paging_options() { let first_page = QueryExecutionOptions { max_rows: Some(100_000), page_size: Some(100), ..Default::default() }; let next_page = QueryExecutionOptions { max_rows: Some(100_000), page_size: Some(100), result_session_id: Some("session-1".to_string()), ..Default::default() }; let single = options_for_sequential_statements(&first_page, 1, Some(DatabaseType::Kingbase)); let existing_cursor = options_for_sequential_statements(&next_page, 2, Some(DatabaseType::Kingbase)); assert_eq!(single.page_size, Some(100)); assert_eq!(single.max_rows, Some(100_000)); assert_eq!(existing_cursor.page_size, Some(100)); assert_eq!(existing_cursor.result_session_id.as_deref(), Some("session-1")); } #[test] fn other_databases_keep_multi_statement_cursor_options() { let options = QueryExecutionOptions { max_rows: Some(100_000), page_size: Some(100), ..Default::default() }; let adjusted = options_for_sequential_statements(&options, 2, Some(DatabaseType::Oracle)); assert_eq!(adjusted.page_size, Some(100)); assert_eq!(adjusted.max_rows, Some(100_000)); } #[test] fn agent_timeout_discards_pool_but_does_not_retry_same_query() { assert!(should_discard_agent_pool_after_error("Query timed out after 30 seconds")); assert!(should_discard_agent_pool_after_error("Agent RPC call timed out (30s)")); assert!(!is_connection_error("Agent RPC call timed out (30s)")); assert_eq!( pool_error_action(Some(DatabaseType::Oracle), "Agent RPC call timed out (30s)"), PoolErrorAction::Discard ); } #[test] fn structured_agent_disposition_controls_pool_recovery() { let quarantined = "Agent RPC error (-1): lost\nDBX_AGENT_ERROR_DATA:{\"category\":\"connection\",\"sessionDisposition\":\"quarantine\"}"; let replace_runtime = "Agent RPC error (-1): saturated\nDBX_AGENT_ERROR_DATA:{\"category\":\"resource\",\"sessionDisposition\":\"replace_runtime\"}"; assert!(should_discard_agent_pool_after_error(quarantined)); assert!(should_discard_agent_pool_after_error(replace_runtime)); assert_eq!(pool_error_action(Some(DatabaseType::Oracle), quarantined), PoolErrorAction::Discard); assert_eq!(pool_error_action(Some(DatabaseType::Oracle), replace_runtime), PoolErrorAction::Discard); } #[test] fn unavailable_agent_pipes_discard_without_replaying_user_operations() { assert!(should_discard_agent_pool_after_error("Agent stdin not available")); assert!(should_discard_agent_pool_after_error("Agent stdout not available")); assert!(is_connection_error("Agent stdin not available")); assert!(is_connection_error("Agent stdout not available")); assert!(is_connection_error("Agent runtime terminated")); assert!(is_connection_error("Agent runtime is unavailable")); assert_eq!( pool_error_action(Some(DatabaseType::Oracle), "Agent stdin not available"), PoolErrorAction::Discard ); assert_eq!(pool_error_action(Some(DatabaseType::Oracle), "Agent runtime terminated"), PoolErrorAction::Discard); assert_eq!( pool_error_action(Some(DatabaseType::Oracle), "Agent runtime is unavailable"), PoolErrorAction::Discard ); } #[test] fn query_results_convert_unsafe_json_integers_to_strings_for_js() { let result = db::QueryResult { columns: vec!["id".to_string(), "nested".to_string()], column_types: Vec::new(), column_sortables: vec![], spatial_columns: vec![], spatial_values: vec![], rows: vec![vec![ serde_json::json!(2_041_797_190_226_354_178_i64), serde_json::json!([1, 2_041_797_190_226_354_178_i64]), ]], affected_rows: 0, execution_time_ms: 0, truncated: false, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: Vec::new(), }; let normalized = normalize_query_result_for_js(result); assert_eq!(normalized.rows[0][0], serde_json::json!("2041797190226354178")); assert_eq!(normalized.rows[0][1], serde_json::json!([1, "2041797190226354178"])); } #[test] fn extracts_server_large_value_markers_and_restores_source_columns() { let mut result = db::QueryResult { columns: vec![ "id".to_string(), "payload".to_string(), format!("{}T_1", crate::sql_dialect::DBX_LARGE_VALUE_BYTES_COLUMN_PREFIX), "note".to_string(), ], column_types: vec!["integer".to_string(), "text".to_string(), "bigint".to_string(), "text".to_string()], column_sortables: vec![true, true, true, true], spatial_columns: Vec::new(), spatial_values: Vec::new(), rows: vec![ vec![ serde_json::json!(1), serde_json::json!("预览文本多"), serde_json::json!("T:4"), serde_json::json!("a"), ], vec![serde_json::json!(2), serde_json::json!("短值"), serde_json::json!("T:4"), serde_json::json!("b")], ], affected_rows: 0, execution_time_ms: 0, truncated: false, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: Vec::new(), }; let cells = extract_server_large_value_markers(&mut result); assert_eq!(result.columns, vec!["id", "payload", "note"]); assert_eq!(result.column_types, vec!["integer", "text", "text"]); assert_eq!( result.rows[0], vec![serde_json::json!(1), serde_json::json!("预览文本..."), serde_json::json!("a")] ); assert_eq!(result.rows[1], vec![serde_json::json!(2), serde_json::json!("短值"), serde_json::json!("b")]); assert_eq!( cells, vec![db::LargeValueCell { row_index: 0, column_index: 1, original_bytes: SERVER_LARGE_VALUE_UNKNOWN_BYTES, }] ); } #[test] fn truncates_server_binary_preview_after_the_configured_byte_count() { let mut result = db::QueryResult { columns: vec![ "raw_value".to_string(), format!("{}B_0", crate::sql_dialect::DBX_LARGE_VALUE_BYTES_COLUMN_PREFIX), ], column_types: vec!["bytea".to_string(), "text".to_string()], column_sortables: vec![true, true], spatial_columns: Vec::new(), spatial_values: Vec::new(), rows: vec![vec![serde_json::json!("0x0102030405"), serde_json::json!("B:4:5")]], affected_rows: 0, execution_time_ms: 0, truncated: false, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: Vec::new(), }; let cells = extract_server_large_value_markers(&mut result); assert_eq!(result.rows, vec![vec![serde_json::json!("0x01020304...")]]); assert_eq!(cells, vec![db::LargeValueCell { row_index: 0, column_index: 0, original_bytes: 5 }]); } #[test] fn extracts_deferred_oracle_clob_markers_without_changing_placeholder() { let mut result = db::QueryResult { columns: vec![ "id".to_string(), "payload".to_string(), format!("{}C_1", crate::sql_dialect::DBX_LARGE_VALUE_BYTES_COLUMN_PREFIX), ], column_types: vec!["number".to_string(), "varchar2".to_string(), "varchar2".to_string()], column_sortables: vec![true; 3], spatial_columns: Vec::new(), spatial_values: Vec::new(), rows: vec![ vec![serde_json::json!(1), serde_json::json!(""), serde_json::json!("D:1")], vec![serde_json::json!(2), serde_json::Value::Null, serde_json::Value::Null], ], affected_rows: 0, execution_time_ms: 0, truncated: false, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: Vec::new(), }; let cells = extract_server_large_value_markers(&mut result); assert_eq!(result.columns, vec!["id", "payload"]); assert_eq!(result.column_types, vec!["number", "clob"]); assert_eq!(result.rows[0], vec![serde_json::json!(1), serde_json::json!("")]); assert_eq!(result.rows[1], vec![serde_json::json!(2), serde_json::Value::Null]); assert_eq!( cells, vec![db::LargeValueCell { row_index: 0, column_index: 1, original_bytes: SERVER_LARGE_VALUE_UNKNOWN_BYTES, }] ); } #[test] fn restores_pgvector_previews_as_arrays_and_marks_only_truncated_values() { let marker = format!("{}V_0", crate::sql_dialect::DBX_LARGE_VALUE_BYTES_COLUMN_PREFIX); let mut result = db::QueryResult { columns: vec!["embedding".to_string(), marker], column_types: vec!["text".to_string(), "text".to_string()], column_sortables: vec![true, true], spatial_columns: Vec::new(), spatial_values: Vec::new(), rows: vec![ vec![serde_json::json!("[0.1,0.2,0.3]"), serde_json::json!("V:9")], vec![serde_json::json!("[0.1,0.2]"), serde_json::json!("V:20")], ], affected_rows: 0, execution_time_ms: 0, truncated: false, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: Vec::new(), }; let cells = extract_server_large_value_markers(&mut result); assert_eq!(result.column_types, vec!["vector"]); assert_eq!(result.rows[0][0], serde_json::json!([0.1, 0.2])); assert_eq!(result.rows[1][0], serde_json::json!([0.1, 0.2])); assert_eq!(cells.len(), 1); assert_eq!(cells[0].row_index, 0); } #[test] fn restores_server_preview_types_without_rows() { let mut result = db::QueryResult { columns: vec![ "embedding".to_string(), format!("{}V_0", crate::sql_dialect::DBX_LARGE_VALUE_BYTES_COLUMN_PREFIX), "metadata".to_string(), format!("{}K_1", crate::sql_dialect::DBX_LARGE_VALUE_BYTES_COLUMN_PREFIX), ], column_types: vec!["text".to_string(), "text".to_string(), "text".to_string(), "text".to_string()], column_sortables: vec![true; 4], spatial_columns: Vec::new(), spatial_values: Vec::new(), rows: Vec::new(), affected_rows: 0, execution_time_ms: 0, truncated: false, session_id: None, has_more: false, elasticsearch_raw_body: None, messages: Vec::new(), }; let cells = extract_server_large_value_markers(&mut result); assert_eq!(result.columns, vec!["embedding", "metadata"]); assert_eq!(result.column_types, vec!["vector", "jsonb"]); assert!(cells.is_empty()); } #[test] fn ordinary_queries_preserve_columns_that_resemble_preview_markers() { let marker = format!("{}0", crate::sql_dialect::DBX_LARGE_VALUE_BYTES_COLUMN_PREFIX); let mut result = empty_query_result(1); result.columns = vec!["payload".to_string(), marker.clone()]; result.column_types = vec!["text".to_string(), "bigint".to_string()]; result.column_sortables = vec![true, true]; result.rows = vec![vec![serde_json::json!("value"), serde_json::json!(123)]]; let ordinary = ExecuteMultiResult::success_with_index_and_optional_server_large_values(result, 0, false); assert_eq!(ordinary.result.columns, vec!["payload".to_string(), marker]); assert_eq!(ordinary.result.rows[0], vec![serde_json::json!("value"), serde_json::json!(123)]); assert!(ordinary.large_value_cells.is_empty()); } #[test] fn single_statement_preview_preserves_absent_statement_index() { let result = single_statement_multi_result(Ok(empty_query_result(1)), true).unwrap(); assert_eq!(result.len(), 1); assert_eq!(result[0].statement_index, None); } #[test] fn query_execution_options_default_use_transaction_is_none() { let opts = QueryExecutionOptions::default(); assert_eq!(opts.use_transaction, None); } #[test] fn query_execution_options_use_transaction_some_true_is_preserved() { let opts = QueryExecutionOptions { use_transaction: Some(true), ..Default::default() }; assert_eq!(opts.use_transaction, Some(true)); } #[test] fn query_execution_options_use_transaction_some_false_is_preserved() { let opts = QueryExecutionOptions { use_transaction: Some(false), ..Default::default() }; assert_eq!(opts.use_transaction, Some(false)); } #[test] fn database_backup_transactions_request_consistent_snapshots() { assert_eq!(postgres_transaction_begin_sql(true), "BEGIN TRANSACTION ISOLATION LEVEL REPEATABLE READ READ ONLY"); assert_eq!(mysql_transaction_isolation_sql(true), Some("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ")); assert_eq!( mysql_transaction_begin_sql_candidates(true), [ "START TRANSACTION WITH CONSISTENT SNAPSHOT, READ ONLY", "START TRANSACTION WITH CONSISTENT SNAPSHOT", "START TRANSACTION", ] ); assert_eq!(mysql_transaction_isolation_sql(false), None); assert_eq!(mysql_transaction_begin_sql_candidates(false), ["START TRANSACTION"]); } #[test] fn mysql_backup_transaction_only_falls_back_for_syntax_errors() { let syntax_error = mysql_async::Error::Server(mysql_async::ServerError { code: 1064, message: "unsupported transaction characteristic".to_string(), state: "42000".to_string(), }); let permission_error = mysql_async::Error::Server(mysql_async::ServerError { code: 1044, message: "access denied".to_string(), state: "42000".to_string(), }); assert!(mysql_error_is_syntax_error(&syntax_error)); assert!(!mysql_error_is_syntax_error(&permission_error)); } }