xiaoxue_admin
/
dbx
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: scope agent metadata by database

This commit is contained in:
t8y2 2026-05-26 11:51:58 +08:00
parent 6d6ba34cac
commit f895bb4de5
5 changed files with 192 additions and 81 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

97
crates/dbx-core/src/connection.rs
View File

@ -148,18 +148,10 @@ impl AppState {
) -> Result<String, String> {
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<String, String> {
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<DatabaseType>,
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"));

76
crates/dbx-core/src/db/agent_driver.rs
View File

@ -382,61 +382,76 @@ impl AgentDriverClient {
self.call_method(AgentMethod::ListSchemas, serde_json::json!({ "database": database })).await
}
pub async fn list_tables<T: DeserializeOwned + Send + 'static>(&mut self, schema: &str) -> Result<T, String> {
self.call_method(AgentMethod::ListTables, agent_schema_params(schema)).await
pub async fn list_tables<T: DeserializeOwned + Send + 'static>(
&mut self,
database: &str,
schema: &str,
) -> Result<T, String> {
self.call_method(AgentMethod::ListTables, agent_schema_params(database, schema)).await
}
pub async fn list_objects<T: DeserializeOwned + Send + 'static>(&mut self, schema: &str) -> Result<T, String> {
self.call_method(AgentMethod::ListObjects, agent_schema_params(schema)).await
pub async fn list_objects<T: DeserializeOwned + Send + 'static>(
&mut self,
database: &str,
schema: &str,
) -> Result<T, String> {
self.call_method(AgentMethod::ListObjects, agent_schema_params(database, schema)).await
}
pub async fn get_object_source<T: DeserializeOwned + Send + 'static, K: Serialize>(
&mut self,
database: &str,
schema: &str,
name: &str,
object_type: &K,
) -> Result<T, String> {
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<T: DeserializeOwned + Send + 'static>(
&mut self,
database: &str,
schema: &str,
table: &str,
) -> Result<T, String> {
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<T: DeserializeOwned + Send + 'static>(
&mut self,
database: &str,
schema: &str,
table: &str,
) -> Result<T, String> {
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<T: DeserializeOwned + Send + 'static>(
&mut self,
database: &str,
schema: &str,
table: &str,
) -> Result<T, String> {
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<T: DeserializeOwned + Send + 'static>(
&mut self,
database: &str,
schema: &str,
table: &str,
) -> Result<T, String> {
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<T: DeserializeOwned + Send + 'static>(
&mut self,
database: &str,
schema: &str,
table: &str,
) -> Result<T, String> {
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<T: DeserializeOwned + Send + 'static>(&mut self, params: Value) -> Result<T, String> {
@ -463,10 +478,11 @@ impl AgentDriverClient {
pub async fn execute_transaction<T: DeserializeOwned + Send + 'static>(
&mut self,
database: Option<&str>,
statements: &[String],
schema: Option<&str>,
) -> Result<T, String> {
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<T: DeserializeOwned + Send + 'static>(
@ -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<K: Serialize>(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<K: Serialize>(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" })
);
}

41
crates/dbx-core/src/query.rs
View File

@ -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);

16
crates/dbx-core/src/schema.rs
View File

@ -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")? {

43
crates/dbx-core/src/transfer.rs
View File

@ -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<RwLock<HashSet<String>>> =
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<DatabaseType, String> {
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<_>>(), 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);
}
}