From f895bb4de5cd898bf47503f7d5b811fa22db172d Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Tue, 26 May 2026 11:51:58 +0800 Subject: [PATCH] fix: scope agent metadata by database --- crates/dbx-core/src/connection.rs | 97 ++++++++++++++++---------- crates/dbx-core/src/db/agent_driver.rs | 76 +++++++++++++------- crates/dbx-core/src/query.rs | 41 ++++++++--- crates/dbx-core/src/schema.rs | 16 ++--- crates/dbx-core/src/transfer.rs | 43 +++++++++++- 5 files changed, 192 insertions(+), 81 deletions(-) diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index dd1bb7815..c2b866cce 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -148,18 +148,10 @@ impl AppState { ) -> Result { let db_type = { let configs = self.configs.read().await; - configs.get(connection_id).map(|c| c.db_type.clone()) + configs.get(connection_id).map(|c| c.db_type) }; - let is_single_conn = db_type.as_ref().is_some_and(database_capabilities::is_single_connection_pool); - let base_pool_key = if is_single_conn { - connection_id.to_string() - } else { - match database { - Some(db) => format!("{connection_id}:{db}"), - None => connection_id.to_string(), - } - }; + let base_pool_key = base_pool_key_for(db_type, connection_id, database, false); let pool_key = session_scoped_pool_key(base_pool_key, client_session_id); let conns = self.connections.read().await; @@ -454,24 +446,11 @@ impl AppState { database: Option<&str>, client_session_id: Option<&str>, ) -> Result { - let is_single_conn = { + let db_type = { let configs = self.configs.read().await; - configs - .get(connection_id) - .map(|c| { - database_capabilities::is_single_connection_pool(&c.db_type) - || c.db_type == DatabaseType::Elasticsearch - }) - .unwrap_or(false) - }; - let base_pool_key = if is_single_conn { - connection_id.to_string() - } else { - match database { - Some(db) => format!("{connection_id}:{db}"), - None => connection_id.to_string(), - } + configs.get(connection_id).map(|c| c.db_type) }; + let base_pool_key = base_pool_key_for(db_type, connection_id, database, true); let pool_key = session_scoped_pool_key(base_pool_key, client_session_id); if self.uses_forwarded_transport(connection_id).await { self.remove_connection_pools(connection_id).await; @@ -494,17 +473,9 @@ impl AppState { }; let db_type = { let configs = self.configs.read().await; - configs.get(connection_id).map(|c| c.db_type.clone()) - }; - let is_single_conn = db_type.as_ref().is_some_and(database_capabilities::is_single_connection_pool); - let base_pool_key = if is_single_conn { - connection_id.to_string() - } else { - match database { - Some(db) => format!("{connection_id}:{db}"), - None => connection_id.to_string(), - } + configs.get(connection_id).map(|c| c.db_type) }; + let base_pool_key = base_pool_key_for(db_type, connection_id, database, false); let pool_key = session_scoped_pool_key(base_pool_key, Some(&session)); Ok(self.connections.write().await.remove(&pool_key).is_some()) } @@ -577,6 +548,28 @@ fn session_scoped_pool_key(base_pool_key: String, client_session_id: Option<&str .unwrap_or(base_pool_key) } +fn base_pool_key_for( + db_type: Option, + connection_id: &str, + database: Option<&str>, + include_elasticsearch_single_pool: bool, +) -> String { + let is_single_connection_pool = db_type.as_ref().is_some_and(|db_type| { + let is_single = database_capabilities::is_single_connection_pool(db_type) + || (include_elasticsearch_single_pool && *db_type == DatabaseType::Elasticsearch); + is_single && !database_capabilities::is_agent_type(db_type) + }); + + if is_single_connection_pool { + connection_id.to_string() + } else { + match database.map(str::trim).filter(|db| !db.is_empty()) { + Some(db) => format!("{connection_id}:{db}"), + None => connection_id.to_string(), + } + } +} + fn default_plugin_dir() -> PathBuf { default_dbx_dir().join("plugins") } @@ -1199,6 +1192,38 @@ mod tests { assert_eq!(scoped.database.as_deref(), Some("ORCL")); } + #[test] + fn agent_single_connection_types_keep_database_scoped_pool_keys() { + assert_eq!( + super::base_pool_key_for(Some(DatabaseType::Kingbase), "kingbase-conn", Some("app1"), false), + "kingbase-conn:app1" + ); + assert_eq!( + super::base_pool_key_for(Some(DatabaseType::Oracle), "oracle-conn", Some("ORCLPDB1"), false), + "oracle-conn:ORCLPDB1" + ); + assert_eq!( + super::base_pool_key_for(Some(DatabaseType::MongoDb), "mongo-conn", Some("shop"), false), + "mongo-conn:shop" + ); + } + + #[test] + fn non_agent_single_connection_types_still_share_pool_keys() { + assert_eq!( + super::base_pool_key_for(Some(DatabaseType::Sqlite), "sqlite-conn", Some("main"), false), + "sqlite-conn" + ); + assert_eq!( + super::base_pool_key_for(Some(DatabaseType::DuckDb), "duckdb-conn", Some("analytics"), false), + "duckdb-conn" + ); + assert_eq!( + super::base_pool_key_for(Some(DatabaseType::Jdbc), "jdbc-conn", Some("analytics"), false), + "jdbc-conn" + ); + } + #[test] fn mysql_hostname_connections_skip_tcp_probe() { let mut config = mysql_config(Some("app")); diff --git a/crates/dbx-core/src/db/agent_driver.rs b/crates/dbx-core/src/db/agent_driver.rs index 67bbd17cd..bf0398283 100644 --- a/crates/dbx-core/src/db/agent_driver.rs +++ b/crates/dbx-core/src/db/agent_driver.rs @@ -382,61 +382,76 @@ impl AgentDriverClient { self.call_method(AgentMethod::ListSchemas, serde_json::json!({ "database": database })).await } - pub async fn list_tables(&mut self, schema: &str) -> Result { - self.call_method(AgentMethod::ListTables, agent_schema_params(schema)).await + pub async fn list_tables( + &mut self, + database: &str, + schema: &str, + ) -> Result { + self.call_method(AgentMethod::ListTables, agent_schema_params(database, schema)).await } - pub async fn list_objects(&mut self, schema: &str) -> Result { - self.call_method(AgentMethod::ListObjects, agent_schema_params(schema)).await + pub async fn list_objects( + &mut self, + database: &str, + schema: &str, + ) -> Result { + self.call_method(AgentMethod::ListObjects, agent_schema_params(database, schema)).await } pub async fn get_object_source( &mut self, + database: &str, schema: &str, name: &str, object_type: &K, ) -> Result { - self.call_method(AgentMethod::GetObjectSource, agent_object_source_params(schema, name, object_type)).await + self.call_method(AgentMethod::GetObjectSource, agent_object_source_params(database, schema, name, object_type)) + .await } pub async fn get_columns( &mut self, + database: &str, schema: &str, table: &str, ) -> Result { - self.call_method(AgentMethod::GetColumns, agent_schema_table_params(schema, table)).await + self.call_method(AgentMethod::GetColumns, agent_schema_table_params(database, schema, table)).await } pub async fn list_indexes( &mut self, + database: &str, schema: &str, table: &str, ) -> Result { - self.call_method(AgentMethod::ListIndexes, agent_schema_table_params(schema, table)).await + self.call_method(AgentMethod::ListIndexes, agent_schema_table_params(database, schema, table)).await } pub async fn list_foreign_keys( &mut self, + database: &str, schema: &str, table: &str, ) -> Result { - self.call_method(AgentMethod::ListForeignKeys, agent_schema_table_params(schema, table)).await + self.call_method(AgentMethod::ListForeignKeys, agent_schema_table_params(database, schema, table)).await } pub async fn list_triggers( &mut self, + database: &str, schema: &str, table: &str, ) -> Result { - self.call_method(AgentMethod::ListTriggers, agent_schema_table_params(schema, table)).await + self.call_method(AgentMethod::ListTriggers, agent_schema_table_params(database, schema, table)).await } pub async fn get_table_ddl( &mut self, + database: &str, schema: &str, table: &str, ) -> Result { - self.call_method(AgentMethod::GetTableDdl, agent_schema_table_params(schema, table)).await + self.call_method(AgentMethod::GetTableDdl, agent_schema_table_params(database, schema, table)).await } pub async fn execute_query(&mut self, params: Value) -> Result { @@ -463,10 +478,11 @@ impl AgentDriverClient { pub async fn execute_transaction( &mut self, + database: Option<&str>, statements: &[String], schema: Option<&str>, ) -> Result { - self.call_method(AgentMethod::ExecuteTransaction, agent_transaction_params(statements, schema)).await + self.call_method(AgentMethod::ExecuteTransaction, agent_transaction_params(database, statements, schema)).await } pub async fn call_mongo_method( @@ -594,24 +610,26 @@ pub fn agent_supports_capability(handshake: Option<&AgentHandshake>, capability: handshake.map(|value| value.supports(capability)).unwrap_or(true) } -pub fn agent_schema_params(schema: &str) -> Value { - serde_json::json!({ "schema": schema }) +pub fn agent_schema_params(database: &str, schema: &str) -> Value { + serde_json::json!({ "database": database, "schema": schema }) } -pub fn agent_schema_table_params(schema: &str, table: &str) -> Value { - serde_json::json!({ "schema": schema, "table": table }) +pub fn agent_schema_table_params(database: &str, schema: &str, table: &str) -> Value { + serde_json::json!({ "database": database, "schema": schema, "table": table }) } -pub fn agent_object_source_params(schema: &str, name: &str, object_type: &K) -> Value { - serde_json::json!({ "schema": schema, "name": name, "object_type": object_type }) +pub fn agent_object_source_params(database: &str, schema: &str, name: &str, object_type: &K) -> Value { + serde_json::json!({ "database": database, "schema": schema, "name": name, "object_type": object_type }) } pub fn agent_close_query_session_params(session_id: &str) -> Value { serde_json::json!({ "sessionId": session_id }) } -pub fn agent_transaction_params(statements: &[String], schema: Option<&str>) -> Value { +pub fn agent_transaction_params(database: Option<&str>, statements: &[String], schema: Option<&str>) -> Value { + let database = database.map(str::trim).filter(|database| !database.is_empty()); serde_json::json!({ + "database": database, "statements": statements, "schema": schema, }) @@ -964,19 +982,27 @@ mod tests { #[test] fn builds_schema_table_and_transaction_params() { - assert_eq!(agent_schema_params("public"), serde_json::json!({ "schema": "public" })); assert_eq!( - agent_schema_table_params("public", "orders"), - serde_json::json!({ "schema": "public", "table": "orders" }) + agent_schema_params("sales", "public"), + serde_json::json!({ "database": "sales", "schema": "public" }) ); assert_eq!( - agent_object_source_params("public", "active_users", &"VIEW"), - serde_json::json!({ "schema": "public", "name": "active_users", "object_type": "VIEW" }) + agent_schema_table_params("sales", "public", "orders"), + serde_json::json!({ "database": "sales", "schema": "public", "table": "orders" }) + ); + assert_eq!( + agent_object_source_params("sales", "public", "active_users", &"VIEW"), + serde_json::json!({ + "database": "sales", + "schema": "public", + "name": "active_users", + "object_type": "VIEW", + }) ); assert_eq!(agent_close_query_session_params("session-1"), serde_json::json!({ "sessionId": "session-1" })); assert_eq!( - agent_transaction_params(&["BEGIN".to_string(), "COMMIT".to_string()], Some("public")), - serde_json::json!({ "statements": ["BEGIN", "COMMIT"], "schema": "public" }) + agent_transaction_params(Some("sales"), &["BEGIN".to_string(), "COMMIT".to_string()], Some("public")), + serde_json::json!({ "database": "sales", "statements": ["BEGIN", "COMMIT"], "schema": "public" }) ); } diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index 7b8ca7663..16733e3c3 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -291,6 +291,7 @@ fn json_value_for_js(value: serde_json::Value) -> serde_json::Value { pub fn agent_execute_query_params( sql: &str, + database: Option<&str>, schema: Option<&str>, options: QueryExecutionOptions, ) -> serde_json::Value { @@ -298,6 +299,9 @@ pub fn agent_execute_query_params( "sql": sql, "maxRows": options.max_rows.unwrap_or(MAX_ROWS), }); + if let Some(database) = database.map(str::trim).filter(|database| !database.is_empty()) { + params["database"] = serde_json::json!(database); + } if let Some(schema) = schema { params["schema"] = serde_json::json!(schema); } @@ -309,6 +313,7 @@ pub fn agent_execute_query_params( pub fn agent_execute_query_page_params( sql: &str, + database: Option<&str>, schema: Option<&str>, options: QueryExecutionOptions, ) -> serde_json::Value { @@ -317,6 +322,9 @@ pub fn agent_execute_query_page_params( "pageSize": options.page_size.unwrap_or(MAX_ROWS), "maxRows": options.max_rows.unwrap_or(MAX_ROWS), }); + if let Some(database) = database.map(str::trim).filter(|database| !database.is_empty()) { + params["database"] = serde_json::json!(database); + } if let Some(schema) = schema { params["schema"] = serde_json::json!(schema); } @@ -559,6 +567,7 @@ pub async fn do_execute( PoolKind::Agent(client) => { let client = client.clone(); let sql = sql.to_string(); + let database = database.map(|s| s.to_string()); let schema = schema.map(|s| s.to_string()); let max_rows = options.max_rows; drop(connections); @@ -568,10 +577,10 @@ pub async fn do_execute( let params = agent_fetch_query_page_params(session_id, options.page_size.unwrap_or(MAX_ROWS)); client.fetch_query_page(params).await } else if options.page_size.is_some() { - let params = agent_execute_query_page_params(&sql, schema.as_deref(), options); + let params = agent_execute_query_page_params(&sql, database.as_deref(), schema.as_deref(), options); client.execute_query_page(params).await } else { - let params = agent_execute_query_params(&sql, schema.as_deref(), options); + let params = agent_execute_query_params(&sql, database.as_deref(), schema.as_deref(), options); client.execute_query(params).await } }) @@ -986,8 +995,10 @@ pub async fn execute_statements_in_transaction( Some(TxPath::Pg(pool)) => exec_tx_pg_inner(pool, statements, schema, start).await, Some(TxPath::Mysql(pool, _bare)) => exec_tx_mysql_inner(pool, statements, start).await, Some(TxPath::Sqlite(pool)) => exec_tx_sqlite_inner(pool, statements, start).await, - Some(TxPath::Explicit) => exec_tx_explicit_inner(state, &pool_key, statements, schema, start).await, - Some(TxPath::None) => exec_tx_none_inner(state, &pool_key, statements, schema, start).await, + Some(TxPath::Explicit) => { + exec_tx_explicit_inner(state, &pool_key, Some(database), statements, schema, start).await + } + Some(TxPath::None) => exec_tx_none_inner(state, &pool_key, Some(database), statements, schema, start).await, None => Err("Connection not found for transaction".to_string()), } } @@ -1122,6 +1133,7 @@ async fn exec_tx_sqlite_inner( async fn exec_tx_explicit_inner( state: &AppState, pool_key: &str, + database: Option<&str>, statements: &[String], schema: Option<&str>, start: std::time::Instant, @@ -1129,24 +1141,25 @@ async fn exec_tx_explicit_inner( let conns = state.connections.read().await; if let Some(crate::connection::PoolKind::Agent(client)) = conns.get(pool_key) { let mut client = client.lock().await; - let result: db::QueryResult = client.execute_transaction(statements, schema).await?; + let result: db::QueryResult = client.execute_transaction(database, statements, schema).await?; return Ok(db::QueryResult { execution_time_ms: start.elapsed().as_millis(), ..result }); } drop(conns); - do_execute(state, pool_key, None, "BEGIN TRANSACTION", schema, None, QueryExecutionOptions::default()) + do_execute(state, pool_key, database, "BEGIN TRANSACTION", schema, None, QueryExecutionOptions::default()) .await .map_err(|e| format!("Failed to begin transaction: {}", e))?; let mut total_affected: u64 = 0; for (i, sql) in statements.iter().enumerate() { - match do_execute(state, pool_key, None, sql, schema, None, QueryExecutionOptions::default()).await { + match do_execute(state, pool_key, database, sql, schema, None, QueryExecutionOptions::default()).await { Ok(result) => { total_affected += result.affected_rows; } Err(e) => { if let Err(rb_err) = - do_execute(state, pool_key, None, "ROLLBACK", schema, None, QueryExecutionOptions::default()).await + do_execute(state, pool_key, database, "ROLLBACK", schema, None, QueryExecutionOptions::default()) + .await { log::error!("ROLLBACK failed after statement {} error: {}", i + 1, rb_err); } @@ -1155,7 +1168,7 @@ async fn exec_tx_explicit_inner( } } - do_execute(state, pool_key, None, "COMMIT", schema, None, QueryExecutionOptions::default()) + do_execute(state, pool_key, database, "COMMIT", schema, None, QueryExecutionOptions::default()) .await .map_err(|e| format!("COMMIT failed: {}", e))?; @@ -1173,6 +1186,7 @@ async fn exec_tx_explicit_inner( async fn exec_tx_none_inner( state: &AppState, pool_key: &str, + database: Option<&str>, statements: &[String], schema: Option<&str>, start: std::time::Instant, @@ -1180,7 +1194,7 @@ async fn exec_tx_none_inner( let mut total_affected: u64 = 0; for (i, sql) in statements.iter().enumerate() { log::info!("[query][tx-none:statement:start] index={} sql={}", i + 1, sql); - match do_execute(state, pool_key, None, sql, schema, None, QueryExecutionOptions::default()).await { + match do_execute(state, pool_key, database, sql, schema, None, QueryExecutionOptions::default()).await { Ok(result) => { total_affected += result.affected_rows; log::info!("[query][tx-none:statement:done] index={} affected_rows={}", i + 1, result.affected_rows); @@ -1421,11 +1435,13 @@ mod tests { fn agent_execute_query_params_include_row_and_fetch_limits() { let params = agent_execute_query_params( "SELECT * FROM events", + Some("analytics"), Some("app"), QueryExecutionOptions { max_rows: Some(500), fetch_size: Some(250), ..Default::default() }, ); assert_eq!(params["sql"], "SELECT * FROM events"); + assert_eq!(params["database"], "analytics"); assert_eq!(params["schema"], "app"); assert_eq!(params["maxRows"], 500); assert_eq!(params["fetchSize"], 250); @@ -1433,9 +1449,10 @@ mod tests { #[test] fn agent_execute_query_params_default_to_safety_row_limit() { - let params = agent_execute_query_params("SELECT * FROM events", None, QueryExecutionOptions::default()); + let params = agent_execute_query_params("SELECT * FROM events", None, None, QueryExecutionOptions::default()); assert_eq!(params["sql"], "SELECT * FROM events"); + assert!(params.get("database").is_none()); assert!(params.get("schema").is_none()); assert_eq!(params["maxRows"], MAX_ROWS); assert!(params.get("fetchSize").is_none()); @@ -1445,11 +1462,13 @@ mod tests { fn agent_execute_query_page_params_include_page_fetch_and_safety_limits() { let params = agent_execute_query_page_params( "SELECT * FROM events", + Some("analytics"), Some("app"), QueryExecutionOptions { page_size: Some(500), fetch_size: Some(250), ..Default::default() }, ); assert_eq!(params["sql"], "SELECT * FROM events"); + assert_eq!(params["database"], "analytics"); assert_eq!(params["schema"], "app"); assert_eq!(params["pageSize"], 500); assert_eq!(params["fetchSize"], 250); diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index 69adb60b9..8f614b89a 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -369,7 +369,7 @@ pub async fn list_tables_core( return db::clickhouse_driver::list_tables(&client, clickhouse_metadata_database(database, schema)).await; } try_sqlserver!(connections, &pool_key, list_tables, schema, filter, limit); - try_agent!(connections, &pool_key, list_tables, schema); + try_agent!(connections, &pool_key, list_tables, database, schema); } let connections = state.connections.read().await; @@ -494,7 +494,7 @@ pub async fn list_objects_core( .await; } try_sqlserver!(connections, &pool_key, list_objects, schema); - try_agent!(connections, &pool_key, list_objects, schema); + try_agent!(connections, &pool_key, list_objects, database, schema); } let connections = state.connections.read().await; @@ -584,7 +584,7 @@ pub async fn get_columns_core( .await; } try_sqlserver!(connections, &pool_key, get_columns, schema, table); - try_agent!(connections, &pool_key, get_columns, schema, table); + try_agent!(connections, &pool_key, get_columns, database, schema, table); } let connections = state.connections.read().await; @@ -612,7 +612,7 @@ pub async fn list_indexes_core( { let connections = state.connections.read().await; try_sqlserver!(connections, &pool_key, list_indexes, schema, table); - try_agent!(connections, &pool_key, list_indexes, schema, table); + try_agent!(connections, &pool_key, list_indexes, database, schema, table); } let connections = state.connections.read().await; @@ -640,7 +640,7 @@ pub async fn list_foreign_keys_core( { let connections = state.connections.read().await; try_sqlserver!(connections, &pool_key, list_foreign_keys, schema, table); - try_agent!(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; @@ -668,7 +668,7 @@ pub async fn list_triggers_core( { let connections = state.connections.read().await; try_sqlserver!(connections, &pool_key, list_triggers, schema, table); - try_agent!(connections, &pool_key, list_triggers, schema, table); + try_agent!(connections, &pool_key, list_triggers, database, schema, table); } let connections = state.connections.read().await; @@ -730,7 +730,7 @@ pub async fn get_table_ddl_core( let mut client = client.lock().await; return build_sqlserver_ddl(&mut client, schema, table).await; } - try_agent!(connections, &pool_key, get_table_ddl, schema, table); + try_agent!(connections, &pool_key, get_table_ddl, database, schema, table); } let connections = state.connections.read().await; @@ -886,7 +886,7 @@ pub async fn get_object_source_core( } else if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { drop(connections); let mut client = client.lock().await; - let result: db::ObjectSource = client.get_object_source(schema, name, &object_type).await?; + let result: db::ObjectSource = client.get_object_source(database, schema, name, &object_type).await?; return Ok(result); } else { match connections.get(&pool_key).ok_or("Pool not found")? { diff --git a/crates/dbx-core/src/transfer.rs b/crates/dbx-core/src/transfer.rs index 6b6d35bad..e368120ca 100644 --- a/crates/dbx-core/src/transfer.rs +++ b/crates/dbx-core/src/transfer.rs @@ -5,6 +5,7 @@ use tokio::sync::RwLock; use crate::connection::{AppState, PoolKind}; use crate::db; use crate::models::connection::DatabaseType; +use crate::query::{agent_execute_query_params, QueryExecutionOptions}; static CANCELLED: std::sync::LazyLock>> = std::sync::LazyLock::new(|| RwLock::new(HashSet::new())); @@ -789,7 +790,7 @@ pub async fn execute_on_pool(state: &AppState, pool_key: &str, sql: &str) -> Res } PoolKind::ClickHouse(client) => { let client = client.clone(); - let database = pool_key.split(':').nth(1).unwrap_or("default").to_string(); + let database = database_from_pool_key(pool_key).unwrap_or("default").to_string(); drop(connections); db::clickhouse_driver::execute_query(&client, &database, sql).await } @@ -799,6 +800,20 @@ pub async fn execute_on_pool(state: &AppState, pool_key: &str, sql: &str) -> Res let mut client = client.lock().await; db::sqlserver::execute_query(&mut client, sql).await } + PoolKind::Agent(client) => { + let client = client.clone(); + let database = database_from_pool_key(pool_key).map(str::to_string); + let sql = sql.to_string(); + drop(connections); + let mut client = client.lock().await; + let params = agent_execute_query_params( + &sql, + database.as_deref(), + None, + QueryExecutionOptions { max_rows: None, ..QueryExecutionOptions::default() }, + ); + client.execute_query(params).await + } PoolKind::DuckDb(con) => { let con = con.clone(); let sql = sql.to_string(); @@ -880,6 +895,16 @@ pub async fn execute_on_pool(state: &AppState, pool_key: &str, sql: &str) -> Res } } +fn database_from_pool_key(pool_key: &str) -> Option<&str> { + pool_key + .split_once(":session:") + .map(|(base, _)| base) + .unwrap_or(pool_key) + .split_once(':') + .map(|(_, database)| database) + .filter(|database| !database.is_empty()) +} + pub async fn get_db_type(state: &AppState, connection_id: &str) -> Result { let configs = state.configs.read().await; configs @@ -939,6 +964,15 @@ pub async fn get_columns_for_transfer( let mut client = client.lock().await; return db::sqlserver::get_columns(&mut client, &schema, &table).await; } + if let Some(PoolKind::Agent(client)) = connections.get(pool_key) { + let client = client.clone(); + let database = database.to_string(); + let schema = schema.to_string(); + let table = table.to_string(); + drop(connections); + let mut client = client.lock().await; + return client.get_columns(&database, &schema, &table).await; + } let pool = connections.get(pool_key).ok_or("Pool not found")?; let schema = schema.to_string(); let table = table.to_string(); @@ -1407,4 +1441,11 @@ mod tests { assert_eq!(columns.iter().map(|c| c.name.as_str()).collect::>(), vec!["id"]); } + + #[test] + fn database_from_pool_key_handles_session_scoped_keys() { + assert_eq!(database_from_pool_key("conn:analytics"), Some("analytics")); + assert_eq!(database_from_pool_key("conn:analytics:session:editor-1"), Some("analytics")); + assert_eq!(database_from_pool_key("conn"), None); + } }