t8y2 committed
1 parent
5f8cc0b08b
commit
a9299e947f
2 files changed
+324
-4
No files matched your search
@@ -24,7 +24,7 @@ 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;
|
||||
use crate::sql::{split_sql_batches, split_sql_statements};
|
||||
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};
|
||||
|
||||
@@ -2280,7 +2280,14 @@ pub async fn execute_multi_core_with_options_for_client_and_progress_typed(
|
||||
}
|
||||
};
|
||||
|
||||
if statements.len() <= 1 && !(mysql_pool.is_some() && options.max_result_bytes.is_some_and(|value| value > 0)) {
|
||||
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();
|
||||
return single_statement_multi_result(
|
||||
execute_sql_statement_with_options_typed(
|
||||
@@ -2383,6 +2390,18 @@ fn single_statement_multi_result(
|
||||
result.map(|result| vec![result.into()])
|
||||
}
|
||||
|
||||
fn mysql_single_statement_uses_batch_route(
|
||||
db_type: Option<DatabaseType>,
|
||||
has_mysql_pool: bool,
|
||||
sql: &str,
|
||||
max_result_bytes: Option<usize>,
|
||||
) -> 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 {
|
||||
async fn execute_statement(&mut self, statement: &str) -> Result<Vec<db::mysql::MySqlQueryResult>, String>;
|
||||
|
||||
@@ -2494,6 +2513,14 @@ fn mysql_non_result_batch_end(
|
||||
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<E>(
|
||||
executor: &mut E,
|
||||
statements: &[String],
|
||||
@@ -2622,7 +2649,8 @@ async fn execute_multi_mysql(
|
||||
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 = !options.continue_on_error && mode == crate::connection::MysqlMode::Normal;
|
||||
let pipeline_non_result_statements =
|
||||
mysql_non_result_pipeline_enabled(statements.len(), options.continue_on_error, mode);
|
||||
let mut conn = match db::mysql::get_conn_with_health_check_with_cancel(
|
||||
pool,
|
||||
operation_budget.checkout_timeout,
|
||||
@@ -5400,6 +5428,53 @@ for line in sys.stdin:
|
||||
);
|
||||
}
|
||||
|
||||
#[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()];
|
||||
|
||||
@@ -4,7 +4,9 @@ use std::time::{Duration, Instant};
|
||||
|
||||
use dbx_core::connection::AppState;
|
||||
use dbx_core::models::connection::{ConnectionConfig, DatabaseType};
|
||||
use dbx_core::query::{execute_multi_core, execute_sql_statement};
|
||||
use dbx_core::query::{
|
||||
execute_multi_core, execute_multi_core_with_options_for_client_typed, execute_sql_statement, QueryExecutionOptions,
|
||||
};
|
||||
use dbx_core::query_result_export::{export_query_result_core, ExportStatus, QueryResultExportRequest};
|
||||
use dbx_core::sql::{split_sql_statements_for_database, SqlFileRequest};
|
||||
use dbx_core::sql_file_import::execute_sql_file_path;
|
||||
@@ -132,6 +134,249 @@ async fn live_mysql_stored_procedure_preserves_all_result_sets() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore = "requires a writable MySQL endpoint for stored procedure creation"]
|
||||
async fn live_mysql_single_call_public_route_preserves_all_result_sets() {
|
||||
let suffix = uuid::Uuid::new_v4().simple().to_string();
|
||||
let connection_id = format!("live-mysql-multi-result-{suffix}");
|
||||
let database = std::env::var("DBX_LIVE_SQL_FILE_MYSQL_DATABASE").expect("DBX_LIVE_SQL_FILE_MYSQL_DATABASE");
|
||||
let config = live_mysql_sql_file_config(&connection_id);
|
||||
let (state, db_path) = app_state_with_config(config).await;
|
||||
let procedure = format!("dbx_issue_5560_{}", &suffix[..8]);
|
||||
let empty_procedure = format!("dbx_issue_5560_empty_{}", &suffix[..8]);
|
||||
let slow_procedure = format!("dbx_issue_5560_slow_{}", &suffix[..8]);
|
||||
let table = format!("dbx_issue_5560_{}", &suffix[..8]);
|
||||
let session_id = format!("dbx-issue-5560-{suffix}");
|
||||
let create_sql = format!(
|
||||
"CREATE PROCEDURE `{procedure}`() BEGIN \
|
||||
SET @dbx_issue_5560_session = 'kept'; \
|
||||
SELECT 11 AS value UNION ALL SELECT 12 AS value; \
|
||||
SELECT 21 AS value UNION ALL SELECT 22 AS value; \
|
||||
SELECT 31 AS value UNION ALL SELECT 32 AS value; \
|
||||
END"
|
||||
);
|
||||
|
||||
execute_sql_statement(&state, &connection_id, &database, &create_sql, None, None)
|
||||
.await
|
||||
.expect("create live multi-result procedure");
|
||||
execute_sql_statement(
|
||||
&state,
|
||||
&connection_id,
|
||||
&database,
|
||||
&format!("CREATE PROCEDURE `{empty_procedure}`() BEGIN SET @dbx_issue_5560_empty = 1; END"),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.expect("create live no-result procedure");
|
||||
execute_sql_statement(
|
||||
&state,
|
||||
&connection_id,
|
||||
&database,
|
||||
&format!("CREATE PROCEDURE `{slow_procedure}`() BEGIN SELECT SLEEP(5) AS slept; END"),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.expect("create live slow procedure");
|
||||
execute_sql_statement(
|
||||
&state,
|
||||
&connection_id,
|
||||
&database,
|
||||
&format!("CREATE TABLE `{table}` (id INT PRIMARY KEY, value INT NOT NULL)"),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.expect("create live DML table");
|
||||
execute_sql_statement(
|
||||
&state,
|
||||
&connection_id,
|
||||
&database,
|
||||
&format!("INSERT INTO `{table}` (id, value) VALUES (1, 1)"),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.expect("seed live DML table");
|
||||
|
||||
let options =
|
||||
QueryExecutionOptions { max_rows: Some(1), client_session_id: Some(session_id.clone()), ..Default::default() };
|
||||
let results = execute_multi_core_with_options_for_client_typed(
|
||||
&state,
|
||||
&connection_id,
|
||||
&database,
|
||||
&format!("CALL `{procedure}`()"),
|
||||
None,
|
||||
None,
|
||||
options.clone(),
|
||||
)
|
||||
.await
|
||||
.expect("execute single CALL through the public multi-result route");
|
||||
let ordinary_select = execute_multi_core_with_options_for_client_typed(
|
||||
&state,
|
||||
&connection_id,
|
||||
&database,
|
||||
"SELECT 42 AS value",
|
||||
None,
|
||||
None,
|
||||
options.clone(),
|
||||
)
|
||||
.await
|
||||
.expect("execute ordinary SELECT through the public route");
|
||||
let ordinary_dml = execute_multi_core_with_options_for_client_typed(
|
||||
&state,
|
||||
&connection_id,
|
||||
&database,
|
||||
&format!("UPDATE `{table}` SET value = value + 1 WHERE id = 1"),
|
||||
None,
|
||||
None,
|
||||
options.clone(),
|
||||
)
|
||||
.await
|
||||
.expect("execute ordinary DML through the public route");
|
||||
let session_result = execute_multi_core_with_options_for_client_typed(
|
||||
&state,
|
||||
&connection_id,
|
||||
&database,
|
||||
"SELECT @dbx_issue_5560_session AS session_value",
|
||||
None,
|
||||
None,
|
||||
options.clone(),
|
||||
)
|
||||
.await
|
||||
.expect("reuse the CALL client session");
|
||||
let empty_results = execute_multi_core_with_options_for_client_typed(
|
||||
&state,
|
||||
&connection_id,
|
||||
&database,
|
||||
&format!("CALL `{empty_procedure}`()"),
|
||||
None,
|
||||
None,
|
||||
options.clone(),
|
||||
)
|
||||
.await
|
||||
.expect("execute no-result CALL through the public route");
|
||||
let sql_error = execute_multi_core_with_options_for_client_typed(
|
||||
&state,
|
||||
&connection_id,
|
||||
&database,
|
||||
&format!("CALL `dbx_issue_5560_missing_{}`()", &suffix[..8]),
|
||||
None,
|
||||
None,
|
||||
options.clone(),
|
||||
)
|
||||
.await
|
||||
.expect("return a structured execution result for a missing procedure");
|
||||
|
||||
state.configs.write().await.get_mut(&connection_id).expect("live config").read_only = true;
|
||||
let read_only_error = execute_multi_core_with_options_for_client_typed(
|
||||
&state,
|
||||
&connection_id,
|
||||
&database,
|
||||
&format!("CALL `{procedure}`()"),
|
||||
None,
|
||||
None,
|
||||
options.clone(),
|
||||
)
|
||||
.await
|
||||
.expect_err("read-only mode must reject CALL before dispatch");
|
||||
state.configs.write().await.get_mut(&connection_id).expect("live config").read_only = false;
|
||||
|
||||
let timeout_results = execute_multi_core_with_options_for_client_typed(
|
||||
&state,
|
||||
&connection_id,
|
||||
&database,
|
||||
&format!("CALL `{slow_procedure}`()"),
|
||||
None,
|
||||
None,
|
||||
QueryExecutionOptions { timeout_secs: Some(1), ..options.clone() },
|
||||
)
|
||||
.await
|
||||
.expect("return a structured timeout result");
|
||||
let cancel_token = CancellationToken::new();
|
||||
let cancel_task = {
|
||||
let cancel_token = cancel_token.clone();
|
||||
tokio::spawn(async move {
|
||||
tokio::time::sleep(Duration::from_millis(200)).await;
|
||||
cancel_token.cancel();
|
||||
})
|
||||
};
|
||||
let canceled_results = execute_multi_core_with_options_for_client_typed(
|
||||
&state,
|
||||
&connection_id,
|
||||
&database,
|
||||
&format!("CALL `{slow_procedure}`()"),
|
||||
None,
|
||||
Some(cancel_token),
|
||||
options.clone(),
|
||||
)
|
||||
.await
|
||||
.expect("return a structured cancellation result");
|
||||
cancel_task.await.unwrap();
|
||||
let recovery_result = execute_multi_core_with_options_for_client_typed(
|
||||
&state,
|
||||
&connection_id,
|
||||
&database,
|
||||
"SELECT 7 AS recovered",
|
||||
None,
|
||||
None,
|
||||
options,
|
||||
)
|
||||
.await
|
||||
.expect("recreate the discarded session pool after timeout and cancellation");
|
||||
|
||||
for cleanup_sql in [
|
||||
format!("DROP PROCEDURE `{procedure}`"),
|
||||
format!("DROP PROCEDURE `{empty_procedure}`"),
|
||||
format!("DROP PROCEDURE `{slow_procedure}`"),
|
||||
format!("DROP TABLE `{table}`"),
|
||||
] {
|
||||
execute_sql_statement(&state, &connection_id, &database, &cleanup_sql, None, None)
|
||||
.await
|
||||
.expect("clean up live issue fixture");
|
||||
}
|
||||
let _ = std::fs::remove_file(db_path);
|
||||
|
||||
assert_eq!(results.len(), 3);
|
||||
assert_eq!(results.iter().map(|result| result.statement_index).collect::<Vec<_>>(), vec![Some(0); 3]);
|
||||
assert_eq!(
|
||||
results.iter().map(|result| result.result.rows[0][0].clone()).collect::<Vec<_>>(),
|
||||
vec![serde_json::json!("11"), serde_json::json!("21"), serde_json::json!("31")]
|
||||
);
|
||||
assert!(results.iter().all(|result| result.result.truncated));
|
||||
assert_eq!(ordinary_select.len(), 1);
|
||||
assert_eq!(ordinary_select[0].statement_index, None);
|
||||
assert_eq!(ordinary_select[0].result.rows, vec![vec![serde_json::json!("42")]]);
|
||||
assert_eq!(ordinary_dml.len(), 1);
|
||||
assert_eq!(ordinary_dml[0].statement_index, None);
|
||||
assert_eq!(ordinary_dml[0].result.affected_rows, 1);
|
||||
assert_eq!(session_result[0].result.rows, vec![vec![serde_json::json!("kept")]]);
|
||||
assert_eq!(empty_results.len(), 1);
|
||||
assert_eq!(empty_results[0].statement_index, Some(0));
|
||||
assert!(empty_results[0].result.columns.is_empty());
|
||||
assert!(empty_results[0].result.rows.is_empty());
|
||||
assert_eq!(sql_error.len(), 1);
|
||||
assert_eq!(sql_error[0].statement_index, Some(0));
|
||||
assert!(sql_error[0].execution_error);
|
||||
let sql_error_message = sql_error[0].result.rows[0][0].as_str().expect("missing procedure error message");
|
||||
assert!(sql_error_message.contains("does not exist"), "unexpected SQL error: {sql_error_message}");
|
||||
assert_eq!(sql_error[0].error.as_ref().map(|error| error.code()), Some("DBX-LEGACY-0001"));
|
||||
assert!(read_only_error.into_legacy_string().to_ascii_lowercase().contains("read-only"));
|
||||
assert_eq!(timeout_results.len(), 1);
|
||||
assert!(timeout_results[0].execution_error);
|
||||
assert_eq!(timeout_results[0].statement_index, Some(0));
|
||||
let timeout_message = timeout_results[0].result.rows[0][0].as_str().expect("timeout error message");
|
||||
assert_eq!(timeout_message, "Query timed out after 1 seconds");
|
||||
assert_eq!(timeout_results[0].error.as_ref().map(|error| error.code()), Some("DBX-LEGACY-0001"));
|
||||
assert_eq!(canceled_results.len(), 1);
|
||||
assert!(canceled_results[0].execution_error);
|
||||
assert_eq!(canceled_results[0].statement_index, Some(0));
|
||||
assert_eq!(canceled_results[0].result.rows[0][0], serde_json::json!("Query canceled"));
|
||||
assert_eq!(canceled_results[0].error.as_ref().map(|error| error.code()), Some("DBX-JDBC-2003"));
|
||||
assert_eq!(recovery_result[0].result.rows, vec![vec![serde_json::json!("7")]]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore = "requires a remote MySQL-compatible endpoint with a limited result-set query"]
|
||||
async fn live_mysql_compatible_limited_text_protocol_query_succeeds() {
|
||||
|
||||
Reference in new issue
Block a user