diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index d9fb76e58..dbd28f2ef 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -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 { - 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 { + 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 { + 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")); diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index 544258a2b..a1ee6d5c6 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -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) -> Option { + 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, +) -> 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 { 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();