github/dbx
4
0
Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-10-19 16:16:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(mysql): honor connection timeout when listing databases

Browse Source 
Closes #4244
This commit is contained in:
t8y2
2026-08-14 15:46:46 +00:00
parent 2d526db91b
commit 3a637de44c
2 changed files with 95 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/dbx-core/src/db/mysql.rs
+59 -14
View File
@@ -9,6 +9,7 @@ use sqlparser::dialect::MySqlDialect;
use sqlparser::parser::Parser;
use std::borrow::Cow;
use std::collections::{HashMap, HashSet};
use std::future::Future;
use std::path::PathBuf;
use std::sync::atomic::AtomicBool;
use std::time::Duration;
@@ -2135,12 +2136,16 @@ const DATABASE_LIST_QUERY_PLAN: [(&str, bool); 2] =
[(SHOW_DATABASES_SQL, true), (INFORMATION_SCHEMA_DATABASES_SQL, false)];
pub async fn list_databases(pool: &MySqlPool) -> Result<Vec<DatabaseInfo>, String> {
list_databases_with_timeout(pool, super::connection_timeout()).await
}
pub async fn list_databases_with_timeout(pool: &MySqlPool, timeout: Duration) -> Result<Vec<DatabaseInfo>, String> {
let [(primary_sql, primary_catalogless), (fallback_sql, fallback_catalogless)] = DATABASE_LIST_QUERY_PLAN;
match list_databases_with_query(pool, primary_sql, primary_catalogless).await {
match list_databases_with_query(pool, primary_sql, primary_catalogless, timeout).await {
Ok(databases) => Ok(databases),
Err(err) => {
log::debug!("Falling back to information_schema.SCHEMATA after SHOW DATABASES failed: {err}");
list_databases_with_query(pool, fallback_sql, fallback_catalogless).await
list_databases_with_query(pool, fallback_sql, fallback_catalogless, timeout).await
}
}
}
@@ -2183,15 +2188,23 @@ async fn list_databases_with_query(
pool: &MySqlPool,
sql: &str,
include_catalogless_when_blank: bool,
timeout: Duration,
) -> Result<Vec<DatabaseInfo>, String> {
let mut conn = get_conn_with_timeout(pool, super::connection_timeout()).await?;
let mut conn = get_conn_with_timeout(pool, timeout).await?;
let result = conn.query_iter(sql).await.map_err(|e| e.to_string())?;
let rows: Vec<mysql_async::Row> = result.collect_and_drop().await.map_err(|e| e.to_string())?;
Ok(database_infos_from_names(rows.iter().map(|row| get_str(row, 0)), include_catalogless_when_blank))
}
pub async fn list_databases_show(pool: &MySqlPool) -> Result<Vec<DatabaseInfo>, String> {
list_databases_with_query(pool, SHOW_DATABASES_SQL, true).await
list_databases_show_with_timeout(pool, super::connection_timeout()).await
}
pub async fn list_databases_show_with_timeout(
pool: &MySqlPool,
timeout: Duration,
) -> Result<Vec<DatabaseInfo>, String> {
list_databases_with_query(pool, SHOW_DATABASES_SQL, true, timeout).await
}
pub(super) fn database_infos_from_names(
@@ -3876,12 +3889,8 @@ async fn get_conn_with_timeout_and_cancel(
timeout: Duration,
cancel_token: Option<&CancellationToken>,
) -> Result<mysql_async::Conn, String> {
let get_future = async {
tokio::time::timeout(timeout, pool.get_conn())
.await
.map_err(|_| "MySQL get connection timed out".to_string())?
.map_err(|e| e.to_string())
};
let get_future =
connection_result_with_timeout(timeout, async { pool.get_conn().await.map_err(|error| error.to_string()) });
match cancel_token {
Some(token) => {
@@ -3896,10 +3905,14 @@ async fn get_conn_with_timeout_and_cancel(
}
pub async fn get_conn_with_timeout(pool: &MySqlPool, timeout: Duration) -> Result<mysql_async::Conn, String> {
tokio::time::timeout(timeout, pool.get_conn())
.await
.map_err(|_| "MySQL get connection timed out".to_string())?
.map_err(|e| e.to_string())
connection_result_with_timeout(timeout, async { pool.get_conn().await.map_err(|error| error.to_string()) }).await
}
async fn connection_result_with_timeout<T, F>(timeout: Duration, future: F) -> Result<T, String>
where
F: Future<Output = Result<T, String>>,
{
tokio::time::timeout(timeout, future).await.map_err(|_| "MySQL get connection timed out".to_string())?
}
async fn ping_conn_with_timeout_and_cancel(
@@ -5302,6 +5315,38 @@ mod tests {
assert_eq!(first_column_value::<String>(&row), Some(String::new()));
}
#[tokio::test]
async fn mysql_database_list_timeout_controls_checkout_deadline() {
let exact = connection_result_with_timeout(Duration::from_millis(100), async {
tokio::time::sleep(Duration::from_millis(60)).await;
Ok::<_, String>("connection")
});
let adjacent = connection_result_with_timeout(Duration::from_millis(50), async {
tokio::time::sleep(Duration::from_millis(10)).await;
Ok::<_, String>("connection")
});
let configured_five = connection_result_with_timeout(Duration::from_millis(50), async {
tokio::time::sleep(Duration::from_millis(60)).await;
Ok::<_, String>("connection")
});
let (exact, adjacent, configured_five) = tokio::join!(exact, adjacent, configured_five);
assert_eq!(exact, Ok("connection"));
assert_eq!(adjacent, Ok("connection"));
assert_eq!(configured_five, Err("MySQL get connection timed out".to_string()));
}
#[tokio::test]
async fn mysql_database_list_timeout_preserves_immediate_errors_and_fallback_plan() {
let error =
connection_result_with_timeout(Duration::from_secs(10), async { Err::<(), _>("driver error".to_string()) })
.await;
assert_eq!(error, Err("driver error".to_string()));
assert_eq!(DATABASE_LIST_QUERY_PLAN, [(SHOW_DATABASES_SQL, true), (INFORMATION_SCHEMA_DATABASES_SQL, false)]);
}
#[test]
fn mysql_first_column_value_handles_null_and_conversion_failures() {
let null_row = mysql_test_row(vec![Value::NULL, Value::Bytes(b"ignored".to_vec())]);
crates/dbx-core/src/schema.rs
+36 -9
View File
@@ -170,6 +170,12 @@ fn agent_metadata_timeout(config: Option<&ConnectionConfig>) -> Option<Duration>
}
}
fn mysql_database_list_timeout(config: Option<&ConnectionConfig>) -> Duration {
config
.map(|config| Duration::from_secs(config.effective_connect_timeout_secs()))
.unwrap_or_else(db::connection_timeout)
}
pub async fn list_databases_core(state: &AppState, connection_id: &str) -> Result<Vec<db::DatabaseInfo>, String> {
retry_metadata_connection(state, connection_id, None, || list_databases_once(state, connection_id)).await
}
@@ -611,6 +617,7 @@ async fn list_databases_once(state: &AppState, connection_id: &str) -> Result<Ve
}
let db_config = connection_config(state, connection_id).await;
let mysql_database_list_timeout = mysql_database_list_timeout(db_config.as_ref());
let connections = state.connections.read().await;
let pool = connections.get(connection_id).ok_or("Connection not found")?;
@@ -621,11 +628,12 @@ async fn list_databases_once(state: &AppState, connection_id: &str) -> Result<Ve
db::dolt::list_databases(p).await
}
PoolKind::Mysql(p, _) if db_config.as_ref().is_some_and(db::mysql_compatible::uses_show_metadata) => {
db::mysql::list_databases_show(p)
db::mysql::list_databases_show_with_timeout(p, mysql_database_list_timeout)
.await
.map(|databases| filter_mysql_system_databases_for_config(databases, db_config.as_ref()))
}
PoolKind::Mysql(p, mode) => dispatch_mysql!(p, mode, db::mysql::list_databases, db::ob_oracle::list_databases),
PoolKind::Mysql(p, mode) if *mode == MysqlMode::OceanBaseOracle => db::ob_oracle::list_databases(p).await,
PoolKind::Mysql(p, _) => db::mysql::list_databases_with_timeout(p, mysql_database_list_timeout).await,
PoolKind::Postgres(p) => db::postgres::list_databases(p).await,
PoolKind::Sqlite(p) => db::sqlite::list_databases(p).await,
PoolKind::Rqlite(client) => db::rqlite_driver::list_databases(client).await,
@@ -2671,13 +2679,14 @@ mod tests {
filter_mysql_system_databases_for_config, filter_object_infos, filter_table_infos, filter_visible_schema_names,
gaussdb_m_view_object_source_sql, gbase8a_object_statistics_sql, is_agent_postgres_metadata_fallback_config,
is_mysql_external_driver_config, is_retryable_metadata_error, metadata_error_action,
metadata_name_or_comment_matches, mysql_external_driver_ddl_from_query_result, mysql_external_driver_ddl_sql,
mysql_object_source_ddl_column_index, mysql_object_source_sql, mysql_table_list_source_for_config,
mysql_table_metadata_catalog, normalize_information_schema_table_type, oracle_columns_from_query_result,
oracle_columns_sql, oracle_object_statistics_dba_segments_sql, oracle_object_statistics_from_query_result,
oracle_object_statistics_rows_only_sql, oracle_object_statistics_sql,
oracle_object_statistics_user_segments_sql, oracle_table_comment_from_query_result, oracle_table_comment_sql,
oracle_table_comments_sql, presto_like_columns_from_query_result, presto_like_information_schema_columns_sql,
metadata_name_or_comment_matches, mysql_database_list_timeout, mysql_external_driver_ddl_from_query_result,
mysql_external_driver_ddl_sql, mysql_object_source_ddl_column_index, mysql_object_source_sql,
mysql_table_list_source_for_config, mysql_table_metadata_catalog, normalize_information_schema_table_type,
oracle_columns_from_query_result, oracle_columns_sql, oracle_object_statistics_dba_segments_sql,
oracle_object_statistics_from_query_result, oracle_object_statistics_rows_only_sql,
oracle_object_statistics_sql, oracle_object_statistics_user_segments_sql,
oracle_table_comment_from_query_result, oracle_table_comment_sql, oracle_table_comments_sql,
presto_like_columns_from_query_result, presto_like_information_schema_columns_sql,
presto_like_information_schema_tables_sql, presto_like_tables_from_query_result, replace_metadata_runtime,
should_query_oracle_columns_via_sql_first, table_comments_from_query_result, table_name_filter_matches,
tdengine_table_comment_like_pattern, tdengine_table_comment_sql, tdengine_table_comments_sql,
@@ -2844,6 +2853,24 @@ mod tests {
}
}
#[test]
fn mysql_database_list_timeout_uses_configured_and_effective_bounds() {
let mut config = test_connection_config(DatabaseType::Mysql);
config.connect_timeout_secs = 10;
assert_eq!(mysql_database_list_timeout(Some(&config)), Duration::from_secs(10));
config.connect_timeout_secs = 0;
assert_eq!(
mysql_database_list_timeout(Some(&config)),
Duration::from_secs(crate::models::connection::default_connect_timeout_secs())
);
config.connect_timeout_secs = 500;
assert_eq!(mysql_database_list_timeout(Some(&config)), Duration::from_secs(300));
assert_eq!(mysql_database_list_timeout(None), db::connection_timeout());
}
#[tokio::test]
async fn turso_schema_dispatch_lists_databases_and_tables() {
let (base_url, server) = spawn_turso_table_server().await;