diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index b09b91034d..317ce8a77b 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -911,6 +911,32 @@ impl AppState { } } } + PoolKind::Postgres(pool) => { + let pool = pool.clone(); + drop(connections); + let timeout = crate::db::connection_timeout(); + match tokio::time::timeout(timeout, pool.get()).await { + Ok(Ok(client)) => match tokio::time::timeout(timeout, client.simple_query("SELECT 1")).await { + Ok(Ok(_)) => false, + Ok(Err(err)) => { + log::warn!("PostgreSQL connection pool '{pool_key}' is stale: {err}"); + true + } + Err(_) => { + log::warn!("PostgreSQL connection pool '{pool_key}' is stale: health check timed out"); + true + } + }, + Ok(Err(err)) => { + log::warn!("PostgreSQL connection pool '{pool_key}' is stale: {err}"); + true + } + Err(_) => { + log::warn!("PostgreSQL connection pool '{pool_key}' is stale: get connection timed out"); + true + } + } + } PoolKind::SqlServer(client) => { let client = client.clone(); drop(connections); @@ -951,7 +977,95 @@ impl AppState { } } } - _ => false, + PoolKind::ClickHouse(client) => { + let client = client.clone(); + drop(connections); + let timeout = crate::db::connection_timeout(); + match db::clickhouse_driver::test_connection(&client, timeout).await { + Ok(()) => false, + Err(err) => { + log::warn!("ClickHouse connection pool '{pool_key}' is stale: {err}"); + true + } + } + } + PoolKind::Elasticsearch(client) => { + let mut client = client.clone(); + drop(connections); + let timeout = crate::db::connection_timeout(); + match db::elasticsearch_driver::test_connection(&mut client, timeout).await { + Ok(()) => false, + Err(err) => { + log::warn!("Elasticsearch connection pool '{pool_key}' is stale: {err}"); + true + } + } + } + PoolKind::VectorDb(client) => { + let client = client.clone(); + drop(connections); + let timeout = crate::db::connection_timeout(); + match db::vector_driver::test_connection(&client, timeout).await { + Ok(()) => false, + Err(err) => { + log::warn!("VectorDB connection pool '{pool_key}' is stale: {err}"); + true + } + } + } + PoolKind::InfluxDb(client) => { + let client = client.clone(); + drop(connections); + let timeout = crate::db::connection_timeout(); + match db::influxdb_driver::test_connection(&client, timeout).await { + Ok(()) => false, + Err(err) => { + log::warn!("InfluxDB connection pool '{pool_key}' is stale: {err}"); + true + } + } + } + PoolKind::Rqlite(client) => { + let client = client.clone(); + drop(connections); + let timeout = crate::db::connection_timeout(); + match db::rqlite_driver::test_connection(&client, timeout).await { + Ok(()) => false, + Err(err) => { + log::warn!("rqlite connection pool '{pool_key}' is stale: {err}"); + true + } + } + } + PoolKind::Turso(client) => { + let client = client.clone(); + drop(connections); + let timeout = crate::db::connection_timeout(); + match db::turso_driver::test_connection(&client, timeout).await { + Ok(()) => false, + Err(err) => { + log::warn!("Turso connection pool '{pool_key}' is stale: {err}"); + true + } + } + } + PoolKind::Agent(client) => { + let client = client.clone(); + drop(connections); + let mut agent = client.lock().await; + match agent.test_connection(serde_json::json!({})).await { + Ok(_) => false, + Err(err) => { + log::warn!("Agent connection pool '{pool_key}' is stale: {err}"); + true + } + } + } + PoolKind::Sqlite(_) + | PoolKind::DuckDb(_) + | PoolKind::ExternalTabular(_) + | PoolKind::ExternalDriver { .. } + | PoolKind::MessageQueue => false, } }; @@ -1170,16 +1284,24 @@ impl AppState { pub async fn refresh_connections(&self) { // Clone pool handles under a short-lived read lock, then release it // before performing I/O-heavy health checks to avoid blocking writers. - let checks: Vec<(String, PoolKind)> = { + // Redis pools are handled separately because RedisConnection cannot be cloned. + let (checks, redis_keys): (Vec<(String, PoolKind)>, Vec) = { let conns = self.connections.read().await; - conns - .iter() - .filter(|(_, pool)| matches!(pool, PoolKind::Mysql(..) | PoolKind::Postgres(..))) - .map(|(key, pool)| (key.clone(), clone_pool_kind(pool))) - .collect() + let mut checks = Vec::with_capacity(conns.len()); + let mut redis_keys = Vec::new(); + for (key, pool) in conns.iter() { + match pool { + PoolKind::Redis(_) => redis_keys.push(key.clone()), + _ => checks.push((key.clone(), clone_pool_kind(pool))), + } + } + (checks, redis_keys) }; let mut dead_keys = Vec::new(); + let timeout = crate::db::connection_timeout(); + + // Check cloned pools (async I/O, no lock held) for (key, pool) in &checks { let healthy = match pool { PoolKind::Mysql(p, _) => match db::mysql::get_conn_with_health_check(p).await { @@ -1189,26 +1311,127 @@ impl AppState { false } }, - PoolKind::Postgres(p) => match p.get().await { - Ok(client) => match client.simple_query("SELECT 1").await { - Ok(_) => true, - Err(e) => { + PoolKind::Postgres(p) => match tokio::time::timeout(timeout, p.get()).await { + Ok(Ok(client)) => match tokio::time::timeout(timeout, client.simple_query("SELECT 1")).await { + Ok(Ok(_)) => true, + Ok(Err(e)) => { log::warn!("PostgreSQL connection pool '{key}' is unhealthy: {e}"); false } + Err(_) => { + log::warn!("PostgreSQL connection pool '{key}' is unhealthy: health check timed out"); + false + } }, - Err(e) => { + Ok(Err(e)) => { log::warn!("PostgreSQL connection pool '{key}' is unhealthy: {e}"); false } + Err(_) => { + log::warn!("PostgreSQL connection pool '{key}' is unhealthy: get connection timed out"); + false + } }, - _ => true, + PoolKind::SqlServer(client) => { + let mut client = client.lock().await; + match db::sqlserver::test_connection(&mut client).await { + Ok(()) => true, + Err(e) => { + log::warn!("SQL Server connection pool '{key}' is unhealthy: {e}"); + false + } + } + } + PoolKind::MongoDb(client) => match db::mongo_driver::test_connection(client, timeout, None).await { + Ok(()) => true, + Err(e) => { + log::warn!("MongoDB connection pool '{key}' is unhealthy: {e}"); + false + } + }, + PoolKind::ClickHouse(client) => match db::clickhouse_driver::test_connection(client, timeout).await { + Ok(()) => true, + Err(e) => { + log::warn!("ClickHouse connection pool '{key}' is unhealthy: {e}"); + false + } + }, + PoolKind::Elasticsearch(client) => { + let mut client = client.clone(); + match db::elasticsearch_driver::test_connection(&mut client, timeout).await { + Ok(()) => true, + Err(e) => { + log::warn!("Elasticsearch connection pool '{key}' is unhealthy: {e}"); + false + } + } + } + PoolKind::VectorDb(client) => match db::vector_driver::test_connection(client, timeout).await { + Ok(()) => true, + Err(e) => { + log::warn!("VectorDB connection pool '{key}' is unhealthy: {e}"); + false + } + }, + PoolKind::InfluxDb(client) => match db::influxdb_driver::test_connection(client, timeout).await { + Ok(()) => true, + Err(e) => { + log::warn!("InfluxDB connection pool '{key}' is unhealthy: {e}"); + false + } + }, + PoolKind::Rqlite(client) => match db::rqlite_driver::test_connection(client, timeout).await { + Ok(()) => true, + Err(e) => { + log::warn!("rqlite connection pool '{key}' is unhealthy: {e}"); + false + } + }, + PoolKind::Turso(client) => match db::turso_driver::test_connection(client, timeout).await { + Ok(()) => true, + Err(e) => { + log::warn!("Turso connection pool '{key}' is unhealthy: {e}"); + false + } + }, + PoolKind::Agent(client) => { + let mut agent = client.lock().await; + match agent.test_connection(serde_json::json!({})).await { + Ok(_) => true, + Err(e) => { + log::warn!("Agent connection pool '{key}' is unhealthy: {e}"); + false + } + } + } + PoolKind::Sqlite(_) + | PoolKind::DuckDb(_) + | PoolKind::ExternalTabular(_) + | PoolKind::ExternalDriver { .. } + | PoolKind::MessageQueue => true, + PoolKind::Redis(_) => unreachable!("Redis handled separately"), }; if !healthy { dead_keys.push(key.clone()); } } + // Check Redis pools (read lock held briefly, no cloning needed) + { + let conns = self.connections.read().await; + for key in &redis_keys { + if let Some(PoolKind::Redis(redis)) = conns.get(key) { + match db::redis_driver::test_connection(redis).await { + Ok(()) => {} + Err(e) => { + log::warn!("Redis connection pool '{key}' is unhealthy: {e}"); + dead_keys.push(key.clone()); + } + } + } + } + } + // Remove dead pools if !dead_keys.is_empty() { self.stop_keepalive_tasks(&dead_keys).await; @@ -1449,7 +1672,26 @@ fn clone_pool_kind(pool: &PoolKind) -> PoolKind { match pool { PoolKind::Mysql(p, mode) => PoolKind::Mysql(p.clone(), *mode), PoolKind::Postgres(p) => PoolKind::Postgres(p.clone()), - other => panic!("clone_pool_kind not supported for {:?}", std::mem::discriminant(other)), + PoolKind::Sqlite(p) => PoolKind::Sqlite(p.clone()), + PoolKind::Rqlite(client) => PoolKind::Rqlite(client.clone()), + PoolKind::Turso(client) => PoolKind::Turso(client.clone()), + #[cfg(feature = "duckdb-bundled")] + PoolKind::DuckDb(con) => PoolKind::DuckDb(con.clone()), + #[cfg(not(feature = "duckdb-bundled"))] + PoolKind::DuckDb(con) => PoolKind::DuckDb(con.clone()), + PoolKind::MongoDb(client) => PoolKind::MongoDb(client.clone()), + PoolKind::ClickHouse(client) => PoolKind::ClickHouse(client.clone()), + PoolKind::SqlServer(client) => PoolKind::SqlServer(client.clone()), + PoolKind::Elasticsearch(client) => PoolKind::Elasticsearch(client.clone()), + PoolKind::VectorDb(client) => PoolKind::VectorDb(client.clone()), + PoolKind::InfluxDb(client) => PoolKind::InfluxDb(client.clone()), + PoolKind::Agent(client) => PoolKind::Agent(client.clone()), + PoolKind::ExternalTabular(ext) => PoolKind::ExternalTabular(ext.clone()), + PoolKind::ExternalDriver { driver_id, config, session } => { + PoolKind::ExternalDriver { driver_id: driver_id.clone(), config: config.clone(), session: session.clone() } + } + PoolKind::MessageQueue => PoolKind::MessageQueue, + PoolKind::Redis(_) => panic!("clone_pool_kind not supported for Redis — handled separately"), } } @@ -1462,19 +1704,33 @@ pub async fn close_pool_kind(pool: PoolKind) { PoolKind::Sqlite(_) => {} PoolKind::Rqlite(_) => {} PoolKind::Turso(_) => {} - PoolKind::Redis(_) => {} + PoolKind::Redis(conn) => { + drop(conn); + } #[cfg(feature = "duckdb-bundled")] PoolKind::DuckDb(con) => { crate::db::duckdb_driver::close_connection(con); } #[cfg(not(feature = "duckdb-bundled"))] PoolKind::DuckDb(_) => {} - PoolKind::MongoDb(_) => {} - PoolKind::ClickHouse(_) => {} - PoolKind::SqlServer(_) => {} - PoolKind::Elasticsearch(_) => {} - PoolKind::VectorDb(_) => {} - PoolKind::InfluxDb(_) => {} + PoolKind::MongoDb(client) => { + drop(client); + } + PoolKind::ClickHouse(client) => { + drop(client); + } + PoolKind::SqlServer(client) => { + drop(client); + } + PoolKind::Elasticsearch(client) => { + drop(client); + } + PoolKind::VectorDb(client) => { + drop(client); + } + PoolKind::InfluxDb(client) => { + drop(client); + } PoolKind::Agent(client) => { let mut client = client.lock().await; let _ = client.disconnect().await; diff --git a/crates/dbx-core/src/db/mongo_driver.rs b/crates/dbx-core/src/db/mongo_driver.rs index d6cbf8cf46..abb9382882 100644 --- a/crates/dbx-core/src/db/mongo_driver.rs +++ b/crates/dbx-core/src/db/mongo_driver.rs @@ -59,14 +59,19 @@ fn is_multi_host_mongo_uri(url: &str) -> bool { host_section.contains(',') } -pub async fn test_connection(client: &Client, _timeout: Duration, database: Option<&str>) -> Result<(), String> { +pub async fn test_connection(client: &Client, timeout: Duration, database: Option<&str>) -> Result<(), String> { let database = database.map(str::trim).filter(|value| !value.is_empty()).unwrap_or("admin"); - client - .database(database) - .run_command(doc! { "ping": 1 }) - .await - .map(|_| ()) - .map_err(|e| format!("MongoDB connection failed: {e}")) + let client = client.clone(); + let database = database.to_string(); + with_connection_timeout("MongoDB", timeout, async move { + client + .database(&database) + .run_command(doc! { "ping": 1 }) + .await + .map(|_| ()) + .map_err(|e| format!("MongoDB connection failed: {e}")) + }) + .await } pub async fn list_databases(client: &Client) -> Result, String> { diff --git a/crates/dbx-core/src/db/mysql.rs b/crates/dbx-core/src/db/mysql.rs index c3a255105c..d228f6db1f 100644 --- a/crates/dbx-core/src/db/mysql.rs +++ b/crates/dbx-core/src/db/mysql.rs @@ -1590,9 +1590,9 @@ fn skip_mysql_quoted(sql: &str, start: usize, quote: u8) -> usize { /// (e.g. after app was backgrounded), it tries again with a fresh connection. pub async fn get_conn_with_health_check(pool: &MySqlPool) -> Result { let mut conn = pool.get_conn().await.map_err(|e| e.to_string())?; - match conn.ping().await { - Ok(()) => Ok(conn), - Err(_) => { + match tokio::time::timeout(crate::db::connection_timeout(), conn.ping()).await { + Ok(Ok(())) => Ok(conn), + _ => { let _ = conn.disconnect().await; pool.get_conn().await.map_err(|e| e.to_string()) } diff --git a/crates/dbx-core/src/db/sqlserver.rs b/crates/dbx-core/src/db/sqlserver.rs index bec8d2de1b..79258da0f8 100644 --- a/crates/dbx-core/src/db/sqlserver.rs +++ b/crates/dbx-core/src/db/sqlserver.rs @@ -580,9 +580,12 @@ pub async fn list_databases(client: &mut SqlServerClient) -> Result Result<(), String> { - let stream = client.simple_query("SELECT 1").await.map_err(|e| e.to_string())?; - let _ = stream.into_first_result().await.map_err(|e| e.to_string())?; - Ok(()) + crate::db::with_connection_timeout("SQL Server", crate::db::connection_timeout(), async { + let stream = client.simple_query("SELECT 1").await.map_err(|e| e.to_string())?; + let _ = stream.into_first_result().await.map_err(|e| e.to_string())?; + Ok(()) + }) + .await } pub async fn list_linked_servers(client: &mut SqlServerClient) -> Result, String> { diff --git a/crates/dbx-core/src/models/connection.rs b/crates/dbx-core/src/models/connection.rs index 288f288290..33386592be 100644 --- a/crates/dbx-core/src/models/connection.rs +++ b/crates/dbx-core/src/models/connection.rs @@ -206,7 +206,7 @@ pub fn default_idle_timeout_secs() -> u64 { } pub fn default_keepalive_interval_secs() -> u64 { - 0 + 60 } fn default_proxy_port() -> u16 { diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index adb29198f9..9f6407f89b 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -915,13 +915,17 @@ pub async fn do_execute( let database = pool_key.split(':').nth(1).unwrap_or("default").to_string(); let max_rows = options.max_rows; drop(connections); - wait_for_query_opt( + let result = wait_for_query_opt( cancel_token, query_timeout, db::clickhouse_driver::execute_query_with_max_rows(&client, &database, sql, max_rows), ) .await - .map(|result| truncate_result_with_max_rows(result, max_rows)) + .map(|result| truncate_result_with_max_rows(result, max_rows)); + if matches!(result.as_ref(), Err(err) if should_discard_pool_after_error(pool_db_type, err)) { + state.remove_pool_by_key(pool_key).await; + } + result } PoolKind::SqlServer(client) => { let client = client.clone(); @@ -953,18 +957,31 @@ pub async fn do_execute( let sql = sql.to_string(); let max_rows = options.max_rows; drop(connections); - wait_for_query_opt(cancel_token, query_timeout, db::elasticsearch_driver::execute_rest_query(&client, &sql)) - .await - .map(|result| truncate_result_with_max_rows(result, max_rows)) + let result = wait_for_query_opt( + cancel_token, + query_timeout, + db::elasticsearch_driver::execute_rest_query(&client, &sql), + ) + .await + .map(|result| truncate_result_with_max_rows(result, max_rows)); + if matches!(result.as_ref(), Err(err) if should_discard_pool_after_error(pool_db_type, err)) { + state.remove_pool_by_key(pool_key).await; + } + result } PoolKind::VectorDb(client) => { let client = client.clone(); let sql = sql.to_string(); let max_rows = options.max_rows; drop(connections); - wait_for_query_opt(cancel_token, query_timeout, db::vector_driver::execute_rest_query(&client, &sql)) - .await - .map(|result| truncate_result_with_max_rows(result, max_rows)) + let result = + wait_for_query_opt(cancel_token, query_timeout, db::vector_driver::execute_rest_query(&client, &sql)) + .await + .map(|result| truncate_result_with_max_rows(result, max_rows)); + if matches!(result.as_ref(), Err(err) if should_discard_pool_after_error(pool_db_type, err)) { + state.remove_pool_by_key(pool_key).await; + } + result } PoolKind::Redis(_) => Err("Use Redis-specific commands".to_string()), PoolKind::MongoDb(_) => Err("Use MongoDB-specific commands".to_string()), @@ -974,9 +991,17 @@ pub async fn do_execute( let database = pool_key.split(':').nth(1).unwrap_or("default").to_string(); let max_rows = options.max_rows; drop(connections); - wait_for_query_opt(cancel_token, query_timeout, db::influxdb_driver::execute_query(&client, &database, sql)) - .await - .map(|result| truncate_result_with_max_rows(result, max_rows)) + let result = wait_for_query_opt( + cancel_token, + query_timeout, + db::influxdb_driver::execute_query(&client, &database, sql), + ) + .await + .map(|result| truncate_result_with_max_rows(result, max_rows)); + if matches!(result.as_ref(), Err(err) if should_discard_pool_after_error(pool_db_type, err)) { + state.remove_pool_by_key(pool_key).await; + } + result } PoolKind::Agent(client) => { let client = client.clone(); diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index 91621e9b50..0d0aa8502d 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -1339,156 +1339,159 @@ pub async fn get_columns_core( schema: &str, table: &str, ) -> Result, String> { - let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; - #[cfg(feature = "duckdb-bundled")] - let duckdb_attached_names = duckdb_attached_database_names(state, connection_id).await; - let db_config = connection_config(state, connection_id).await; + retry_metadata_connection(state, connection_id, Some(database), || async { + let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; + #[cfg(feature = "duckdb-bundled")] + let duckdb_attached_names = duckdb_attached_database_names(state, connection_id).await; + let db_config = connection_config(state, connection_id).await; - { - let connections = state.connections.read().await; - #[cfg(feature = "duckdb-bundled")] - if let Some(ext_pool) = extract_pool!(&connections, &pool_key, ExternalTabular) { - drop(connections); - let cache = ext_pool.cache.clone(); - let table = table.to_string(); - return tokio::task::spawn_blocking(move || { - let con = cache.lock().map_err(|e| e.to_string())?; - duckdb_query_columns(&con, &table) - }) - .await - .map_err(|e| e.to_string())?; - } - if let Some(PoolKind::ExternalDriver { config, session, .. }) = connections.get(&pool_key) { - let config = config.clone(); - let session = session.clone(); - drop(connections); - let columns = session - .invoke::>( - "getColumns", - serde_json::json!({ - "connection": config.as_ref(), - "database": database, - "schema": schema, - "table": table, - }), - ) - .await?; - return Ok(deduplicate_column_infos(columns)); - } - #[cfg(feature = "duckdb-bundled")] - if let Some(con) = extract_pool!(&connections, &pool_key, DuckDb) { - drop(connections); - let con = con.lock().map_err(|e| e.to_string())?; - return duckdb_query_columns_in_database_with_attached( - &con, - database, - schema, - table, - &duckdb_attached_names, - ); - } - if let Some(client) = extract_pool!(&connections, &pool_key, ClickHouse) { - drop(connections); - return db::clickhouse_driver::get_columns(&client, clickhouse_metadata_database(database, schema), table) + { + let connections = state.connections.read().await; + #[cfg(feature = "duckdb-bundled")] + if let Some(ext_pool) = extract_pool!(&connections, &pool_key, ExternalTabular) { + drop(connections); + let cache = ext_pool.cache.clone(); + let table = table.to_string(); + return tokio::task::spawn_blocking(move || { + let con = cache.lock().map_err(|e| e.to_string())?; + duckdb_query_columns(&con, &table) + }) .await - .map(deduplicate_column_infos); - } - if let Some(client) = extract_pool!(&connections, &pool_key, InfluxDb) { - drop(connections); - return db::influxdb_driver::get_columns(&client, database, table).await.map(deduplicate_column_infos); - } - if let Some(linked) = crate::sql_dialect::parse_sqlserver_linked_schema_ref(schema) { - if let Some(client) = extract_pool!(&connections, &pool_key, SqlServer) { + .map_err(|e| e.to_string())?; + } + if let Some(PoolKind::ExternalDriver { config, session, .. }) = connections.get(&pool_key) { + let config = config.clone(); + let session = session.clone(); + drop(connections); + let columns = session + .invoke::>( + "getColumns", + serde_json::json!({ + "connection": config.as_ref(), + "database": database, + "schema": schema, + "table": table, + }), + ) + .await?; + return Ok(deduplicate_column_infos(columns)); + } + #[cfg(feature = "duckdb-bundled")] + if let Some(con) = extract_pool!(&connections, &pool_key, DuckDb) { + drop(connections); + let con = con.lock().map_err(|e| e.to_string())?; + return duckdb_query_columns_in_database_with_attached( + &con, + database, + schema, + table, + &duckdb_attached_names, + ); + } + if let Some(client) = extract_pool!(&connections, &pool_key, ClickHouse) { + drop(connections); + return db::clickhouse_driver::get_columns(&client, clickhouse_metadata_database(database, schema), table) + .await + .map(deduplicate_column_infos); + } + if let Some(client) = extract_pool!(&connections, &pool_key, InfluxDb) { + drop(connections); + return db::influxdb_driver::get_columns(&client, database, table).await.map(deduplicate_column_infos); + } + if let Some(linked) = crate::sql_dialect::parse_sqlserver_linked_schema_ref(schema) { + if let Some(client) = extract_pool!(&connections, &pool_key, SqlServer) { + drop(connections); + let mut client = client.lock().await; + return db::sqlserver::get_linked_server_columns( + &mut client, + &linked.server, + &linked.catalog, + &linked.schema, + table, + ) + .await + .map(deduplicate_column_infos); + } + } + try_sqlserver!(connections, &pool_key, get_columns, schema, table); + if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { + let fallback_config = db_config.clone(); drop(connections); let mut client = client.lock().await; - return db::sqlserver::get_linked_server_columns( - &mut client, - &linked.server, - &linked.catalog, - &linked.schema, - table, - ) - .await - .map(deduplicate_column_infos); - } - } - try_sqlserver!(connections, &pool_key, get_columns, schema, table); - if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { - let fallback_config = db_config.clone(); - drop(connections); - let mut client = client.lock().await; - match client.get_columns::>(database, schema, table).await { - Ok(columns) if !columns.is_empty() => return Ok(deduplicate_column_infos(columns)), - Ok(columns) => { - if let Some(config) = fallback_config.as_ref() { - match native_postgres_metadata_pool(state, connection_id, database, config).await { - Ok(Some(pool)) => { + match client.get_columns::>(database, schema, table).await { + Ok(columns) if !columns.is_empty() => return Ok(deduplicate_column_infos(columns)), + Ok(columns) => { + if let Some(config) = fallback_config.as_ref() { + match native_postgres_metadata_pool(state, connection_id, database, config).await { + Ok(Some(pool)) => { + return db::postgres::get_columns(&pool, schema, table) + .await + .map(deduplicate_column_infos); + } + Ok(None) => return Ok(deduplicate_column_infos(columns)), + Err(error) => { + log::warn!( + "[schema][agent:get_columns:fallback-failed] connection_id={} database={} schema={} table={} error={}", + connection_id, + database, + schema, + table, + error + ); + } + } + } + return Ok(deduplicate_column_infos(columns)); + } + Err(agent_error) => { + if let Some(config) = fallback_config.as_ref() { + if let Some(pool) = + native_postgres_metadata_pool(state, connection_id, database, config).await? + { return db::postgres::get_columns(&pool, schema, table) .await - .map(deduplicate_column_infos); - } - Ok(None) => return Ok(deduplicate_column_infos(columns)), - Err(error) => { - log::warn!( - "[schema][agent:get_columns:fallback-failed] connection_id={} database={} schema={} table={} error={}", - connection_id, - database, - schema, - table, - error - ); + .map(deduplicate_column_infos) + .map_err(|fallback_error| { + format!( + "{agent_error}\n\nNative PostgreSQL metadata fallback failed: {fallback_error}" + ) + }); } } + return Err(agent_error); } - return Ok(deduplicate_column_infos(columns)); - } - Err(agent_error) => { - if let Some(config) = fallback_config.as_ref() { - if let Some(pool) = - native_postgres_metadata_pool(state, connection_id, database, config).await? - { - return db::postgres::get_columns(&pool, schema, table) - .await - .map(deduplicate_column_infos) - .map_err(|fallback_error| { - format!( - "{agent_error}\n\nNative PostgreSQL metadata fallback failed: {fallback_error}" - ) - }); - } - } - return Err(agent_error); } } } - } - let connections = state.connections.read().await; - let pool = connections.get(&pool_key).ok_or("Pool not found")?; + let connections = state.connections.read().await; + let pool = connections.get(&pool_key).ok_or("Pool not found")?; - match pool { - PoolKind::Mysql(p, _) if db_config.as_ref().is_some_and(is_manticoresearch_config) => { - let metadata_database = mysql_show_metadata_database_for_config(db_config.as_ref(), database); - db::manticoresearch::get_columns(p, metadata_database, table).await.map(deduplicate_column_infos) + match pool { + PoolKind::Mysql(p, _) if db_config.as_ref().is_some_and(is_manticoresearch_config) => { + let metadata_database = mysql_show_metadata_database_for_config(db_config.as_ref(), database); + db::manticoresearch::get_columns(p, metadata_database, table).await.map(deduplicate_column_infos) + } + PoolKind::Mysql(p, _) if db_config.as_ref().is_some_and(is_doris_family_config) => { + let metadata_database = mysql_show_metadata_database_for_config(db_config.as_ref(), database); + db::mysql::get_columns_show(p, metadata_database, table).await.map(deduplicate_column_infos) + } + PoolKind::Mysql(p, mode) => { + dispatch_mysql!(p, mode, db::mysql::get_columns, db::ob_oracle::get_columns, database, table) + .map(deduplicate_column_infos) + } + PoolKind::Postgres(p) if db_config.as_ref().is_some_and(is_questdb_config) => { + db::questdb::get_columns(p, schema, table).await.map(deduplicate_column_infos) + } + PoolKind::Postgres(p) => db::postgres::get_columns(p, schema, table).await.map(deduplicate_column_infos), + PoolKind::Sqlite(p) => db::sqlite::get_columns(p, schema, table).await.map(deduplicate_column_infos), + PoolKind::Rqlite(client) => { + db::rqlite_driver::get_columns(client, schema, table).await.map(deduplicate_column_infos) + } + _ => Ok(vec![]), } - PoolKind::Mysql(p, _) if db_config.as_ref().is_some_and(is_doris_family_config) => { - let metadata_database = mysql_show_metadata_database_for_config(db_config.as_ref(), database); - db::mysql::get_columns_show(p, metadata_database, table).await.map(deduplicate_column_infos) - } - PoolKind::Mysql(p, mode) => { - dispatch_mysql!(p, mode, db::mysql::get_columns, db::ob_oracle::get_columns, database, table) - .map(deduplicate_column_infos) - } - PoolKind::Postgres(p) if db_config.as_ref().is_some_and(is_questdb_config) => { - db::questdb::get_columns(p, schema, table).await.map(deduplicate_column_infos) - } - PoolKind::Postgres(p) => db::postgres::get_columns(p, schema, table).await.map(deduplicate_column_infos), - PoolKind::Sqlite(p) => db::sqlite::get_columns(p, schema, table).await.map(deduplicate_column_infos), - PoolKind::Rqlite(client) => { - db::rqlite_driver::get_columns(client, schema, table).await.map(deduplicate_column_infos) - } - _ => Ok(vec![]), - } + }) + .await } fn deduplicate_column_infos(columns: Vec) -> Vec { @@ -1544,33 +1547,36 @@ pub async fn list_indexes_core( if crate::sql_dialect::parse_sqlserver_linked_schema_ref(schema).is_some() { return Ok(vec![]); } - let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; - let db_config = connection_config(state, connection_id).await; + retry_metadata_connection(state, connection_id, Some(database), || async { + let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; + let db_config = connection_config(state, connection_id).await; + + { + let connections = state.connections.read().await; + try_sqlserver!(connections, &pool_key, list_indexes, schema, table); + try_agent!(connections, &pool_key, list_indexes, database, schema, table); + } - { let connections = state.connections.read().await; - try_sqlserver!(connections, &pool_key, list_indexes, schema, table); - try_agent!(connections, &pool_key, list_indexes, database, schema, table); - } + let pool = connections.get(&pool_key).ok_or("Pool not found")?; - let connections = state.connections.read().await; - let pool = connections.get(&pool_key).ok_or("Pool not found")?; - - match pool { - PoolKind::Mysql(p, mode) => { - if db_config.as_ref().is_some_and(is_manticoresearch_config) { - return db::manticoresearch::list_indexes(p, table).await; + match pool { + PoolKind::Mysql(p, mode) => { + if db_config.as_ref().is_some_and(is_manticoresearch_config) { + return db::manticoresearch::list_indexes(p, table).await; + } + dispatch_mysql!(p, mode, db::mysql::list_indexes, db::ob_oracle::list_indexes, schema, table) } - dispatch_mysql!(p, mode, db::mysql::list_indexes, db::ob_oracle::list_indexes, schema, table) + PoolKind::Postgres(p) if db_config.as_ref().is_some_and(is_questdb_config) => { + db::questdb::list_indexes(p, schema, table).await + } + PoolKind::Postgres(p) => db::postgres::list_indexes(p, schema, table).await, + PoolKind::Sqlite(p) => db::sqlite::list_indexes(p, schema, table).await, + PoolKind::Rqlite(client) => db::rqlite_driver::list_indexes(client, schema, table).await, + _ => Ok(vec![]), } - PoolKind::Postgres(p) if db_config.as_ref().is_some_and(is_questdb_config) => { - db::questdb::list_indexes(p, schema, table).await - } - PoolKind::Postgres(p) => db::postgres::list_indexes(p, schema, table).await, - PoolKind::Sqlite(p) => db::sqlite::list_indexes(p, schema, table).await, - PoolKind::Rqlite(client) => db::rqlite_driver::list_indexes(client, schema, table).await, - _ => Ok(vec![]), - } + }) + .await } pub async fn list_foreign_keys_core( @@ -1583,26 +1589,29 @@ pub async fn list_foreign_keys_core( if crate::sql_dialect::parse_sqlserver_linked_schema_ref(schema).is_some() { return Ok(vec![]); } - let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; + retry_metadata_connection(state, connection_id, Some(database), || async { + let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; - { - let connections = state.connections.read().await; - try_sqlserver!(connections, &pool_key, list_foreign_keys, schema, table); - try_agent!(connections, &pool_key, list_foreign_keys, database, schema, table); - } - - let connections = state.connections.read().await; - let pool = connections.get(&pool_key).ok_or("Pool not found")?; - - match pool { - PoolKind::Mysql(p, mode) => { - dispatch_mysql!(p, mode, db::mysql::list_foreign_keys, db::ob_oracle::list_foreign_keys, schema, table) + { + let connections = state.connections.read().await; + try_sqlserver!(connections, &pool_key, list_foreign_keys, schema, table); + try_agent!(connections, &pool_key, list_foreign_keys, database, schema, table); } - PoolKind::Postgres(p) => db::postgres::list_foreign_keys(p, schema, table).await, - PoolKind::Sqlite(p) => db::sqlite::list_foreign_keys(p, schema, table).await, - PoolKind::Rqlite(client) => db::rqlite_driver::list_foreign_keys(client, schema, table).await, - _ => Ok(vec![]), - } + + let connections = state.connections.read().await; + let pool = connections.get(&pool_key).ok_or("Pool not found")?; + + match pool { + PoolKind::Mysql(p, mode) => { + dispatch_mysql!(p, mode, db::mysql::list_foreign_keys, db::ob_oracle::list_foreign_keys, schema, table) + } + PoolKind::Postgres(p) => db::postgres::list_foreign_keys(p, schema, table).await, + PoolKind::Sqlite(p) => db::sqlite::list_foreign_keys(p, schema, table).await, + PoolKind::Rqlite(client) => db::rqlite_driver::list_foreign_keys(client, schema, table).await, + _ => Ok(vec![]), + } + }) + .await } pub async fn list_triggers_core( @@ -1615,26 +1624,29 @@ pub async fn list_triggers_core( if crate::sql_dialect::parse_sqlserver_linked_schema_ref(schema).is_some() { return Ok(vec![]); } - let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; + retry_metadata_connection(state, connection_id, Some(database), || async { + let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; - { - let connections = state.connections.read().await; - try_sqlserver!(connections, &pool_key, list_triggers, schema, table); - try_agent!(connections, &pool_key, list_triggers, database, schema, table); - } - - let connections = state.connections.read().await; - let pool = connections.get(&pool_key).ok_or("Pool not found")?; - - match pool { - PoolKind::Mysql(p, mode) => { - dispatch_mysql!(p, mode, db::mysql::list_triggers, db::ob_oracle::list_triggers, schema, table) + { + let connections = state.connections.read().await; + try_sqlserver!(connections, &pool_key, list_triggers, schema, table); + try_agent!(connections, &pool_key, list_triggers, database, schema, table); } - PoolKind::Postgres(p) => db::postgres::list_triggers(p, schema, table).await, - PoolKind::Sqlite(p) => db::sqlite::list_triggers(p, schema, table).await, - PoolKind::Rqlite(client) => db::rqlite_driver::list_triggers(client, schema, table).await, - _ => Ok(vec![]), - } + + let connections = state.connections.read().await; + let pool = connections.get(&pool_key).ok_or("Pool not found")?; + + match pool { + PoolKind::Mysql(p, mode) => { + dispatch_mysql!(p, mode, db::mysql::list_triggers, db::ob_oracle::list_triggers, schema, table) + } + PoolKind::Postgres(p) => db::postgres::list_triggers(p, schema, table).await, + PoolKind::Sqlite(p) => db::sqlite::list_triggers(p, schema, table).await, + PoolKind::Rqlite(client) => db::rqlite_driver::list_triggers(client, schema, table).await, + _ => Ok(vec![]), + } + }) + .await } pub async fn list_functions_core( @@ -1643,14 +1655,17 @@ pub async fn list_functions_core( database: &str, schema: &str, ) -> Result, String> { - let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; - let connections = state.connections.read().await; - let pool = connections.get(&pool_key).ok_or("Pool not found")?; + retry_metadata_connection(state, connection_id, Some(database), || async { + let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; + let connections = state.connections.read().await; + let pool = connections.get(&pool_key).ok_or("Pool not found")?; - match pool { - PoolKind::Postgres(p) => db::postgres::list_functions(p, schema).await, - _ => Ok(vec![]), - } + match pool { + PoolKind::Postgres(p) => db::postgres::list_functions(p, schema).await, + _ => Ok(vec![]), + } + }) + .await } pub async fn list_sequences_core( @@ -1660,14 +1675,17 @@ pub async fn list_sequences_core( schema: &str, with_last_values: bool, ) -> Result, String> { - let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; - let connections = state.connections.read().await; - let pool = connections.get(&pool_key).ok_or("Pool not found")?; + retry_metadata_connection(state, connection_id, Some(database), || async { + let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; + let connections = state.connections.read().await; + let pool = connections.get(&pool_key).ok_or("Pool not found")?; - match pool { - PoolKind::Postgres(p) => db::postgres::list_sequences(p, schema, with_last_values).await, - _ => Ok(vec![]), - } + match pool { + PoolKind::Postgres(p) => db::postgres::list_sequences(p, schema, with_last_values).await, + _ => Ok(vec![]), + } + }) + .await } pub async fn list_rules_core( @@ -1676,14 +1694,17 @@ pub async fn list_rules_core( database: &str, schema: &str, ) -> Result, String> { - let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; - let connections = state.connections.read().await; - let pool = connections.get(&pool_key).ok_or("Pool not found")?; + retry_metadata_connection(state, connection_id, Some(database), || async { + let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; + let connections = state.connections.read().await; + let pool = connections.get(&pool_key).ok_or("Pool not found")?; - match pool { - PoolKind::Postgres(p) => db::postgres::list_rules(p, schema).await, - _ => Ok(vec![]), - } + match pool { + PoolKind::Postgres(p) => db::postgres::list_rules(p, schema).await, + _ => Ok(vec![]), + } + }) + .await } pub async fn list_owners_core( @@ -1692,14 +1713,17 @@ pub async fn list_owners_core( database: &str, schema: &str, ) -> Result, String> { - let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; - let connections = state.connections.read().await; - let pool = connections.get(&pool_key).ok_or("Pool not found")?; + retry_metadata_connection(state, connection_id, Some(database), || async { + let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; + let connections = state.connections.read().await; + let pool = connections.get(&pool_key).ok_or("Pool not found")?; - match pool { - PoolKind::Postgres(p) => db::postgres::list_owners(p, schema).await, - _ => Ok(vec![]), - } + match pool { + PoolKind::Postgres(p) => db::postgres::list_owners(p, schema).await, + _ => Ok(vec![]), + } + }) + .await } pub async fn get_table_ddl_core( diff --git a/packages/mcp-server/src/index.ts b/packages/mcp-server/src/index.ts index fa8538f4e5..24dad69d13 100644 --- a/packages/mcp-server/src/index.ts +++ b/packages/mcp-server/src/index.ts @@ -46,7 +46,7 @@ function formatQueryToolResult(result: QueryResult, title?: string) { } export const DBX_CONNECTION_TYPE_DESCRIPTION = - "Database type: postgres, mysql, sqlite, rqlite, redis, duckdb, clickhouse, sqlserver, mongodb, oracle, elasticsearch, etcd, doris, starrocks, manticoresearch, redshift, dameng, kingbase, highgo, vastbase, goldendb, databend, gaussdb, kwdb, yashandb, databricks, saphana, teradata, vertica, firebird, exasol, opengauss, oceanbase-oracle, questdb, gbase, h2, snowflake, trino, hive, db2, informix, influxdb, iris, neo4j, cassandra, bigquery, kylin, sundb, tdengine, iotdb, xugu, jdbc, access, mq"; + "Database type: postgres, mysql, sqlite, rqlite, redis, duckdb, clickhouse, sqlserver, mongodb, oracle, elasticsearch, etcd, doris, starrocks, manticoresearch, milvus, qdrant, redshift, dameng, kingbase, highgo, vastbase, goldendb, databend, gaussdb, kwdb, yashandb, databricks, saphana, teradata, vertica, firebird, exasol, opengauss, oceanbase-oracle, questdb, gbase, h2, snowflake, trino, hive, db2, informix, influxdb, iris, neo4j, cassandra, bigquery, kylin, sundb, tdengine, iotdb, xugu, jdbc, access, mq"; const FILE_CAPABLE_CONNECTION_TYPES = new Set(["sqlite", "duckdb", "access", "h2"]); interface McpScope { diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 061704b4c0..c5847cd20f 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -8,7 +8,6 @@ use commands::connection::AppState; use dbx_core::storage::{DesktopIconTheme, DesktopSettings, Storage}; use std::sync::Arc; use std::time::Instant; -#[cfg(target_os = "macos")] use tauri::RunEvent; use tauri::{ menu::MenuBuilder, @@ -700,5 +699,14 @@ pub fn run() { } }); } + + if let RunEvent::Resumed = &event { + let app_handle = app_handle.clone(); + tauri::async_runtime::spawn(async move { + if let Some(state) = app_handle.try_state::() { + state.refresh_connections().await; + } + }); + } }); }