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

feat: fix connection cascade failure, add auto-recovery for all drivers

Browse Source
This commit is contained in:
dbin0123
2026-06-19 23:06:21 +08:00
committed by GitHub
parent c9ab2fb60a
commit 92d3f64eae
9 changed files with 590 additions and 269 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/dbx-core/src/connection.rs
+277 -21
View File
@@ -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<String>) = {
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;
crates/dbx-core/src/db/mongo_driver.rs
+12 -7
View File
@@ -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<Vec<String>, String> {
crates/dbx-core/src/db/mysql.rs
+3 -3
View File
@@ -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<mysql_async::Conn, String> {
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())
}
crates/dbx-core/src/db/sqlserver.rs
+6 -3
View File
@@ -580,9 +580,12 @@ pub async fn list_databases(client: &mut SqlServerClient) -> Result<Vec<Database
}
pub async fn test_connection(client: &mut SqlServerClient) -> 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<Vec<LinkedServerInfo>, String> {
crates/dbx-core/src/models/connection.rs
+1 -1
View File
@@ -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 {
crates/dbx-core/src/query.rs
+36 -11
View File
@@ -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();
crates/dbx-core/src/schema.rs
+245 -221
View File
@@ -1339,156 +1339,159 @@ pub async fn get_columns_core(
schema: &str,
table: &str,
) -> Result<Vec<db::ColumnInfo>, 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::<Vec<db::ColumnInfo>>(
"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::<Vec<db::ColumnInfo>>(
"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::<Vec<db::ColumnInfo>>(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::<Vec<db::ColumnInfo>>(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<db::ColumnInfo>) -> Vec<db::ColumnInfo> {
@@ -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<Vec<db::FunctionInfo>, 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<Vec<db::SequenceInfo>, 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<Vec<db::RuleInfo>, 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<Vec<db::OwnerInfo>, 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(
packages/mcp-server/src/index.ts
+1 -1
View File
@@ -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 {
src-tauri/src/lib.rs
+9 -1
View File
@@ -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::<AppState>() {
state.refresh_connections().await;
}
});
}
});
}