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

fix(sqlserver): isolate legacy metadata pool from work pools

Browse Source 
Closes #9183
This commit is contained in:
zipg
2026-09-16 22:40:46 +08:00
committed by GitHub
parent 81a0ed922a
commit 81292760e2
1 changed files with 152 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/dbx-core/src/connection.rs
+152 -15
View File
@@ -941,6 +941,17 @@ pub fn sqlserver_uses_legacy_driver(config: &ConnectionConfig) -> bool {
.is_some_and(|profile| profile.eq_ignore_ascii_case(db::sqlserver::SQLSERVER_LEGACY_DRIVER_PROFILE))
}
fn metadata_pool_database<'a>(config: Option<&ConnectionConfig>, database: Option<&'a str>) -> Option<&'a str> {
if config.is_some_and(sqlserver_uses_legacy_driver) {
// The legacy SQL Server Agent switches catalogs on the borrowed JDBC connection for
// each metadata request. Reuse the connection-level pool so expanding a database does
// not create another physical login session on SQL Server 2000.
None
} else {
database
}
}
pub fn sqlserver_legacy_driver_error(agent_error: &str) -> String {
// This mapper handles both AgentManager launch strings and Agent call errors, so context
// must remain before any structured-error compatibility marker.
@@ -2143,9 +2154,14 @@ impl AppState {
database: Option<&str>,
client_session_id: Option<&str>,
) -> Result<String, String> {
let config = {
let configs = self.configs.read().await;
configs.get(connection_id).cloned()
};
let pool_database = metadata_pool_database(config.as_ref(), database);
self.get_or_create_pool_for_session_inner(
connection_id,
database,
pool_database,
None,
client_session_id,
AgentSessionRole::Metadata,
@@ -4008,7 +4024,12 @@ impl AppState {
};
let db_type = config.as_ref().map(|config| config.db_type);
let catalog = catalog.map(str::trim).filter(|value| !value.is_empty());
let base_pool_key = base_pool_key_for_with_catalog(db_type, connection_id, database, catalog, true);
let pool_database = if session_role == AgentSessionRole::Metadata {
metadata_pool_database(config.as_ref(), database)
} else {
database
};
let base_pool_key = base_pool_key_for_with_catalog(db_type, connection_id, pool_database, catalog, true);
let pool_key = pool_key_for_session_role(config.as_ref(), base_pool_key, client_session_id, session_role);
if self.uses_forwarded_transport(connection_id).await {
self.remove_connection_pools(connection_id).await;
@@ -4024,7 +4045,7 @@ impl AppState {
}
self.get_or_create_pool_for_session_inner(
connection_id,
database,
pool_database,
catalog,
client_session_id,
session_role,
@@ -4055,8 +4076,13 @@ impl AppState {
database: Option<&str>,
client_session_id: &str,
) -> Result<bool, String> {
let config = {
let configs = self.configs.read().await;
configs.get(connection_id).cloned()
};
let pool_database = metadata_pool_database(config.as_ref(), database);
let Some((pool_key, pool)) = self
.take_client_session_pool(connection_id, database, client_session_id, AgentSessionRole::Metadata)
.take_client_session_pool(connection_id, pool_database, client_session_id, AgentSessionRole::Metadata)
.await?
else {
return Ok(false);
@@ -4107,7 +4133,12 @@ impl AppState {
configs.get(connection_id).cloned()
};
let db_type = config.as_ref().map(|config| config.db_type);
let base_pool_key = base_pool_key_for(db_type, connection_id, database, false);
let pool_database = if session_role == AgentSessionRole::Metadata {
metadata_pool_database(config.as_ref(), database)
} else {
database
};
let base_pool_key = base_pool_key_for(db_type, connection_id, pool_database, false);
let pool_key =
pool_key_for_session_role(config.as_ref(), base_pool_key.clone(), Some(client_session_id), session_role);
if pool_key == base_pool_key {
@@ -4146,7 +4177,8 @@ impl AppState {
configs.get(connection_id).cloned()
};
let db_type = config.as_ref().map(|config| config.db_type);
let base_pool_key = base_pool_key_for(db_type, connection_id, database, false);
let pool_database = metadata_pool_database(config.as_ref(), database);
let base_pool_key = base_pool_key_for(db_type, connection_id, pool_database, false);
let pool_key =
pool_key_for_session_role(config.as_ref(), base_pool_key, client_session_id, AgentSessionRole::Metadata);
self.detach_pool_by_key(&pool_key, true).await
@@ -4165,7 +4197,8 @@ impl AppState {
configs.get(connection_id).cloned()
};
let db_type = config.as_ref().map(|config| config.db_type);
let base_pool_key = base_pool_key_for(db_type, connection_id, database, false);
let pool_database = metadata_pool_database(config.as_ref(), database);
let base_pool_key = base_pool_key_for(db_type, connection_id, pool_database, false);
let pool_key =
pool_key_for_session_role(config.as_ref(), base_pool_key, client_session_id, AgentSessionRole::Metadata);
if let Some(session_id) = agent_session_id {
@@ -5682,8 +5715,13 @@ fn pool_key_for_session_role(
) -> String {
let pool_key = session_scoped_pool_key_for(config, base_pool_key, client_session_id);
if session_role == AgentSessionRole::Metadata
&& config.is_some_and(|config| database_capabilities::is_agent_type(&config.db_type))
&& config.is_some_and(|config| {
database_capabilities::is_agent_type(&config.db_type) || sqlserver_uses_legacy_driver(config)
})
{
// The legacy SQL Server Agent borrows one connection-level pool for metadata across
// databases and switches catalogs per request. The role suffix keeps that shared
// pool from colliding with workload pools on the bare connection id.
format!("{pool_key}:role:metadata")
} else {
pool_key
@@ -6136,13 +6174,14 @@ mod tests {
connection_probe_endpoints, connection_remote_endpoint, connection_url_for_endpoint,
database_connection_config, database_connection_config_with_catalog,
gaussdb_identifier_quote_from_query_result, gaussdb_m_jdbc_config_for_endpoint, gaussdb_uses_m_jdbc_driver,
kafka_single_loopback_bootstrap_endpoint, metadata_connection_config, mysql_metadata_fallback_url,
mysql_pool_setup_queries, oceanbase_mysql_setup_queries, prestosql_jdbc_config_for_endpoint,
redacted_connection_url_for_endpoint, redis_sentinel_transport_id, redis_sentinel_transport_prefix,
sqlserver_legacy_agent_config, sqlserver_legacy_driver_error, sqlserver_uses_legacy_driver,
task_client_session_id, transport_layers_through_last_ssh, upsert_connection_url_param, uses_bare_mysql_pool,
uses_tcp_probe, validate_connection_url_params, validate_h2_database_path, AppState, MysqlMode, PoolKind,
TxnConnection, GAUSSDB_M_JDBC_DRIVER_CLASS, GAUSSDB_M_JDBC_DRIVER_PROFILE, PRESTOSQL_JDBC_DRIVER_CLASS,
kafka_single_loopback_bootstrap_endpoint, metadata_connection_config, metadata_pool_database,
mysql_metadata_fallback_url, mysql_pool_setup_queries, oceanbase_mysql_setup_queries,
prestosql_jdbc_config_for_endpoint, redacted_connection_url_for_endpoint, redis_sentinel_transport_id,
redis_sentinel_transport_prefix, sqlserver_legacy_agent_config, sqlserver_legacy_driver_error,
sqlserver_uses_legacy_driver, task_client_session_id, transport_layers_through_last_ssh,
upsert_connection_url_param, uses_bare_mysql_pool, uses_tcp_probe, validate_connection_url_params,
validate_h2_database_path, AppState, MysqlMode, PoolKind, TxnConnection, GAUSSDB_M_JDBC_DRIVER_CLASS,
GAUSSDB_M_JDBC_DRIVER_PROFILE, PRESTOSQL_JDBC_DRIVER_CLASS,
};
use crate::agent_connection::{
agent_connect_params, mongo_legacy_error_with_auth_hint, mongo_uses_legacy_driver,
@@ -6664,6 +6703,16 @@ mod tests {
assert!(sqlserver_uses_legacy_driver(&legacy));
}
#[test]
fn legacy_sqlserver_metadata_reuses_connection_pool_across_databases() {
let mut config = mysql_config(Some("master"));
config.db_type = DatabaseType::SqlServer;
let legacy = sqlserver_legacy_agent_config(&config);
assert_eq!(metadata_pool_database(Some(&legacy), Some("app")), None);
assert_eq!(metadata_pool_database(Some(&config), Some("app")), Some("app"));
}
#[test]
fn sqlserver_legacy_url_param_requires_canonicalization_before_agent_driver_selection() {
let mut config = mysql_config(Some("master"));
@@ -7888,6 +7937,49 @@ mod tests {
assert_eq!(metadata, "conn:analytics:session:task_1");
}
#[test]
fn legacy_sqlserver_metadata_pool_keys_share_role_isolated_key() {
let mut config = mysql_config(Some("master"));
config.db_type = DatabaseType::SqlServer;
let legacy = sqlserver_legacy_agent_config(&config);
let metadata_key = |database: Option<&str>, client_session_id: Option<&str>| {
let pool_database = super::metadata_pool_database(Some(&legacy), database);
let base_pool_key = super::base_pool_key_for(Some(legacy.db_type), "conn", pool_database, false);
super::pool_key_for_session_role(
Some(&legacy),
base_pool_key,
client_session_id,
crate::agent_connection::AgentSessionRole::Metadata,
)
};
// Every database resolves to one shared, role-isolated metadata pool key that never
// collides with the bare connection-level workload pool.
assert_eq!(metadata_key(Some("a"), Some("task:1")), "conn:session:task_1:role:metadata");
assert_eq!(metadata_key(Some("b"), Some("task:1")), "conn:session:task_1:role:metadata");
assert_eq!(metadata_key(Some("a"), None), "conn:role:metadata");
assert_ne!(metadata_key(Some("a"), None), "conn");
let workload = super::pool_key_for_session_role(
Some(&legacy),
"conn".to_string(),
Some("task:1"),
crate::agent_connection::AgentSessionRole::Workload,
);
assert_eq!(workload, "conn:session:task_1");
assert_ne!(metadata_key(Some("a"), Some("task:1")), workload);
// Without the legacy driver profile the metadata role keeps sharing workload keys.
let shared = super::pool_key_for_session_role(
Some(&config),
"conn".to_string(),
Some("task:1"),
crate::agent_connection::AgentSessionRole::Metadata,
);
assert_eq!(shared, "conn:session:task_1");
}
#[test]
fn redis_sentinel_transport_ids_are_connection_scoped_by_role_and_endpoint() {
let endpoint = db::redis_driver::RedisNodeEndpoint { host: "10.0.0.8".to_string(), port: 6379 };
@@ -8693,6 +8785,51 @@ for line in sys.stdin:
let _ = std::fs::remove_dir_all(dir);
}
#[tokio::test]
async fn legacy_sqlserver_metadata_close_finds_role_isolated_pool() {
let (state, dir) = test_app_state().await;
let mut config = mysql_config(Some("master"));
config.id = "conn".to_string();
config.db_type = DatabaseType::SqlServer;
let legacy = sqlserver_legacy_agent_config(&config);
state.configs.write().await.insert(legacy.id.clone(), legacy.clone());
let metadata_pool_key = {
let pool_database = super::metadata_pool_database(Some(&legacy), Some("a"));
let base_pool_key = super::base_pool_key_for(Some(legacy.db_type), "conn", pool_database, false);
super::pool_key_for_session_role(
Some(&legacy),
base_pool_key,
Some("metadata-session"),
crate::agent_connection::AgentSessionRole::Metadata,
)
};
assert_eq!(metadata_pool_key, "conn:session:metadata-session:role:metadata");
let workload_pool_key = "conn".to_string();
{
let mut connections = state.connections.write().await;
connections.insert(metadata_pool_key.clone(), agent_pool_stub());
connections.insert(workload_pool_key.clone(), agent_pool_stub());
}
{
let mut activity = state.pool_activity.write().await;
activity.insert(metadata_pool_key.clone(), super::PoolActivity::now());
activity.insert(workload_pool_key.clone(), super::PoolActivity::now());
}
// Closing through any database resolves the same shared metadata pool and leaves the
// connection-level workload pool untouched.
assert!(state.close_metadata_session_pool("conn", Some("b"), "metadata-session").await.unwrap());
assert!(!state.connections.read().await.contains_key(&metadata_pool_key));
assert!(!state.pool_activity.read().await.contains_key(&metadata_pool_key));
assert!(state.connections.read().await.contains_key(&workload_pool_key));
assert!(state.pool_activity.read().await.contains_key(&workload_pool_key));
state.shutdown(Duration::from_secs(1)).await;
let _ = std::fs::remove_dir_all(dir);
}
#[tokio::test]
async fn reclaim_close_replace_runtime_never_restores_failed_pool() {
let (state, dir) = test_app_state().await;