xiaoxue_admin
/
dbx
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(mysql): preserve session state for multi-statement queries

This commit is contained in:
t8y2 2026-05-30 17:56:28 +08:00
parent 7e3a34f1e1
commit 0eacda465e
2 changed files with 84 additions and 31 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

33
crates/dbx-core/src/db/mysql.rs
View File

@ -849,13 +849,12 @@ pub async fn get_conn_with_health_check(pool: &MySqlPool) -> Result<mysql_async:
}
}
async fn execute_result_set_with_text_protocol(
pool: &MySqlPool,
async fn execute_result_set_with_text_protocol_on_conn(
conn: &mut mysql_async::Conn,
sql: &str,
row_limit: usize,
start: Instant,
) -> Result<QueryResult, String> {
let mut conn = get_conn_with_health_check(pool).await?;
let mut result = conn.query_iter(sql).await.map_err(|e| e.to_string())?;
let columns: Vec<String> = result.columns_ref().iter().map(|c| c.name_str().to_string()).collect();
@ -891,13 +890,12 @@ async fn execute_result_set_with_text_protocol(
})
}
async fn execute_result_set_with_prepared_protocol(
pool: &MySqlPool,
async fn execute_result_set_with_prepared_protocol_on_conn(
conn: &mut mysql_async::Conn,
sql: &str,
row_limit: usize,
start: Instant,
) -> Result<QueryResult, String> {
let mut conn = get_conn_with_health_check(pool).await?;
let mut result = conn.exec_iter(sql, ()).await.map_err(|e| e.to_string())?;
let columns: Vec<String> = result.columns_ref().iter().map(|c| c.name_str().to_string()).collect();
@ -942,35 +940,44 @@ pub async fn execute_query_with_max_rows(
sql: &str,
bare: bool,
max_rows: Option<usize>,
) -> Result<QueryResult, String> {
let mut conn = get_conn_with_health_check(pool).await?;
execute_query_on_conn_with_max_rows(&mut conn, sql, bare, max_rows).await
}
pub async fn execute_query_on_conn_with_max_rows(
conn: &mut mysql_async::Conn,
sql: &str,
bare: bool,
max_rows: Option<usize>,
) -> Result<QueryResult, String> {
let start = Instant::now();
let row_limit = query_result_row_limit(max_rows);
if is_result_set_query(sql) {
if bare || requires_text_protocol_query(sql) {
execute_result_set_with_text_protocol(pool, sql, row_limit, start).await
execute_result_set_with_text_protocol_on_conn(conn, sql, row_limit, start).await
} else {
match execute_result_set_with_prepared_protocol(pool, sql, row_limit, start).await {
match execute_result_set_with_prepared_protocol_on_conn(conn, sql, row_limit, start).await {
Ok(result) => Ok(result),
Err(err) if mysql_error_should_retry_with_text_protocol(&err) => {
execute_result_set_with_text_protocol(pool, sql, row_limit, start).await
execute_result_set_with_text_protocol_on_conn(conn, sql, row_limit, start).await
}
Err(err) => Err(err),
}
}
} else {
let mut conn = get_conn_with_health_check(pool).await?;
let previous_explicit_timestamp_defaults = enable_explicit_timestamp_defaults_for_query(&mut conn, sql).await;
let previous_explicit_timestamp_defaults = enable_explicit_timestamp_defaults_for_query(conn, sql).await;
let result = match conn.query_iter(sql).await {
Ok(result) => result,
Err(err) => {
restore_explicit_timestamp_defaults_for_query(&mut conn, previous_explicit_timestamp_defaults).await;
restore_explicit_timestamp_defaults_for_query(conn, previous_explicit_timestamp_defaults).await;
return Err(err.to_string());
}
};
let affected_rows = result.affected_rows();
let drop_result = result.drop_result().await;
restore_explicit_timestamp_defaults_for_query(&mut conn, previous_explicit_timestamp_defaults).await;
restore_explicit_timestamp_defaults_for_query(conn, previous_explicit_timestamp_defaults).await;
drop_result.map_err(|e| e.to_string())?;
Ok(QueryResult {

82
crates/dbx-core/src/query.rs
View File

@ -774,6 +774,15 @@ pub async fn execute_multi_core_with_options(
|| split_sql_statements(sql),
|db_type| crate::sql::split_sql_statements_for_database(sql, db_type),
);
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 {
let single_sql = statements.into_iter().next().unwrap_or_default();
let result = execute_sql_statement_with_options(
@ -789,18 +798,14 @@ pub async fn execute_multi_core_with_options(
return Ok(vec![result]);
}
if let Some((pool, mode)) = mysql_pool {
return execute_multi_mysql(&pool, mode, &statements, cancel_token, options).await;
}
let mut results = Vec::with_capacity(statements.len());
for stmt in &statements {
if is_canceled(&cancel_token) {
results.push(db::QueryResult {
columns: vec!["Error".to_string()],
rows: vec![vec![serde_json::Value::String(canceled_error())]],
affected_rows: 0,
execution_time_ms: 0,
truncated: false,
session_id: None,
has_more: false,
});
results.push(error_query_result(canceled_error()));
break;
}
match execute_sql_statement_with_options(
@ -816,15 +821,7 @@ pub async fn execute_multi_core_with_options(
{
Ok(r) => results.push(r),
Err(e) => {
results.push(db::QueryResult {
columns: vec!["Error".to_string()],
rows: vec![vec![serde_json::Value::String(e)]],
affected_rows: 0,
execution_time_ms: 0,
truncated: false,
session_id: None,
has_more: false,
});
results.push(error_query_result(e));
}
}
}
@ -832,6 +829,55 @@ pub async fn execute_multi_core_with_options(
Ok(results)
}
async fn execute_multi_mysql(
pool: &db::mysql::MySqlPool,
mode: crate::connection::MysqlMode,
statements: &[String],
cancel_token: Option<CancellationToken>,
options: QueryExecutionOptions,
) -> Result<Vec<db::QueryResult>, String> {
let query_timeout = resolve_query_timeout(options.timeout_secs);
let bare = mode == crate::connection::MysqlMode::Bare;
let max_rows = options.max_rows;
let mut conn = match db::mysql::get_conn_with_health_check(pool).await {
Ok(conn) => conn,
Err(err) => return Ok(vec![error_query_result(err)]),
};
let mut results = Vec::with_capacity(statements.len());
for stmt in statements {
if is_canceled(&cancel_token) {
results.push(error_query_result(canceled_error()));
break;
}
match wait_for_query_opt(
cancel_token.clone(),
query_timeout,
db::mysql::execute_query_on_conn_with_max_rows(&mut conn, stmt, bare, max_rows),
)
.await
{
Ok(result) => results.push(result),
Err(err) => results.push(error_query_result(err)),
}
}
Ok(results)
}
fn error_query_result(message: String) -> db::QueryResult {
db::QueryResult {
columns: vec!["Error".to_string()],
rows: vec![vec![serde_json::Value::String(message)]],
affected_rows: 0,
execution_time_ms: 0,
truncated: false,
session_id: None,
has_more: false,
}
}
async fn execute_multi_sqlserver(
state: &AppState,
pool_key: &str,