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

fix(oceanbase): sync query timeout before execution

This commit is contained in:
t8y2 2026-07-06 13:40:20 +08:00
parent ae050dadd0
commit 0c965eb836
2 changed files with 86 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

33
crates/dbx-core/src/connection.rs
View File

@ -64,12 +64,16 @@ fn is_oceanbase_mysql_config(config: &ConnectionConfig) -> bool {
&& config.driver_profile.as_deref().is_some_and(|profile| profile.eq_ignore_ascii_case("oceanbase"))
}
fn oceanbase_mysql_setup_queries(config: &ConnectionConfig) -> Vec<String> {
if !is_oceanbase_mysql_config(config) || config.query_timeout_secs == 0 {
return Vec::new();
pub(crate) fn oceanbase_mysql_query_timeout_sql(config: &ConnectionConfig, timeout_secs: u64) -> Option<String> {
if !is_oceanbase_mysql_config(config) || timeout_secs == 0 {
return None;
}
let timeout_us = config.query_timeout_secs.saturating_mul(1_000_000);
vec![format!("SET ob_query_timeout = {timeout_us}")]
let timeout_us = timeout_secs.saturating_mul(1_000_000);
Some(format!("SET ob_query_timeout = {timeout_us}"))
}
fn oceanbase_mysql_setup_queries(config: &ConnectionConfig) -> Vec<String> {
oceanbase_mysql_query_timeout_sql(config, config.query_timeout_secs).into_iter().collect()
}
pub enum PoolKind {
@ -2967,10 +2971,10 @@ async fn detect_ob_oracle_mode(config: &ConnectionConfig, pool: &db::mysql::MySq
mod tests {
use super::{
agent_connect_timeout, connection_remote_endpoint, connection_url_for_endpoint, database_connection_config,
metadata_connection_config, mysql_metadata_fallback_url, oceanbase_mysql_setup_queries,
prestosql_jdbc_config_for_endpoint, redacted_connection_url_for_endpoint, redis_sentinel_transport_id,
redis_sentinel_transport_prefix, uses_bare_mysql_pool, uses_tcp_probe, validate_h2_database_path, AppState,
PoolKind, PRESTOSQL_JDBC_DRIVER_CLASS,
metadata_connection_config, mysql_metadata_fallback_url, oceanbase_mysql_query_timeout_sql,
oceanbase_mysql_setup_queries, prestosql_jdbc_config_for_endpoint, redacted_connection_url_for_endpoint,
redis_sentinel_transport_id, redis_sentinel_transport_prefix, uses_bare_mysql_pool, uses_tcp_probe,
validate_h2_database_path, AppState, PoolKind, PRESTOSQL_JDBC_DRIVER_CLASS,
};
use crate::agent_connection::{
agent_connect_params, mongo_legacy_error_with_auth_hint, mongo_uses_legacy_driver,
@ -3181,6 +3185,17 @@ mod tests {
assert_eq!(oceanbase_mysql_setup_queries(&config), vec!["SET ob_query_timeout = 30000000"]);
}
#[test]
fn oceanbase_mysql_query_timeout_sql_accepts_large_timeout() {
let mut config = mysql_config(Some("dbx"));
config.driver_profile = Some("oceanbase".to_string());
assert_eq!(
oceanbase_mysql_query_timeout_sql(&config, 300_000),
Some("SET ob_query_timeout = 300000000000".to_string())
);
}
#[test]
fn oceanbase_mysql_setup_queries_skip_disabled_timeout() {
let mut config = mysql_config(Some("dbx"));

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

@ -937,6 +937,30 @@ async fn configured_operation_budget_for_pool_key(state: &AppState, pool_key: &s
.unwrap_or_else(DbOperationBudget::with_defaults)
}
fn oceanbase_mysql_session_timeout_sql(config: Option<&ConnectionConfig>, timeout_secs: Option<u64>) -> Option<String> {
let config = config?;
let timeout_secs = timeout_secs.unwrap_or(config.query_timeout_secs);
crate::connection::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<u64>,
) -> 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)]
pub async fn do_execute(
state: &AppState,
@ -1075,6 +1099,7 @@ pub async fn do_execute(
});
});
}
apply_oceanbase_mysql_session_timeout(state, pool_key, &mut conn, options.timeout_secs).await?;
wait_for_query_opt(
cancel_token,
query_timeout,
@ -1764,6 +1789,8 @@ async fn execute_multi_mysql(
};
let mut results = Vec::with_capacity(statements.len());
apply_oceanbase_mysql_session_timeout(state, pool_key, &mut conn, options.timeout_secs).await?;
for stmt in statements {
if is_canceled(&cancel_token) {
results.push(error_query_result(canceled_error()));
@ -2096,7 +2123,7 @@ pub async fn execute_statements_in_transaction(
exec_tx_pg_inner(pool, statements, schema, start, operation_budget.clone(), cancel_context).await
}
Some(TxPath::Mysql(pool, _bare)) => {
exec_tx_mysql_inner(pool, statements, start, operation_budget.clone()).await
exec_tx_mysql_inner(state, &pool_key, pool, statements, start, operation_budget.clone()).await
}
Some(TxPath::Sqlite(pool)) => exec_tx_sqlite_inner(pool, statements, start).await,
Some(TxPath::Explicit) => {
@ -2221,12 +2248,15 @@ async fn exec_tx_pg_statements(
}
async fn exec_tx_mysql_inner(
state: &AppState,
pool_key: &str,
pool: mysql_async::Pool,
statements: &[String],
start: std::time::Instant,
budget: DbOperationBudget,
) -> Result<db::QueryResult, String> {
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?;
mysql_query_drop_with_timeout(
&mut conn,
"START TRANSACTION",
@ -2914,6 +2944,37 @@ mod tests {
assert!(!is_agent_execute_batch_unsupported("Agent RPC error (-1): unknown method: execute_query"));
}
#[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();