mirror of
https://github.com/t8y2/dbx.git
synced 2026-10-19 16:16:45 +00:00
9916 lines
419 KiB
Rust
9916 lines
419 KiB
Rust
use std::collections::{HashMap, HashSet};
|
||
use std::path::PathBuf;
|
||
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
|
||
use std::sync::Arc;
|
||
use std::time::{Duration, Instant};
|
||
use tokio::sync::{watch, Mutex, OwnedSemaphorePermit, RwLock, Semaphore};
|
||
|
||
use mysql_async::prelude::Queryable;
|
||
use mysql_async::Row as MysqlRow;
|
||
|
||
use crate::agent_connection::{
|
||
agent_connect_params, agent_connect_params_with_role, h2_file_path_from_jdbc_url, hive_uses_zookeeper_discovery,
|
||
is_h2_file_connection, mongo_legacy_error_with_auth_hint, mongo_uses_legacy_driver,
|
||
oracle_alternate_connect_config_labels, oracle_alternate_connect_configs, oracle_error_with_driver_hint,
|
||
should_retry_mongo_with_legacy_driver, trino_like_jdbc_connection_string, AgentSessionRole,
|
||
};
|
||
use crate::agent_manager::{AgentManager, JavaRuntimeMode, DEFAULT_JRE_KEY};
|
||
use crate::agent_recovery::{RecoveryDecision, RecoveryPolicy, RecoveryScope};
|
||
use crate::database_capabilities;
|
||
use crate::db;
|
||
use crate::db::agent_driver::{AgentCallError, AgentMethod};
|
||
use crate::db::http_tunnel::HttpTunnelManager;
|
||
use crate::db::proxy_tunnel::ProxyTunnelManager;
|
||
use crate::db::ssh_tunnel::TunnelManager;
|
||
use crate::models::connection::{
|
||
database_info_from_protocol_value, parse_jdbc_host_port, parse_mongo_first_host, rewrite_jdbc_url_host,
|
||
ConnectionConfig, ConnectionTestResult, DatabaseConnectionInfo, DatabaseType, TransportLayerConfig,
|
||
};
|
||
use crate::mongo_oidc::MongoOidcBrowserOpener;
|
||
use crate::nacos::config::{NACOS_CONSOLE_SESSION_PASSWORD, NACOS_PRIMARY_SESSION_PASSWORD};
|
||
use crate::path_utils::expand_tilde;
|
||
use crate::plugins::{PluginConnectionHandle, PluginDriverSession, PluginHost, PluginRegistry, PluginRuntimeEnv};
|
||
use crate::query_cancel::RunningQueries;
|
||
use crate::session_credentials::SessionCredentialStore;
|
||
use crate::storage::{normalize_duckdb_worker_max_processes, Storage, DUCKDB_WORKER_MAX_PROCESSES_DEFAULT};
|
||
use crate::task_supervisor::TaskSupervisor;
|
||
|
||
pub const JDBC_PLUGIN_NOT_INSTALLED: &str =
|
||
"JDBC plugin is not installed. Install the optional JDBC plugin to use this connection.";
|
||
pub const PRESTOSQL_JDBC_DRIVER_CLASS: &str = "io.prestosql.jdbc.PrestoDriver";
|
||
pub const GAUSSDB_M_JDBC_DRIVER_PROFILE: &str = "gaussdb-m";
|
||
pub const GAUSSDB_M_JDBC_DRIVER_CLASS: &str = "com.huawei.gaussdb.jdbc.Driver";
|
||
const SQLSERVER_LEGACY_DRIVER_INSTALL_HINT: &str =
|
||
"Install the SQL Server legacy compatibility component from Driver Manager, or open the connection settings and enable SQL Server legacy compatibility mode again.";
|
||
const DEFAULT_AGENT_CONNECT_TIMEOUT_SECS: u64 = 30;
|
||
const ACCESS_AGENT_CONNECT_TIMEOUT_SECS: u64 = 30;
|
||
const POOL_CLOSE_TIMEOUT_SECS: u64 = 3;
|
||
const HEALTH_CHECK_POOL_ACQUIRE_TIMEOUT: Duration = Duration::from_millis(500);
|
||
pub(crate) const METADATA_POOL_ACQUIRE_TIMEOUT: Duration = Duration::from_secs(5);
|
||
pub(crate) const METADATA_POOL_DEFAULT_LIMIT: usize = 6;
|
||
pub(crate) const METADATA_POOL_SQLSERVER_LIMIT: usize = 1;
|
||
|
||
mod duckdb_types {
|
||
#[cfg(feature = "duckdb-sidecar")]
|
||
use std::sync::Arc;
|
||
#[cfg(feature = "duckdb-sidecar")]
|
||
pub type DuckDbWorkerHandle = Arc<crate::db::duckdb_worker_process::DuckDbWorkerClient>;
|
||
#[cfg(not(feature = "duckdb-sidecar"))]
|
||
pub type DuckDbWorkerHandle = ();
|
||
}
|
||
|
||
use duckdb_types::DuckDbWorkerHandle;
|
||
|
||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||
pub enum MysqlMode {
|
||
Normal,
|
||
Bare,
|
||
OceanBaseOracle,
|
||
}
|
||
|
||
fn oceanbase_mysql_setup_queries(config: &ConnectionConfig) -> Vec<String> {
|
||
db::oceanbase_mysql::query_timeout_sql(config, config.query_timeout_secs).into_iter().collect()
|
||
}
|
||
|
||
fn mysql_pool_setup_queries(config: &ConnectionConfig, url: &str) -> Vec<String> {
|
||
let mut queries = oceanbase_mysql_setup_queries(config);
|
||
if let Some(dialect) = db::mysql::mysql_catalog_dialect(config.db_type, config.driver_profile.as_deref()) {
|
||
if let Some(query) = db::mysql::catalog_setup_query_for_url(dialect, url) {
|
||
queries.push(query);
|
||
}
|
||
}
|
||
queries
|
||
}
|
||
|
||
#[derive(Clone)]
|
||
pub enum PoolKind {
|
||
Mysql(db::mysql::MySqlPool, MysqlMode),
|
||
Postgres(deadpool_postgres::Pool),
|
||
Sqlite(db::sqlite::SqliteHandle),
|
||
Rqlite(db::rqlite_driver::RqliteClient),
|
||
Turso(db::turso_driver::TursoClient),
|
||
CloudflareD1(db::cloudflare_d1_driver::CloudflareD1Client),
|
||
Redis(Arc<db::redis_driver::RedisConnection>),
|
||
DuckDbWorker(DuckDbWorkerHandle),
|
||
MongoDb(mongodb::Client),
|
||
DynamoDb(db::dynamodb_driver::DynamoDbClient),
|
||
ClickHouse(db::clickhouse_driver::ChClient),
|
||
SqlServer(Arc<tokio::sync::Mutex<db::sqlserver::SqlServerClient>>),
|
||
Elasticsearch(db::elasticsearch_driver::EsClient),
|
||
Easysearch(db::easysearch_driver::EasysearchClient),
|
||
Meilisearch(db::meilisearch_driver::MeilisearchClient),
|
||
HBase(db::hbase_driver::HBaseClient),
|
||
VectorDb(db::vector_driver::VectorClient),
|
||
InfluxDb(db::influxdb_driver::InfluxdbClient),
|
||
InfluxDb3(db::influxdb3_driver::Influxdb3Client),
|
||
VictoriaMetrics(db::victoriametrics_driver::VictoriaMetricsClient),
|
||
Agent(Arc<db::agent_driver::PooledAgentClient>),
|
||
ExternalDriver {
|
||
driver_id: String,
|
||
config: Arc<ConnectionConfig>,
|
||
session: Arc<PluginDriverSession>,
|
||
},
|
||
PluginConnection(PluginConnectionHandle),
|
||
/// Message queue admin connection (not a data query pool; serves as a
|
||
/// marker that this connection_id is a valid MQ admin connection).
|
||
MessageQueue,
|
||
/// Nacos admin connection marker.
|
||
Nacos,
|
||
Consul(crate::consul::ConsulClient),
|
||
/// MQTT broker connection with an active client.
|
||
#[cfg(feature = "mq-admin")]
|
||
Mqtt(Arc<super::mqtt::client::MqttClient>),
|
||
}
|
||
|
||
impl PoolKind {
|
||
pub fn agent(client: db::agent_driver::AgentDriverClient) -> Self {
|
||
Self::Agent(Arc::new(db::agent_driver::PooledAgentClient::new(client)))
|
||
}
|
||
|
||
fn is_available_for_routing(&self) -> bool {
|
||
match self {
|
||
Self::Agent(client) => client.is_runtime_available(),
|
||
_ => true,
|
||
}
|
||
}
|
||
}
|
||
|
||
#[derive(Clone)]
|
||
struct PoolPublication(Arc<()>);
|
||
|
||
impl PoolPublication {
|
||
fn new() -> Self {
|
||
Self(Arc::new(()))
|
||
}
|
||
|
||
fn is_same(&self, other: &Self) -> bool {
|
||
Arc::ptr_eq(&self.0, &other.0)
|
||
}
|
||
}
|
||
|
||
#[derive(Clone)]
|
||
struct PoolPublicationSnapshot {
|
||
pool: PoolKind,
|
||
publication: PoolPublication,
|
||
}
|
||
|
||
#[cfg(test)]
|
||
#[derive(Default)]
|
||
struct StalePoolCleanupBarriers {
|
||
before_removal: Option<(Arc<tokio::sync::Barrier>, Arc<tokio::sync::Barrier>)>,
|
||
after_removal: Option<(Arc<tokio::sync::Barrier>, Arc<tokio::sync::Barrier>)>,
|
||
}
|
||
|
||
/// Internal pool registry that assigns an opaque identity to every published
|
||
/// entry. The identity belongs to the registry publication rather than to the
|
||
/// driver handle, so replacing an entry always creates a new generation even
|
||
/// when the replacement reuses a cloned handle.
|
||
#[doc(hidden)]
|
||
pub struct ConnectionPoolRegistry {
|
||
pools: HashMap<String, PoolKind>,
|
||
publications: HashMap<String, PoolPublication>,
|
||
}
|
||
|
||
impl ConnectionPoolRegistry {
|
||
fn new() -> Self {
|
||
Self { pools: HashMap::new(), publications: HashMap::new() }
|
||
}
|
||
|
||
pub fn insert(&mut self, pool_key: String, pool: PoolKind) -> Option<PoolKind> {
|
||
self.publications.insert(pool_key.clone(), PoolPublication::new());
|
||
self.pools.insert(pool_key, pool)
|
||
}
|
||
|
||
pub fn remove(&mut self, pool_key: &str) -> Option<PoolKind> {
|
||
self.publications.remove(pool_key);
|
||
self.pools.remove(pool_key)
|
||
}
|
||
|
||
#[cfg(test)]
|
||
fn clear(&mut self) {
|
||
self.publications.clear();
|
||
self.pools.clear();
|
||
}
|
||
|
||
fn drain(&mut self) -> std::collections::hash_map::Drain<'_, String, PoolKind> {
|
||
self.publications.clear();
|
||
self.pools.drain()
|
||
}
|
||
|
||
fn snapshot(&self, pool_key: &str) -> Option<PoolPublicationSnapshot> {
|
||
Some(PoolPublicationSnapshot {
|
||
pool: self.pools.get(pool_key)?.clone(),
|
||
publication: self.publications.get(pool_key)?.clone(),
|
||
})
|
||
}
|
||
|
||
fn remove_if_publication(&mut self, pool_key: &str, expected: &PoolPublication) -> Option<PoolKind> {
|
||
let is_current = self.publications.get(pool_key).is_some_and(|current| current.is_same(expected));
|
||
is_current.then(|| self.remove(pool_key)).flatten()
|
||
}
|
||
}
|
||
|
||
impl std::ops::Deref for ConnectionPoolRegistry {
|
||
type Target = HashMap<String, PoolKind>;
|
||
|
||
fn deref(&self) -> &Self::Target {
|
||
&self.pools
|
||
}
|
||
}
|
||
|
||
enum ConnectionDatabaseInfoSource {
|
||
Agent(Arc<db::agent_driver::PooledAgentClient>),
|
||
MongoAgent(Arc<db::agent_driver::PooledAgentClient>, Option<String>),
|
||
ExternalDriver {
|
||
config: Arc<ConnectionConfig>,
|
||
session: Arc<PluginDriverSession>,
|
||
},
|
||
NativeMysql(db::mysql::MySqlPool),
|
||
NativeHBase(db::hbase_driver::HBaseClient),
|
||
NativeMongo(mongodb::Client, Option<String>),
|
||
Meilisearch(db::meilisearch_driver::MeilisearchClient),
|
||
VictoriaMetrics(db::victoriametrics_driver::VictoriaMetricsClient),
|
||
Redis(String),
|
||
Nacos,
|
||
Consul(Box<crate::consul::ConsulClient>),
|
||
#[cfg(feature = "mq-admin")]
|
||
MessageQueue,
|
||
}
|
||
|
||
/// Held connection for a manual transaction session
|
||
pub enum TxnConnection {
|
||
Postgres(Box<deadpool_postgres::Object>),
|
||
/// A dedicated MySQL connection. Cancellation may consume and discard this
|
||
/// connection instead of trying to reuse it after an interrupted result set.
|
||
Mysql(Option<mysql_async::Conn>),
|
||
/// Dedicated agent multi_session workload client with an open sticky TX.
|
||
Agent {
|
||
client: Arc<db::agent_driver::PooledAgentClient>,
|
||
/// Client-session id used when opening the dedicated agent pool.
|
||
client_session_id: String,
|
||
database: Option<String>,
|
||
cleanup_guard: ClientSessionPoolCleanupGuard,
|
||
},
|
||
/// Dedicated external-driver process whose shared JDBC connection owns the TX.
|
||
ExternalDriver {
|
||
session: Arc<PluginDriverSession>,
|
||
config: Arc<ConnectionConfig>,
|
||
client_session_id: String,
|
||
database: Option<String>,
|
||
cleanup_guard: ClientSessionPoolCleanupGuard,
|
||
},
|
||
}
|
||
|
||
pub struct TransactionSession {
|
||
pub connection: Arc<Mutex<TxnConnection>>,
|
||
pub pool_key: String,
|
||
pub last_activity: std::time::Instant,
|
||
pub busy: bool,
|
||
pub connection_id: String,
|
||
pub database: String,
|
||
pub schema: Option<String>,
|
||
}
|
||
|
||
macro_rules! agent_connection_pool_database_type {
|
||
() => {
|
||
DatabaseType::Dameng
|
||
| DatabaseType::Kingbase
|
||
| DatabaseType::Highgo
|
||
| DatabaseType::Uxdb
|
||
| DatabaseType::Vastbase
|
||
| DatabaseType::Goldendb
|
||
| DatabaseType::Databend
|
||
| DatabaseType::Yashandb
|
||
| DatabaseType::Databricks
|
||
| DatabaseType::SapHana
|
||
| DatabaseType::Teradata
|
||
| DatabaseType::Vertica
|
||
| DatabaseType::Firebird
|
||
| DatabaseType::Exasol
|
||
| DatabaseType::OceanbaseOracle
|
||
| DatabaseType::Gbase
|
||
| DatabaseType::Oracle
|
||
| DatabaseType::H2
|
||
| DatabaseType::Snowflake
|
||
| DatabaseType::Trino
|
||
| DatabaseType::Hive
|
||
| DatabaseType::Kyuubi
|
||
| DatabaseType::Impala
|
||
| DatabaseType::Argo
|
||
| DatabaseType::Spark
|
||
| DatabaseType::Db2
|
||
| DatabaseType::Informix
|
||
| DatabaseType::Neo4j
|
||
| DatabaseType::Cassandra
|
||
| DatabaseType::Bigquery
|
||
| DatabaseType::Spanner
|
||
| DatabaseType::Kylin
|
||
| DatabaseType::Ignite
|
||
| DatabaseType::Ignite3
|
||
| DatabaseType::Sundb
|
||
| DatabaseType::Oscar
|
||
| DatabaseType::Tdengine
|
||
| DatabaseType::Xugu
|
||
| DatabaseType::Iotdb
|
||
| DatabaseType::Etcd
|
||
| DatabaseType::ZooKeeper
|
||
| DatabaseType::Iris
|
||
| DatabaseType::Access
|
||
};
|
||
}
|
||
|
||
pub struct AppState {
|
||
connections: Arc<RwLock<ConnectionPoolRegistry>>,
|
||
task_supervisor: TaskSupervisor,
|
||
pool_activity: Arc<RwLock<HashMap<String, PoolActivity>>>,
|
||
draining_pools: Arc<std::sync::Mutex<HashMap<String, watch::Sender<bool>>>>,
|
||
connection_attempts: RwLock<HashMap<String, ConnectionAttemptState>>,
|
||
pub configs: RwLock<HashMap<String, ConnectionConfig>>,
|
||
pub running_queries: RunningQueries,
|
||
pub tunnels: TunnelManager,
|
||
pub proxy_tunnels: ProxyTunnelManager,
|
||
pub http_tunnels: HttpTunnelManager,
|
||
pub storage: Storage,
|
||
pub plugins: PluginRegistry,
|
||
pub plugin_host: PluginHost,
|
||
pub agent_manager: crate::agent_manager::AgentManager,
|
||
pub nacos_registry: crate::nacos::NacosAdminRegistry,
|
||
duckdb_worker_process_isolation: AtomicBool,
|
||
duckdb_worker_max_processes: AtomicUsize,
|
||
/// PostgreSQL TLS cancel context, keyed by pool_key.
|
||
/// Used to reconstruct a TLS connector compatible with the original connection when cancelling.
|
||
postgres_cancel_contexts: Arc<RwLock<HashMap<String, db::postgres::PostgresCancelContext>>>,
|
||
pub transaction_sessions: Arc<RwLock<HashMap<String, TransactionSession>>>,
|
||
/// `save_password=false` 连接本次运行期的临时密码(内存,进程退出即丢,
|
||
/// 绝不落盘)。键为 `(owner_scope, connection_id)`:桌面端 owner 为空串,
|
||
/// Web 端 owner 为已认证会话 token,不同登录会话互不可见。建池/池重建/
|
||
/// AI/元数据从它读取,前端通过状态接口查询。
|
||
pub session_credentials: SessionCredentialStore,
|
||
/// In-memory, never-persisted 1/5 minute write overrides for read-only connections.
|
||
pub write_unlock_windows: crate::write_unlock::WriteUnlockWindows,
|
||
metadata_gates: Arc<Mutex<HashMap<String, Arc<Semaphore>>>>,
|
||
mongo_oidc_browser_opener: std::sync::RwLock<Option<MongoOidcBrowserOpener>>,
|
||
#[cfg(feature = "mq-admin")]
|
||
pub mq_registry: crate::mq::MqAdminRegistry,
|
||
}
|
||
|
||
fn transport_layers_through_last_ssh(layers: &[TransportLayerConfig]) -> Result<&[TransportLayerConfig], String> {
|
||
let Some(last_ssh_index) = layers.iter().rposition(|layer| matches!(layer, TransportLayerConfig::Ssh(_))) else {
|
||
return Err("Connection has no enabled SSH tunnel layer".to_string());
|
||
};
|
||
Ok(&layers[..=last_ssh_index])
|
||
}
|
||
|
||
/// 活跃时间以进程内单调时钟的相对毫秒存储(AtomicU64):热路径每条查询都要
|
||
/// 更新它,读锁 + 原子写让并发查询不再在全局写锁上串行化。
|
||
/// 基准偏移让测试能构造"过去"的时间点(否则进程刚启动时相对毫秒接近 0)。
|
||
const POOL_ACTIVITY_BASE_MS: u64 = 86_400_000;
|
||
|
||
fn pool_activity_epoch() -> Instant {
|
||
static EPOCH: std::sync::OnceLock<Instant> = std::sync::OnceLock::new();
|
||
*EPOCH.get_or_init(Instant::now)
|
||
}
|
||
|
||
fn pool_activity_now_ms() -> u64 {
|
||
POOL_ACTIVITY_BASE_MS + pool_activity_epoch().elapsed().as_millis() as u64
|
||
}
|
||
|
||
#[cfg_attr(not(test), allow(dead_code))]
|
||
struct PoolActivity {
|
||
last_used_at_ms: std::sync::atomic::AtomicU64,
|
||
}
|
||
|
||
#[derive(Clone, Copy)]
|
||
struct ConnectionAttemptState {
|
||
server_attempt: u64,
|
||
client_attempt: Option<u64>,
|
||
}
|
||
|
||
pub(crate) fn metadata_concurrency_limit(db_type: DatabaseType, max_connections: usize) -> usize {
|
||
if db_type == DatabaseType::SqlServer {
|
||
METADATA_POOL_SQLSERVER_LIMIT
|
||
} else {
|
||
max_connections.saturating_sub(2).clamp(1, METADATA_POOL_DEFAULT_LIMIT)
|
||
}
|
||
}
|
||
|
||
pub(crate) fn uses_metadata_gate(db_type: DatabaseType) -> bool {
|
||
matches!(db_type, DatabaseType::Mysql | DatabaseType::Postgres | DatabaseType::SqlServer)
|
||
}
|
||
|
||
/// Session-scoped metadata gate allowance for a database type.
|
||
///
|
||
/// A non-empty `client_session_id` normally means the caller opened a
|
||
/// dedicated single-connection session pool (for example MySQL tab sessions),
|
||
/// so the gate allows only one in-flight metadata operation. PostgreSQL is the
|
||
/// exception: `get_or_create_pool_for_session` builds a fresh full-size
|
||
/// (10-connection) pool for each export session, so an allowance of 1 would
|
||
/// serialize the database-export metadata prefetch (4 concurrent DDL/column
|
||
/// lookups) and make three of them time out after METADATA_POOL_ACQUIRE_TIMEOUT
|
||
/// with "DBX metadata pool is busy". Returning the pool's real capacity here
|
||
/// flows into metadata_concurrency_limit, which caps the effective gate at
|
||
/// METADATA_POOL_DEFAULT_LIMIT (6) -- still enough for the 4-wide prefetch and
|
||
/// still isolated from the UI/base pool.
|
||
fn metadata_gate_session_allowance(db_type: DatabaseType) -> usize {
|
||
if db_type == DatabaseType::Postgres {
|
||
// Matches db::postgres::connect's Pool::builder().max_size(10).
|
||
10
|
||
} else {
|
||
// MySQL and other session-scoped pools are single-connection.
|
||
1
|
||
}
|
||
}
|
||
|
||
fn metadata_gate_key(
|
||
connection_id: &str,
|
||
database: Option<&str>,
|
||
db_type: DatabaseType,
|
||
client_session_id: Option<&str>,
|
||
) -> String {
|
||
let session = if db_type == DatabaseType::SqlServer {
|
||
""
|
||
} else {
|
||
client_session_id.map(str::trim).filter(|session| !session.is_empty()).unwrap_or_default()
|
||
};
|
||
format!("{connection_id}\0{}\0{}", database.unwrap_or_default(), session)
|
||
}
|
||
|
||
struct PoolDrainGuard {
|
||
pool_key: String,
|
||
draining_pools: Arc<std::sync::Mutex<HashMap<String, watch::Sender<bool>>>>,
|
||
signal: watch::Sender<bool>,
|
||
}
|
||
|
||
impl Drop for PoolDrainGuard {
|
||
fn drop(&mut self) {
|
||
self.draining_pools.lock().unwrap_or_else(|error| error.into_inner()).remove(&self.pool_key);
|
||
let _ = self.signal.send(false);
|
||
}
|
||
}
|
||
|
||
impl PoolActivity {
|
||
fn now() -> Self {
|
||
Self { last_used_at_ms: std::sync::atomic::AtomicU64::new(pool_activity_now_ms()) }
|
||
}
|
||
|
||
fn touch(&self) {
|
||
// fetch_max 防止"先算后写"交错导致时间戳倒退(A 算得 100 被挂起,
|
||
// B 写入 200 后 A 再写 100)
|
||
self.last_used_at_ms.fetch_max(pool_activity_now_ms(), std::sync::atomic::Ordering::Relaxed);
|
||
}
|
||
|
||
#[cfg_attr(not(test), allow(dead_code))]
|
||
fn elapsed(&self) -> Duration {
|
||
Duration::from_millis(
|
||
pool_activity_now_ms().saturating_sub(self.last_used_at_ms.load(std::sync::atomic::Ordering::Relaxed)),
|
||
)
|
||
}
|
||
|
||
#[cfg(test)]
|
||
fn idle_for(idle: Duration) -> Self {
|
||
Self {
|
||
last_used_at_ms: std::sync::atomic::AtomicU64::new(
|
||
pool_activity_now_ms().saturating_sub(idle.as_millis() as u64),
|
||
),
|
||
}
|
||
}
|
||
}
|
||
|
||
pub struct PoolActivityTouch {
|
||
pool_key: String,
|
||
connections: Arc<RwLock<ConnectionPoolRegistry>>,
|
||
pool_activity: Arc<RwLock<HashMap<String, PoolActivity>>>,
|
||
task_supervisor: TaskSupervisor,
|
||
}
|
||
|
||
#[derive(Clone)]
|
||
struct PoolRoutingControl {
|
||
connections: Arc<RwLock<ConnectionPoolRegistry>>,
|
||
pool_activity: Arc<RwLock<HashMap<String, PoolActivity>>>,
|
||
postgres_cancel_contexts: Arc<RwLock<HashMap<String, db::postgres::PostgresCancelContext>>>,
|
||
task_supervisor: TaskSupervisor,
|
||
}
|
||
|
||
pub struct ClientSessionPoolCleanupGuard {
|
||
pool_key: String,
|
||
routing: PoolRoutingControl,
|
||
armed: bool,
|
||
}
|
||
|
||
impl Drop for PoolActivityTouch {
|
||
fn drop(&mut self) {
|
||
let pool_key = self.pool_key.clone();
|
||
let connections = self.connections.clone();
|
||
let pool_activity = self.pool_activity.clone();
|
||
self.task_supervisor.spawn_replace(format!("pool-activity:{pool_key}"), move |_| async move {
|
||
if !connections.read().await.contains_key(&pool_key) {
|
||
return;
|
||
}
|
||
if let Some(activity) = pool_activity.read().await.get(&pool_key) {
|
||
activity.touch();
|
||
return;
|
||
}
|
||
pool_activity.write().await.insert(pool_key, PoolActivity::now());
|
||
});
|
||
}
|
||
}
|
||
|
||
impl ClientSessionPoolCleanupGuard {
|
||
pub(crate) fn disarm(&mut self) {
|
||
self.armed = false;
|
||
}
|
||
}
|
||
|
||
impl Drop for ClientSessionPoolCleanupGuard {
|
||
fn drop(&mut self) {
|
||
if !self.armed {
|
||
return;
|
||
}
|
||
let pool_key = self.pool_key.clone();
|
||
let routing = self.routing.clone();
|
||
routing.stop_keepalive(&pool_key);
|
||
let supervisor = routing.task_supervisor.clone();
|
||
supervisor.spawn_once(format!("client-session-cleanup:{pool_key}"), move |_| async move {
|
||
routing.detach_pool_by_key(&pool_key, false).await;
|
||
});
|
||
}
|
||
}
|
||
|
||
impl PoolRoutingControl {
|
||
fn stop_keepalive(&self, pool_key: &str) {
|
||
self.task_supervisor.stop(&format!("keepalive:{pool_key}"));
|
||
}
|
||
|
||
async fn detach_pool_by_key(&self, pool_key: &str, replace_agent_runtime: bool) -> bool {
|
||
let removed = {
|
||
let mut connections = self.connections.write().await;
|
||
let Some(pool) = connections.remove(pool_key) else {
|
||
return false;
|
||
};
|
||
let mut removed = vec![(pool_key.to_string(), pool)];
|
||
if replace_agent_runtime {
|
||
let sibling_keys = shared_runtime_sibling_keys(&connections, &removed[0].1);
|
||
let protects_manual_txn = sibling_keys.iter().any(|key| is_manual_transaction_pool_key(key))
|
||
|| is_manual_transaction_pool_key(pool_key);
|
||
if protects_manual_txn {
|
||
log::warn!(
|
||
"Skipping shared Agent runtime kill for '{pool_key}' because a manual-transaction session is active on the same runtime"
|
||
);
|
||
} else {
|
||
for key in sibling_keys {
|
||
if let Some(pool) = connections.remove(&key) {
|
||
removed.push((key, pool));
|
||
}
|
||
}
|
||
fail_stop_removed_agent_pool(pool_key, &removed[0].1);
|
||
}
|
||
}
|
||
removed
|
||
};
|
||
|
||
self.finish_detach(removed).await;
|
||
true
|
||
}
|
||
|
||
async fn detach_agent_pool_if_current(
|
||
&self,
|
||
pool_key: &str,
|
||
expected_client: &Arc<db::agent_driver::PooledAgentClient>,
|
||
expected_publication: Option<&PoolPublication>,
|
||
replace_agent_runtime: bool,
|
||
) -> bool {
|
||
let removed = {
|
||
let mut connections = self.connections.write().await;
|
||
let publication_is_current = expected_publication.is_none_or(|expected| {
|
||
connections.publications.get(pool_key).is_some_and(|current| current.is_same(expected))
|
||
});
|
||
if !publication_is_current {
|
||
return false;
|
||
}
|
||
let is_current = matches!(
|
||
connections.get(pool_key),
|
||
Some(PoolKind::Agent(current)) if Arc::ptr_eq(current, expected_client)
|
||
);
|
||
if !is_current && !replace_agent_runtime {
|
||
return false;
|
||
}
|
||
if replace_agent_runtime {
|
||
let runtime_keys = connections
|
||
.iter()
|
||
.filter_map(|(key, pool)| match pool {
|
||
PoolKind::Agent(client) if expected_client.shares_runtime_with(client) => Some(key.clone()),
|
||
_ => None,
|
||
})
|
||
.collect::<Vec<_>>();
|
||
// Sticky manual-transaction sessions pin work on the shared runtime.
|
||
// Prefer quarantining only the failed route over killing every sibling TX.
|
||
let protects_manual_txn = runtime_keys.iter().any(|key| is_manual_transaction_pool_key(key))
|
||
|| is_manual_transaction_pool_key(pool_key);
|
||
if protects_manual_txn {
|
||
log::warn!(
|
||
"Agent pool '{pool_key}' requested runtime replacement, but a manual-transaction session shares the runtime; detaching only this pool"
|
||
);
|
||
if is_current {
|
||
let pool = connections.remove(pool_key).expect("current Agent pool must still exist");
|
||
vec![(pool_key.to_string(), pool)]
|
||
} else {
|
||
Vec::new()
|
||
}
|
||
} else {
|
||
let removed = runtime_keys
|
||
.into_iter()
|
||
.filter_map(|key| connections.remove(&key).map(|pool| (key, pool)))
|
||
.collect::<Vec<_>>();
|
||
if !expected_client.fail_stop() {
|
||
log::warn!("Failed to terminate the shared Agent runtime while detaching pool '{pool_key}'");
|
||
}
|
||
removed
|
||
}
|
||
} else {
|
||
let pool = connections.remove(pool_key).expect("current Agent pool must still exist");
|
||
vec![(pool_key.to_string(), pool)]
|
||
}
|
||
};
|
||
|
||
let detached = !removed.is_empty();
|
||
self.finish_detach(removed).await;
|
||
detached
|
||
}
|
||
|
||
async fn finish_detach(&self, removed: Vec<(String, PoolKind)>) {
|
||
for (key, _) in &removed {
|
||
self.stop_keepalive(key);
|
||
}
|
||
{
|
||
let mut activity = self.pool_activity.write().await;
|
||
let mut cancel_contexts = self.postgres_cancel_contexts.write().await;
|
||
for (key, _) in &removed {
|
||
activity.remove(key);
|
||
cancel_contexts.remove(key);
|
||
}
|
||
}
|
||
self.close_removed_in_background(removed);
|
||
}
|
||
|
||
async fn close_pool_with_timeout(&self, pool_key: String, pool: PoolKind) {
|
||
let agent_client = match &pool {
|
||
PoolKind::Agent(client) => Some(client.clone()),
|
||
_ => None,
|
||
};
|
||
let protects_manual_txn = match agent_client.as_ref() {
|
||
Some(client) if client.uses_shared_runtime() => {
|
||
let connections = {
|
||
let registry = self.connections.read().await;
|
||
registry.pools.clone()
|
||
};
|
||
connections.iter().any(|(key, pool)| {
|
||
is_manual_transaction_pool_key(key)
|
||
&& matches!(pool, PoolKind::Agent(sibling) if client.shares_runtime_with(sibling))
|
||
})
|
||
}
|
||
_ => false,
|
||
};
|
||
match tokio::time::timeout(Duration::from_secs(POOL_CLOSE_TIMEOUT_SECS), close_pool_kind(pool)).await {
|
||
Ok(Ok(())) => {}
|
||
Ok(Err(error)) => {
|
||
log::warn!("Failed to close connection pool '{pool_key}': {error}");
|
||
if protects_manual_txn {
|
||
log::warn!(
|
||
"Leaving shared Agent runtime running after close failure for '{pool_key}' to protect sibling sessions"
|
||
);
|
||
} else if let Some(client) = agent_client.filter(|_| should_replace_agent_runtime(&error)) {
|
||
self.replace_runtime_after_close_failure(&pool_key, &client).await;
|
||
}
|
||
}
|
||
Err(_) => {
|
||
if protects_manual_txn {
|
||
log::warn!(
|
||
"Timed out closing shared Agent pool '{pool_key}' after {POOL_CLOSE_TIMEOUT_SECS}s; leaving the runtime running so sibling sessions (e.g. manual transactions) stay alive"
|
||
);
|
||
} else {
|
||
log::warn!(
|
||
"Timed out closing connection pool '{pool_key}' after {POOL_CLOSE_TIMEOUT_SECS}s; replacing a shared Agent runtime when present."
|
||
);
|
||
if let Some(client) = agent_client {
|
||
self.replace_runtime_after_close_failure(&pool_key, &client).await;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
async fn replace_runtime_after_close_failure(
|
||
&self,
|
||
closed_pool_key: &str,
|
||
failed_client: &Arc<db::agent_driver::PooledAgentClient>,
|
||
) {
|
||
let removed = {
|
||
let mut connections = self.connections.write().await;
|
||
let sibling_keys = connections
|
||
.iter()
|
||
.filter_map(|(key, pool)| match pool {
|
||
PoolKind::Agent(client) if failed_client.shares_runtime_with(client) => Some(key.clone()),
|
||
_ => None,
|
||
})
|
||
.collect::<Vec<_>>();
|
||
let removed = sibling_keys
|
||
.into_iter()
|
||
.filter_map(|key| connections.remove(&key).map(|pool| (key, pool)))
|
||
.collect::<Vec<_>>();
|
||
if !failed_client.fail_stop() {
|
||
log::warn!(
|
||
"Detached Agent pool '{closed_pool_key}' after close failure, but its legacy process could not be terminated without the client lock"
|
||
);
|
||
}
|
||
removed
|
||
};
|
||
self.finish_detach(removed).await;
|
||
}
|
||
|
||
async fn replace_runtime_after_open_failure(&self, failed_runtime: &Arc<db::agent_driver::AgentRuntimeClient>) {
|
||
let removed = {
|
||
let mut connections = self.connections.write().await;
|
||
let keys = connections
|
||
.iter()
|
||
.filter_map(|(key, pool)| match pool {
|
||
PoolKind::Agent(client) if client.uses_runtime(failed_runtime) => Some(key.clone()),
|
||
_ => None,
|
||
})
|
||
.collect::<Vec<_>>();
|
||
let removed =
|
||
keys.into_iter().filter_map(|key| connections.remove(&key).map(|pool| (key, pool))).collect::<Vec<_>>();
|
||
failed_runtime.kill();
|
||
removed
|
||
};
|
||
self.finish_detach(removed).await;
|
||
}
|
||
|
||
async fn close_removed(&self, removed: Vec<(String, PoolKind)>) {
|
||
for (pool_key, pool) in removed {
|
||
self.close_pool_with_timeout(pool_key, pool).await;
|
||
}
|
||
}
|
||
|
||
fn close_removed_in_background(&self, removed: Vec<(String, PoolKind)>) {
|
||
if removed.is_empty() {
|
||
return;
|
||
}
|
||
let pool_count = removed.len();
|
||
let routing = self.clone();
|
||
let task_key = format!("pool-close:{}", uuid::Uuid::new_v4());
|
||
if !self.task_supervisor.spawn_once(task_key, move |_| async move {
|
||
for (pool_key, pool) in removed {
|
||
routing.close_pool_with_timeout(pool_key, pool).await;
|
||
}
|
||
}) {
|
||
log::debug!("Dropped {pool_count} detached pool handle(s) during application shutdown");
|
||
}
|
||
}
|
||
}
|
||
|
||
fn shared_runtime_sibling_keys(connections: &HashMap<String, PoolKind>, source_pool: &PoolKind) -> Vec<String> {
|
||
let PoolKind::Agent(source_client) = source_pool else {
|
||
return Vec::new();
|
||
};
|
||
connections
|
||
.iter()
|
||
.filter_map(|(key, pool)| match pool {
|
||
PoolKind::Agent(client) if source_client.shares_runtime_with(client) => Some(key.clone()),
|
||
_ => None,
|
||
})
|
||
.collect()
|
||
}
|
||
|
||
fn fail_stop_removed_agent_pool(pool_key: &str, pool: &PoolKind) {
|
||
let PoolKind::Agent(client) = pool else {
|
||
return;
|
||
};
|
||
if !client.fail_stop() {
|
||
log::warn!(
|
||
"Detached busy legacy Agent pool '{pool_key}', but its process cannot be terminated without the client lock"
|
||
);
|
||
}
|
||
}
|
||
|
||
fn should_replace_agent_runtime(error: &str) -> bool {
|
||
crate::db::agent_driver::agent_recovery_decision(error, RecoveryScope::ConnectionOpen).replaces_runtime()
|
||
}
|
||
|
||
pub fn metadata_connection_config(config: &ConnectionConfig) -> ConnectionConfig {
|
||
let mut db_config = config.canonicalized();
|
||
if database_capabilities::is_metadata_connection_scoped(&db_config.db_type) {
|
||
db_config.database = None;
|
||
}
|
||
db_config
|
||
}
|
||
|
||
pub fn database_connection_config(config: &ConnectionConfig, database: Option<&str>) -> ConnectionConfig {
|
||
database_connection_config_with_catalog(config, database, None)
|
||
}
|
||
|
||
/// Like [`database_connection_config`], but optionally injects a Doris/StarRocks
|
||
/// `catalog=<name>` URL parameter so mysql_async emits `SET catalog` during
|
||
/// connection setup (before any `USE <database>`).
|
||
pub fn database_connection_config_with_catalog(
|
||
config: &ConnectionConfig,
|
||
database: Option<&str>,
|
||
catalog: Option<&str>,
|
||
) -> ConnectionConfig {
|
||
let mut db_config = if database.is_some() { config.clone() } else { metadata_connection_config(config) };
|
||
if let Some(db) = database {
|
||
if !matches!(
|
||
db_config.db_type,
|
||
DatabaseType::Oracle
|
||
| DatabaseType::Dameng
|
||
| DatabaseType::MongoDb
|
||
| DatabaseType::OceanbaseOracle
|
||
| DatabaseType::CloudflareD1
|
||
) {
|
||
db_config.database = Some(db.to_string());
|
||
}
|
||
}
|
||
if let Some(catalog) = catalog.map(str::trim).filter(|catalog| !catalog.is_empty()) {
|
||
db_config.url_params = Some(upsert_connection_url_param(db_config.url_params.as_deref(), "catalog", catalog));
|
||
}
|
||
db_config
|
||
}
|
||
|
||
/// Insert or replace a single `key=value` entry in a connection URL-params string.
|
||
pub fn upsert_connection_url_param(params: Option<&str>, key: &str, value: &str) -> String {
|
||
let key = key.trim();
|
||
let value = value.trim();
|
||
let key_lower = key.to_ascii_lowercase();
|
||
let encoded_value = percent_encoding::utf8_percent_encode(value, percent_encoding::NON_ALPHANUMERIC).to_string();
|
||
let mut parts: Vec<String> = params
|
||
.unwrap_or("")
|
||
.trim()
|
||
.trim_start_matches('?')
|
||
.split('&')
|
||
.filter(|part| !part.trim().is_empty())
|
||
.filter(|part| {
|
||
part.split_once('=')
|
||
.map(|(existing_key, _)| existing_key.trim().to_ascii_lowercase() != key_lower)
|
||
.unwrap_or(true)
|
||
})
|
||
.map(str::to_string)
|
||
.collect();
|
||
parts.push(format!("{key}={encoded_value}"));
|
||
parts.join("&")
|
||
}
|
||
|
||
pub fn prestosql_jdbc_config_for_endpoint(
|
||
config: &ConnectionConfig,
|
||
host: &str,
|
||
port: u16,
|
||
) -> Result<ConnectionConfig, String> {
|
||
let mut jdbc_config = config.clone();
|
||
jdbc_config.connection_string =
|
||
Some(trino_like_jdbc_connection_string(config, host, port, config.effective_database().unwrap_or(""))?);
|
||
jdbc_config.url_params = None;
|
||
if jdbc_config.jdbc_driver_class.as_deref().is_none_or(|value| value.trim().is_empty()) {
|
||
jdbc_config.jdbc_driver_class = Some(PRESTOSQL_JDBC_DRIVER_CLASS.to_string());
|
||
}
|
||
Ok(jdbc_config)
|
||
}
|
||
|
||
pub fn gaussdb_uses_m_jdbc_driver(config: &ConnectionConfig) -> bool {
|
||
config.db_type == DatabaseType::Gaussdb
|
||
&& config
|
||
.driver_profile
|
||
.as_deref()
|
||
.is_some_and(|profile| profile.eq_ignore_ascii_case(GAUSSDB_M_JDBC_DRIVER_PROFILE))
|
||
}
|
||
|
||
pub fn gaussdb_m_jdbc_config_for_endpoint(config: &ConnectionConfig, host: &str, port: u16) -> ConnectionConfig {
|
||
let mut jdbc_config = config.clone();
|
||
let mut jdbc_url = format!("jdbc:{}", config.redacted_connection_url_with_host(host, port));
|
||
let raw_params = config.url_params.as_deref().unwrap_or("").trim().trim_start_matches('?');
|
||
let explicit_sslmode = raw_params.split('&').find_map(|part| {
|
||
let (key, value) = part.split_once('=')?;
|
||
key.trim().eq_ignore_ascii_case("sslmode").then(|| value.trim().to_ascii_lowercase())
|
||
});
|
||
let explicit_ssl = raw_params.split('&').find_map(|part| {
|
||
let (key, value) = part.split_once('=')?;
|
||
key.trim().eq_ignore_ascii_case("ssl").then(|| value.trim().eq_ignore_ascii_case("true"))
|
||
});
|
||
let sslmode = explicit_sslmode.unwrap_or_else(|| {
|
||
if explicit_ssl == Some(false) {
|
||
"disable".to_string()
|
||
} else if config.ssl {
|
||
"require".to_string()
|
||
} else {
|
||
"prefer".to_string()
|
||
}
|
||
});
|
||
let params = upsert_connection_url_param(Some(raw_params), "sslmode", &sslmode);
|
||
let params = upsert_connection_url_param(Some(¶ms), "ssl", if sslmode == "disable" { "false" } else { "true" });
|
||
let target_server_type = config
|
||
.external_config
|
||
.as_ref()
|
||
.and_then(|ext| ext.get("gaussdbTargetServerType"))
|
||
.and_then(|v| v.as_str())
|
||
.filter(|v| *v == "master" || *v == "slave" || *v == "any");
|
||
let params = if let Some(value) = target_server_type {
|
||
upsert_connection_url_param(Some(¶ms), "targetServerType", value)
|
||
} else {
|
||
params
|
||
};
|
||
if !params.is_empty() {
|
||
jdbc_url.push('?');
|
||
jdbc_url.push_str(¶ms);
|
||
}
|
||
jdbc_config.connection_string = Some(jdbc_url);
|
||
jdbc_config.jdbc_driver_class = Some(GAUSSDB_M_JDBC_DRIVER_CLASS.to_string());
|
||
jdbc_config
|
||
}
|
||
|
||
pub fn sqlserver_legacy_agent_config(config: &ConnectionConfig) -> ConnectionConfig {
|
||
let mut legacy_config = config.clone();
|
||
legacy_config.driver_profile = Some(db::sqlserver::SQLSERVER_LEGACY_DRIVER_PROFILE.to_string());
|
||
legacy_config.driver_label = Some(db::sqlserver::SQLSERVER_LEGACY_DRIVER_LABEL.to_string());
|
||
legacy_config
|
||
}
|
||
|
||
pub fn sqlserver_uses_legacy_driver(config: &ConnectionConfig) -> bool {
|
||
config
|
||
.driver_profile
|
||
.as_deref()
|
||
.is_some_and(|profile| profile.eq_ignore_ascii_case(db::sqlserver::SQLSERVER_LEGACY_DRIVER_PROFILE))
|
||
}
|
||
|
||
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.
|
||
if agent_error.contains("driver is not installed") {
|
||
crate::db::agent_driver::append_legacy_error_context(agent_error, SQLSERVER_LEGACY_DRIVER_INSTALL_HINT)
|
||
} else {
|
||
agent_error.to_string()
|
||
}
|
||
}
|
||
|
||
pub async fn connect_mysql_metadata_pool(
|
||
config: &ConnectionConfig,
|
||
db_config: &ConnectionConfig,
|
||
host: &str,
|
||
port: u16,
|
||
connect_timeout: std::time::Duration,
|
||
max_connections: usize,
|
||
) -> Result<(db::mysql::MySqlPool, MysqlMode), String> {
|
||
let url = connection_url_for_endpoint(db_config, host, port);
|
||
let idle_timeout_secs = Some(db_config.idle_timeout_secs);
|
||
let extra_setup_queries = mysql_pool_setup_queries(db_config, &url);
|
||
if db_config.needs_bare_mysql() {
|
||
return match connect_bare_mysql_pool_with_setup(
|
||
db_config,
|
||
&url,
|
||
connect_timeout,
|
||
max_connections,
|
||
&extra_setup_queries,
|
||
)
|
||
.await
|
||
{
|
||
Ok(pool) => Ok((pool, MysqlMode::Bare)),
|
||
Err(err) => {
|
||
let fallback_url = mysql_metadata_fallback_url(config, db_config, host, port);
|
||
if let Some(fallback_url) = fallback_url {
|
||
log::info!(
|
||
"MySQL metadata connection without a default database failed ({err}); retrying with configured default database."
|
||
);
|
||
connect_bare_mysql_pool_with_setup(
|
||
db_config,
|
||
&fallback_url,
|
||
connect_timeout,
|
||
max_connections,
|
||
&extra_setup_queries,
|
||
)
|
||
.await
|
||
.map(|pool| (pool, MysqlMode::Bare))
|
||
} else if let Some(db) = db_config.effective_database() {
|
||
let mut unscoped_config = db_config.clone();
|
||
unscoped_config.database = None;
|
||
let unscoped_url = connection_url_for_endpoint(&unscoped_config, host, port);
|
||
log::info!("MySQL connection with database in URL failed ({err}); retrying without database in URL and using USE statement.");
|
||
connect_bare_mysql_pool_with_setup_database(
|
||
&unscoped_config,
|
||
&unscoped_url,
|
||
connect_timeout,
|
||
max_connections,
|
||
db,
|
||
&extra_setup_queries,
|
||
)
|
||
.await
|
||
.map(|pool| (pool, MysqlMode::Bare))
|
||
} else {
|
||
Err(err)
|
||
}
|
||
}
|
||
};
|
||
}
|
||
|
||
match db::mysql::connect_with_ca_cert_pool_limit_idle_and_setup(
|
||
&url,
|
||
Some(&db_config.ca_cert_path),
|
||
connect_timeout,
|
||
max_connections,
|
||
idle_timeout_secs,
|
||
&extra_setup_queries,
|
||
)
|
||
.await
|
||
{
|
||
Ok(pool) => {
|
||
let mode = detect_ob_oracle_mode(config, &pool).await;
|
||
Ok((pool, mode))
|
||
}
|
||
Err(err) => {
|
||
let fallback_url = mysql_metadata_fallback_url(config, db_config, host, port);
|
||
if let Some(fallback_url) = fallback_url {
|
||
log::info!(
|
||
"MySQL metadata connection without a default database failed ({err}); retrying with configured default database."
|
||
);
|
||
let pool = db::mysql::connect_with_ca_cert_pool_limit_idle_and_setup(
|
||
&fallback_url,
|
||
Some(&config.ca_cert_path),
|
||
connect_timeout,
|
||
max_connections,
|
||
idle_timeout_secs,
|
||
&extra_setup_queries,
|
||
)
|
||
.await?;
|
||
let mode = detect_ob_oracle_mode(config, &pool).await;
|
||
Ok((pool, mode))
|
||
} else if let Some(db) = db_config.effective_database() {
|
||
let mut unscoped_config = db_config.clone();
|
||
unscoped_config.database = None;
|
||
let unscoped_url = connection_url_for_endpoint(&unscoped_config, host, port);
|
||
log::info!("MySQL connection with database in URL failed ({err}); retrying without database in URL and using USE statement.");
|
||
let pool = db::mysql::connect_with_ca_cert_pool_limit_idle_and_setup_database(
|
||
&unscoped_url,
|
||
Some(&config.ca_cert_path),
|
||
connect_timeout,
|
||
max_connections,
|
||
idle_timeout_secs,
|
||
Some(db),
|
||
&extra_setup_queries,
|
||
)
|
||
.await?;
|
||
let mode = detect_ob_oracle_mode(config, &pool).await;
|
||
Ok((pool, mode))
|
||
} else {
|
||
Err(err)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
pub async fn connect_bare_metadata_pool(
|
||
db_config: &ConnectionConfig,
|
||
host: &str,
|
||
port: u16,
|
||
connect_timeout: std::time::Duration,
|
||
max_connections: usize,
|
||
) -> Result<db::mysql::MySqlPool, String> {
|
||
let url = connection_url_for_endpoint(db_config, host, port);
|
||
let extra_setup_queries = mysql_pool_setup_queries(db_config, &url);
|
||
if db_config.effective_database().is_none() {
|
||
return connect_bare_mysql_pool_with_setup(
|
||
db_config,
|
||
&url,
|
||
connect_timeout,
|
||
max_connections,
|
||
&extra_setup_queries,
|
||
)
|
||
.await;
|
||
}
|
||
|
||
let mut unscoped_config = db_config.clone();
|
||
unscoped_config.database = None;
|
||
let unscoped_url = connection_url_for_endpoint(&unscoped_config, host, port);
|
||
if unscoped_url == url {
|
||
return connect_bare_mysql_pool_with_setup(
|
||
db_config,
|
||
&url,
|
||
connect_timeout,
|
||
max_connections,
|
||
&extra_setup_queries,
|
||
)
|
||
.await;
|
||
}
|
||
|
||
let preferred =
|
||
connect_bare_mysql_pool_with_setup(db_config, &url, connect_timeout, max_connections, &extra_setup_queries);
|
||
let unscoped = connect_bare_mysql_pool_with_setup(
|
||
db_config,
|
||
&unscoped_url,
|
||
connect_timeout,
|
||
max_connections,
|
||
&extra_setup_queries,
|
||
);
|
||
tokio::pin!(preferred);
|
||
tokio::pin!(unscoped);
|
||
|
||
tokio::select! {
|
||
result = &mut preferred => match result {
|
||
Ok(pool) => Ok(pool),
|
||
Err(preferred_err) => match (&mut unscoped).await {
|
||
Ok(pool) => Ok(pool),
|
||
Err(unscoped_err) => Err(format!(
|
||
"Connection with the configured database failed: {preferred_err}\n\nConnection without a default database also failed: {unscoped_err}"
|
||
)),
|
||
},
|
||
},
|
||
result = &mut unscoped => match result {
|
||
Ok(pool) => Ok(pool),
|
||
Err(unscoped_err) => match (&mut preferred).await {
|
||
Ok(pool) => Ok(pool),
|
||
Err(preferred_err) => Err(format!(
|
||
"Connection with the configured database failed: {preferred_err}\n\nConnection without a default database also failed: {unscoped_err}"
|
||
)),
|
||
},
|
||
},
|
||
}
|
||
}
|
||
|
||
async fn connect_bare_mysql_pool_with_setup(
|
||
db_config: &ConnectionConfig,
|
||
url: &str,
|
||
connect_timeout: std::time::Duration,
|
||
max_connections: usize,
|
||
extra_setup_queries: &[String],
|
||
) -> Result<db::mysql::MySqlPool, String> {
|
||
if db_config.bare_mysql_uses_tls() {
|
||
let idle_timeout_secs = Some(db_config.idle_timeout_secs);
|
||
db::mysql::connect_compatible_with_ca_cert_pool_limit_idle_and_setup(
|
||
url,
|
||
Some(&db_config.ca_cert_path),
|
||
connect_timeout,
|
||
max_connections,
|
||
idle_timeout_secs,
|
||
extra_setup_queries,
|
||
)
|
||
.await
|
||
} else {
|
||
db::mysql::connect_bare_with_pool_limit_and_setup(url, connect_timeout, max_connections, extra_setup_queries)
|
||
.await
|
||
}
|
||
}
|
||
|
||
async fn connect_bare_mysql_pool_with_setup_database(
|
||
db_config: &ConnectionConfig,
|
||
url: &str,
|
||
connect_timeout: std::time::Duration,
|
||
max_connections: usize,
|
||
setup_database: &str,
|
||
extra_setup_queries: &[String],
|
||
) -> Result<db::mysql::MySqlPool, String> {
|
||
// Some MySQL proxies reject the default database in the handshake; pass it
|
||
// separately so DB-layer setup keeps the normal charset/catalog/USE order.
|
||
if db_config.bare_mysql_uses_tls() {
|
||
let idle_timeout_secs = Some(db_config.idle_timeout_secs);
|
||
db::mysql::connect_compatible_with_ca_cert_pool_limit_idle_and_setup_database(
|
||
url,
|
||
Some(&db_config.ca_cert_path),
|
||
connect_timeout,
|
||
max_connections,
|
||
idle_timeout_secs,
|
||
Some(setup_database),
|
||
extra_setup_queries,
|
||
)
|
||
.await
|
||
} else {
|
||
db::mysql::connect_bare_with_pool_limit_and_setup_database(
|
||
url,
|
||
connect_timeout,
|
||
max_connections,
|
||
Some(setup_database),
|
||
extra_setup_queries,
|
||
)
|
||
.await
|
||
}
|
||
}
|
||
|
||
fn mysql_metadata_fallback_url(
|
||
config: &ConnectionConfig,
|
||
db_config: &ConnectionConfig,
|
||
host: &str,
|
||
port: u16,
|
||
) -> Option<String> {
|
||
if db_config.db_type != DatabaseType::Mysql || db_config.effective_database().is_some() {
|
||
return None;
|
||
}
|
||
config.effective_database()?;
|
||
Some(connection_url_for_endpoint(config, host, port))
|
||
}
|
||
|
||
impl AppState {
|
||
/// Return an owned pool handle. The registry read lock is released before
|
||
/// the caller can perform any asynchronous database operation.
|
||
pub async fn pool_handle(&self, pool_key: &str) -> Option<PoolKind> {
|
||
self.connections.read().await.get(pool_key).cloned()
|
||
}
|
||
|
||
async fn pool_publication_snapshot(&self, pool_key: &str) -> Option<PoolPublicationSnapshot> {
|
||
self.connections.read().await.snapshot(pool_key)
|
||
}
|
||
|
||
async fn connection_pool_publication_snapshots(&self) -> Vec<(String, PoolPublicationSnapshot)> {
|
||
let connections = self.connections.read().await;
|
||
connections
|
||
.pools
|
||
.keys()
|
||
.filter_map(|pool_key| connections.snapshot(pool_key).map(|snapshot| (pool_key.clone(), snapshot)))
|
||
.collect()
|
||
}
|
||
|
||
/// Return an owned snapshot for operations that need to inspect multiple
|
||
/// entries. Cloning handles is cheap and prevents registry guards from
|
||
/// leaking into asynchronous database work.
|
||
pub async fn connection_pools_snapshot(&self) -> HashMap<String, PoolKind> {
|
||
self.connections.read().await.pools.clone()
|
||
}
|
||
|
||
/// Inspect the registry while holding its read lock. The callback is
|
||
/// deliberately synchronous so no database I/O can run under the lock.
|
||
pub async fn with_connection_pools<R>(&self, inspect: impl FnOnce(&HashMap<String, PoolKind>) -> R) -> R {
|
||
let connections = self.connections.read().await;
|
||
inspect(&connections.pools)
|
||
}
|
||
|
||
/// Mutate the registry atomically. The callback is deliberately
|
||
/// synchronous; asynchronous cleanup must use values returned from it.
|
||
pub async fn update_connection_pools<R>(&self, update: impl FnOnce(&mut ConnectionPoolRegistry) -> R) -> R {
|
||
let mut connections = self.connections.write().await;
|
||
update(&mut connections)
|
||
}
|
||
|
||
fn pool_routing_control(&self) -> PoolRoutingControl {
|
||
PoolRoutingControl {
|
||
connections: self.connections.clone(),
|
||
pool_activity: self.pool_activity.clone(),
|
||
postgres_cancel_contexts: self.postgres_cancel_contexts.clone(),
|
||
task_supervisor: self.task_supervisor.clone(),
|
||
}
|
||
}
|
||
|
||
/// save_password=false 连接:持久化/运行态 config 的 password 恒为空,从本次
|
||
/// 运行期会话凭据仓库补主密码,使手动/编辑器/AI/元数据/池重建复用首次输入,
|
||
/// 不再反复弹窗。会话凭据只在内存中,进程退出即丢;"断开并忘记"后仓库为空,
|
||
/// 此处自然回到空密码(需重新输入)。
|
||
///
|
||
/// owner 作用域取自 [`crate::session_credentials::current_credential_owner`]:
|
||
/// Web 请求由鉴权中间件注入会话 token,桌面端/后台任务为空串。这样不同登录
|
||
/// 会话只会读取各自输入的密码,不会跨会话复用。
|
||
pub fn apply_session_credential(
|
||
&self,
|
||
config: &ConnectionConfig,
|
||
db_config: &mut ConnectionConfig,
|
||
connection_id: &str,
|
||
) {
|
||
if !config.save_password && db_config.password.is_empty() {
|
||
let owner = crate::session_credentials::current_credential_owner().unwrap_or_default();
|
||
if let Some(session_password) = self.session_credentials.get(&owner, connection_id) {
|
||
db_config.password = session_password;
|
||
}
|
||
}
|
||
}
|
||
|
||
/// no-save 连接复用已有池时,若该池由"另一个 owner 会话"创建,则不应复用
|
||
/// (否则会话 B 会直接使用会话 A 输入的临时密码建立的连接)。返回 `true`
|
||
/// 表示需要销毁旧池、以当前 owner 的凭据重建。
|
||
async fn pool_credential_owner_mismatch(&self, config: &ConnectionConfig, pool_key: &str) -> bool {
|
||
if config.save_password {
|
||
return false;
|
||
}
|
||
let owner = crate::session_credentials::current_credential_owner().unwrap_or_default();
|
||
if owner.is_empty() {
|
||
// 桌面端/未注入 owner:单用户场景,按既有逻辑复用。
|
||
return false;
|
||
}
|
||
self.session_credentials.pool_owner_mismatch(pool_key, &owner)
|
||
}
|
||
|
||
async fn handle_shared_connection_open_error(
|
||
&self,
|
||
error: crate::agent_runtime::SharedConnectionOpenError,
|
||
) -> String {
|
||
if let Some(runtime) = error.runtime.as_ref() {
|
||
self.pool_routing_control().replace_runtime_after_open_failure(runtime).await;
|
||
}
|
||
error.message
|
||
}
|
||
|
||
async fn spawn_routed_shared_agent_client(
|
||
&self,
|
||
db_type: &DatabaseType,
|
||
driver_profile: Option<&str>,
|
||
extra_java_args: &[String],
|
||
agent_session_id: String,
|
||
connect_params: serde_json::Value,
|
||
connect_timeout: Duration,
|
||
) -> Result<db::agent_driver::AgentDriverClient, String> {
|
||
match self
|
||
.agent_manager
|
||
.spawn_shared_connection_client(
|
||
db_type,
|
||
driver_profile,
|
||
extra_java_args,
|
||
agent_session_id,
|
||
connect_params,
|
||
connect_timeout,
|
||
)
|
||
.await
|
||
{
|
||
Ok(client) => Ok(client),
|
||
Err(error) => Err(self.handle_shared_connection_open_error(error).await),
|
||
}
|
||
}
|
||
|
||
pub fn new(storage: Storage) -> Self {
|
||
Self::new_with_plugin_dir(storage, default_plugin_dir())
|
||
}
|
||
|
||
pub fn new_with_plugin_dir(storage: Storage, plugin_dir: PathBuf) -> Self {
|
||
Self::new_with_plugin_dir_and_app_version(storage, plugin_dir, env!("CARGO_PKG_VERSION"))
|
||
}
|
||
|
||
pub fn new_with_plugin_dir_and_app_version(
|
||
storage: Storage,
|
||
plugin_dir: PathBuf,
|
||
app_version: impl Into<String>,
|
||
) -> Self {
|
||
Self::new_with_plugin_and_agent_dir_and_app_version(storage, plugin_dir, default_agent_dir(), app_version)
|
||
}
|
||
|
||
pub fn new_with_plugin_and_agent_dir_and_app_version(
|
||
storage: Storage,
|
||
plugin_dir: PathBuf,
|
||
agent_dir: PathBuf,
|
||
app_version: impl Into<String>,
|
||
) -> Self {
|
||
let data_dir = storage.data_dir().to_path_buf();
|
||
let app_version = app_version.into();
|
||
Self {
|
||
connections: Arc::new(RwLock::new(ConnectionPoolRegistry::new())),
|
||
task_supervisor: TaskSupervisor::new(),
|
||
pool_activity: Arc::new(RwLock::new(HashMap::new())),
|
||
draining_pools: Arc::new(std::sync::Mutex::new(HashMap::new())),
|
||
connection_attempts: RwLock::new(HashMap::new()),
|
||
configs: RwLock::new(HashMap::new()),
|
||
running_queries: RunningQueries::default(),
|
||
tunnels: TunnelManager::new(data_dir),
|
||
proxy_tunnels: ProxyTunnelManager::new(),
|
||
http_tunnels: HttpTunnelManager::new(),
|
||
storage,
|
||
plugins: PluginRegistry::new_with_app_version(plugin_dir.clone(), app_version.clone()),
|
||
plugin_host: PluginHost::new(PluginRegistry::new_with_app_version(plugin_dir, app_version.clone())),
|
||
agent_manager: crate::agent_manager::AgentManager::new_with_base_dir_and_app_version(
|
||
agent_dir,
|
||
app_version,
|
||
),
|
||
nacos_registry: crate::nacos::NacosAdminRegistry::new(),
|
||
duckdb_worker_process_isolation: AtomicBool::new(false),
|
||
duckdb_worker_max_processes: AtomicUsize::new(DUCKDB_WORKER_MAX_PROCESSES_DEFAULT),
|
||
postgres_cancel_contexts: Arc::new(RwLock::new(HashMap::new())),
|
||
transaction_sessions: Arc::new(RwLock::new(HashMap::new())),
|
||
session_credentials: SessionCredentialStore::new(),
|
||
write_unlock_windows: crate::write_unlock::WriteUnlockWindows::default(),
|
||
metadata_gates: Arc::new(Mutex::new(HashMap::new())),
|
||
mongo_oidc_browser_opener: std::sync::RwLock::new(None),
|
||
#[cfg(feature = "mq-admin")]
|
||
mq_registry: crate::mq::MqAdminRegistry::new(),
|
||
}
|
||
}
|
||
|
||
pub fn set_mongo_oidc_browser_opener(&self, opener: MongoOidcBrowserOpener) {
|
||
*self.mongo_oidc_browser_opener.write().expect("MongoDB OIDC browser opener lock poisoned") = Some(opener);
|
||
}
|
||
|
||
pub fn mongo_oidc_browser_opener(&self) -> Option<MongoOidcBrowserOpener> {
|
||
self.mongo_oidc_browser_opener.read().expect("MongoDB OIDC browser opener lock poisoned").clone()
|
||
}
|
||
|
||
pub(crate) async fn acquire_metadata_permit(
|
||
&self,
|
||
connection_id: &str,
|
||
database: Option<&str>,
|
||
db_type: DatabaseType,
|
||
client_session_id: Option<&str>,
|
||
) -> Result<OwnedSemaphorePermit, String> {
|
||
let key = metadata_gate_key(connection_id, database, db_type, client_session_id);
|
||
let has_session = client_session_id.map(str::trim).is_some_and(|session| !session.is_empty());
|
||
let max_connections = if has_session { metadata_gate_session_allowance(db_type) } else { 10 };
|
||
let limit = metadata_concurrency_limit(db_type, max_connections);
|
||
let gate = {
|
||
let mut gates = self.metadata_gates.lock().await;
|
||
gates.entry(key.clone()).or_insert_with(|| Arc::new(Semaphore::new(limit))).clone()
|
||
};
|
||
let started = Instant::now();
|
||
let queued = gate.available_permits() == 0;
|
||
let permit = match tokio::time::timeout(METADATA_POOL_ACQUIRE_TIMEOUT, gate.acquire_owned()).await {
|
||
Ok(Ok(permit)) => permit,
|
||
Ok(Err(_)) => return Err(crate::query::METADATA_POOL_BUSY_ERROR.to_string()),
|
||
Err(_) => {
|
||
log::warn!(
|
||
"[metadata:pool:busy] connection_id={} database={} wait_ms={} limit={}",
|
||
connection_id,
|
||
database.unwrap_or_default(),
|
||
started.elapsed().as_millis(),
|
||
limit
|
||
);
|
||
return Err(crate::query::METADATA_POOL_BUSY_ERROR.to_string());
|
||
}
|
||
};
|
||
if queued {
|
||
log::debug!(
|
||
"[metadata:pool:acquired] connection_id={} database={} wait_ms={} limit={}",
|
||
connection_id,
|
||
database.unwrap_or_default(),
|
||
started.elapsed().as_millis(),
|
||
limit
|
||
);
|
||
}
|
||
Ok(permit)
|
||
}
|
||
|
||
async fn clear_metadata_gates_for_connection(&self, connection_id: &str) {
|
||
let prefix = format!("{connection_id}\0");
|
||
self.metadata_gates.lock().await.retain(|key, _| !key.starts_with(&prefix));
|
||
}
|
||
|
||
async fn clear_metadata_gate_for_database(&self, connection_id: &str, database: Option<&str>) {
|
||
let prefix = format!("{connection_id}\0{}\0", database.unwrap_or_default());
|
||
self.metadata_gates.lock().await.retain(|key, _| !key.starts_with(&prefix));
|
||
}
|
||
|
||
pub fn jdbc_unavailable_error(&self) -> String {
|
||
match self.plugins.find_driver("jdbc") {
|
||
Ok(Some(_)) => "JDBC plugin is installed, but the connection could not be opened.".to_string(),
|
||
Ok(None) => JDBC_PLUGIN_NOT_INSTALLED.to_string(),
|
||
Err(err) => format!("Failed to inspect JDBC plugin: {err}"),
|
||
}
|
||
}
|
||
|
||
pub fn set_duckdb_worker_process_isolation_enabled(&self, enabled: bool) {
|
||
self.duckdb_worker_process_isolation.store(enabled, Ordering::Relaxed);
|
||
}
|
||
|
||
pub fn set_duckdb_worker_max_processes(&self, max_processes: usize) {
|
||
self.duckdb_worker_max_processes.store(normalize_duckdb_worker_max_processes(max_processes), Ordering::Relaxed);
|
||
}
|
||
|
||
pub async fn apply_duckdb_worker_process_isolation(&self, enabled: bool) {
|
||
let previous = self.duckdb_worker_process_isolation.swap(enabled, Ordering::Relaxed);
|
||
if previous != enabled {
|
||
self.remove_duckdb_pools_detached().await;
|
||
}
|
||
}
|
||
|
||
#[cfg(feature = "duckdb-sidecar")]
|
||
pub async fn create_duckdb_pool(&self, config: &ConnectionConfig) -> Result<PoolKind, String> {
|
||
let attached_databases = config
|
||
.attached_databases
|
||
.iter()
|
||
.map(|attached| crate::models::connection::AttachedDatabaseConfig {
|
||
name: attached.name.clone(),
|
||
path: expand_tilde(&attached.path),
|
||
})
|
||
.collect();
|
||
let path = expand_tilde(&config.host);
|
||
let init_script = config.init_script.clone();
|
||
let process_limit = self.duckdb_worker_max_processes.load(Ordering::Relaxed);
|
||
let installed_driver = self.agent_manager.driver_native_path("duckdb");
|
||
let has_driver_override =
|
||
std::env::var_os(db::duckdb_worker_process::DUCKDB_DRIVER_PATH_ENV).is_some_and(|value| !value.is_empty());
|
||
let client = if installed_driver.is_file() && !has_driver_override {
|
||
db::duckdb_worker_process::DuckDbWorkerClient::open_with_executable_and_process_limit(
|
||
installed_driver,
|
||
path,
|
||
attached_databases,
|
||
init_script,
|
||
process_limit,
|
||
)
|
||
.await?
|
||
} else {
|
||
db::duckdb_worker_process::DuckDbWorkerClient::open_with_process_limit(
|
||
path,
|
||
attached_databases,
|
||
init_script,
|
||
process_limit,
|
||
)
|
||
.await?
|
||
};
|
||
Ok(PoolKind::DuckDbWorker(Arc::new(client)))
|
||
}
|
||
|
||
#[cfg(feature = "duckdb-sidecar")]
|
||
pub async fn test_duckdb_connection_config(&self, config: &ConnectionConfig) -> Result<(), String> {
|
||
// Test the submitted form as a fresh session so unsaved ATTACH/init
|
||
// changes cannot be masked by a pool created from older settings.
|
||
let pool = self.create_duckdb_pool(config).await?;
|
||
close_pool_kind(pool).await?;
|
||
Ok(())
|
||
}
|
||
|
||
pub async fn test_external_driver(&self, driver_id: &str, config: &ConnectionConfig) -> Result<String, String> {
|
||
self.test_external_driver_with_info(driver_id, config).await.map(|result| result.message)
|
||
}
|
||
|
||
pub async fn test_external_driver_with_info(
|
||
&self,
|
||
driver_id: &str,
|
||
config: &ConnectionConfig,
|
||
) -> Result<ConnectionTestResult, String> {
|
||
let params = serde_json::json!({ "connection": config });
|
||
let env = self.external_driver_runtime_env(driver_id)?;
|
||
let response = self
|
||
.plugins
|
||
.invoke_driver_with_env_and_timeout::<serde_json::Value>(
|
||
driver_id,
|
||
"testConnection",
|
||
params,
|
||
env,
|
||
Some(external_driver_connect_timeout(config)),
|
||
)
|
||
.await?;
|
||
Ok(ConnectionTestResult::success("Connection successful")
|
||
.with_database_info(database_info_from_protocol_value(&response)))
|
||
}
|
||
|
||
pub async fn external_driver_pool(&self, driver_id: &str, config: &ConnectionConfig) -> Result<PoolKind, String> {
|
||
let env = self.external_driver_runtime_env(driver_id)?;
|
||
let session = self.plugins.start_driver_session_with_env(driver_id, env).await?;
|
||
let params = serde_json::json!({ "connection": config });
|
||
session
|
||
.invoke_with_timeout::<serde_json::Value>("connect", params, Some(external_driver_connect_timeout(config)))
|
||
.await?;
|
||
Ok(PoolKind::ExternalDriver { driver_id: driver_id.to_string(), config: Arc::new(config.clone()), session })
|
||
}
|
||
|
||
pub async fn test_sqlserver_connection(
|
||
&self,
|
||
config: &ConnectionConfig,
|
||
host: &str,
|
||
port: u16,
|
||
connect_timeout: Duration,
|
||
) -> Result<String, String> {
|
||
self.test_sqlserver_connection_with_info(config, host, port, connect_timeout).await.map(|result| result.message)
|
||
}
|
||
|
||
pub async fn test_sqlserver_connection_with_info(
|
||
&self,
|
||
config: &ConnectionConfig,
|
||
host: &str,
|
||
port: u16,
|
||
connect_timeout: Duration,
|
||
) -> Result<ConnectionTestResult, String> {
|
||
if sqlserver_uses_legacy_driver(config) {
|
||
let legacy_config = sqlserver_legacy_agent_config(config);
|
||
let connect_params =
|
||
agent_connect_params(&legacy_config, host, port, legacy_config.effective_database().unwrap_or(""))?;
|
||
let mut client = self
|
||
.agent_manager
|
||
.spawn(&legacy_config.db_type, legacy_config.driver_profile.as_deref())
|
||
.await
|
||
.map_err(|err| sqlserver_legacy_driver_error(&err))?;
|
||
let response = client
|
||
.call_method_with_timeout::<serde_json::Value>(
|
||
AgentMethod::TestConnection,
|
||
connect_params,
|
||
Some(agent_connect_timeout(&legacy_config)),
|
||
)
|
||
.await
|
||
.map_err(|err| sqlserver_legacy_driver_error(&err))?;
|
||
client.disconnect().await.ok();
|
||
return Ok(ConnectionTestResult::success(
|
||
"Connection successful (via SQL Server legacy compatibility driver)",
|
||
)
|
||
.with_database_info(database_info_from_protocol_value(&response)));
|
||
}
|
||
|
||
db::sqlserver::connect_with_port_explicit(
|
||
host,
|
||
port,
|
||
config.sqlserver_port_explicit(),
|
||
&config.username,
|
||
&config.password,
|
||
config.database.as_deref(),
|
||
connect_timeout,
|
||
)
|
||
.await?;
|
||
Ok(ConnectionTestResult::success("Connection successful"))
|
||
}
|
||
|
||
pub async fn connect_sqlserver_pool(
|
||
&self,
|
||
config: &ConnectionConfig,
|
||
host: &str,
|
||
port: u16,
|
||
connect_timeout: Duration,
|
||
) -> Result<PoolKind, String> {
|
||
if sqlserver_uses_legacy_driver(config) {
|
||
let legacy_config = sqlserver_legacy_agent_config(config);
|
||
let connect_params =
|
||
agent_connect_params(&legacy_config, host, port, legacy_config.effective_database().unwrap_or(""))?;
|
||
let mut client = self
|
||
.agent_manager
|
||
.spawn(&legacy_config.db_type, legacy_config.driver_profile.as_deref())
|
||
.await
|
||
.map_err(|err| sqlserver_legacy_driver_error(&err))?;
|
||
client
|
||
.call_method_with_timeout::<serde_json::Value>(
|
||
AgentMethod::Connect,
|
||
connect_params,
|
||
Some(agent_connect_timeout(&legacy_config)),
|
||
)
|
||
.await
|
||
.map_err(|err| sqlserver_legacy_driver_error(&err))?;
|
||
return Ok(PoolKind::agent(client));
|
||
}
|
||
|
||
let client = db::sqlserver::connect_with_port_explicit(
|
||
host,
|
||
port,
|
||
config.sqlserver_port_explicit(),
|
||
&config.username,
|
||
&config.password,
|
||
config.database.as_deref(),
|
||
connect_timeout,
|
||
)
|
||
.await?;
|
||
Ok(PoolKind::SqlServer(Arc::new(tokio::sync::Mutex::new(client))))
|
||
}
|
||
|
||
pub fn external_driver_runtime_env(&self, driver_id: &str) -> Result<PluginRuntimeEnv, String> {
|
||
if driver_id != "jdbc" {
|
||
return Ok(PluginRuntimeEnv::default());
|
||
}
|
||
let state = self.agent_manager.load_state();
|
||
if state.java_runtime.mode == JavaRuntimeMode::Managed && !self.agent_manager.is_jre_installed(DEFAULT_JRE_KEY)
|
||
{
|
||
return Ok(PluginRuntimeEnv::default());
|
||
}
|
||
let java = self.agent_manager.resolve_java_runtime(&state, DEFAULT_JRE_KEY)?;
|
||
Ok(PluginRuntimeEnv::default().with_var("DBX_JAVA_BIN", java.to_string_lossy().to_string()))
|
||
}
|
||
|
||
fn begin_pool_drain(&self, pool_key: &str) -> Option<PoolDrainGuard> {
|
||
let mut draining = self.draining_pools.lock().unwrap_or_else(|error| error.into_inner());
|
||
if draining.contains_key(pool_key) {
|
||
return None;
|
||
}
|
||
let (signal, _) = watch::channel(true);
|
||
draining.insert(pool_key.to_string(), signal.clone());
|
||
Some(PoolDrainGuard { pool_key: pool_key.to_string(), draining_pools: self.draining_pools.clone(), signal })
|
||
}
|
||
|
||
async fn wait_for_pool_drain(&self, pool_key: &str) {
|
||
loop {
|
||
let receiver = self
|
||
.draining_pools
|
||
.lock()
|
||
.unwrap_or_else(|error| error.into_inner())
|
||
.get(pool_key)
|
||
.map(watch::Sender::subscribe);
|
||
let Some(mut receiver) = receiver else {
|
||
return;
|
||
};
|
||
if *receiver.borrow_and_update() && receiver.changed().await.is_err() {
|
||
return;
|
||
}
|
||
}
|
||
}
|
||
|
||
async fn insert_connection_pool_inner(
|
||
&self,
|
||
pool_key: String,
|
||
pool: PoolKind,
|
||
config: &ConnectionConfig,
|
||
wait_for_drain: bool,
|
||
) -> Result<(), String> {
|
||
if wait_for_drain {
|
||
self.wait_for_pool_drain(&pool_key).await;
|
||
}
|
||
#[cfg(feature = "mq-admin")]
|
||
let mq_keepalive_adapter = if matches!(&pool, PoolKind::MessageQueue) {
|
||
self.mq_registry.get_cached_adapter(&config.id).await
|
||
} else {
|
||
None
|
||
};
|
||
let routing = self.pool_routing_control();
|
||
let previous = loop {
|
||
let mut connections = self.connections.write().await;
|
||
if !pool.is_available_for_routing() {
|
||
break Err(pool);
|
||
}
|
||
let Ok(mut activity) = self.pool_activity.try_write() else {
|
||
// Idle reclamation reads activity before routing. Never await that lock while
|
||
// holding the routing lock; release and retry to preserve a single lock order.
|
||
drop(connections);
|
||
tokio::task::yield_now().await;
|
||
continue;
|
||
};
|
||
// Abort the old probe while the route cannot change underneath us. A failed
|
||
// candidate never reaches this point, so the existing route keeps its state.
|
||
routing.stop_keepalive(&pool_key);
|
||
activity.insert(pool_key.clone(), PoolActivity::now());
|
||
// Manual-transaction sessions pin a sticky connection/TX. Keepalive
|
||
// detach on these pools was able to tear down the shared agent runtime
|
||
// (and any open TX) when close timed out — skip probes for them.
|
||
let skip_keepalive = pool_key.contains(":session:manual-txn-");
|
||
if !skip_keepalive {
|
||
self.start_keepalive_task(
|
||
&pool_key,
|
||
&pool,
|
||
config,
|
||
#[cfg(feature = "mq-admin")]
|
||
mq_keepalive_adapter.clone(),
|
||
);
|
||
}
|
||
break Ok(connections.insert(pool_key.clone(), pool));
|
||
};
|
||
let previous = match previous {
|
||
Ok(previous) => previous,
|
||
Err(pool) => {
|
||
if let PoolKind::Agent(client) = &pool {
|
||
routing.replace_runtime_after_close_failure(&pool_key, client).await;
|
||
}
|
||
routing.close_pool_with_timeout(pool_key, pool).await;
|
||
return Err("Agent runtime is unavailable while publishing the connection pool".to_string());
|
||
}
|
||
};
|
||
if let Some(pool) = previous {
|
||
routing.close_pool_with_timeout(pool_key.clone(), pool).await;
|
||
}
|
||
let route_is_available = self
|
||
.with_connection_pools(|connections| {
|
||
connections.get(&pool_key).is_some_and(PoolKind::is_available_for_routing)
|
||
})
|
||
.await;
|
||
if !route_is_available {
|
||
routing.detach_pool_by_key(&pool_key, true).await;
|
||
return Err("Agent runtime is unavailable while publishing the connection pool".to_string());
|
||
}
|
||
// 记录 no-save 池由哪个 owner 会话创建,供多会话复用判定使用
|
||
// (见 pool_credential_owner_mismatch)。已保存密码连接始终共享池,不记录。
|
||
if !config.save_password {
|
||
let owner = crate::session_credentials::current_credential_owner().unwrap_or_default();
|
||
if !owner.is_empty() {
|
||
self.session_credentials.record_pool_owner(&pool_key, &owner);
|
||
}
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
pub async fn insert_connection_pool(
|
||
&self,
|
||
pool_key: String,
|
||
pool: PoolKind,
|
||
config: &ConnectionConfig,
|
||
) -> Result<(), String> {
|
||
self.insert_connection_pool_inner(pool_key, pool, config, true).await
|
||
}
|
||
|
||
pub async fn begin_connection_attempt(&self, connection_id: &str) -> u64 {
|
||
self.begin_connection_attempt_with_client_attempt(connection_id, None).await
|
||
}
|
||
|
||
pub async fn begin_connection_attempt_with_client_attempt(
|
||
&self,
|
||
connection_id: &str,
|
||
client_attempt: Option<u64>,
|
||
) -> u64 {
|
||
let mut attempts = self.connection_attempts.write().await;
|
||
let next = attempts.get(connection_id).map(|state| state.server_attempt).unwrap_or(0).wrapping_add(1);
|
||
attempts.insert(connection_id.to_string(), ConnectionAttemptState { server_attempt: next, client_attempt });
|
||
next
|
||
}
|
||
|
||
pub async fn supersede_connection_attempt(&self, connection_id: &str) {
|
||
self.begin_connection_attempt(connection_id).await;
|
||
}
|
||
|
||
pub async fn supersede_connection_attempt_if_client_attempt(
|
||
&self,
|
||
connection_id: &str,
|
||
client_attempt: u64,
|
||
) -> bool {
|
||
let mut attempts = self.connection_attempts.write().await;
|
||
let Some(current) = attempts.get(connection_id).copied() else {
|
||
return false;
|
||
};
|
||
if current.client_attempt != Some(client_attempt) {
|
||
return false;
|
||
}
|
||
attempts.insert(
|
||
connection_id.to_string(),
|
||
ConnectionAttemptState { server_attempt: current.server_attempt.wrapping_add(1), client_attempt: None },
|
||
);
|
||
true
|
||
}
|
||
|
||
async fn connection_attempt_is_current(&self, connection_id: &str, attempt: u64) -> bool {
|
||
self.connection_attempts.read().await.get(connection_id).map(|state| state.server_attempt) == Some(attempt)
|
||
}
|
||
|
||
pub async fn ensure_current_connection_attempt(
|
||
&self,
|
||
connection_id: &str,
|
||
attempt: Option<u64>,
|
||
) -> Result<(), String> {
|
||
let Some(attempt) = attempt else {
|
||
return Ok(());
|
||
};
|
||
if self.connection_attempt_is_current(connection_id, attempt).await {
|
||
Ok(())
|
||
} else {
|
||
Err("Connection attempt was superseded by a newer attempt".to_string())
|
||
}
|
||
}
|
||
|
||
pub async fn insert_connection_pool_for_attempt(
|
||
&self,
|
||
connection_id: &str,
|
||
attempt: u64,
|
||
pool_key: String,
|
||
pool: PoolKind,
|
||
config: &ConnectionConfig,
|
||
) -> Result<(), String> {
|
||
if let Err(err) = self.ensure_current_connection_attempt(connection_id, Some(attempt)).await {
|
||
self.pool_routing_control().close_pool_with_timeout(pool_key, pool).await;
|
||
return Err(err);
|
||
}
|
||
self.insert_connection_pool(pool_key, pool, config).await
|
||
}
|
||
|
||
async fn discard_stale_connection_attempt_pool(
|
||
&self,
|
||
connection_id: &str,
|
||
pool_key: String,
|
||
pool: PoolKind,
|
||
config: &ConnectionConfig,
|
||
) {
|
||
#[cfg(feature = "mq-admin")]
|
||
if matches!(pool, PoolKind::MessageQueue) {
|
||
self.mq_registry.drop_connection(connection_id).await;
|
||
}
|
||
self.reset_connection_transport_for_config(connection_id, config).await;
|
||
self.pool_routing_control().close_pool_with_timeout(pool_key, pool).await;
|
||
}
|
||
|
||
fn start_keepalive_task(
|
||
&self,
|
||
pool_key: &str,
|
||
pool: &PoolKind,
|
||
config: &ConnectionConfig,
|
||
#[cfg(feature = "mq-admin")] mq_adapter: Option<std::sync::Arc<dyn crate::mq::port::MessageQueueAdmin>>,
|
||
) {
|
||
let interval_secs = config.keepalive_interval_secs;
|
||
let mut target = keepalive_target_from_pool(pool, config);
|
||
#[cfg(feature = "mq-admin")]
|
||
if target.is_none() {
|
||
if let Some(adapter) = mq_adapter {
|
||
target = Some(KeepaliveTarget::MessageQueue(adapter));
|
||
}
|
||
}
|
||
if interval_secs == 0 {
|
||
return;
|
||
}
|
||
if interval_secs > 0 && target.is_none() {
|
||
log::debug!(
|
||
"Connection keepalive requested for '{pool_key}', but this database driver does not keep a pingable client handle."
|
||
);
|
||
return;
|
||
};
|
||
|
||
let key = pool_key.to_string();
|
||
let interval = Duration::from_secs(interval_secs.max(1));
|
||
// MQ keepalive runs a full adapter test_connection (agent RPC); use query timeout
|
||
// so slow clusters are not spuriously dropped by the shorter connect timeout.
|
||
#[cfg(feature = "mq-admin")]
|
||
let timeout = if matches!(target, Some(KeepaliveTarget::MessageQueue(_))) {
|
||
match config.effective_query_timeout_secs() {
|
||
0 => Duration::from_secs(300), // UI "unlimited" still needs a keepalive bound
|
||
secs => Duration::from_secs(secs.max(1)),
|
||
}
|
||
} else {
|
||
Duration::from_secs(config.effective_connect_timeout_secs().max(1))
|
||
};
|
||
#[cfg(not(feature = "mq-admin"))]
|
||
let timeout = Duration::from_secs(config.effective_connect_timeout_secs().max(1));
|
||
let routing = self.pool_routing_control();
|
||
let connections = self.connections.clone();
|
||
let running_queries = self.running_queries.clone();
|
||
// MQ pool markers are empty; close_pool_kind is a no-op, so keepalive must
|
||
// drop the registry adapter or reconnect would reuse a dead agent.
|
||
#[cfg(feature = "mq-admin")]
|
||
let mq_registry = self.mq_registry.clone();
|
||
#[cfg(feature = "mq-admin")]
|
||
let mq_connection_id = config.id.clone();
|
||
self.task_supervisor.spawn_replace(format!("keepalive:{pool_key}"), move |shutdown| async move {
|
||
loop {
|
||
tokio::select! {
|
||
_ = shutdown.cancelled() => break,
|
||
_ = tokio::time::sleep(interval) => {}
|
||
}
|
||
|
||
if running_queries.is_pool_active(&key) {
|
||
continue;
|
||
}
|
||
|
||
if let Some(target) = target.as_mut() {
|
||
let result = tokio::time::timeout(timeout, ping_keepalive_target(target, timeout)).await;
|
||
match result {
|
||
Ok(Ok(())) => {}
|
||
Ok(Err(err)) => {
|
||
log::warn!("Connection keepalive failed for '{key}': {err}; invalidating pool");
|
||
let replace_runtime =
|
||
err.recovery_decision().is_some_and(RecoveryDecision::replaces_runtime);
|
||
// PoolKind::MessageQueue is a unit marker, so matches_pool cannot
|
||
// Arc::ptr_eq. Identity-check the registry adapter before teardown.
|
||
#[cfg(feature = "mq-admin")]
|
||
if let KeepaliveTarget::MessageQueue(adapter) = target {
|
||
if !mq_registry.is_current_adapter(&mq_connection_id, adapter).await {
|
||
log::debug!("Skipping stale MQ keepalive result for replaced pool '{key}'");
|
||
break;
|
||
}
|
||
}
|
||
#[cfg(feature = "mq-admin")]
|
||
let drop_mq = matches!(target, KeepaliveTarget::MessageQueue(_));
|
||
if !detach_keepalive_target_if_current(
|
||
&routing,
|
||
&connections,
|
||
&key,
|
||
target,
|
||
replace_runtime,
|
||
)
|
||
.await
|
||
{
|
||
log::debug!("Skipping stale keepalive result for replaced pool '{key}'");
|
||
} else {
|
||
#[cfg(feature = "mq-admin")]
|
||
if drop_mq {
|
||
mq_registry.drop_connection(&mq_connection_id).await;
|
||
}
|
||
}
|
||
break;
|
||
}
|
||
Err(_) => {
|
||
log::warn!(
|
||
"Connection keepalive timed out for '{key}' after {}s; invalidating pool",
|
||
timeout.as_secs()
|
||
);
|
||
#[cfg(feature = "mq-admin")]
|
||
if let KeepaliveTarget::MessageQueue(adapter) = target {
|
||
if !mq_registry.is_current_adapter(&mq_connection_id, adapter).await {
|
||
log::debug!("Skipping stale MQ keepalive timeout for replaced pool '{key}'");
|
||
break;
|
||
}
|
||
}
|
||
#[cfg(feature = "mq-admin")]
|
||
let drop_mq = matches!(target, KeepaliveTarget::MessageQueue(_));
|
||
if !detach_keepalive_target_if_current(&routing, &connections, &key, target, false).await {
|
||
log::debug!("Skipping stale keepalive timeout for replaced pool '{key}'");
|
||
} else {
|
||
#[cfg(feature = "mq-admin")]
|
||
if drop_mq {
|
||
mq_registry.drop_connection(&mq_connection_id).await;
|
||
}
|
||
}
|
||
break;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
});
|
||
}
|
||
|
||
async fn stop_keepalive_task(&self, pool_key: &str) {
|
||
self.task_supervisor.stop(&format!("keepalive:{pool_key}"));
|
||
}
|
||
|
||
async fn stop_keepalive_task_and_wait(&self, pool_key: &str) {
|
||
self.task_supervisor.stop_and_wait(&format!("keepalive:{pool_key}")).await;
|
||
}
|
||
|
||
async fn stop_keepalive_tasks(&self, pool_keys: &[String]) {
|
||
let keys: Vec<String> = pool_keys.iter().map(|pool_key| format!("keepalive:{pool_key}")).collect();
|
||
self.task_supervisor.stop_many(keys.iter().map(String::as_str));
|
||
}
|
||
|
||
pub async fn touch_pool_activity(&self, pool_key: &str) {
|
||
// 热路径:读锁下原子更新;仅条目缺失(首次/已被清理)才退化为写锁插入
|
||
if let Some(activity) = self.pool_activity.read().await.get(pool_key) {
|
||
activity.touch();
|
||
return;
|
||
}
|
||
self.pool_activity.write().await.insert(pool_key.to_string(), PoolActivity::now());
|
||
}
|
||
|
||
/// Get the PostgreSQL TLS cancel context (used to reconstruct the TLS connector when cancelling a query).
|
||
pub async fn get_postgres_cancel_context(&self, pool_key: &str) -> Option<db::postgres::PostgresCancelContext> {
|
||
self.postgres_cancel_contexts.read().await.get(pool_key).cloned()
|
||
}
|
||
|
||
pub fn pool_activity_touch(&self, pool_key: &str) -> PoolActivityTouch {
|
||
PoolActivityTouch {
|
||
pool_key: pool_key.to_string(),
|
||
connections: self.connections.clone(),
|
||
pool_activity: self.pool_activity.clone(),
|
||
task_supervisor: self.task_supervisor.clone(),
|
||
}
|
||
}
|
||
|
||
pub async fn shutdown(&self, deadline: Duration) {
|
||
self.running_queries.cancel_all();
|
||
let removed_pools = self.drain_all_connection_pools().await;
|
||
self.transaction_sessions.write().await.clear();
|
||
|
||
let shutdown = async {
|
||
let routing = self.pool_routing_control();
|
||
tokio::join!(
|
||
self.task_supervisor.shutdown(deadline),
|
||
routing.close_removed(removed_pools),
|
||
self.tunnels.stop_all_tunnels(),
|
||
self.proxy_tunnels.stop_all_tunnels(),
|
||
self.http_tunnels.stop_all_tunnels(),
|
||
self.agent_manager.stop_daemons(),
|
||
);
|
||
};
|
||
if tokio::time::timeout(deadline, shutdown).await.is_err() {
|
||
log::warn!("Timed out shutting down DBX runtime resources after {}ms", deadline.as_millis());
|
||
}
|
||
}
|
||
|
||
/// Cancels the tasks whose frontend consumer lives in the webview renderer
|
||
/// session that is about to be reloaded after a WebView2 renderer process
|
||
/// failure.
|
||
///
|
||
/// A renderer reload keeps the application process alive, so only the tasks
|
||
/// bound to the (now-dying) renderer session must be torn down. All of those
|
||
/// (SQL execution, counts/explains and exports) are registered in
|
||
/// [`Self::running_queries`], so this narrow boundary is exactly
|
||
/// [`RunningQueries::cancel_all`]: it signals their cancellation tokens and
|
||
/// fires any registered interrupts, which is what the underlying drivers
|
||
/// (and `query_result_export`) poll to stop work on the database.
|
||
///
|
||
/// Connection pools, tunnels, transaction sessions and daemons are
|
||
/// *application-scoped* and are reused by the reloaded frontend, so they are
|
||
/// deliberately NOT closed here — closing them is the job of
|
||
/// [`Self::shutdown`] on a full application restart. This keeps a renderer
|
||
/// reload from evicting a pool and tearing down tunnels that the reloaded
|
||
/// page still needs.
|
||
///
|
||
/// Returns the number of tasks signalled for cancellation (0 when there was
|
||
/// nothing running).
|
||
pub fn cancel_webview_reload_session_tasks(&self) -> usize {
|
||
let cancelled = self.running_queries.cancel_all();
|
||
if cancelled > 0 {
|
||
log::info!("cancelled {cancelled} webview-session query/export tasks before renderer reload");
|
||
}
|
||
cancelled
|
||
}
|
||
|
||
#[cfg(test)]
|
||
pub fn supervised_task_count(&self) -> usize {
|
||
self.task_supervisor.active_count()
|
||
}
|
||
|
||
pub async fn get_or_create_pool(&self, connection_id: &str, database: Option<&str>) -> Result<String, String> {
|
||
self.get_or_create_pool_for_session(connection_id, database, None).await
|
||
}
|
||
|
||
pub async fn get_or_create_pool_with_catalog(
|
||
&self,
|
||
connection_id: &str,
|
||
database: Option<&str>,
|
||
catalog: Option<&str>,
|
||
) -> Result<String, String> {
|
||
self.get_or_create_pool_for_session_with_catalog(connection_id, database, catalog, None).await
|
||
}
|
||
|
||
pub async fn get_or_create_pool_for_connection_attempt(
|
||
&self,
|
||
connection_id: &str,
|
||
database: Option<&str>,
|
||
attempt: u64,
|
||
) -> Result<String, String> {
|
||
self.get_or_create_pool_for_session_inner(
|
||
connection_id,
|
||
database,
|
||
None,
|
||
None,
|
||
AgentSessionRole::Workload,
|
||
Some(attempt),
|
||
)
|
||
.await
|
||
}
|
||
|
||
pub async fn get_or_create_pool_for_session(
|
||
&self,
|
||
connection_id: &str,
|
||
database: Option<&str>,
|
||
client_session_id: Option<&str>,
|
||
) -> Result<String, String> {
|
||
self.get_or_create_pool_for_session_with_catalog(connection_id, database, None, client_session_id).await
|
||
}
|
||
|
||
pub async fn get_or_create_pool_for_session_with_catalog(
|
||
&self,
|
||
connection_id: &str,
|
||
database: Option<&str>,
|
||
catalog: Option<&str>,
|
||
client_session_id: Option<&str>,
|
||
) -> Result<String, String> {
|
||
self.get_or_create_pool_for_session_inner(
|
||
connection_id,
|
||
database,
|
||
catalog,
|
||
client_session_id,
|
||
AgentSessionRole::Workload,
|
||
None,
|
||
)
|
||
.await
|
||
}
|
||
|
||
pub(crate) async fn get_or_create_metadata_pool_for_session(
|
||
&self,
|
||
connection_id: &str,
|
||
database: Option<&str>,
|
||
client_session_id: Option<&str>,
|
||
) -> Result<String, String> {
|
||
self.get_or_create_pool_for_session_inner(
|
||
connection_id,
|
||
database,
|
||
None,
|
||
client_session_id,
|
||
AgentSessionRole::Metadata,
|
||
None,
|
||
)
|
||
.await
|
||
}
|
||
|
||
async fn get_or_create_pool_for_session_inner(
|
||
&self,
|
||
connection_id: &str,
|
||
database: Option<&str>,
|
||
catalog: Option<&str>,
|
||
client_session_id: Option<&str>,
|
||
session_role: AgentSessionRole,
|
||
connection_attempt: Option<u64>,
|
||
) -> Result<String, String> {
|
||
let config = {
|
||
let configs = self.configs.read().await;
|
||
configs.get(connection_id).ok_or("Connection config not found")?.clone()
|
||
};
|
||
validate_connection_url_params(&config)?;
|
||
let db_type = Some(config.db_type);
|
||
let validate_existing_pool = should_validate_existing_pool_before_reuse(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, false);
|
||
let pool_key = pool_key_for_session_role(Some(&config), base_pool_key.clone(), client_session_id, session_role);
|
||
|
||
loop {
|
||
self.wait_for_pool_drain(&pool_key).await;
|
||
if self.pool_handle(&pool_key).await.is_some() {
|
||
if self.remove_pool_if_duckdb_isolation_mismatch(&pool_key).await {
|
||
// Recreate below using the current DuckDB isolation mode.
|
||
} else if self.pool_credential_owner_mismatch(&config, &pool_key).await {
|
||
// 创建该 no-save 池的是另一个登录会话:销毁旧池,用当前会话的
|
||
// 凭据重建,避免跨会话复用他人输入的临时密码。
|
||
self.remove_stale_connection_pool(&pool_key).await;
|
||
break;
|
||
} else if !validate_existing_pool || !self.remove_stale_connection_pool(&pool_key).await {
|
||
self.touch_pool_activity(&pool_key).await;
|
||
return Ok(pool_key);
|
||
}
|
||
break;
|
||
}
|
||
// A reclaim may have removed the pool after the first drain check. Wait
|
||
// for its confirmed close or rollback before deciding to create a new one.
|
||
self.wait_for_pool_drain(&pool_key).await;
|
||
if self.pool_handle(&pool_key).await.is_some() {
|
||
continue;
|
||
}
|
||
break;
|
||
}
|
||
|
||
let mut db_config = database_connection_config_with_catalog(&config, database, catalog);
|
||
self.apply_session_credential(&config, &mut db_config, connection_id);
|
||
|
||
validate_h2_file_connection(&db_config)?;
|
||
self.ensure_current_connection_attempt(connection_id, connection_attempt).await?;
|
||
let (host, port) = self.connection_host_port(connection_id, &db_config).await?;
|
||
if let Err(err) = self.ensure_current_connection_attempt(connection_id, connection_attempt).await {
|
||
self.reset_connection_transport_for_config(connection_id, &db_config).await;
|
||
return Err(err);
|
||
}
|
||
probe_connection_endpoint(&db_config, &host, port).await?;
|
||
if let Err(err) = self.ensure_current_connection_attempt(connection_id, connection_attempt).await {
|
||
self.reset_connection_transport_for_config(connection_id, &db_config).await;
|
||
return Err(err);
|
||
}
|
||
let url = connection_url_for_endpoint(&db_config, &host, port);
|
||
let connect_timeout = std::time::Duration::from_secs(db_config.effective_connect_timeout_secs());
|
||
let idle_timeout = std::time::Duration::from_secs(db_config.idle_timeout_secs);
|
||
let mysql_pool_max_connections = mysql_pool_max_connections_for_session(client_session_id);
|
||
let pool = match db_config.db_type {
|
||
DatabaseType::Mysql => {
|
||
let (pool, mode) = connect_mysql_metadata_pool(
|
||
&config,
|
||
&db_config,
|
||
&host,
|
||
port,
|
||
connect_timeout,
|
||
mysql_pool_max_connections,
|
||
)
|
||
.await?;
|
||
PoolKind::Mysql(pool, mode)
|
||
}
|
||
DatabaseType::Doris | DatabaseType::StarRocks | DatabaseType::ManticoreSearch => {
|
||
let pool = if database.is_none() {
|
||
connect_bare_metadata_pool(&db_config, &host, port, connect_timeout, mysql_pool_max_connections)
|
||
.await?
|
||
} else {
|
||
connect_bare_mysql_pool_with_setup(
|
||
&db_config,
|
||
&url,
|
||
connect_timeout,
|
||
mysql_pool_max_connections,
|
||
&mysql_pool_setup_queries(&db_config, &url),
|
||
)
|
||
.await?
|
||
};
|
||
PoolKind::Mysql(pool, MysqlMode::Bare)
|
||
}
|
||
DatabaseType::Gaussdb if gaussdb_uses_m_jdbc_driver(&db_config) => {
|
||
let jdbc_config = gaussdb_m_jdbc_config_for_endpoint(&db_config, &host, port);
|
||
self.external_driver_pool("jdbc", &jdbc_config).await?
|
||
}
|
||
DatabaseType::Postgres
|
||
| DatabaseType::Redshift
|
||
| DatabaseType::Gaussdb
|
||
| DatabaseType::Kwdb
|
||
| DatabaseType::Questdb
|
||
| DatabaseType::OpenGauss => {
|
||
// A session pool must have one physical client: sequential
|
||
// statements can otherwise lose temp tables/SET state when a
|
||
// pool checkout selects another PostgreSQL connection.
|
||
let pg_pool = db::postgres::connect_with_max_connections(
|
||
&url,
|
||
connect_timeout,
|
||
postgres_pool_max_connections_for_session(client_session_id),
|
||
)
|
||
.await?;
|
||
// Build TLS cancel context for reconstructing TLS connection during cancel
|
||
if let Some(ctx) = db::postgres::build_postgres_cancel_context(&url) {
|
||
self.postgres_cancel_contexts.write().await.insert(pool_key.clone(), ctx);
|
||
}
|
||
PoolKind::Postgres(pg_pool)
|
||
}
|
||
DatabaseType::Sqlite => {
|
||
if db::sqlite_worker::sqlite_ssh_worker_requested(&db_config) {
|
||
let transport_layers = self.resolved_transport_layers(&db_config).await?;
|
||
let worker = db::sqlite_worker::connect_sqlite_worker(
|
||
&self.tunnels,
|
||
&self.agent_manager,
|
||
self.storage.data_dir(),
|
||
connection_id,
|
||
&db_config,
|
||
&transport_layers,
|
||
)
|
||
.await?;
|
||
PoolKind::Sqlite(db::sqlite::SqliteHandle::from_worker(worker))
|
||
} else {
|
||
let sqlite_path = expand_tilde(&db_config.host);
|
||
db::sqlite::validate_persistent_attachments(
|
||
&sqlite_path,
|
||
&db_config.password,
|
||
!db_config.attached_databases.is_empty(),
|
||
)?;
|
||
let extensions =
|
||
db::sqlite::sqlite_extension_specs_from_url_params(db_config.url_params.as_deref())
|
||
.into_iter()
|
||
.map(|mut extension| {
|
||
extension.path = expand_tilde(&extension.path);
|
||
extension
|
||
})
|
||
.collect();
|
||
let pool = db::sqlite::connect_path_with_cipher_key_and_extensions(
|
||
&sqlite_path,
|
||
&db_config.password,
|
||
extensions,
|
||
)
|
||
.await?;
|
||
for attached in &db_config.attached_databases {
|
||
db::sqlite::attach_database(&pool, &attached.name, &expand_tilde(&attached.path))?;
|
||
}
|
||
PoolKind::Sqlite(pool)
|
||
}
|
||
}
|
||
DatabaseType::Rqlite => {
|
||
let client = db::rqlite_driver::RqliteClient::new(
|
||
&url,
|
||
db_config.url_params.as_deref(),
|
||
&db_config.username,
|
||
&db_config.password,
|
||
db_config.ssl,
|
||
connect_timeout,
|
||
)?;
|
||
db::rqlite_driver::test_connection(&client, connect_timeout).await?;
|
||
PoolKind::Rqlite(client)
|
||
}
|
||
DatabaseType::Turso => {
|
||
let auth_token = if !db_config.password.is_empty() {
|
||
db_config.password.clone()
|
||
} else {
|
||
db_config.url_params.as_deref().and_then(extract_auth_token_from_params).unwrap_or_default()
|
||
};
|
||
let client = db::turso_driver::TursoClient::new(&url, &auth_token, db_config.ssl, connect_timeout)?;
|
||
db::turso_driver::test_connection(&client, connect_timeout).await?;
|
||
PoolKind::Turso(client)
|
||
}
|
||
DatabaseType::CloudflareD1 => {
|
||
PoolKind::CloudflareD1(db::cloudflare_d1_driver::connect(&db_config, connect_timeout).await?)
|
||
}
|
||
DatabaseType::Redis => {
|
||
let con = if db_config.uses_redis_cluster() {
|
||
db::redis_driver::RedisConnection::Cluster(
|
||
self.connect_redis_cluster(connection_id, &db_config).await?,
|
||
)
|
||
} else if db_config.uses_redis_sentinel() {
|
||
db::redis_driver::RedisConnection::Direct(tokio::sync::Mutex::new(
|
||
self.connect_redis_sentinel(connection_id, &db_config).await?,
|
||
))
|
||
} else {
|
||
db::redis_driver::RedisConnection::Direct(tokio::sync::Mutex::new(
|
||
db::redis_driver::connect_standalone(&db_config, &host, port, connect_timeout).await?,
|
||
))
|
||
};
|
||
PoolKind::Redis(Arc::new(con))
|
||
}
|
||
#[cfg(feature = "duckdb-sidecar")]
|
||
DatabaseType::DuckDb => self.create_duckdb_pool(&db_config).await?,
|
||
#[cfg(not(feature = "duckdb-sidecar"))]
|
||
DatabaseType::DuckDb => {
|
||
return Err("DuckDB support is not compiled in this build.".to_string());
|
||
}
|
||
DatabaseType::MongoDb => {
|
||
let uses_oidc = db::mongo_driver::mongo_uri_uses_oidc(&url);
|
||
if mongo_uses_legacy_driver(&db_config) && !uses_oidc {
|
||
log::info!("Using configured MongoDB legacy driver for connection_id={connection_id}");
|
||
let connect_params = serde_json::json!({ "connection": agent_connect_params(&db_config, &host, port, db_config.effective_database().unwrap_or(""))? });
|
||
let mut client = self.agent_manager.spawn(&DatabaseType::MongoDb, Some("mongodb-legacy")).await?;
|
||
client.connect(connect_params).await.map_err(|err| mongo_legacy_error_with_auth_hint(&err))?;
|
||
PoolKind::agent(client)
|
||
} else {
|
||
let native_err = match db::mongo_driver::connect_with_oidc(
|
||
&url,
|
||
connect_timeout,
|
||
idle_timeout,
|
||
self.mongo_oidc_browser_opener(),
|
||
)
|
||
.await
|
||
{
|
||
Ok(client) => match db::mongo_driver::test_connection_for_url(
|
||
&client,
|
||
&url,
|
||
connect_timeout,
|
||
db_config.effective_database(),
|
||
)
|
||
.await
|
||
{
|
||
Ok(()) => {
|
||
// Re-check: another task may have created the pool while we were connecting.
|
||
if self.pool_handle(&pool_key).await.is_some() {
|
||
self.pool_routing_control()
|
||
.close_pool_with_timeout(pool_key.clone(), PoolKind::MongoDb(client))
|
||
.await;
|
||
return Ok(pool_key);
|
||
}
|
||
if let Err(err) =
|
||
self.ensure_current_connection_attempt(connection_id, connection_attempt).await
|
||
{
|
||
self.discard_stale_connection_attempt_pool(
|
||
connection_id,
|
||
pool_key.clone(),
|
||
PoolKind::MongoDb(client),
|
||
&db_config,
|
||
)
|
||
.await;
|
||
return Err(err);
|
||
}
|
||
self.insert_connection_pool(pool_key.clone(), PoolKind::MongoDb(client), &db_config)
|
||
.await?;
|
||
return Ok(pool_key);
|
||
}
|
||
Err(e) => e,
|
||
},
|
||
Err(e) => e,
|
||
};
|
||
if !uses_oidc && should_retry_mongo_with_legacy_driver(&native_err) {
|
||
log::info!("Native MongoDB driver failed ({native_err}), falling back to agent driver");
|
||
let connect_params = serde_json::json!({ "connection": agent_connect_params(&db_config, &host, port, db_config.effective_database().unwrap_or(""))? });
|
||
let legacy_agent_key =
|
||
AgentManager::db_type_to_agent_key(&DatabaseType::MongoDb, Some("mongodb-legacy"))
|
||
.ok_or_else(|| "MongoDB (Legacy) Agent mapping is unavailable".to_string())?;
|
||
crate::agent_service::ensure_agent_driver_ready(&self.agent_manager, legacy_agent_key)
|
||
.await
|
||
.map_err(|err| {
|
||
format!("{native_err}\n\nFailed to prepare MongoDB (Legacy) fallback driver: {err}")
|
||
})?;
|
||
let mut client =
|
||
self.agent_manager.spawn(&DatabaseType::MongoDb, Some("mongodb-legacy")).await.map_err(
|
||
|err| {
|
||
format!(
|
||
"{native_err}\n\nFallback with MongoDB (Legacy) driver failed: {}",
|
||
mongo_legacy_error_with_auth_hint(&err)
|
||
)
|
||
},
|
||
)?;
|
||
client.connect(connect_params).await.map_err(|err| {
|
||
format!(
|
||
"{native_err}\n\nFallback with MongoDB (Legacy) driver failed: {}",
|
||
mongo_legacy_error_with_auth_hint(&err)
|
||
)
|
||
})?;
|
||
PoolKind::agent(client)
|
||
} else {
|
||
return Err(native_err);
|
||
}
|
||
}
|
||
}
|
||
DatabaseType::DynamoDb => {
|
||
let client = db::dynamodb_driver::connect(&db_config, &host, port)?;
|
||
db::dynamodb_driver::test_connection(&client, connect_timeout).await?;
|
||
PoolKind::DynamoDb(client)
|
||
}
|
||
DatabaseType::ClickHouse => {
|
||
let username = if db_config.username.is_empty() { None } else { Some(db_config.username.clone()) };
|
||
let password = if db_config.password.is_empty() { None } else { Some(db_config.password.clone()) };
|
||
let client = db::clickhouse_driver::ChClient::new_with_ca_cert(
|
||
&url,
|
||
username,
|
||
password,
|
||
Some(&db_config.ca_cert_path),
|
||
db_config.url_params.as_deref(),
|
||
connect_timeout,
|
||
)?;
|
||
db::clickhouse_driver::test_connection(&client, connect_timeout).await?;
|
||
PoolKind::ClickHouse(client)
|
||
}
|
||
DatabaseType::SqlServer => self.connect_sqlserver_pool(&db_config, &host, port, connect_timeout).await?,
|
||
DatabaseType::Elasticsearch => {
|
||
let mut client = db::elasticsearch_driver::EsClient::from_config(
|
||
&url,
|
||
Some(&db_config.username),
|
||
Some(&db_config.password),
|
||
db_config.ssl,
|
||
db_config.url_params.as_deref(),
|
||
db_config.external_config.as_ref(),
|
||
connect_timeout,
|
||
);
|
||
db::elasticsearch_driver::test_connection(&mut client, connect_timeout).await?;
|
||
PoolKind::Elasticsearch(client)
|
||
}
|
||
DatabaseType::Easysearch => {
|
||
let mut client = db::easysearch_driver::EasysearchClient::from_config(
|
||
&url,
|
||
Some(&db_config.username),
|
||
Some(&db_config.password),
|
||
db_config.ssl,
|
||
db_config.url_params.as_deref(),
|
||
db_config.external_config.as_ref(),
|
||
connect_timeout,
|
||
);
|
||
db::easysearch_driver::test_connection(&mut client, connect_timeout).await?;
|
||
PoolKind::Easysearch(client)
|
||
}
|
||
DatabaseType::Meilisearch => {
|
||
let client = db::meilisearch_driver::MeilisearchClient::new_for_config(
|
||
&url,
|
||
Some(&db_config.password),
|
||
db_config.ssl,
|
||
db_config.url_params.as_deref(),
|
||
db_config.external_config.as_ref(),
|
||
connect_timeout,
|
||
)?;
|
||
db::meilisearch_driver::test_connection(&client, connect_timeout).await?;
|
||
PoolKind::Meilisearch(client)
|
||
}
|
||
DatabaseType::Hbase => {
|
||
let client = db::hbase_driver::HBaseClient::new(
|
||
&url,
|
||
Some(&db_config.username),
|
||
Some(&db_config.password),
|
||
false,
|
||
connect_timeout,
|
||
)?;
|
||
db::hbase_driver::test_connection(&client, connect_timeout).await?;
|
||
PoolKind::HBase(client)
|
||
}
|
||
DatabaseType::Qdrant | DatabaseType::Milvus | DatabaseType::Weaviate | DatabaseType::ChromaDb => {
|
||
let kind = match db_config.db_type {
|
||
DatabaseType::Qdrant => db::vector_driver::VectorDbKind::Qdrant,
|
||
DatabaseType::Milvus => db::vector_driver::VectorDbKind::Milvus,
|
||
DatabaseType::Weaviate => db::vector_driver::VectorDbKind::Weaviate,
|
||
DatabaseType::ChromaDb => db::vector_driver::VectorDbKind::ChromaDb,
|
||
_ => unreachable!(),
|
||
};
|
||
let client = db::vector_driver::VectorClient::new(
|
||
kind,
|
||
&url,
|
||
Some(&db_config.username),
|
||
Some(&db_config.password),
|
||
db_config.ssl,
|
||
connect_timeout,
|
||
)
|
||
.with_database(db_config.database.as_deref());
|
||
db::vector_driver::test_connection(&client, connect_timeout).await?;
|
||
PoolKind::VectorDb(client)
|
||
}
|
||
DatabaseType::InfluxDb => {
|
||
let client = db::influxdb_driver::InfluxdbClient::new_for_config(&url, &db_config, connect_timeout)?;
|
||
db::influxdb_driver::test_connection(&client, connect_timeout).await?;
|
||
PoolKind::InfluxDb(client)
|
||
}
|
||
DatabaseType::InfluxDb3 => {
|
||
let client = db::influxdb3_driver::Influxdb3Client::new_for_config(&url, &db_config, connect_timeout)?;
|
||
db::influxdb3_driver::test_connection(&client, connect_timeout).await?;
|
||
PoolKind::InfluxDb3(client)
|
||
}
|
||
DatabaseType::VictoriaMetrics => {
|
||
let client = db::victoriametrics_driver::VictoriaMetricsClient::new_for_config(
|
||
&url,
|
||
&db_config,
|
||
connect_timeout,
|
||
)?;
|
||
db::victoriametrics_driver::test_connection(&client, connect_timeout).await?;
|
||
PoolKind::VictoriaMetrics(client)
|
||
}
|
||
DatabaseType::Nacos => {
|
||
let admin_config = self.nacos_admin_config_for_connection(connection_id, &config).await?;
|
||
let adapter = self.nacos_registry.build_transient_config(admin_config).await?;
|
||
adapter.test_connection().await?;
|
||
PoolKind::Nacos
|
||
}
|
||
DatabaseType::Consul => {
|
||
let mut consul_config = crate::consul::ConsulConfig::from_connection(&db_config)?;
|
||
let original_host = consul_config.base_url.host_str().unwrap_or_default();
|
||
let original_port = consul_config.base_url.port_or_known_default().unwrap_or(db_config.port);
|
||
if host != original_host || port != original_port {
|
||
consul_config = consul_config.with_connect_override(&host, port);
|
||
}
|
||
let client = crate::consul::ConsulClient::new(consul_config).await?;
|
||
client.probe().await?;
|
||
PoolKind::Consul(client)
|
||
}
|
||
agent_connection_pool_database_type!() => {
|
||
let connect_params = agent_connect_params_with_role(
|
||
&db_config,
|
||
&host,
|
||
port,
|
||
db_config.effective_database().unwrap_or(""),
|
||
session_role,
|
||
)?;
|
||
if db_config.db_type != DatabaseType::ZooKeeper {
|
||
let agent_session_id = uuid::Uuid::new_v4().simple().to_string();
|
||
let mut initial_result = self
|
||
.spawn_routed_shared_agent_client(
|
||
&db_config.db_type,
|
||
db_config.driver_profile.as_deref(),
|
||
&db_config.agent_java_options,
|
||
agent_session_id.clone(),
|
||
connect_params,
|
||
agent_connect_timeout(&db_config),
|
||
)
|
||
.await;
|
||
if initial_result.as_ref().is_err_and(|err| {
|
||
normalize_client_session_id(client_session_id).is_some()
|
||
&& is_connection_slot_exhausted_error(err)
|
||
}) && self.reclaim_idle_base_pool_for_session(connection_id, &base_pool_key).await
|
||
{
|
||
log::warn!(
|
||
"Reclaimed an idle metadata pool for '{connection_id}' after the database rejected a new Agent session due to exhausted connection slots"
|
||
);
|
||
for retry_delay_ms in [150, 350] {
|
||
tokio::time::sleep(Duration::from_millis(retry_delay_ms)).await;
|
||
initial_result = self
|
||
.spawn_routed_shared_agent_client(
|
||
&db_config.db_type,
|
||
db_config.driver_profile.as_deref(),
|
||
&db_config.agent_java_options,
|
||
agent_session_id.clone(),
|
||
agent_connect_params_with_role(
|
||
&db_config,
|
||
&host,
|
||
port,
|
||
db_config.effective_database().unwrap_or(""),
|
||
session_role,
|
||
)?,
|
||
agent_connect_timeout(&db_config),
|
||
)
|
||
.await;
|
||
if !initial_result.as_ref().is_err_and(|err| is_connection_slot_exhausted_error(err)) {
|
||
break;
|
||
}
|
||
}
|
||
}
|
||
let mut client = match initial_result {
|
||
Ok(client) => client,
|
||
Err(err) => {
|
||
let alternate_configs = oracle_alternate_connect_configs(&db_config, &err);
|
||
if alternate_configs.is_empty() {
|
||
if err.contains("does not support multi_session protocol v2") {
|
||
let mut client = self
|
||
.agent_manager
|
||
.spawn_with_extra_java_args(
|
||
&db_config.db_type,
|
||
db_config.driver_profile.as_deref(),
|
||
&db_config.agent_java_options,
|
||
)
|
||
.await?;
|
||
client
|
||
.call_method_with_timeout::<serde_json::Value>(
|
||
AgentMethod::Connect,
|
||
agent_connect_params_with_role(
|
||
&db_config,
|
||
&host,
|
||
port,
|
||
db_config.effective_database().unwrap_or(""),
|
||
session_role,
|
||
)?,
|
||
Some(agent_connect_timeout(&db_config)),
|
||
)
|
||
.await?;
|
||
client
|
||
} else {
|
||
return Err(oracle_error_with_driver_hint(&db_config, &err));
|
||
}
|
||
} else {
|
||
let mut fallback_errors = Vec::new();
|
||
let mut connected = None;
|
||
for alternate_config in alternate_configs {
|
||
let label =
|
||
oracle_alternate_connect_config_labels(std::slice::from_ref(&alternate_config))
|
||
.into_iter()
|
||
.next()
|
||
.unwrap_or_else(|| "alternate".to_string());
|
||
let alternate_params = agent_connect_params_with_role(
|
||
&alternate_config,
|
||
&host,
|
||
port,
|
||
alternate_config.effective_database().unwrap_or(""),
|
||
session_role,
|
||
)?;
|
||
match self
|
||
.spawn_routed_shared_agent_client(
|
||
&alternate_config.db_type,
|
||
alternate_config.driver_profile.as_deref(),
|
||
&alternate_config.agent_java_options,
|
||
agent_session_id.clone(),
|
||
alternate_params,
|
||
agent_connect_timeout(&alternate_config),
|
||
)
|
||
.await
|
||
{
|
||
Ok(client) => {
|
||
connected = Some(client);
|
||
break;
|
||
}
|
||
Err(alternate_err) => fallback_errors.push(format!("{label}: {alternate_err}")),
|
||
}
|
||
}
|
||
connected.ok_or_else(|| {
|
||
format!(
|
||
"{err}\n\nFallback with alternate Oracle connection descriptors failed: {}",
|
||
fallback_errors.join("\n")
|
||
)
|
||
})?
|
||
}
|
||
}
|
||
};
|
||
if db_config.db_type == DatabaseType::Kingbase {
|
||
let identifier_quote = client
|
||
.connection_info(Some(agent_connect_timeout(&db_config)))
|
||
.await
|
||
.ok()
|
||
.map(|info| info.identifier_quote);
|
||
client.set_identifier_quote(identifier_quote);
|
||
}
|
||
PoolKind::agent(client)
|
||
} else {
|
||
// ZooKeeper JVM properties are connection-scoped; shared agent daemons must not inherit them.
|
||
let mut client = self
|
||
.agent_manager
|
||
.spawn_with_extra_java_args(
|
||
&db_config.db_type,
|
||
db_config.driver_profile.as_deref(),
|
||
&db_config.agent_java_options,
|
||
)
|
||
.await?;
|
||
let connect_result = client
|
||
.call_method_with_timeout::<serde_json::Value>(
|
||
AgentMethod::Connect,
|
||
connect_params,
|
||
Some(agent_connect_timeout(&db_config)),
|
||
)
|
||
.await;
|
||
if let Err(err) = connect_result {
|
||
let alternate_configs = oracle_alternate_connect_configs(&db_config, &err);
|
||
if !alternate_configs.is_empty() {
|
||
log::warn!(
|
||
"Oracle connect failed with {:?} descriptor: {}. Retrying with Oracle JDBC URL variants: {:?}.",
|
||
db_config.oracle_connection_type,
|
||
err,
|
||
oracle_alternate_connect_config_labels(&alternate_configs)
|
||
);
|
||
let mut fallback_errors = Vec::new();
|
||
let mut connected = false;
|
||
for alternate_config in alternate_configs {
|
||
let label =
|
||
oracle_alternate_connect_config_labels(std::slice::from_ref(&alternate_config))
|
||
.into_iter()
|
||
.next()
|
||
.unwrap_or_else(|| "alternate".to_string());
|
||
match client
|
||
.call_method_with_timeout::<serde_json::Value>(
|
||
AgentMethod::Connect,
|
||
agent_connect_params_with_role(
|
||
&alternate_config,
|
||
&host,
|
||
port,
|
||
alternate_config.effective_database().unwrap_or(""),
|
||
session_role,
|
||
)?,
|
||
Some(agent_connect_timeout(&alternate_config)),
|
||
)
|
||
.await
|
||
{
|
||
Ok(_) => {
|
||
connected = true;
|
||
break;
|
||
}
|
||
Err(alternate_err) => {
|
||
fallback_errors.push(format!("{label}: {alternate_err}"));
|
||
}
|
||
}
|
||
}
|
||
if !connected {
|
||
return Err(format!(
|
||
"{err}\n\nFallback with alternate Oracle JDBC URLs failed: {}",
|
||
fallback_errors.join("\n")
|
||
));
|
||
}
|
||
} else {
|
||
return Err(oracle_error_with_driver_hint(&db_config, &err));
|
||
}
|
||
}
|
||
PoolKind::agent(client)
|
||
}
|
||
}
|
||
DatabaseType::PrestoSql => {
|
||
let jdbc_config = prestosql_jdbc_config_for_endpoint(&db_config, &host, port)?;
|
||
self.external_driver_pool("jdbc", &jdbc_config).await?
|
||
}
|
||
DatabaseType::Jdbc => {
|
||
let mut jdbc_config = db_config.clone();
|
||
if host != config.host || port != config.port {
|
||
if let Some(ref url) = jdbc_config.connection_string {
|
||
jdbc_config.connection_string = Some(rewrite_jdbc_url_host(url, &host, port)?);
|
||
}
|
||
}
|
||
self.external_driver_pool("jdbc", &jdbc_config).await?
|
||
}
|
||
DatabaseType::Plugin => return Err("Plugin-owned connections use the plugin host".to_string()),
|
||
#[cfg(feature = "mq-admin")]
|
||
DatabaseType::MessageQueue => {
|
||
// MQ admin connections don't hold a data query pool. We just test
|
||
// connectivity via the mq_registry and insert a marker so this
|
||
// connection_id is recognized as valid.
|
||
let mqc = self.mq_admin_config_for_connection(connection_id, &config).await?;
|
||
let agent_launch = crate::mq::service::resolve_mq_agent_launch_spec(&mqc, self);
|
||
// Temporary "__test_*" probes must not retain agents in the registry.
|
||
// reconnect fast-path caching only applies to durable connection ids;
|
||
// drain_connection_pools no longer drops MQ adapters (reconnect reuse).
|
||
if connection_id.starts_with(crate::runtime_config::TEST_PROBE_ID_PREFIX) {
|
||
let adapter = self.mq_registry.build_transient_config(mqc, agent_launch).await?;
|
||
adapter.test_connection().await?;
|
||
if let Err(err) = self.ensure_current_connection_attempt(connection_id, connection_attempt).await {
|
||
self.reset_connection_transport_for_config(connection_id, &db_config).await;
|
||
return Err(err);
|
||
}
|
||
// adapter drops here and kills agent-backed processes (RocketMQ/Kafka/RabbitMQ).
|
||
PoolKind::MessageQueue
|
||
} else {
|
||
let build = match self.mq_registry.get_or_build_config(connection_id, mqc, agent_launch).await {
|
||
Ok(build) => build,
|
||
Err(err) => {
|
||
self.mq_registry.drop_connection(connection_id).await;
|
||
return Err(err);
|
||
}
|
||
};
|
||
if let Err(err) = crate::mq::validate_mq_adapter_after_build(&build).await {
|
||
self.mq_registry.drop_connection(connection_id).await;
|
||
return Err(err);
|
||
}
|
||
if let Err(err) = self.ensure_current_connection_attempt(connection_id, connection_attempt).await {
|
||
self.mq_registry.drop_connection(connection_id).await;
|
||
self.reset_connection_transport_for_config(connection_id, &db_config).await;
|
||
return Err(err);
|
||
}
|
||
PoolKind::MessageQueue
|
||
}
|
||
}
|
||
#[cfg(not(feature = "mq-admin"))]
|
||
DatabaseType::MessageQueue => {
|
||
return Err(
|
||
"Message queue admin support is not compiled in this build. Rebuild with the 'mq-admin' feature."
|
||
.to_string(),
|
||
);
|
||
}
|
||
#[cfg(feature = "mq-admin")]
|
||
DatabaseType::Mqtt => {
|
||
let mqtt_config = crate::mqtt::types::MqttConnectionConfig::from_connection(&config)?;
|
||
let client = crate::mqtt::client::MqttClient::connect(mqtt_config).await?;
|
||
PoolKind::Mqtt(client)
|
||
}
|
||
#[cfg(not(feature = "mq-admin"))]
|
||
DatabaseType::Mqtt => {
|
||
return Err(
|
||
"MQTT support is not compiled in this build. Rebuild with the 'mq-admin' feature.".to_string()
|
||
);
|
||
}
|
||
};
|
||
|
||
if let Err(err) = self.ensure_current_connection_attempt(connection_id, connection_attempt).await {
|
||
self.discard_stale_connection_attempt_pool(connection_id, pool_key.clone(), pool, &db_config).await;
|
||
return Err(err);
|
||
}
|
||
self.insert_connection_pool(pool_key.clone(), pool, &db_config).await?;
|
||
Ok(pool_key)
|
||
}
|
||
|
||
/// Returns the enabled transport layers for a connection with tunnel
|
||
/// profile references resolved: a layer carrying a `profile_id` is
|
||
/// replaced by the shared profile from storage (Settings > Tunnels), so
|
||
/// edits to a profile take effect for every connection referencing it.
|
||
/// Fails when a referenced profile no longer exists — connecting without
|
||
/// the intended tunnel would silently bypass it.
|
||
pub async fn resolved_transport_layers(
|
||
&self,
|
||
config: &ConnectionConfig,
|
||
) -> Result<Vec<TransportLayerConfig>, String> {
|
||
let layers = config.effective_transport_layers();
|
||
if layers.iter().all(|layer| layer.profile_id().is_empty()) {
|
||
return Ok(layers);
|
||
}
|
||
|
||
let profiles: HashMap<String, TransportLayerConfig> = self
|
||
.storage
|
||
.load_tunnel_profiles()
|
||
.await?
|
||
.into_iter()
|
||
.map(|profile| (profile.id().to_string(), profile))
|
||
.collect();
|
||
|
||
layers
|
||
.into_iter()
|
||
.map(|layer| {
|
||
let profile_id = layer.profile_id();
|
||
if profile_id.is_empty() {
|
||
return Ok(layer);
|
||
}
|
||
let Some(profile) = profiles.get(profile_id) else {
|
||
let label = if layer.name().is_empty() { profile_id } else { layer.name() };
|
||
return Err(format!(
|
||
"Tunnel profile '{label}' referenced by this connection no longer exists. Re-create it in Settings > Tunnels or edit the connection's tunnel settings."
|
||
));
|
||
};
|
||
// Validate the stored reference again at connect time because synced or
|
||
// externally supplied configs may bypass the editor's type constraints.
|
||
if !layer.same_type_as(profile) {
|
||
return Err(format!(
|
||
"Tunnel profile '{}' has a different type than the referencing transport layer.",
|
||
if layer.name().is_empty() { profile_id } else { layer.name() }
|
||
));
|
||
}
|
||
Ok(layer.resolved_from_profile(profile))
|
||
})
|
||
.collect()
|
||
}
|
||
|
||
/// Tests a shared tunnel profile in isolation (no downstream database), for
|
||
/// the Test button in Settings > Tunnels.
|
||
///
|
||
/// - SSH: starting an SSH tunnel connects and authenticates eagerly, so a
|
||
/// successful start verifies host reachability and credentials.
|
||
/// - Proxy (HTTP CONNECT / SOCKS5): performs a standalone handshake test
|
||
/// against the proxy endpoint to verify reachability and credentials.
|
||
/// - HTTP tunnel: connects lazily (nothing happens until traffic flows), so
|
||
/// there is nothing to verify here without a target to probe.
|
||
pub async fn test_tunnel_profile(&self, profile: &TransportLayerConfig) -> Result<String, String> {
|
||
match profile {
|
||
TransportLayerConfig::Ssh(ssh) => {
|
||
// A `~/.ssh/config` alias declaring `ProxyJump` resolves to more
|
||
// than one hop; every other alias (the common case) resolves to
|
||
// exactly one, matching `resolve_ssh_tunnel_config`.
|
||
let chain = crate::ssh_config::resolve_ssh_tunnel_chain(ssh);
|
||
let leaf = chain.last().expect("resolve_ssh_tunnel_chain always returns at least one hop");
|
||
if leaf.host.trim().is_empty() {
|
||
return Err("SSH host is required.".to_string());
|
||
}
|
||
// A throwaway id so the probe never reuses or evicts a live tunnel, and
|
||
// a sentinel forward target: SSH auth completes on connect, before any
|
||
// channel to this target is opened, so it need not be reachable.
|
||
let probe_id = format!("__tunnel_profile_test__:{}", uuid::Uuid::new_v4());
|
||
let result = if chain.len() == 1 {
|
||
let timeout = if leaf.connect_timeout_secs == 0 {
|
||
crate::models::connection::default_ssh_connect_timeout_secs()
|
||
} else {
|
||
leaf.connect_timeout_secs
|
||
};
|
||
self.tunnels
|
||
.start_tunnel(
|
||
&probe_id,
|
||
&leaf.host,
|
||
leaf.port,
|
||
&leaf.host,
|
||
leaf.port,
|
||
&leaf.user,
|
||
&leaf.password,
|
||
&leaf.key_path,
|
||
&leaf.key_passphrase,
|
||
leaf.use_ssh_agent,
|
||
&leaf.ssh_agent_sock_path,
|
||
&leaf.auth_method,
|
||
timeout,
|
||
"127.0.0.1",
|
||
1,
|
||
false,
|
||
leaf.allow_exec_channel_proxy,
|
||
)
|
||
.await
|
||
.map(|_| ())
|
||
} else {
|
||
self.tunnels.start_chain(&probe_id, &chain, &leaf.host, leaf.port).await.map(|_| ())
|
||
};
|
||
self.tunnels.stop_tunnel(&probe_id).await;
|
||
result.map(|_| "SSH tunnel connection successful".to_string())
|
||
}
|
||
TransportLayerConfig::Proxy(proxy) => {
|
||
if proxy.host.trim().is_empty() {
|
||
return Err("Proxy host is required.".to_string());
|
||
}
|
||
if proxy.port == 0 {
|
||
return Err("Proxy port is required.".to_string());
|
||
}
|
||
crate::db::proxy_tunnel::test_proxy_endpoint(
|
||
proxy.proxy_type,
|
||
&proxy.host,
|
||
proxy.port,
|
||
&proxy.username,
|
||
&proxy.password,
|
||
proxy.test_target.as_deref(),
|
||
)
|
||
.await
|
||
}
|
||
TransportLayerConfig::HttpTunnel(_) => {
|
||
Err("Tunnel test is not supported for HTTP tunnel profiles.".to_string())
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Tests the enabled SSH chain without opening a database connection.
|
||
/// Layers before the final SSH hop are included because they may be
|
||
/// required to reach that hop; layers after it are unrelated to SSH auth.
|
||
pub async fn test_connection_ssh_tunnel(&self, config: &ConnectionConfig) -> Result<String, String> {
|
||
let resolved_layers = self.resolved_transport_layers(config).await?;
|
||
let test_layers = transport_layers_through_last_ssh(&resolved_layers)?;
|
||
let probe_id = format!("__connection_ssh_test__:{}", uuid::Uuid::new_v4());
|
||
let result = db::transport_layer_tunnel::start_transport_layers(
|
||
&probe_id,
|
||
test_layers,
|
||
"127.0.0.1",
|
||
1,
|
||
&self.tunnels,
|
||
&self.proxy_tunnels,
|
||
&self.http_tunnels,
|
||
)
|
||
.await;
|
||
db::transport_layer_tunnel::stop_transport_layers(
|
||
&probe_id,
|
||
test_layers.len(),
|
||
&self.tunnels,
|
||
&self.proxy_tunnels,
|
||
&self.http_tunnels,
|
||
)
|
||
.await;
|
||
result.map(|_| "SSH tunnel connection successful".to_string())
|
||
}
|
||
|
||
pub async fn connection_host_port(
|
||
&self,
|
||
connection_id: &str,
|
||
config: &ConnectionConfig,
|
||
) -> Result<(String, u16), String> {
|
||
let transport_layers = self.resolved_transport_layers(config).await?;
|
||
if transport_layers.is_empty() || db::sqlite_worker::sqlite_ssh_worker_requested(config) {
|
||
return Ok((config.host.clone(), config.port));
|
||
}
|
||
if config.uses_oracle_tns() {
|
||
// A TNS descriptor may contain several failover addresses, so rewriting it
|
||
// through one local tunnel endpoint would silently break Oracle Net routing.
|
||
return Err("Oracle TNS connections cannot be combined with SSH, proxy, or HTTP tunnel layers. Remove the transport layer or use Service Name/SID mode.".to_string());
|
||
}
|
||
if hive_uses_zookeeper_discovery(config) {
|
||
return Err("Hive ZooKeeper service discovery cannot be combined with SSH, proxy, or HTTP tunnel layers because discovered HiveServer2 nodes would bypass the configured transport. Remove the transport layer or use a direct HiveServer2 host and port.".to_string());
|
||
}
|
||
|
||
#[cfg(feature = "mq-admin")]
|
||
if config.db_type == DatabaseType::MessageQueue
|
||
&& crate::mq::config::MqAdminConfig::from_connection(config)?.system_kind
|
||
== crate::mq::types::MqSystemKind::RocketMq
|
||
{
|
||
self.rocketmq_socks_proxy_for_transport_layers(connection_id, &transport_layers).await?;
|
||
return Ok((config.host.clone(), config.port));
|
||
}
|
||
|
||
let (remote_host, remote_port) = connection_remote_endpoint(config);
|
||
let local_port = db::transport_layer_tunnel::start_transport_layers(
|
||
connection_id,
|
||
&transport_layers,
|
||
&remote_host,
|
||
remote_port,
|
||
&self.tunnels,
|
||
&self.proxy_tunnels,
|
||
&self.http_tunnels,
|
||
)
|
||
.await?;
|
||
|
||
Ok(("127.0.0.1".to_string(), local_port))
|
||
}
|
||
|
||
pub async fn connect_redis_sentinel(
|
||
&self,
|
||
connection_id: &str,
|
||
config: &ConnectionConfig,
|
||
) -> Result<redis::aio::MultiplexedConnection, String> {
|
||
let transport_layers = self.resolved_transport_layers(config).await?;
|
||
if transport_layers.is_empty() {
|
||
return db::redis_driver::connect_sentinel(config).await;
|
||
}
|
||
|
||
let result = async {
|
||
let sentinel_nodes = db::redis_driver::redis_sentinel_node_endpoints(config)?;
|
||
let connect_timeout = std::time::Duration::from_secs(config.effective_connect_timeout_secs());
|
||
let layer_count = transport_layers.len();
|
||
let mut last_error = None;
|
||
|
||
for sentinel in sentinel_nodes {
|
||
let sentinel_tunnel_id = redis_sentinel_transport_id(connection_id, "sentinel", &sentinel);
|
||
let sentinel_local_port = match db::transport_layer_tunnel::start_transport_layers(
|
||
&sentinel_tunnel_id,
|
||
&transport_layers,
|
||
&sentinel.host,
|
||
sentinel.port,
|
||
&self.tunnels,
|
||
&self.proxy_tunnels,
|
||
&self.http_tunnels,
|
||
)
|
||
.await
|
||
{
|
||
Ok(port) => port,
|
||
Err(err) => {
|
||
last_error =
|
||
Some(format!("Redis Sentinel {}:{} transport failed: {err}", sentinel.host, sentinel.port));
|
||
continue;
|
||
}
|
||
};
|
||
|
||
let master =
|
||
match db::redis_driver::discover_sentinel_master(config, "127.0.0.1", sentinel_local_port).await {
|
||
Ok(master) => master,
|
||
Err(err) => {
|
||
last_error = Some(format!(
|
||
"Redis Sentinel {}:{} master lookup failed: {err}",
|
||
sentinel.host, sentinel.port
|
||
));
|
||
db::transport_layer_tunnel::stop_transport_layers(
|
||
&sentinel_tunnel_id,
|
||
layer_count,
|
||
&self.tunnels,
|
||
&self.proxy_tunnels,
|
||
&self.http_tunnels,
|
||
)
|
||
.await;
|
||
continue;
|
||
}
|
||
};
|
||
|
||
let master_tunnel_id = redis_sentinel_transport_id(connection_id, "master", &master);
|
||
let master_local_port = match db::transport_layer_tunnel::start_transport_layers(
|
||
&master_tunnel_id,
|
||
&transport_layers,
|
||
&master.host,
|
||
master.port,
|
||
&self.tunnels,
|
||
&self.proxy_tunnels,
|
||
&self.http_tunnels,
|
||
)
|
||
.await
|
||
{
|
||
Ok(port) => port,
|
||
Err(err) => {
|
||
last_error = Some(format!(
|
||
"Redis Sentinel master {}:{} transport failed: {err}",
|
||
master.host, master.port
|
||
));
|
||
db::transport_layer_tunnel::stop_transport_layers(
|
||
&sentinel_tunnel_id,
|
||
layer_count,
|
||
&self.tunnels,
|
||
&self.proxy_tunnels,
|
||
&self.http_tunnels,
|
||
)
|
||
.await;
|
||
continue;
|
||
}
|
||
};
|
||
|
||
match db::redis_driver::connect_standalone(config, "127.0.0.1", master_local_port, connect_timeout)
|
||
.await
|
||
{
|
||
Ok(con) => return Ok(con),
|
||
Err(err) => {
|
||
last_error = Some(format!(
|
||
"Redis Sentinel master {}:{} connection failed: {err}",
|
||
master.host, master.port
|
||
));
|
||
db::transport_layer_tunnel::stop_transport_layers(
|
||
&master_tunnel_id,
|
||
layer_count,
|
||
&self.tunnels,
|
||
&self.proxy_tunnels,
|
||
&self.http_tunnels,
|
||
)
|
||
.await;
|
||
db::transport_layer_tunnel::stop_transport_layers(
|
||
&sentinel_tunnel_id,
|
||
layer_count,
|
||
&self.tunnels,
|
||
&self.proxy_tunnels,
|
||
&self.http_tunnels,
|
||
)
|
||
.await;
|
||
}
|
||
}
|
||
}
|
||
|
||
Err(last_error.unwrap_or_else(|| "Redis Sentinel master discovery failed".to_string()))
|
||
}
|
||
.await;
|
||
|
||
if result.is_err() {
|
||
let redis_sentinel_prefix = redis_sentinel_transport_prefix(connection_id);
|
||
self.tunnels.stop_tunnels_with_prefix(&redis_sentinel_prefix).await;
|
||
self.proxy_tunnels.stop_tunnels_with_prefix(&redis_sentinel_prefix).await;
|
||
self.http_tunnels.stop_tunnels_with_prefix(&redis_sentinel_prefix).await;
|
||
}
|
||
|
||
result
|
||
}
|
||
|
||
pub async fn connect_redis_cluster(
|
||
&self,
|
||
connection_id: &str,
|
||
config: &ConnectionConfig,
|
||
) -> Result<db::redis_driver::RedisClusterPool, String> {
|
||
let transport_layers = self.resolved_transport_layers(config).await?;
|
||
if transport_layers.is_empty() {
|
||
return db::redis_driver::connect_cluster(config).await;
|
||
}
|
||
|
||
let result = async {
|
||
let seed_nodes = db::redis_driver::redis_cluster_seed_nodes(config)?;
|
||
let seed_routes = self.redis_cluster_node_routes(connection_id, &transport_layers, &seed_nodes).await?;
|
||
let (auth, slot_ranges) =
|
||
db::redis_driver::discover_cluster_slot_ranges_from_routes(config, &seed_routes).await?;
|
||
let master_nodes = db::redis_driver::unique_master_nodes(&slot_ranges);
|
||
let node_routes = self.redis_cluster_node_routes(connection_id, &transport_layers, &master_nodes).await?;
|
||
|
||
db::redis_driver::connect_routed_cluster(config, seed_routes, slot_ranges, node_routes, auth).await
|
||
}
|
||
.await;
|
||
|
||
if result.is_err() {
|
||
let redis_cluster_prefix = redis_cluster_transport_prefix(connection_id);
|
||
self.tunnels.stop_tunnels_with_prefix(&redis_cluster_prefix).await;
|
||
self.proxy_tunnels.stop_tunnels_with_prefix(&redis_cluster_prefix).await;
|
||
self.http_tunnels.stop_tunnels_with_prefix(&redis_cluster_prefix).await;
|
||
}
|
||
|
||
result
|
||
}
|
||
|
||
async fn redis_cluster_node_routes(
|
||
&self,
|
||
connection_id: &str,
|
||
transport_layers: &[crate::models::connection::TransportLayerConfig],
|
||
nodes: &[db::redis_driver::RedisNodeEndpoint],
|
||
) -> Result<Vec<db::redis_driver::RedisNodeRoute>, String> {
|
||
let mut routes = Vec::with_capacity(nodes.len());
|
||
for node in nodes {
|
||
let tunnel_id = redis_cluster_transport_id(connection_id, node);
|
||
let local_port = db::transport_layer_tunnel::start_transport_layers(
|
||
&tunnel_id,
|
||
transport_layers,
|
||
&node.host,
|
||
node.port,
|
||
&self.tunnels,
|
||
&self.proxy_tunnels,
|
||
&self.http_tunnels,
|
||
)
|
||
.await?;
|
||
routes.push(db::redis_driver::RedisNodeRoute {
|
||
advertised: node.clone(),
|
||
connect: db::redis_driver::RedisNodeEndpoint { host: "127.0.0.1".to_string(), port: local_port },
|
||
});
|
||
}
|
||
Ok(routes)
|
||
}
|
||
|
||
#[cfg(feature = "mq-admin")]
|
||
async fn rocketmq_socks_proxy_for_transport_layers(
|
||
&self,
|
||
connection_id: &str,
|
||
transport_layers: &[TransportLayerConfig],
|
||
) -> Result<crate::mq::config::MqSocksProxy, String> {
|
||
// ProxyType 仅此 mq-admin 分支用到,局部导入避免在关闭 mq-admin 时
|
||
// 顶层 import 触发 unused 警告。
|
||
use crate::models::connection::ProxyType;
|
||
let final_layer = transport_layers.last().ok_or("No transport layers configured")?;
|
||
match final_layer {
|
||
TransportLayerConfig::Ssh(_) => {
|
||
let local_port = db::transport_layer_tunnel::start_transport_layers_with_final_ssh_socks5(
|
||
connection_id,
|
||
transport_layers,
|
||
&self.tunnels,
|
||
&self.proxy_tunnels,
|
||
&self.http_tunnels,
|
||
)
|
||
.await?;
|
||
Ok(crate::mq::config::MqSocksProxy {
|
||
host: "127.0.0.1".to_string(),
|
||
port: local_port,
|
||
username: String::new(),
|
||
password: String::new(),
|
||
})
|
||
}
|
||
TransportLayerConfig::Proxy(proxy) if proxy.proxy_type == ProxyType::Socks5 => {
|
||
if transport_layers.len() == 1 {
|
||
Ok(crate::mq::config::MqSocksProxy {
|
||
host: proxy.host.clone(),
|
||
port: proxy.port,
|
||
username: proxy.username.clone(),
|
||
password: proxy.password.clone(),
|
||
})
|
||
} else {
|
||
let local_port = db::transport_layer_tunnel::start_transport_layers(
|
||
connection_id,
|
||
&transport_layers[..transport_layers.len() - 1],
|
||
&proxy.host,
|
||
proxy.port,
|
||
&self.tunnels,
|
||
&self.proxy_tunnels,
|
||
&self.http_tunnels,
|
||
)
|
||
.await?;
|
||
Ok(crate::mq::config::MqSocksProxy {
|
||
host: "127.0.0.1".to_string(),
|
||
port: local_port,
|
||
username: proxy.username.clone(),
|
||
password: proxy.password.clone(),
|
||
})
|
||
}
|
||
}
|
||
TransportLayerConfig::Proxy(_) => {
|
||
Err("RocketMQ requires a SOCKS5 proxy as the final proxy layer".to_string())
|
||
}
|
||
TransportLayerConfig::HttpTunnel(_) => {
|
||
Err("RocketMQ does not support an HTTP tunnel as the final transport layer".to_string())
|
||
}
|
||
}
|
||
}
|
||
|
||
#[cfg(feature = "mq-admin")]
|
||
pub async fn mq_admin_config_for_connection(
|
||
&self,
|
||
connection_id: &str,
|
||
config: &ConnectionConfig,
|
||
) -> Result<crate::mq::config::MqAdminConfig, String> {
|
||
let mqc = crate::mq::config::MqAdminConfig::from_connection(config)?;
|
||
if !config.has_effective_transport_layers() {
|
||
return Ok(mqc);
|
||
}
|
||
|
||
if mqc.system_kind == crate::mq::types::MqSystemKind::RocketMq {
|
||
let transport_layers = self.resolved_transport_layers(config).await?;
|
||
let proxy = self.rocketmq_socks_proxy_for_transport_layers(connection_id, &transport_layers).await?;
|
||
return Ok(mqc.with_socks_proxy(&proxy.host, proxy.port, &proxy.username, &proxy.password));
|
||
}
|
||
|
||
if mqc.system_kind == crate::mq::types::MqSystemKind::RabbitMq {
|
||
let transport_layers = self.resolved_transport_layers(config).await?;
|
||
let (amqp_host, amqp_port) = crate::mq::adapters::rabbitmq::primary_amqp_endpoint(&mqc)?;
|
||
let management_endpoint = crate::mq::adapters::rabbitmq::management_endpoint(&mqc)?;
|
||
let amqp_local_port = db::transport_layer_tunnel::start_transport_layers(
|
||
connection_id,
|
||
&transport_layers,
|
||
&amqp_host,
|
||
amqp_port,
|
||
&self.tunnels,
|
||
&self.proxy_tunnels,
|
||
&self.http_tunnels,
|
||
)
|
||
.await?;
|
||
let mut mqc = mqc.with_connect_override("127.0.0.1", amqp_local_port);
|
||
if let Some((management_host, management_port)) = management_endpoint {
|
||
let management_transport_id = rabbitmq_management_transport_id(connection_id);
|
||
let management_local_port = match db::transport_layer_tunnel::start_transport_layers(
|
||
&management_transport_id,
|
||
&transport_layers,
|
||
&management_host,
|
||
management_port,
|
||
&self.tunnels,
|
||
&self.proxy_tunnels,
|
||
&self.http_tunnels,
|
||
)
|
||
.await
|
||
{
|
||
Ok(port) => port,
|
||
Err(error) => {
|
||
db::transport_layer_tunnel::stop_transport_layers(
|
||
&management_transport_id,
|
||
transport_layers.len(),
|
||
&self.tunnels,
|
||
&self.proxy_tunnels,
|
||
&self.http_tunnels,
|
||
)
|
||
.await;
|
||
db::transport_layer_tunnel::stop_transport_layers(
|
||
connection_id,
|
||
transport_layers.len(),
|
||
&self.tunnels,
|
||
&self.proxy_tunnels,
|
||
&self.http_tunnels,
|
||
)
|
||
.await;
|
||
return Err(error);
|
||
}
|
||
};
|
||
mqc = mqc.with_management_connect_override("127.0.0.1", management_local_port);
|
||
}
|
||
return Ok(mqc);
|
||
}
|
||
|
||
if mqc.system_kind == crate::mq::types::MqSystemKind::Kafka {
|
||
if let Some((bootstrap_host, bootstrap_port)) = kafka_single_loopback_bootstrap_endpoint(&mqc.extra) {
|
||
let transport_layers = self.resolved_transport_layers(config).await?;
|
||
if matches!(transport_layers.last(), Some(TransportLayerConfig::Ssh(_))) {
|
||
let local_port = db::transport_layer_tunnel::start_transport_layers_with_final_ssh_local_port(
|
||
connection_id,
|
||
&transport_layers,
|
||
&bootstrap_host,
|
||
bootstrap_port,
|
||
Some(bootstrap_port),
|
||
&self.tunnels,
|
||
&self.proxy_tunnels,
|
||
&self.http_tunnels,
|
||
)
|
||
.await?;
|
||
return Ok(mqc.with_connect_override("127.0.0.1", local_port));
|
||
}
|
||
}
|
||
}
|
||
|
||
let (host, port) = self.connection_host_port(connection_id, config).await?;
|
||
Ok(mqc.with_connect_override(&host, port))
|
||
}
|
||
|
||
pub async fn nacos_admin_config_for_connection(
|
||
&self,
|
||
connection_id: &str,
|
||
config: &ConnectionConfig,
|
||
) -> Result<crate::nacos::config::NacosAdminConfig, String> {
|
||
let mut nacos_config = crate::nacos::config::NacosAdminConfig::from_connection(config)?;
|
||
if !config.save_password {
|
||
let owner = crate::session_credentials::current_credential_owner().unwrap_or_default();
|
||
let primary_password = self
|
||
.session_credentials
|
||
.get_for_purpose(&owner, connection_id, NACOS_PRIMARY_SESSION_PASSWORD)
|
||
.or_else(|| self.session_credentials.get(&owner, connection_id));
|
||
if let (Some(password), crate::nacos::config::NacosAuthConfig::UsernamePassword { password: target, .. }) =
|
||
(primary_password, &mut nacos_config.auth)
|
||
{
|
||
*target = password;
|
||
}
|
||
if let (
|
||
Some(password),
|
||
crate::nacos::config::NacosRNacosConsoleAuth::UsernamePassword { password: target, .. },
|
||
) = (
|
||
self.session_credentials.get_for_purpose(&owner, connection_id, NACOS_CONSOLE_SESSION_PASSWORD),
|
||
&mut nacos_config.rnacos_console_auth,
|
||
) {
|
||
*target = password;
|
||
}
|
||
}
|
||
if !config.has_effective_transport_layers() {
|
||
return Ok(nacos_config);
|
||
}
|
||
|
||
let (host, port) = self.connection_host_port(connection_id, config).await?;
|
||
let nacos_config = nacos_config.with_server_endpoint(&host, port)?;
|
||
let transport_layers = self.resolved_transport_layers(config).await?;
|
||
if transport_layers.is_empty() {
|
||
return Ok(nacos_config);
|
||
}
|
||
|
||
if nacos_config.rnacos_console_addr.is_empty() {
|
||
return Ok(nacos_config);
|
||
}
|
||
|
||
let console_url = reqwest::Url::parse(&nacos_config.rnacos_console_addr)
|
||
.map_err(|error| format!("r-nacos console address is invalid: {error}"))?;
|
||
let console_host = console_url
|
||
.host_str()
|
||
.filter(|host| !host.is_empty())
|
||
.ok_or_else(|| "r-nacos console address does not include a host".to_string())?;
|
||
let console_port = console_url
|
||
.port_or_known_default()
|
||
.ok_or_else(|| "r-nacos console address does not include a port".to_string())?;
|
||
let console_transport_id = rnacos_console_transport_id(connection_id);
|
||
let local_port = match db::transport_layer_tunnel::start_transport_layers(
|
||
&console_transport_id,
|
||
&transport_layers,
|
||
console_host,
|
||
console_port,
|
||
&self.tunnels,
|
||
&self.proxy_tunnels,
|
||
&self.http_tunnels,
|
||
)
|
||
.await
|
||
{
|
||
Ok(port) => port,
|
||
Err(error) => {
|
||
db::transport_layer_tunnel::stop_transport_layers(
|
||
&console_transport_id,
|
||
transport_layers.len(),
|
||
&self.tunnels,
|
||
&self.proxy_tunnels,
|
||
&self.http_tunnels,
|
||
)
|
||
.await;
|
||
return Err(format!("r-nacos console transport failed: {error}"));
|
||
}
|
||
};
|
||
match nacos_config.with_rnacos_console_endpoint("127.0.0.1", local_port) {
|
||
Ok(nacos_config) => Ok(nacos_config),
|
||
Err(error) => {
|
||
db::transport_layer_tunnel::stop_transport_layers(
|
||
&console_transport_id,
|
||
transport_layers.len(),
|
||
&self.tunnels,
|
||
&self.proxy_tunnels,
|
||
&self.http_tunnels,
|
||
)
|
||
.await;
|
||
Err(error)
|
||
}
|
||
}
|
||
}
|
||
|
||
pub async fn replace_nacos_session_credential(
|
||
&self,
|
||
owner_scope: &str,
|
||
connection_id: &str,
|
||
username: &str,
|
||
password: &str,
|
||
) -> Result<(), String> {
|
||
let username = username.trim();
|
||
if username.is_empty() || password.is_empty() {
|
||
return Err("Nacos username and replacement password are required".to_string());
|
||
}
|
||
let config = self.configs.read().await.get(connection_id).cloned().ok_or("Connection not found")?;
|
||
if config.db_type != DatabaseType::Nacos {
|
||
return Err("Connection is not a Nacos connection".to_string());
|
||
}
|
||
if config.save_password {
|
||
return Err("Connection saves its password; update the persisted connection instead".to_string());
|
||
}
|
||
let nacos_config = crate::nacos::config::NacosAdminConfig::from_connection(&config)?;
|
||
let primary_matches = matches!(
|
||
&nacos_config.auth,
|
||
crate::nacos::config::NacosAuthConfig::UsernamePassword { username: current, .. } if current == username
|
||
);
|
||
let console_matches = matches!(
|
||
&nacos_config.rnacos_console_auth,
|
||
crate::nacos::config::NacosRNacosConsoleAuth::UsernamePassword { username: current, .. } if current == username
|
||
);
|
||
if !primary_matches && !console_matches {
|
||
return Err(format!("Nacos user {username} is not used by this connection"));
|
||
}
|
||
if primary_matches {
|
||
let _ = self.session_credentials.set(owner_scope, connection_id, password);
|
||
self.session_credentials.set_for_purpose(
|
||
owner_scope,
|
||
connection_id,
|
||
NACOS_PRIMARY_SESSION_PASSWORD,
|
||
password,
|
||
);
|
||
}
|
||
if console_matches {
|
||
self.session_credentials.set_for_purpose(
|
||
owner_scope,
|
||
connection_id,
|
||
NACOS_CONSOLE_SESSION_PASSWORD,
|
||
password,
|
||
);
|
||
}
|
||
self.nacos_registry.drop_connection(connection_id).await;
|
||
Ok(())
|
||
}
|
||
|
||
async fn remove_stale_connection_pool(&self, pool_key: &str) -> bool {
|
||
if self.running_queries.is_pool_active(pool_key) {
|
||
return false;
|
||
}
|
||
|
||
let Some(checked) = self.pool_publication_snapshot(pool_key).await else {
|
||
return false;
|
||
};
|
||
let stale = {
|
||
match &checked.pool {
|
||
PoolKind::Mysql(pool, _) => {
|
||
let pool = pool.clone();
|
||
match db::mysql::checkout_mysql_conn(&pool, HEALTH_CHECK_POOL_ACQUIRE_TIMEOUT).await {
|
||
// The 500 ms probe budget is intentionally shorter than a foreground checkout. A timeout
|
||
// while waiting, creating, or recycling is inconclusive: slow remote handshakes and active
|
||
// metadata exports can legitimately exceed it. Removing the pool here would start competing
|
||
// reconnects while useful work is still running.
|
||
Err(err @ db::PoolCheckoutError::Timeout { .. }) => {
|
||
log::debug!(
|
||
"MySQL connection pool '{pool_key}' did not finish a health checkout; keeping pool: {err}"
|
||
);
|
||
false
|
||
}
|
||
Err(err) => {
|
||
log::warn!("MySQL connection pool '{pool_key}' is stale: {err}");
|
||
true
|
||
}
|
||
Ok(mut conn) => {
|
||
let timeout = crate::db::connection_timeout();
|
||
match tokio::time::timeout(timeout, conn.ping()).await {
|
||
Ok(Ok(())) => false,
|
||
Ok(Err(err)) => {
|
||
log::warn!("MySQL connection pool '{pool_key}' is stale: {err}");
|
||
true
|
||
}
|
||
Err(_) => {
|
||
log::warn!("MySQL connection pool '{pool_key}' is stale: health check timed out");
|
||
true
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
PoolKind::Postgres(pool) => {
|
||
let pool = pool.clone();
|
||
let timeout = crate::db::connection_timeout();
|
||
match db::postgres::checkout_postgres_client_classified(
|
||
&pool,
|
||
None,
|
||
HEALTH_CHECK_POOL_ACQUIRE_TIMEOUT,
|
||
)
|
||
.await
|
||
{
|
||
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
|
||
}
|
||
},
|
||
Err(err) if err.is_pool_saturation() => {
|
||
log::debug!(
|
||
"PostgreSQL connection pool '{pool_key}' is busy; skipping health probe: {err}"
|
||
);
|
||
false
|
||
}
|
||
Err(err) => {
|
||
log::warn!("PostgreSQL connection pool '{pool_key}' is stale: {err}");
|
||
true
|
||
}
|
||
}
|
||
}
|
||
PoolKind::SqlServer(client) => {
|
||
let client = client.clone();
|
||
let mut client = client.lock().await;
|
||
match db::sqlserver::test_connection(&mut client).await {
|
||
Ok(()) => false,
|
||
Err(err) => {
|
||
log::warn!("SQL Server connection pool '{pool_key}' is stale: {err}");
|
||
true
|
||
}
|
||
}
|
||
}
|
||
PoolKind::Redis(redis) => match db::redis_driver::test_connection(redis).await {
|
||
Ok(()) => false,
|
||
Err(err) => {
|
||
log::warn!("Redis connection pool '{pool_key}' is stale: {err}");
|
||
true
|
||
}
|
||
},
|
||
PoolKind::MongoDb(client) => {
|
||
let client = client.clone();
|
||
let (connect_timeout, database) = {
|
||
let configs = self.configs.read().await;
|
||
let config = config_for_pool_key(pool_key, &configs);
|
||
(
|
||
config
|
||
.map(|config| Duration::from_secs(config.effective_connect_timeout_secs().max(1)))
|
||
.unwrap_or_else(|| Duration::from_secs(1)),
|
||
config.and_then(|config| config.effective_database().map(str::to_string)),
|
||
)
|
||
};
|
||
match db::mongo_driver::test_connection(&client, connect_timeout, database.as_deref()).await {
|
||
Ok(()) => false,
|
||
Err(err) => {
|
||
log::warn!("MongoDB connection pool '{pool_key}' is stale: {err}");
|
||
true
|
||
}
|
||
}
|
||
}
|
||
PoolKind::DynamoDb(client) => {
|
||
let client = client.clone();
|
||
let timeout = crate::db::connection_timeout();
|
||
match db::dynamodb_driver::test_connection(&client, timeout).await {
|
||
Ok(()) => false,
|
||
Err(err) => {
|
||
log::warn!("DynamoDB connection pool '{pool_key}' is stale: {err}");
|
||
true
|
||
}
|
||
}
|
||
}
|
||
PoolKind::ClickHouse(client) => {
|
||
let client = client.clone();
|
||
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();
|
||
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::Easysearch(client) => {
|
||
let mut client = client.clone();
|
||
let timeout = crate::db::connection_timeout();
|
||
match db::easysearch_driver::test_connection(&mut client, timeout).await {
|
||
Ok(()) => false,
|
||
Err(err) => {
|
||
log::warn!("Easysearch connection pool '{pool_key}' is stale: {err}");
|
||
true
|
||
}
|
||
}
|
||
}
|
||
PoolKind::Meilisearch(client) => {
|
||
let client = client.clone();
|
||
let timeout = crate::db::connection_timeout();
|
||
match db::meilisearch_driver::test_connection(&client, timeout).await {
|
||
Ok(()) => false,
|
||
Err(err) => {
|
||
log::warn!("Meilisearch connection pool '{pool_key}' is stale: {err}");
|
||
true
|
||
}
|
||
}
|
||
}
|
||
PoolKind::HBase(client) => {
|
||
let client = client.clone();
|
||
let timeout = crate::db::connection_timeout();
|
||
match db::hbase_driver::test_connection(&client, timeout).await {
|
||
Ok(_) => false,
|
||
Err(err) => {
|
||
log::warn!("HBase connection pool '{pool_key}' is stale: {err}");
|
||
true
|
||
}
|
||
}
|
||
}
|
||
PoolKind::VectorDb(client) => {
|
||
let client = client.clone();
|
||
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();
|
||
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::InfluxDb3(client) => {
|
||
let client = client.clone();
|
||
let timeout = crate::db::connection_timeout();
|
||
match db::influxdb3_driver::test_connection(&client, timeout).await {
|
||
Ok(()) => false,
|
||
Err(err) => {
|
||
log::warn!("InfluxDB 3 connection pool '{pool_key}' is stale: {err}");
|
||
true
|
||
}
|
||
}
|
||
}
|
||
PoolKind::VictoriaMetrics(client) => {
|
||
let client = client.clone();
|
||
let timeout = crate::db::connection_timeout();
|
||
match db::victoriametrics_driver::test_connection(&client, timeout).await {
|
||
Ok(()) => false,
|
||
Err(err) => {
|
||
log::warn!("VictoriaMetrics connection pool '{pool_key}' is stale: {err}");
|
||
true
|
||
}
|
||
}
|
||
}
|
||
PoolKind::Rqlite(client) => {
|
||
let client = client.clone();
|
||
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();
|
||
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::CloudflareD1(client) => {
|
||
let client = client.clone();
|
||
let timeout = crate::db::connection_timeout();
|
||
match db::cloudflare_d1_driver::test_connection(&client, timeout).await {
|
||
Ok(()) => false,
|
||
Err(err) => {
|
||
log::warn!("Cloudflare D1 connection pool '{pool_key}' is stale: {err}");
|
||
true
|
||
}
|
||
}
|
||
}
|
||
PoolKind::Agent(client) => {
|
||
let client = client.clone();
|
||
let Ok(mut agent) = client.try_lock() else {
|
||
log::debug!("Agent connection pool '{pool_key}' is busy; skipping health probe");
|
||
return false;
|
||
};
|
||
let timeout = crate::db::connection_timeout();
|
||
match agent.validate_connection_typed(Some(timeout)).await {
|
||
Ok(_) => false,
|
||
Err(err) if is_agent_validate_connection_unsupported(&err.to_string()) => {
|
||
log::debug!(
|
||
"Agent connection pool '{pool_key}' does not support validate_connection; keeping pool"
|
||
);
|
||
false
|
||
}
|
||
Err(err) if RecoveryPolicy::decide(&err, RecoveryScope::Keepalive).replaces_runtime() => {
|
||
log::warn!(
|
||
"Agent connection pool '{pool_key}' requested runtime replacement during health probe: {err}"
|
||
);
|
||
drop(agent);
|
||
return self
|
||
.pool_routing_control()
|
||
.detach_agent_pool_if_current(pool_key, &client, Some(&checked.publication), true)
|
||
.await;
|
||
}
|
||
Err(err) => {
|
||
log::warn!("Agent connection pool '{pool_key}' is stale: {err}");
|
||
drop(agent);
|
||
return self
|
||
.pool_routing_control()
|
||
.detach_agent_pool_if_current(pool_key, &client, Some(&checked.publication), false)
|
||
.await;
|
||
}
|
||
}
|
||
}
|
||
PoolKind::Sqlite(_)
|
||
| PoolKind::DuckDbWorker(_)
|
||
| PoolKind::ExternalDriver { .. }
|
||
| PoolKind::MessageQueue
|
||
| PoolKind::Nacos
|
||
| PoolKind::Consul(_) => false,
|
||
PoolKind::PluginConnection(handle) => !handle.is_running(),
|
||
#[cfg(feature = "mq-admin")]
|
||
PoolKind::Mqtt(_) => false,
|
||
}
|
||
};
|
||
|
||
if !stale {
|
||
return false;
|
||
}
|
||
|
||
self.remove_stale_pool_if_current(pool_key, &checked.publication).await
|
||
}
|
||
async fn remove_stale_pool_if_current(&self, pool_key: &str, checked_publication: &PoolPublication) -> bool {
|
||
self.remove_stale_pool_if_current_inner(
|
||
pool_key,
|
||
checked_publication,
|
||
#[cfg(test)]
|
||
None,
|
||
)
|
||
.await
|
||
}
|
||
|
||
async fn remove_stale_pool_if_current_inner(
|
||
&self,
|
||
pool_key: &str,
|
||
checked_publication: &PoolPublication,
|
||
#[cfg(test)] cleanup_barriers: Option<StalePoolCleanupBarriers>,
|
||
) -> bool {
|
||
#[cfg(test)]
|
||
if let Some((cleanup_ready, continue_cleanup)) =
|
||
cleanup_barriers.as_ref().and_then(|barriers| barriers.before_removal.as_ref())
|
||
{
|
||
cleanup_ready.wait().await;
|
||
continue_cleanup.wait().await;
|
||
}
|
||
|
||
let routing = self.pool_routing_control();
|
||
let removed = loop {
|
||
let mut connections = self.connections.write().await;
|
||
let is_current =
|
||
connections.publications.get(pool_key).is_some_and(|current| current.is_same(checked_publication));
|
||
if !is_current {
|
||
log::debug!(
|
||
"Connection pool '{pool_key}' was replaced while its health check was running; keeping the current route"
|
||
);
|
||
return false;
|
||
}
|
||
let Ok(mut activity) = self.pool_activity.try_write() else {
|
||
drop(connections);
|
||
tokio::task::yield_now().await;
|
||
continue;
|
||
};
|
||
let removed = connections
|
||
.remove_if_publication(pool_key, checked_publication)
|
||
.expect("checked pool publication must remain current while routing is locked");
|
||
#[cfg(test)]
|
||
if let Some((route_removed, continue_cleanup)) =
|
||
cleanup_barriers.as_ref().and_then(|barriers| barriers.after_removal.as_ref())
|
||
{
|
||
route_removed.wait().await;
|
||
continue_cleanup.wait().await;
|
||
}
|
||
routing.stop_keepalive(pool_key);
|
||
activity.remove(pool_key);
|
||
break removed;
|
||
};
|
||
|
||
self.postgres_cancel_contexts.write().await.remove(pool_key);
|
||
routing.close_pool_with_timeout(pool_key.to_string(), removed).await;
|
||
true
|
||
}
|
||
|
||
pub async fn reconnect_pool(&self, connection_id: &str, database: Option<&str>) -> Result<String, String> {
|
||
self.reconnect_pool_for_session(connection_id, database, None).await
|
||
}
|
||
|
||
pub async fn reconnect_pool_for_session(
|
||
&self,
|
||
connection_id: &str,
|
||
database: Option<&str>,
|
||
client_session_id: Option<&str>,
|
||
) -> Result<String, String> {
|
||
self.reconnect_pool_for_session_with_catalog(connection_id, database, None, client_session_id).await
|
||
}
|
||
|
||
pub async fn reconnect_pool_for_session_with_catalog(
|
||
&self,
|
||
connection_id: &str,
|
||
database: Option<&str>,
|
||
catalog: Option<&str>,
|
||
client_session_id: Option<&str>,
|
||
) -> Result<String, String> {
|
||
self.reconnect_pool_for_session_with_catalog_and_role(
|
||
connection_id,
|
||
database,
|
||
catalog,
|
||
client_session_id,
|
||
AgentSessionRole::Workload,
|
||
)
|
||
.await
|
||
}
|
||
|
||
pub(crate) async fn reconnect_metadata_pool_for_session(
|
||
&self,
|
||
connection_id: &str,
|
||
database: Option<&str>,
|
||
client_session_id: Option<&str>,
|
||
) -> Result<String, String> {
|
||
self.reconnect_pool_for_session_with_catalog_and_role(
|
||
connection_id,
|
||
database,
|
||
None,
|
||
client_session_id,
|
||
AgentSessionRole::Metadata,
|
||
)
|
||
.await
|
||
}
|
||
|
||
async fn reconnect_pool_for_session_with_catalog_and_role(
|
||
&self,
|
||
connection_id: &str,
|
||
database: Option<&str>,
|
||
catalog: Option<&str>,
|
||
client_session_id: Option<&str>,
|
||
session_role: AgentSessionRole,
|
||
) -> Result<String, String> {
|
||
let config = {
|
||
let configs = self.configs.read().await;
|
||
configs.get(connection_id).cloned()
|
||
};
|
||
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_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;
|
||
self.reset_connection_transport(connection_id).await;
|
||
} else {
|
||
self.stop_keepalive_task(&pool_key).await;
|
||
self.pool_activity.write().await.remove(&pool_key);
|
||
self.postgres_cancel_contexts.write().await.remove(&pool_key);
|
||
let removed = self.connections.write().await.remove(&pool_key);
|
||
if let Some(pool) = removed {
|
||
self.pool_routing_control().close_pool_with_timeout(pool_key.clone(), pool).await;
|
||
}
|
||
}
|
||
self.get_or_create_pool_for_session_inner(
|
||
connection_id,
|
||
database,
|
||
catalog,
|
||
client_session_id,
|
||
session_role,
|
||
None,
|
||
)
|
||
.await
|
||
}
|
||
|
||
pub async fn close_client_session_pool(
|
||
&self,
|
||
connection_id: &str,
|
||
database: Option<&str>,
|
||
client_session_id: &str,
|
||
) -> Result<bool, String> {
|
||
let Some((pool_key, pool)) = self
|
||
.take_client_session_pool(connection_id, database, client_session_id, AgentSessionRole::Workload)
|
||
.await?
|
||
else {
|
||
return Ok(false);
|
||
};
|
||
self.pool_routing_control().close_pool_with_timeout(pool_key, pool).await;
|
||
Ok(true)
|
||
}
|
||
|
||
pub(crate) async fn close_metadata_session_pool(
|
||
&self,
|
||
connection_id: &str,
|
||
database: Option<&str>,
|
||
client_session_id: &str,
|
||
) -> Result<bool, String> {
|
||
let Some((pool_key, pool)) = self
|
||
.take_client_session_pool(connection_id, database, client_session_id, AgentSessionRole::Metadata)
|
||
.await?
|
||
else {
|
||
return Ok(false);
|
||
};
|
||
self.pool_routing_control().close_pool_with_timeout(pool_key, pool).await;
|
||
Ok(true)
|
||
}
|
||
|
||
pub(crate) async fn metadata_session_pool_cleanup_guard(
|
||
&self,
|
||
connection_id: &str,
|
||
database: Option<&str>,
|
||
client_session_id: &str,
|
||
) -> Option<ClientSessionPoolCleanupGuard> {
|
||
self.client_session_pool_cleanup_guard_for_role(
|
||
connection_id,
|
||
database,
|
||
client_session_id,
|
||
AgentSessionRole::Metadata,
|
||
)
|
||
.await
|
||
}
|
||
|
||
pub(crate) async fn workload_session_pool_cleanup_guard(
|
||
&self,
|
||
connection_id: &str,
|
||
database: Option<&str>,
|
||
client_session_id: &str,
|
||
) -> Option<ClientSessionPoolCleanupGuard> {
|
||
self.client_session_pool_cleanup_guard_for_role(
|
||
connection_id,
|
||
database,
|
||
client_session_id,
|
||
AgentSessionRole::Workload,
|
||
)
|
||
.await
|
||
}
|
||
|
||
async fn client_session_pool_cleanup_guard_for_role(
|
||
&self,
|
||
connection_id: &str,
|
||
database: Option<&str>,
|
||
client_session_id: &str,
|
||
session_role: AgentSessionRole,
|
||
) -> Option<ClientSessionPoolCleanupGuard> {
|
||
let config = {
|
||
let configs = self.configs.read().await;
|
||
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_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 {
|
||
return None;
|
||
}
|
||
Some(ClientSessionPoolCleanupGuard { pool_key, routing: self.pool_routing_control(), armed: true })
|
||
}
|
||
|
||
/// Removes a session-scoped pool immediately and schedules the potentially slow driver
|
||
/// shutdown on the supervised background task set.
|
||
pub async fn detach_client_session_pool(
|
||
&self,
|
||
connection_id: &str,
|
||
database: Option<&str>,
|
||
client_session_id: &str,
|
||
) -> Result<bool, String> {
|
||
let Some(removed) = self
|
||
.take_client_session_pool(connection_id, database, client_session_id, AgentSessionRole::Workload)
|
||
.await?
|
||
else {
|
||
return Ok(false);
|
||
};
|
||
self.pool_routing_control().close_removed_in_background(vec![removed]);
|
||
Ok(true)
|
||
}
|
||
|
||
#[cfg(test)]
|
||
pub(crate) async fn replace_runtime_for_metadata_pool(
|
||
&self,
|
||
connection_id: &str,
|
||
database: Option<&str>,
|
||
client_session_id: Option<&str>,
|
||
) -> bool {
|
||
let config = {
|
||
let configs = self.configs.read().await;
|
||
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_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
|
||
}
|
||
|
||
pub(crate) async fn detach_metadata_pool_after_recovery(
|
||
&self,
|
||
connection_id: &str,
|
||
database: Option<&str>,
|
||
client_session_id: Option<&str>,
|
||
agent_session_id: Option<&str>,
|
||
replace_agent_runtime: bool,
|
||
) -> bool {
|
||
let config = {
|
||
let configs = self.configs.read().await;
|
||
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_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 {
|
||
let expected_client = {
|
||
match self.pool_handle(&pool_key).await.as_ref() {
|
||
Some(PoolKind::Agent(client)) if client.matches_session_id(session_id) => Some(client.clone()),
|
||
Some(PoolKind::Agent(_)) => return false,
|
||
_ => None,
|
||
}
|
||
};
|
||
if let Some(client) = expected_client {
|
||
return self.detach_agent_pool_if_current(&pool_key, &client, replace_agent_runtime).await;
|
||
}
|
||
}
|
||
self.detach_pool_by_key(&pool_key, replace_agent_runtime).await
|
||
}
|
||
|
||
/// Detaches a pool before cleanup so a stuck Agent close cannot delay replacement.
|
||
pub async fn detach_pool_by_key(&self, pool_key: &str, replace_agent_runtime: bool) -> bool {
|
||
self.pool_routing_control().detach_pool_by_key(pool_key, replace_agent_runtime).await
|
||
}
|
||
|
||
pub(crate) async fn detach_agent_pool_if_current(
|
||
&self,
|
||
pool_key: &str,
|
||
expected_client: &Arc<db::agent_driver::PooledAgentClient>,
|
||
replace_agent_runtime: bool,
|
||
) -> bool {
|
||
self.pool_routing_control()
|
||
.detach_agent_pool_if_current(pool_key, expected_client, None, replace_agent_runtime)
|
||
.await
|
||
}
|
||
|
||
async fn take_client_session_pool(
|
||
&self,
|
||
connection_id: &str,
|
||
database: Option<&str>,
|
||
client_session_id: &str,
|
||
session_role: AgentSessionRole,
|
||
) -> Result<Option<(String, PoolKind)>, String> {
|
||
let session = normalize_client_session_id(Some(client_session_id));
|
||
let Some(session) = session else {
|
||
return Ok(None);
|
||
};
|
||
let config = {
|
||
let configs = self.configs.read().await;
|
||
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_key = pool_key_for_session_role(config.as_ref(), base_pool_key.clone(), Some(&session), session_role);
|
||
if pool_key == base_pool_key {
|
||
return Ok(None);
|
||
}
|
||
self.stop_keepalive_task(&pool_key).await;
|
||
self.pool_activity.write().await.remove(&pool_key);
|
||
self.postgres_cancel_contexts.write().await.remove(&pool_key);
|
||
let removed = self.update_connection_pools(|connections| connections.remove(&pool_key)).await;
|
||
Ok(removed.map(|pool| (pool_key, pool)))
|
||
}
|
||
|
||
pub async fn remove_pool_by_key(&self, pool_key: &str) -> bool {
|
||
self.stop_keepalive_task(pool_key).await;
|
||
self.pool_activity.write().await.remove(pool_key);
|
||
self.postgres_cancel_contexts.write().await.remove(pool_key);
|
||
let removed = self.connections.write().await.remove(pool_key);
|
||
if let Some(pool) = removed {
|
||
self.pool_routing_control().close_pool_with_timeout(pool_key.to_string(), pool).await;
|
||
true
|
||
} else {
|
||
false
|
||
}
|
||
}
|
||
|
||
async fn reclaim_idle_base_pool_for_session(&self, connection_id: &str, preferred_base_pool_key: &str) -> bool {
|
||
let pool_prefix = format!("{connection_id}:");
|
||
let activity = self.pool_activity.read().await;
|
||
let connections = self.connection_pools_snapshot().await;
|
||
let mut candidates: Vec<(String, (usize, u64))> = connections
|
||
.iter()
|
||
.filter_map(|(key, pool)| {
|
||
if !matches!(pool, PoolKind::Agent(_))
|
||
|| (key != connection_id && !key.starts_with(&pool_prefix))
|
||
|| is_session_scoped_pool_key(key)
|
||
|| self.running_queries.is_pool_active(key)
|
||
{
|
||
return None;
|
||
}
|
||
let preferred_rank = usize::from(key != preferred_base_pool_key);
|
||
let last_used = activity
|
||
.get(key)
|
||
.map(|value| value.last_used_at_ms.load(std::sync::atomic::Ordering::Relaxed))
|
||
.unwrap_or(0);
|
||
Some((key.clone(), (preferred_rank, last_used)))
|
||
})
|
||
.collect();
|
||
drop(connections);
|
||
drop(activity);
|
||
candidates.sort_by_key(|(_, rank)| *rank);
|
||
|
||
for (pool_key, _) in candidates {
|
||
let Some(_drain) = self.begin_pool_drain(&pool_key) else {
|
||
continue;
|
||
};
|
||
if self.try_reclaim_idle_agent_pool(&pool_key).await {
|
||
return true;
|
||
}
|
||
}
|
||
false
|
||
}
|
||
|
||
async fn try_reclaim_idle_agent_pool(&self, pool_key: &str) -> bool {
|
||
self.stop_keepalive_task_and_wait(pool_key).await;
|
||
if self.running_queries.is_pool_active(pool_key) {
|
||
self.restart_agent_keepalive(pool_key).await;
|
||
return false;
|
||
}
|
||
|
||
let removed = {
|
||
let mut connections = self.connections.write().await;
|
||
let exclusively_idle = match connections.get(pool_key) {
|
||
Some(PoolKind::Agent(client)) => Arc::strong_count(client) == 1 && client.try_lock().is_ok(),
|
||
_ => false,
|
||
};
|
||
exclusively_idle.then(|| connections.remove(pool_key)).flatten()
|
||
};
|
||
let Some(pool) = removed else {
|
||
self.restart_agent_keepalive(pool_key).await;
|
||
return false;
|
||
};
|
||
|
||
self.pool_activity.write().await.remove(pool_key);
|
||
self.postgres_cancel_contexts.write().await.remove(pool_key);
|
||
match close_reclaimed_agent_pool(pool).await {
|
||
Ok(()) => true,
|
||
Err((PoolKind::Agent(client), error)) if should_replace_agent_runtime(&error) => {
|
||
log::warn!("Reclaimed Agent pool '{pool_key}' requested runtime replacement while closing: {error}");
|
||
self.pool_routing_control().replace_runtime_after_close_failure(pool_key, &client).await;
|
||
true
|
||
}
|
||
Err((pool, error)) => {
|
||
log::warn!("Failed to close reclaimed Agent pool '{pool_key}': {error}; restoring the pool");
|
||
self.restore_reclaimed_agent_pool(pool_key, pool).await;
|
||
false
|
||
}
|
||
}
|
||
}
|
||
|
||
async fn restart_agent_keepalive(&self, pool_key: &str) {
|
||
let config = {
|
||
let configs = self.configs.read().await;
|
||
config_for_pool_key(pool_key, &configs).cloned()
|
||
};
|
||
let client = {
|
||
match self.pool_handle(pool_key).await.as_ref() {
|
||
Some(PoolKind::Agent(client)) => Some(client.clone()),
|
||
_ => None,
|
||
}
|
||
};
|
||
if let (Some(config), Some(client)) = (config, client) {
|
||
self.start_keepalive_task(
|
||
pool_key,
|
||
&PoolKind::Agent(client),
|
||
&config,
|
||
#[cfg(feature = "mq-admin")]
|
||
None,
|
||
);
|
||
}
|
||
}
|
||
|
||
async fn restore_reclaimed_agent_pool(&self, pool_key: &str, pool: PoolKind) {
|
||
let config = {
|
||
let configs = self.configs.read().await;
|
||
config_for_pool_key(pool_key, &configs).cloned()
|
||
};
|
||
if let Some(config) = config {
|
||
let _ = self.insert_connection_pool_inner(pool_key.to_string(), pool, &config, false).await;
|
||
} else {
|
||
self.pool_activity.write().await.insert(pool_key.to_string(), PoolActivity::now());
|
||
self.connections.write().await.insert(pool_key.to_string(), pool);
|
||
}
|
||
}
|
||
|
||
async fn remove_pool_if_duckdb_isolation_mismatch(&self, _pool_key: &str) -> bool {
|
||
false
|
||
}
|
||
|
||
pub async fn close_database_pool(&self, connection_id: &str, database: Option<&str>) -> Result<bool, String> {
|
||
let db_type = {
|
||
let configs = self.configs.read().await;
|
||
configs.get(connection_id).map(|c| c.db_type)
|
||
};
|
||
if database.is_some() && db_type.is_some_and(|db_type| shares_database_pool_with_connection(&db_type)) {
|
||
return Ok(false);
|
||
}
|
||
let base_pool_key = base_pool_key_for(db_type, connection_id, database, false);
|
||
let session_prefix = format!("{base_pool_key}:session:");
|
||
let metadata_role_key = format!("{base_pool_key}:role:metadata");
|
||
let keys_to_remove: Vec<String> = self
|
||
.connections
|
||
.read()
|
||
.await
|
||
.keys()
|
||
.filter(|key| *key == &base_pool_key || *key == &metadata_role_key || key.starts_with(&session_prefix))
|
||
.cloned()
|
||
.collect();
|
||
self.stop_keepalive_tasks(&keys_to_remove).await;
|
||
{
|
||
let mut activity = self.pool_activity.write().await;
|
||
let mut cancel_contexts = self.postgres_cancel_contexts.write().await;
|
||
for key in &keys_to_remove {
|
||
activity.remove(key);
|
||
cancel_contexts.remove(key);
|
||
}
|
||
}
|
||
let mut conns = self.connections.write().await;
|
||
let mut removed = Vec::with_capacity(keys_to_remove.len());
|
||
for key in keys_to_remove {
|
||
if let Some(pool) = conns.remove(&key) {
|
||
removed.push((key, pool));
|
||
}
|
||
}
|
||
drop(conns);
|
||
let closed = !removed.is_empty();
|
||
self.clear_metadata_gate_for_database(connection_id, database).await;
|
||
for (key, pool) in removed {
|
||
self.pool_routing_control().close_pool_with_timeout(key, pool).await;
|
||
}
|
||
Ok(closed)
|
||
}
|
||
|
||
pub async fn active_agent_connection_driver_keys(&self) -> HashSet<String> {
|
||
let configs = self.configs.read().await;
|
||
let connections = self.connection_pools_snapshot().await;
|
||
let mut keys = HashSet::new();
|
||
|
||
for (pool_key, pool) in connections.iter() {
|
||
#[cfg(feature = "duckdb-sidecar")]
|
||
if matches!(pool, PoolKind::DuckDbWorker(_)) {
|
||
keys.insert("duckdb".to_string());
|
||
continue;
|
||
}
|
||
if !matches!(pool, PoolKind::Agent(_)) {
|
||
continue;
|
||
}
|
||
let Some(config) = config_for_pool_key(pool_key, &configs) else {
|
||
continue;
|
||
};
|
||
if let Some(agent_key) = crate::agent_manager::AgentManager::db_type_to_agent_key(
|
||
&config.db_type,
|
||
config.driver_profile.as_deref(),
|
||
) {
|
||
keys.insert(agent_key.to_string());
|
||
}
|
||
}
|
||
|
||
keys
|
||
}
|
||
|
||
pub async fn prepare_agent_driver_updates(&self, driver_keys: &[String]) -> HashSet<String> {
|
||
let candidates = driver_keys.iter().cloned().collect::<HashSet<_>>();
|
||
if candidates.is_empty() {
|
||
return HashSet::new();
|
||
}
|
||
|
||
let blockers = self
|
||
.active_agent_connection_driver_keys()
|
||
.await
|
||
.into_iter()
|
||
.filter(|key| candidates.contains(key))
|
||
.collect::<HashSet<_>>();
|
||
if !blockers.is_empty() {
|
||
return blockers;
|
||
}
|
||
|
||
for key in &candidates {
|
||
self.agent_manager.stop_daemon_by_key(key).await;
|
||
}
|
||
|
||
// A connection may have started while idle runtimes were stopping.
|
||
self.active_agent_connection_driver_keys().await.into_iter().filter(|key| candidates.contains(key)).collect()
|
||
}
|
||
|
||
pub async fn connection_identifier_quote(
|
||
&self,
|
||
connection_id: &str,
|
||
database: Option<&str>,
|
||
) -> Result<Option<String>, String> {
|
||
let config = self
|
||
.configs
|
||
.read()
|
||
.await
|
||
.get(connection_id)
|
||
.cloned()
|
||
.ok_or_else(|| format!("Connection config not found: {connection_id}"))?;
|
||
let pool_key = self.get_or_create_pool(connection_id, database).await?;
|
||
enum IdentifierQuoteSource {
|
||
NativeGaussdb(deadpool_postgres::Pool),
|
||
Agent(Arc<db::agent_driver::PooledAgentClient>),
|
||
ExternalDriver { config: Arc<ConnectionConfig>, session: Arc<PluginDriverSession> },
|
||
}
|
||
let source = {
|
||
match self.pool_handle(&pool_key).await.as_ref() {
|
||
Some(PoolKind::Postgres(pool)) if config.db_type == DatabaseType::Gaussdb => {
|
||
Some(IdentifierQuoteSource::NativeGaussdb(pool.clone()))
|
||
}
|
||
Some(PoolKind::Agent(client)) if database_capabilities::is_agent_type(&config.db_type) => {
|
||
Some(IdentifierQuoteSource::Agent(client.clone()))
|
||
}
|
||
Some(PoolKind::ExternalDriver { config, session, .. }) => {
|
||
Some(IdentifierQuoteSource::ExternalDriver { config: config.clone(), session: session.clone() })
|
||
}
|
||
_ => None,
|
||
}
|
||
};
|
||
match source {
|
||
Some(IdentifierQuoteSource::NativeGaussdb(pool)) => Ok(db::postgres::gaussdb_identifier_quote(&pool).await),
|
||
Some(IdentifierQuoteSource::Agent(client)) => {
|
||
let mut agent = client.lock().await;
|
||
let info = agent.connection_info(Some(db::connection_timeout())).await?;
|
||
Ok(Some(info.identifier_quote))
|
||
}
|
||
Some(IdentifierQuoteSource::ExternalDriver { config, session }) => {
|
||
let response = session
|
||
.invoke_with_timeout::<db::QueryResult>(
|
||
"executeQuery",
|
||
serde_json::json!({
|
||
"connection": config.as_ref(),
|
||
"sql": db::postgres::GAUSSDB_COMPATIBILITY_SQL,
|
||
"database": config.effective_database().unwrap_or(""),
|
||
"schema": null,
|
||
"maxRows": 1,
|
||
"timeoutSecs": 5,
|
||
}),
|
||
Some(db::connection_timeout()),
|
||
)
|
||
.await;
|
||
Ok(response.ok().and_then(|result| gaussdb_identifier_quote_from_query_result(&result)))
|
||
}
|
||
None => Ok(None),
|
||
}
|
||
}
|
||
|
||
pub async fn connection_database_info(
|
||
&self,
|
||
connection_id: &str,
|
||
database: Option<&str>,
|
||
) -> Result<Option<DatabaseConnectionInfo>, String> {
|
||
let config = self
|
||
.configs
|
||
.read()
|
||
.await
|
||
.get(connection_id)
|
||
.cloned()
|
||
.ok_or_else(|| format!("Connection config not found: {connection_id}"))?;
|
||
let pool_key = self.get_or_create_pool(connection_id, database).await?;
|
||
let source = {
|
||
match self.pool_handle(&pool_key).await.as_ref() {
|
||
Some(PoolKind::Agent(client)) if config.db_type == DatabaseType::MongoDb => {
|
||
Some(ConnectionDatabaseInfoSource::MongoAgent(client.clone(), database.map(str::to_string)))
|
||
}
|
||
Some(PoolKind::Agent(client)) => Some(ConnectionDatabaseInfoSource::Agent(client.clone())),
|
||
Some(PoolKind::ExternalDriver { config, session, .. }) => {
|
||
Some(ConnectionDatabaseInfoSource::ExternalDriver {
|
||
config: config.clone(),
|
||
session: session.clone(),
|
||
})
|
||
}
|
||
Some(PoolKind::Mysql(pool, _)) => Some(ConnectionDatabaseInfoSource::NativeMysql(pool.clone())),
|
||
Some(PoolKind::HBase(client)) => Some(ConnectionDatabaseInfoSource::NativeHBase(client.clone())),
|
||
Some(PoolKind::MongoDb(client)) => {
|
||
Some(ConnectionDatabaseInfoSource::NativeMongo(client.clone(), database.map(str::to_string)))
|
||
}
|
||
Some(PoolKind::Meilisearch(client)) => Some(ConnectionDatabaseInfoSource::Meilisearch(client.clone())),
|
||
Some(PoolKind::VictoriaMetrics(client)) => {
|
||
Some(ConnectionDatabaseInfoSource::VictoriaMetrics(client.clone()))
|
||
}
|
||
Some(PoolKind::Redis(_)) => Some(ConnectionDatabaseInfoSource::Redis(pool_key.clone())),
|
||
Some(PoolKind::Nacos) => Some(ConnectionDatabaseInfoSource::Nacos),
|
||
Some(PoolKind::Consul(client)) => Some(ConnectionDatabaseInfoSource::Consul(Box::new(client.clone()))),
|
||
#[cfg(feature = "mq-admin")]
|
||
Some(PoolKind::MessageQueue) => Some(ConnectionDatabaseInfoSource::MessageQueue),
|
||
_ => None,
|
||
}
|
||
};
|
||
|
||
match source {
|
||
Some(ConnectionDatabaseInfoSource::Agent(client)) => {
|
||
let mut agent = client.lock().await;
|
||
Ok(agent.connection_info(Some(db::connection_timeout())).await?.database_info)
|
||
}
|
||
Some(ConnectionDatabaseInfoSource::MongoAgent(client, database)) => {
|
||
let mut agent = client.lock().await;
|
||
let version = agent.mongo_server_version::<String>(database.as_deref().unwrap_or("admin")).await?;
|
||
Ok(Some(DatabaseConnectionInfo {
|
||
product_name: Some("MongoDB".to_string()),
|
||
product_version: Some(version),
|
||
current_database: database,
|
||
driver_name: Some("MongoDB legacy Agent".to_string()),
|
||
..Default::default()
|
||
}))
|
||
}
|
||
Some(ConnectionDatabaseInfoSource::ExternalDriver { config, session }) => {
|
||
let response = session
|
||
.invoke_with_timeout::<serde_json::Value>(
|
||
"connectionInfo",
|
||
serde_json::json!({ "connection": config.as_ref() }),
|
||
Some(db::connection_timeout()),
|
||
)
|
||
.await?;
|
||
Ok(database_info_from_protocol_value(&response))
|
||
}
|
||
Some(ConnectionDatabaseInfoSource::NativeMysql(pool)) => {
|
||
db::mysql::database_connection_info(&pool, db::mysql::protocol_product_name(&config)).await.map(Some)
|
||
}
|
||
Some(ConnectionDatabaseInfoSource::NativeHBase(client)) => {
|
||
db::hbase_driver::database_connection_info(&client).await
|
||
}
|
||
Some(ConnectionDatabaseInfoSource::NativeMongo(client, database)) => {
|
||
db::mongo_driver::database_connection_info(&client, database.as_deref()).await.map(Some)
|
||
}
|
||
Some(ConnectionDatabaseInfoSource::Meilisearch(client)) => {
|
||
db::meilisearch_driver::database_connection_info(&client).await.map(Some)
|
||
}
|
||
Some(ConnectionDatabaseInfoSource::VictoriaMetrics(client)) => {
|
||
db::victoriametrics_driver::database_connection_info(&client, db::connection_timeout()).await.map(Some)
|
||
}
|
||
Some(ConnectionDatabaseInfoSource::Redis(pool_key)) => match self.pool_handle(&pool_key).await.as_ref() {
|
||
Some(PoolKind::Redis(redis)) => db::redis_driver::database_connection_info(redis).await.map(Some),
|
||
_ => Ok(None),
|
||
},
|
||
Some(ConnectionDatabaseInfoSource::Nacos) => {
|
||
let admin_config = self.nacos_admin_config_for_connection(connection_id, &config).await?;
|
||
let admin = self.nacos_registry.get_or_build_config(connection_id, admin_config).await?;
|
||
Ok(crate::nacos::service::database_info_from_connection(&admin.test_connection().await?))
|
||
}
|
||
Some(ConnectionDatabaseInfoSource::Consul(client)) => {
|
||
let identity = client.agent_self().await?;
|
||
Ok(Some(DatabaseConnectionInfo {
|
||
product_name: Some("Consul".to_string()),
|
||
product_version: identity.version,
|
||
server_comment: Some(format!("Agent {}", identity.node)),
|
||
driver_name: Some("Consul HTTP API".to_string()),
|
||
..Default::default()
|
||
}))
|
||
}
|
||
#[cfg(feature = "mq-admin")]
|
||
Some(ConnectionDatabaseInfoSource::MessageQueue) => {
|
||
Ok(crate::mq::service::mq_database_connection_info(self, connection_id).await?)
|
||
}
|
||
None => Ok(None),
|
||
}
|
||
}
|
||
|
||
pub async fn save_connection_database_info(
|
||
&self,
|
||
connection_id: &str,
|
||
database_info: Option<DatabaseConnectionInfo>,
|
||
) -> Result<(), String> {
|
||
self.storage.save_connection_database_info(connection_id, database_info.clone()).await?;
|
||
if let Some(config) = self.configs.write().await.get_mut(connection_id) {
|
||
config.database_info = database_info;
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
pub async fn reset_connection_transport(&self, connection_id: &str) {
|
||
let layer_count = {
|
||
let configs = self.configs.read().await;
|
||
configs.get(connection_id).map(|config| config.effective_transport_layers().len()).unwrap_or(0)
|
||
};
|
||
self.reset_connection_transport_layers(connection_id, layer_count).await;
|
||
}
|
||
|
||
pub async fn reset_connection_transport_for_config(&self, connection_id: &str, config: &ConnectionConfig) {
|
||
let existing_layer_count = {
|
||
let configs = self.configs.read().await;
|
||
configs.get(connection_id).map(|config| config.effective_transport_layers().len()).unwrap_or(0)
|
||
};
|
||
let layer_count = existing_layer_count.max(config.effective_transport_layers().len());
|
||
self.reset_connection_transport_layers(connection_id, layer_count).await;
|
||
}
|
||
|
||
async fn reset_connection_transport_layers(&self, connection_id: &str, layer_count: usize) {
|
||
let redis_cluster_prefix = redis_cluster_transport_prefix(connection_id);
|
||
self.tunnels.stop_tunnels_with_prefix(&redis_cluster_prefix).await;
|
||
self.proxy_tunnels.stop_tunnels_with_prefix(&redis_cluster_prefix).await;
|
||
self.http_tunnels.stop_tunnels_with_prefix(&redis_cluster_prefix).await;
|
||
let redis_sentinel_prefix = redis_sentinel_transport_prefix(connection_id);
|
||
self.tunnels.stop_tunnels_with_prefix(&redis_sentinel_prefix).await;
|
||
self.proxy_tunnels.stop_tunnels_with_prefix(&redis_sentinel_prefix).await;
|
||
self.http_tunnels.stop_tunnels_with_prefix(&redis_sentinel_prefix).await;
|
||
let sqlite_worker_prefix = db::sqlite_worker::sqlite_worker_chain_id(connection_id);
|
||
self.tunnels.stop_tunnels_with_prefix(&sqlite_worker_prefix).await;
|
||
db::transport_layer_tunnel::stop_transport_layers(
|
||
connection_id,
|
||
layer_count,
|
||
&self.tunnels,
|
||
&self.proxy_tunnels,
|
||
&self.http_tunnels,
|
||
)
|
||
.await;
|
||
db::transport_layer_tunnel::stop_transport_layers(
|
||
&rnacos_console_transport_id(connection_id),
|
||
layer_count,
|
||
&self.tunnels,
|
||
&self.proxy_tunnels,
|
||
&self.http_tunnels,
|
||
)
|
||
.await;
|
||
db::transport_layer_tunnel::stop_transport_layers(
|
||
&rabbitmq_management_transport_id(connection_id),
|
||
layer_count,
|
||
&self.tunnels,
|
||
&self.proxy_tunnels,
|
||
&self.http_tunnels,
|
||
)
|
||
.await;
|
||
self.tunnels.stop_tunnel(connection_id).await;
|
||
self.proxy_tunnels.stop_tunnel(connection_id).await;
|
||
self.http_tunnels.stop_tunnel(connection_id).await;
|
||
}
|
||
|
||
/// Health-check the base connection pool for a given connection_id.
|
||
/// Returns `Ok(())` if the pool exists and is healthy, `Err` otherwise.
|
||
/// If the pool is unhealthy it is removed from the map so subsequent
|
||
/// `get_or_create_pool` calls will transparently recreate it.
|
||
pub async fn check_connection_health(&self, connection_id: &str) -> Result<(), String> {
|
||
let db_type = {
|
||
let configs = self.configs.read().await;
|
||
configs.get(connection_id).map(|c| c.db_type)
|
||
};
|
||
let pool_key = base_pool_key_for(db_type, connection_id, None, false);
|
||
|
||
// Check if pool exists first
|
||
if self.pool_handle(&pool_key).await.is_none() {
|
||
return Err("No active connection pool found".to_string());
|
||
}
|
||
|
||
// `remove_stale_connection_pool` returns true if the pool was stale (and removed)
|
||
if self.remove_stale_connection_pool(&pool_key).await {
|
||
return Err("Connection pool is unhealthy".to_string());
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
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 = self.connection_pool_publication_snapshots().await;
|
||
|
||
let mut failed_agent_checks = Vec::new();
|
||
let mut dead_pools = Vec::new();
|
||
let timeout = crate::db::connection_timeout();
|
||
|
||
// Check cloned pools (async I/O, no lock held)
|
||
for (key, checked) in &checks {
|
||
let pool = &checked.pool;
|
||
let healthy = match pool {
|
||
PoolKind::Mysql(p, _) => match db::mysql::get_conn_with_health_check(p).await {
|
||
Ok(_) => true,
|
||
Err(e) if crate::query::is_pool_saturation_error(&e) => {
|
||
log::debug!("MySQL connection pool '{key}' is busy; skipping health probe: {e}");
|
||
true
|
||
}
|
||
Err(e) => {
|
||
log::warn!("MySQL connection pool '{key}' is unhealthy: {e}");
|
||
false
|
||
}
|
||
},
|
||
PoolKind::Postgres(p) => {
|
||
match db::postgres::checkout_postgres_client_classified(p, None, timeout).await {
|
||
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(error) if error.is_pool_saturation() => {
|
||
log::debug!("PostgreSQL connection pool '{key}' is busy; skipping health probe: {error}");
|
||
true
|
||
}
|
||
Err(error) => {
|
||
log::warn!("PostgreSQL connection pool '{key}' is unhealthy: {error}");
|
||
false
|
||
}
|
||
}
|
||
}
|
||
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::DynamoDb(client) => match db::dynamodb_driver::test_connection(client, timeout).await {
|
||
Ok(()) => true,
|
||
Err(e) => {
|
||
log::warn!("DynamoDB 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::Easysearch(client) => {
|
||
let mut client = client.clone();
|
||
match db::easysearch_driver::test_connection(&mut client, timeout).await {
|
||
Ok(()) => true,
|
||
Err(e) => {
|
||
log::warn!("Easysearch connection pool '{key}' is unhealthy: {e}");
|
||
false
|
||
}
|
||
}
|
||
}
|
||
PoolKind::Meilisearch(client) => {
|
||
let client = client.clone();
|
||
match db::meilisearch_driver::test_connection(&client, timeout).await {
|
||
Ok(()) => true,
|
||
Err(e) => {
|
||
log::warn!("Meilisearch connection pool '{key}' is unhealthy: {e}");
|
||
false
|
||
}
|
||
}
|
||
}
|
||
PoolKind::HBase(client) => match db::hbase_driver::test_connection(client, timeout).await {
|
||
Ok(_) => true,
|
||
Err(e) => {
|
||
log::warn!("HBase 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::InfluxDb3(client) => match db::influxdb3_driver::test_connection(client, timeout).await {
|
||
Ok(()) => true,
|
||
Err(e) => {
|
||
log::warn!("InfluxDB 3 connection pool '{key}' is unhealthy: {e}");
|
||
false
|
||
}
|
||
},
|
||
PoolKind::VictoriaMetrics(client) => {
|
||
match db::victoriametrics_driver::test_connection(client, timeout).await {
|
||
Ok(()) => true,
|
||
Err(e) => {
|
||
log::warn!("VictoriaMetrics 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::CloudflareD1(client) => {
|
||
match db::cloudflare_d1_driver::test_connection(client, timeout).await {
|
||
Ok(()) => true,
|
||
Err(e) => {
|
||
log::warn!("Cloudflare D1 connection pool '{key}' is unhealthy: {e}");
|
||
false
|
||
}
|
||
}
|
||
}
|
||
PoolKind::Agent(client) => {
|
||
let Ok(mut agent) = client.try_lock() else {
|
||
log::debug!("Agent connection pool '{key}' is busy; skipping resume health probe");
|
||
continue;
|
||
};
|
||
match agent.validate_connection_typed(Some(timeout)).await {
|
||
Ok(_) => true,
|
||
Err(err) if is_agent_validate_connection_unsupported(&err.to_string()) => {
|
||
log::debug!(
|
||
"Agent connection pool '{key}' does not support validate_connection; keeping pool"
|
||
);
|
||
true
|
||
}
|
||
Err(e) => {
|
||
log::warn!("Agent connection pool '{key}' is unhealthy: {e}");
|
||
failed_agent_checks.push((
|
||
key.clone(),
|
||
client.clone(),
|
||
checked.publication.clone(),
|
||
RecoveryPolicy::decide(&e, RecoveryScope::Keepalive).replaces_runtime(),
|
||
));
|
||
false
|
||
}
|
||
}
|
||
}
|
||
PoolKind::Sqlite(_)
|
||
| PoolKind::DuckDbWorker(_)
|
||
| PoolKind::ExternalDriver { .. }
|
||
| PoolKind::MessageQueue
|
||
| PoolKind::Nacos
|
||
| PoolKind::Consul(_) => true,
|
||
PoolKind::PluginConnection(handle) => handle.is_running(),
|
||
#[cfg(feature = "mq-admin")]
|
||
PoolKind::Mqtt(_) => true,
|
||
PoolKind::Redis(redis) => match db::redis_driver::test_connection(redis).await {
|
||
Ok(()) => true,
|
||
Err(e) => {
|
||
log::warn!("Redis connection pool '{key}' is unhealthy: {e}");
|
||
false
|
||
}
|
||
},
|
||
};
|
||
if !healthy && !matches!(pool, PoolKind::Agent(_)) {
|
||
dead_pools.push((key.clone(), checked.publication.clone()));
|
||
}
|
||
}
|
||
|
||
let mut detached_pool_keys = Vec::new();
|
||
for (key, client, publication, replace_runtime) in failed_agent_checks {
|
||
if self
|
||
.pool_routing_control()
|
||
.detach_agent_pool_if_current(&key, &client, Some(&publication), replace_runtime)
|
||
.await
|
||
{
|
||
detached_pool_keys.push(key);
|
||
}
|
||
}
|
||
|
||
// Remove dead pools
|
||
if !dead_pools.is_empty() {
|
||
let mut conns = self.connections.write().await;
|
||
let mut removed = Vec::with_capacity(dead_pools.len());
|
||
for (key, publication) in &dead_pools {
|
||
if let Some(pool) = conns.remove_if_publication(key, publication) {
|
||
removed.push((key.clone(), pool));
|
||
} else {
|
||
log::debug!("Skipping stale refresh health result for replaced pool '{key}'");
|
||
}
|
||
}
|
||
drop(conns);
|
||
detached_pool_keys.extend(removed.iter().map(|(key, _)| key.clone()));
|
||
self.pool_routing_control().finish_detach(removed).await;
|
||
}
|
||
|
||
// Only failed pools may require a fresh transport. Healthy tunnel listeners must keep
|
||
// their local ports stable across app resume and visibility-triggered health checks.
|
||
let tunnel_connection_ids: HashSet<String> = {
|
||
let configs = self.configs.read().await;
|
||
detached_pool_keys
|
||
.iter()
|
||
.filter_map(|pool_key| config_for_pool_key(pool_key, &configs))
|
||
.filter(|config| config.has_effective_transport_layers())
|
||
.map(|config| config.id.clone())
|
||
.collect()
|
||
};
|
||
for connection_id in tunnel_connection_ids {
|
||
self.reset_connection_transport(&connection_id).await;
|
||
// Tunnels will be re-created on next pool access via connection_host_port
|
||
}
|
||
}
|
||
|
||
pub async fn remove_connection_pools(&self, connection_id: &str) {
|
||
let removed = self.drain_connection_pools(connection_id).await;
|
||
self.clear_metadata_gates_for_connection(connection_id).await;
|
||
self.pool_routing_control().close_removed(removed).await;
|
||
}
|
||
|
||
pub async fn remove_connection_pools_detached(&self, connection_id: &str) {
|
||
let removed = self.drain_connection_pools(connection_id).await;
|
||
self.clear_metadata_gates_for_connection(connection_id).await;
|
||
self.pool_routing_control().close_removed_in_background(removed);
|
||
}
|
||
|
||
pub async fn invalidate_agent_pool_if_current(
|
||
&self,
|
||
pool_key: &str,
|
||
expected: &Arc<db::agent_driver::PooledAgentClient>,
|
||
) -> bool {
|
||
self.detach_agent_pool_if_current(pool_key, expected, false).await
|
||
}
|
||
|
||
async fn drain_all_connection_pools(&self) -> Vec<(String, PoolKind)> {
|
||
let pool_keys = self.connection_pools_snapshot().await.keys().cloned().collect::<Vec<_>>();
|
||
self.stop_keepalive_tasks(&pool_keys).await;
|
||
self.pool_activity.write().await.clear();
|
||
self.session_credentials.clear_pool_owners();
|
||
self.postgres_cancel_contexts.write().await.clear();
|
||
self.draining_pools.lock().unwrap_or_else(|error| error.into_inner()).clear();
|
||
self.metadata_gates.lock().await.clear();
|
||
self.connections.write().await.drain().collect()
|
||
}
|
||
|
||
#[cfg(feature = "duckdb-sidecar")]
|
||
async fn remove_duckdb_pools_detached(&self) {
|
||
let removed = self.drain_duckdb_pools().await;
|
||
self.pool_routing_control().close_removed_in_background(removed);
|
||
}
|
||
|
||
#[cfg(not(feature = "duckdb-sidecar"))]
|
||
async fn remove_duckdb_pools_detached(&self) {}
|
||
|
||
pub async fn remove_external_driver_pools(&self, driver_id: &str) {
|
||
let removed = self.drain_external_driver_pools(driver_id).await;
|
||
self.pool_routing_control().close_removed(removed).await;
|
||
}
|
||
|
||
pub async fn remove_plugin_connection_pools(&self, plugin_id: &str) {
|
||
let removed = self.drain_plugin_connection_pools(plugin_id).await;
|
||
self.pool_routing_control().close_removed(removed).await;
|
||
}
|
||
|
||
pub async fn invoke_plugin_connection_action(
|
||
&self,
|
||
config: ConnectionConfig,
|
||
action_id: &str,
|
||
) -> Result<crate::plugins::PluginConnectionActionResult, String> {
|
||
let (host, port) = self.connection_host_port(&config.id, &config).await?;
|
||
self.plugin_host.invoke_connection_action(&config, action_id, &host, port).await
|
||
}
|
||
|
||
async fn drain_connection_pools(&self, connection_id: &str) -> Vec<(String, PoolKind)> {
|
||
let pool_prefix = format!("{connection_id}:");
|
||
let keys_to_remove: Vec<String> = self
|
||
.connections
|
||
.read()
|
||
.await
|
||
.keys()
|
||
.filter(|k| *k == connection_id || k.starts_with(&pool_prefix))
|
||
.cloned()
|
||
.collect();
|
||
self.stop_keepalive_tasks(&keys_to_remove).await;
|
||
{
|
||
let mut activity = self.pool_activity.write().await;
|
||
let mut cancel_contexts = self.postgres_cancel_contexts.write().await;
|
||
for key in &keys_to_remove {
|
||
activity.remove(key);
|
||
cancel_contexts.remove(key);
|
||
}
|
||
}
|
||
self.session_credentials.remove_pool_owners(&keys_to_remove);
|
||
let mut conns = self.connections.write().await;
|
||
let mut removed = Vec::with_capacity(keys_to_remove.len());
|
||
for key in keys_to_remove {
|
||
if let Some(pool) = conns.remove(&key) {
|
||
removed.push((key, pool));
|
||
}
|
||
}
|
||
drop(conns);
|
||
removed
|
||
}
|
||
|
||
#[cfg(feature = "duckdb-sidecar")]
|
||
async fn drain_duckdb_pools(&self) -> Vec<(String, PoolKind)> {
|
||
let keys_to_remove: Vec<String> = self
|
||
.connections
|
||
.read()
|
||
.await
|
||
.iter()
|
||
.filter_map(|(key, pool)| {
|
||
let is_duckdb = matches!(pool, PoolKind::DuckDbWorker(_));
|
||
is_duckdb.then(|| key.clone())
|
||
})
|
||
.collect();
|
||
self.stop_keepalive_tasks(&keys_to_remove).await;
|
||
{
|
||
let mut activity = self.pool_activity.write().await;
|
||
let mut cancel_contexts = self.postgres_cancel_contexts.write().await;
|
||
for key in &keys_to_remove {
|
||
activity.remove(key);
|
||
cancel_contexts.remove(key);
|
||
}
|
||
}
|
||
let mut conns = self.connections.write().await;
|
||
let mut removed = Vec::with_capacity(keys_to_remove.len());
|
||
for key in keys_to_remove {
|
||
if let Some(pool) = conns.remove(&key) {
|
||
removed.push((key, pool));
|
||
}
|
||
}
|
||
removed
|
||
}
|
||
|
||
async fn drain_external_driver_pools(&self, driver_id: &str) -> Vec<(String, PoolKind)> {
|
||
let keys_to_remove: Vec<String> = self
|
||
.connections
|
||
.read()
|
||
.await
|
||
.iter()
|
||
.filter_map(|(key, pool)| match pool {
|
||
PoolKind::ExternalDriver { driver_id: pool_driver_id, .. } if pool_driver_id == driver_id => {
|
||
Some(key.clone())
|
||
}
|
||
_ => None,
|
||
})
|
||
.collect();
|
||
self.stop_keepalive_tasks(&keys_to_remove).await;
|
||
{
|
||
let mut activity = self.pool_activity.write().await;
|
||
for key in &keys_to_remove {
|
||
activity.remove(key);
|
||
}
|
||
}
|
||
let mut conns = self.connections.write().await;
|
||
let mut removed = Vec::with_capacity(keys_to_remove.len());
|
||
for key in keys_to_remove {
|
||
if let Some(pool) = conns.remove(&key) {
|
||
removed.push((key, pool));
|
||
}
|
||
}
|
||
removed
|
||
}
|
||
|
||
async fn drain_plugin_connection_pools(&self, plugin_id: &str) -> Vec<(String, PoolKind)> {
|
||
let keys_to_remove: Vec<String> = self
|
||
.connections
|
||
.read()
|
||
.await
|
||
.iter()
|
||
.filter_map(|(key, pool)| match pool {
|
||
PoolKind::PluginConnection(handle) if handle.plugin_id == plugin_id => Some(key.clone()),
|
||
_ => None,
|
||
})
|
||
.collect();
|
||
self.stop_keepalive_tasks(&keys_to_remove).await;
|
||
{
|
||
let mut activity = self.pool_activity.write().await;
|
||
for key in &keys_to_remove {
|
||
activity.remove(key);
|
||
}
|
||
}
|
||
let mut conns = self.connections.write().await;
|
||
let mut removed = Vec::with_capacity(keys_to_remove.len());
|
||
for key in keys_to_remove {
|
||
if let Some(pool) = conns.remove(&key) {
|
||
removed.push((key, pool));
|
||
}
|
||
}
|
||
removed
|
||
}
|
||
|
||
async fn uses_forwarded_transport(&self, connection_id: &str) -> bool {
|
||
let configs = self.configs.read().await;
|
||
configs.get(connection_id).is_some_and(|config| config.has_effective_transport_layers())
|
||
}
|
||
}
|
||
|
||
fn gaussdb_identifier_quote_from_query_result(result: &db::QueryResult) -> Option<String> {
|
||
let compatibility_mode = result.rows.first()?.first()?.as_str()?;
|
||
db::postgres::gaussdb_identifier_quote_for_compatibility_mode(compatibility_mode).map(str::to_string)
|
||
}
|
||
|
||
enum KeepaliveTarget {
|
||
Mysql(db::mysql::MySqlPool),
|
||
Postgres(deadpool_postgres::Pool),
|
||
Rqlite(db::rqlite_driver::RqliteClient),
|
||
Turso(db::turso_driver::TursoClient),
|
||
MongoDb {
|
||
client: mongodb::Client,
|
||
database: Option<String>,
|
||
},
|
||
ClickHouse(db::clickhouse_driver::ChClient),
|
||
SqlServer(Arc<tokio::sync::Mutex<db::sqlserver::SqlServerClient>>),
|
||
Elasticsearch(db::elasticsearch_driver::EsClient),
|
||
Easysearch(db::easysearch_driver::EasysearchClient),
|
||
HBase(db::hbase_driver::HBaseClient),
|
||
VectorDb(db::vector_driver::VectorClient),
|
||
InfluxDb(db::influxdb_driver::InfluxdbClient),
|
||
InfluxDb3(db::influxdb3_driver::Influxdb3Client),
|
||
VictoriaMetrics(db::victoriametrics_driver::VictoriaMetricsClient),
|
||
Agent(Arc<db::agent_driver::PooledAgentClient>),
|
||
#[cfg(feature = "mq-admin")]
|
||
MessageQueue(Arc<dyn crate::mq::port::MessageQueueAdmin>),
|
||
}
|
||
|
||
#[derive(Debug)]
|
||
enum KeepaliveError {
|
||
Agent(AgentCallError),
|
||
Legacy(String),
|
||
}
|
||
|
||
impl KeepaliveError {
|
||
fn recovery_decision(&self) -> Option<RecoveryDecision> {
|
||
match self {
|
||
Self::Agent(error) => Some(RecoveryPolicy::decide(error, RecoveryScope::Keepalive)),
|
||
Self::Legacy(_) => None,
|
||
}
|
||
}
|
||
}
|
||
|
||
impl std::fmt::Display for KeepaliveError {
|
||
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||
match self {
|
||
Self::Agent(error) => error.fmt(formatter),
|
||
Self::Legacy(error) => formatter.write_str(error),
|
||
}
|
||
}
|
||
}
|
||
|
||
impl From<String> for KeepaliveError {
|
||
fn from(error: String) -> Self {
|
||
Self::Legacy(error)
|
||
}
|
||
}
|
||
|
||
impl KeepaliveTarget {
|
||
fn matches_pool(&self, pool: &PoolKind) -> bool {
|
||
match (self, pool) {
|
||
(Self::Agent(expected), PoolKind::Agent(current)) => Arc::ptr_eq(expected, current),
|
||
(Self::SqlServer(expected), PoolKind::SqlServer(current)) => Arc::ptr_eq(expected, current),
|
||
#[cfg(feature = "mq-admin")]
|
||
(Self::MessageQueue(_), PoolKind::MessageQueue) => true,
|
||
#[cfg(feature = "mq-admin")]
|
||
(Self::MessageQueue(_), _) | (_, PoolKind::MessageQueue) => false,
|
||
(Self::Agent(_), _) | (_, PoolKind::Agent(_)) | (Self::SqlServer(_), _) | (_, PoolKind::SqlServer(_)) => {
|
||
false
|
||
}
|
||
_ => true,
|
||
}
|
||
}
|
||
}
|
||
|
||
async fn remove_keepalive_pool_if_current(
|
||
connections: &Arc<RwLock<ConnectionPoolRegistry>>,
|
||
pool_key: &str,
|
||
target: &KeepaliveTarget,
|
||
) -> Option<PoolKind> {
|
||
let mut pools = connections.write().await;
|
||
if pools.get(pool_key).is_some_and(|pool| target.matches_pool(pool)) {
|
||
pools.remove(pool_key)
|
||
} else {
|
||
None
|
||
}
|
||
}
|
||
|
||
async fn detach_keepalive_target_if_current(
|
||
routing: &PoolRoutingControl,
|
||
connections: &Arc<RwLock<ConnectionPoolRegistry>>,
|
||
pool_key: &str,
|
||
target: &KeepaliveTarget,
|
||
replace_agent_runtime: bool,
|
||
) -> bool {
|
||
if let KeepaliveTarget::Agent(expected) = target {
|
||
return routing.detach_agent_pool_if_current(pool_key, expected, None, replace_agent_runtime).await;
|
||
}
|
||
let Some(pool) = remove_keepalive_pool_if_current(connections, pool_key, target).await else {
|
||
return false;
|
||
};
|
||
routing.finish_detach(vec![(pool_key.to_string(), pool)]).await;
|
||
true
|
||
}
|
||
|
||
fn keepalive_target_from_pool(pool: &PoolKind, config: &ConnectionConfig) -> Option<KeepaliveTarget> {
|
||
match pool {
|
||
PoolKind::Mysql(pool, _) => Some(KeepaliveTarget::Mysql(pool.clone())),
|
||
PoolKind::Postgres(pool) => Some(KeepaliveTarget::Postgres(pool.clone())),
|
||
PoolKind::Rqlite(client) => Some(KeepaliveTarget::Rqlite(client.clone())),
|
||
PoolKind::Turso(client) => Some(KeepaliveTarget::Turso(client.clone())),
|
||
PoolKind::MongoDb(client) => Some(KeepaliveTarget::MongoDb {
|
||
client: client.clone(),
|
||
database: config.effective_database().map(str::to_string),
|
||
}),
|
||
PoolKind::ClickHouse(client) => Some(KeepaliveTarget::ClickHouse(client.clone())),
|
||
PoolKind::SqlServer(client) => Some(KeepaliveTarget::SqlServer(client.clone())),
|
||
PoolKind::Elasticsearch(client) => Some(KeepaliveTarget::Elasticsearch(client.clone())),
|
||
PoolKind::Easysearch(client) => Some(KeepaliveTarget::Easysearch(client.clone())),
|
||
PoolKind::HBase(client) => Some(KeepaliveTarget::HBase(client.clone())),
|
||
PoolKind::VectorDb(client) => Some(KeepaliveTarget::VectorDb(client.clone())),
|
||
PoolKind::InfluxDb(client) => Some(KeepaliveTarget::InfluxDb(client.clone())),
|
||
PoolKind::InfluxDb3(client) => Some(KeepaliveTarget::InfluxDb3(client.clone())),
|
||
PoolKind::VictoriaMetrics(client) => Some(KeepaliveTarget::VictoriaMetrics(client.clone())),
|
||
PoolKind::Agent(client) => Some(KeepaliveTarget::Agent(client.clone())),
|
||
_ => None,
|
||
}
|
||
}
|
||
|
||
async fn ping_keepalive_target(target: &mut KeepaliveTarget, timeout: Duration) -> Result<(), KeepaliveError> {
|
||
match target {
|
||
KeepaliveTarget::Mysql(pool) => {
|
||
let mut conn = db::mysql::get_conn_with_health_check(pool).await?;
|
||
conn.ping().await.map_err(|error| KeepaliveError::Legacy(error.to_string()))
|
||
}
|
||
KeepaliveTarget::Postgres(pool) => {
|
||
let client = pool.get().await.map_err(|e| format!("PostgreSQL pool error: {e}"))?;
|
||
client.simple_query("SELECT 1").await.map(|_| ()).map_err(|error| KeepaliveError::Legacy(error.to_string()))
|
||
}
|
||
KeepaliveTarget::Rqlite(client) => {
|
||
db::rqlite_driver::test_connection(client, timeout).await.map_err(Into::into)
|
||
}
|
||
KeepaliveTarget::Turso(client) => db::turso_driver::test_connection(client, timeout).await.map_err(Into::into),
|
||
KeepaliveTarget::MongoDb { client, database } => {
|
||
db::mongo_driver::test_connection(client, timeout, database.as_deref()).await.map_err(Into::into)
|
||
}
|
||
KeepaliveTarget::ClickHouse(client) => {
|
||
db::clickhouse_driver::test_connection(client, timeout).await.map_err(Into::into)
|
||
}
|
||
KeepaliveTarget::SqlServer(client) => {
|
||
let Ok(mut client) = client.try_lock() else {
|
||
return Ok(());
|
||
};
|
||
db::sqlserver::test_connection(&mut client).await.map_err(Into::into)
|
||
}
|
||
KeepaliveTarget::Elasticsearch(client) => {
|
||
db::elasticsearch_driver::test_connection(client, timeout).await.map_err(Into::into)
|
||
}
|
||
KeepaliveTarget::Easysearch(client) => {
|
||
db::easysearch_driver::test_connection(client, timeout).await.map_err(Into::into)
|
||
}
|
||
KeepaliveTarget::HBase(client) => {
|
||
db::hbase_driver::test_connection(client, timeout).await.map(|_| ()).map_err(Into::into)
|
||
}
|
||
KeepaliveTarget::VectorDb(client) => {
|
||
db::vector_driver::test_connection(client, timeout).await.map_err(Into::into)
|
||
}
|
||
KeepaliveTarget::InfluxDb(client) => {
|
||
db::influxdb_driver::test_connection(client, timeout).await.map_err(Into::into)
|
||
}
|
||
KeepaliveTarget::InfluxDb3(client) => {
|
||
db::influxdb3_driver::test_connection(client, timeout).await.map_err(Into::into)
|
||
}
|
||
KeepaliveTarget::VictoriaMetrics(client) => {
|
||
db::victoriametrics_driver::test_connection(client, timeout).await.map_err(Into::into)
|
||
}
|
||
KeepaliveTarget::Agent(client) => {
|
||
let Ok(mut client) = client.try_lock() else {
|
||
return Ok(());
|
||
};
|
||
match client.validate_connection_typed(Some(timeout)).await {
|
||
Ok(_) => Ok(()),
|
||
Err(error) if is_agent_validate_connection_unsupported(&error.to_string()) => Ok(()),
|
||
Err(error) => Err(KeepaliveError::Agent(error)),
|
||
}
|
||
}
|
||
#[cfg(feature = "mq-admin")]
|
||
KeepaliveTarget::MessageQueue(adapter) => {
|
||
adapter.test_connection().await.map(|_| ()).map_err(KeepaliveError::from)
|
||
}
|
||
}
|
||
}
|
||
|
||
fn is_agent_validate_connection_unsupported(err: &str) -> bool {
|
||
let lower = err.to_ascii_lowercase();
|
||
lower.contains("validate_connection") && (lower.contains("unknown method") || lower.contains("method not found"))
|
||
}
|
||
|
||
fn connection_remote_endpoint(config: &ConnectionConfig) -> (String, u16) {
|
||
if config.db_type == DatabaseType::MongoDb {
|
||
config
|
||
.connection_string
|
||
.as_deref()
|
||
.filter(|s| !s.is_empty())
|
||
.and_then(parse_mongo_first_host)
|
||
.unwrap_or_else(|| (config.host.clone(), config.port))
|
||
} else if config.db_type == DatabaseType::Jdbc {
|
||
config
|
||
.connection_string
|
||
.as_deref()
|
||
.filter(|s| !s.is_empty())
|
||
.and_then(parse_jdbc_host_port)
|
||
.unwrap_or_else(|| (config.host.clone(), config.port))
|
||
} else if config.db_type == DatabaseType::MessageQueue {
|
||
#[cfg(feature = "mq-admin")]
|
||
if let Some(endpoint) = parse_rabbitmq_amqp_host_port(config) {
|
||
return endpoint;
|
||
}
|
||
parse_mq_admin_host_port(config).unwrap_or_else(|| (config.host.clone(), config.port))
|
||
} else if config.db_type == DatabaseType::Mqtt {
|
||
parse_mqtt_broker_host_port(config).unwrap_or_else(|| (config.host.clone(), config.port))
|
||
} else if config.db_type == DatabaseType::Nacos {
|
||
parse_nacos_server_host_port(config).unwrap_or_else(|| (config.host.clone(), config.port))
|
||
} else if config.db_type == DatabaseType::Consul {
|
||
parse_consul_server_host_port(config).unwrap_or_else(|| (config.host.clone(), config.port))
|
||
} else {
|
||
(config.host.clone(), config.port)
|
||
}
|
||
}
|
||
|
||
fn rnacos_console_transport_id(connection_id: &str) -> String {
|
||
format!("{connection_id}:rnacos-console")
|
||
}
|
||
|
||
fn rabbitmq_management_transport_id(connection_id: &str) -> String {
|
||
format!("{connection_id}:rabbitmq-management")
|
||
}
|
||
|
||
#[cfg(feature = "mq-admin")]
|
||
fn parse_rabbitmq_amqp_host_port(config: &ConnectionConfig) -> Option<(String, u16)> {
|
||
let mqc = crate::mq::config::MqAdminConfig::from_connection(config).ok()?;
|
||
if mqc.system_kind != crate::mq::types::MqSystemKind::RabbitMq {
|
||
return None;
|
||
}
|
||
crate::mq::adapters::rabbitmq::primary_amqp_endpoint(&mqc).ok()
|
||
}
|
||
|
||
fn parse_mq_admin_host_port(config: &ConnectionConfig) -> Option<(String, u16)> {
|
||
let value = config
|
||
.external_config
|
||
.as_ref()?
|
||
.get("adminUrl")
|
||
.or_else(|| config.external_config.as_ref()?.get("admin_url"))?
|
||
.as_str()?
|
||
.trim();
|
||
if value.is_empty() {
|
||
return None;
|
||
}
|
||
let url = reqwest::Url::parse(value).ok()?;
|
||
let host = url.host_str()?.to_string();
|
||
let port = url.port_or_known_default()?;
|
||
Some((host, port))
|
||
}
|
||
|
||
#[cfg(any(feature = "mq-admin", test))]
|
||
fn kafka_single_loopback_bootstrap_endpoint(extra: &serde_json::Value) -> Option<(String, u16)> {
|
||
let value = extra.get("bootstrapServers").or_else(|| extra.get("bootstrap_servers"))?.as_str()?.trim();
|
||
let mut endpoints = value
|
||
.split(|ch: char| ch.is_whitespace() || matches!(ch, ',' | ';' | ',' | ';'))
|
||
.filter(|endpoint| !endpoint.is_empty());
|
||
let endpoint = endpoints.next()?;
|
||
if endpoints.next().is_some() {
|
||
return None;
|
||
}
|
||
let address = endpoint.rsplit_once("://").map_or(endpoint, |(_, address)| address);
|
||
let url = reqwest::Url::parse(&format!("kafka://{address}")).ok()?;
|
||
if url.host_str()? != "127.0.0.1"
|
||
|| url.username() != ""
|
||
|| url.password().is_some()
|
||
|| url.query().is_some()
|
||
|| url.fragment().is_some()
|
||
|| !matches!(url.path(), "" | "/")
|
||
{
|
||
return None;
|
||
}
|
||
Some(("127.0.0.1".to_string(), url.port()?))
|
||
}
|
||
|
||
fn parse_nacos_server_host_port(config: &ConnectionConfig) -> Option<(String, u16)> {
|
||
let value = config
|
||
.external_config
|
||
.as_ref()?
|
||
.get("serverAddr")
|
||
.or_else(|| config.external_config.as_ref()?.get("server_addr"))?
|
||
.as_str()?
|
||
.trim();
|
||
if value.is_empty() {
|
||
return None;
|
||
}
|
||
let url = reqwest::Url::parse(value).ok()?;
|
||
let host = url.host_str()?.to_string();
|
||
let port = url.port_or_known_default()?;
|
||
Some((host, port))
|
||
}
|
||
|
||
fn parse_consul_server_host_port(config: &ConnectionConfig) -> Option<(String, u16)> {
|
||
let value = config
|
||
.external_config
|
||
.as_ref()?
|
||
.get("serverAddr")
|
||
.or_else(|| config.external_config.as_ref()?.get("server_addr"))?
|
||
.as_str()?
|
||
.trim();
|
||
if value.is_empty() {
|
||
return None;
|
||
}
|
||
let url = reqwest::Url::parse(value).ok()?;
|
||
Some((url.host_str()?.to_string(), url.port_or_known_default()?))
|
||
}
|
||
|
||
fn parse_mqtt_broker_host_port(config: &ConnectionConfig) -> Option<(String, u16)> {
|
||
let value = config
|
||
.external_config
|
||
.as_ref()?
|
||
.get("host")
|
||
.or_else(|| config.external_config.as_ref()?.get("host"))?
|
||
.as_str()?
|
||
.trim();
|
||
if value.is_empty() {
|
||
return None;
|
||
}
|
||
let port = config
|
||
.external_config
|
||
.as_ref()?
|
||
.get("port")
|
||
.or_else(|| config.external_config.as_ref()?.get("port"))
|
||
.and_then(|v| v.as_u64())
|
||
.unwrap_or(1883) as u16;
|
||
Some((value.to_string(), port))
|
||
}
|
||
|
||
fn normalize_client_session_id(client_session_id: Option<&str>) -> Option<String> {
|
||
client_session_id.map(str::trim).filter(|session| !session.is_empty()).map(|session| session.replace(':', "_"))
|
||
}
|
||
|
||
pub fn task_client_session_id(task_kind: &str, task_id: &str) -> String {
|
||
format!("{task_kind}:{task_id}")
|
||
}
|
||
|
||
fn mysql_pool_max_connections_for_session(client_session_id: Option<&str>) -> usize {
|
||
if normalize_client_session_id(client_session_id).is_some() {
|
||
1
|
||
} else {
|
||
10
|
||
}
|
||
}
|
||
|
||
fn postgres_pool_max_connections_for_session(client_session_id: Option<&str>) -> usize {
|
||
if normalize_client_session_id(client_session_id).is_some() {
|
||
1
|
||
} else {
|
||
10
|
||
}
|
||
}
|
||
|
||
fn redis_cluster_transport_prefix(connection_id: &str) -> String {
|
||
format!("{connection_id}:redis-cluster:")
|
||
}
|
||
|
||
fn redis_cluster_transport_id(connection_id: &str, endpoint: &db::redis_driver::RedisNodeEndpoint) -> String {
|
||
format!(
|
||
"{}{host}:{port}",
|
||
redis_cluster_transport_prefix(connection_id),
|
||
host = endpoint.host,
|
||
port = endpoint.port
|
||
)
|
||
}
|
||
|
||
fn redis_sentinel_transport_prefix(connection_id: &str) -> String {
|
||
format!("{connection_id}:redis-sentinel:")
|
||
}
|
||
|
||
fn redis_sentinel_transport_id(
|
||
connection_id: &str,
|
||
role: &str,
|
||
endpoint: &db::redis_driver::RedisNodeEndpoint,
|
||
) -> String {
|
||
format!(
|
||
"{}{role}:{host}:{port}",
|
||
redis_sentinel_transport_prefix(connection_id),
|
||
host = endpoint.host,
|
||
port = endpoint.port
|
||
)
|
||
}
|
||
|
||
fn session_scoped_pool_key(base_pool_key: String, client_session_id: Option<&str>) -> String {
|
||
normalize_client_session_id(client_session_id)
|
||
.map(|session| format!("{base_pool_key}:session:{session}"))
|
||
.unwrap_or(base_pool_key)
|
||
}
|
||
|
||
fn is_session_scoped_pool_key(pool_key: &str) -> bool {
|
||
pool_key.contains(":session:")
|
||
}
|
||
|
||
fn is_manual_transaction_pool_key(pool_key: &str) -> bool {
|
||
pool_key.contains(":session:manual-txn-")
|
||
}
|
||
|
||
pub(crate) fn config_for_pool_key<'a>(
|
||
pool_key: &str,
|
||
configs: &'a HashMap<String, ConnectionConfig>,
|
||
) -> Option<&'a ConnectionConfig> {
|
||
configs
|
||
.iter()
|
||
.filter(|(connection_id, _)| {
|
||
pool_key.strip_prefix(connection_id.as_str()).is_some_and(|rest| rest.is_empty() || rest.starts_with(':'))
|
||
})
|
||
.max_by_key(|(connection_id, _)| connection_id.len())
|
||
.map(|(_, config)| config)
|
||
}
|
||
|
||
fn session_scoped_pool_key_for(
|
||
config: Option<&ConnectionConfig>,
|
||
base_pool_key: String,
|
||
client_session_id: Option<&str>,
|
||
) -> String {
|
||
let shares_base_pool = config.is_some_and(|config| {
|
||
matches!(config.db_type, DatabaseType::DuckDb | DatabaseType::CloudflareD1)
|
||
|| (config.db_type == DatabaseType::Sqlite && db::sqlite::is_memory_database_path(&config.host))
|
||
});
|
||
if shares_base_pool {
|
||
// DuckDB and D1 already use connection-scoped handles. In-memory SQLite databases
|
||
// only exist inside one connection, so a session-scoped handle would point tabs at
|
||
// a different empty database.
|
||
return base_pool_key;
|
||
}
|
||
session_scoped_pool_key(base_pool_key, client_session_id)
|
||
}
|
||
|
||
/// 两个运行态连接配置是否视为同一连接(用于决定是否销毁连接池)。
|
||
///
|
||
/// 仅当双方都是 `save_password=false` 时才忽略 `password` 字段的差异:这类连接
|
||
/// 在 connect 时运行态配置可能携带会话密码,持久化同步后为空,这种空值差异不应
|
||
/// 触发池重建。若任一方 `save_password=true`,密码是真实的连接参数,任何密码变更
|
||
/// (包括用户保存了新密码)都必须销毁旧池,否则旧池会继续用旧密码认证。
|
||
pub fn connection_configs_pool_equivalent(a: &ConnectionConfig, b: &ConnectionConfig) -> bool {
|
||
if a == b {
|
||
return true;
|
||
}
|
||
if !a.save_password && !b.save_password {
|
||
let mut a = a.clone();
|
||
a.password.clear();
|
||
let mut b = b.clone();
|
||
b.password.clear();
|
||
return a == b;
|
||
}
|
||
false
|
||
}
|
||
|
||
/// Whether transient credentials can safely survive a persisted config update.
|
||
///
|
||
/// This comparison is intentionally conservative: only presentation, local
|
||
/// visibility, and runtime-policy fields are ignored. Endpoint, account,
|
||
/// transport, TLS, driver, and unclassified external-config changes still
|
||
/// invalidate credentials. Nacos managed namespaces are a local discovery
|
||
/// scope and therefore do not change which account a transient password
|
||
/// belongs to.
|
||
pub fn connection_configs_session_credentials_compatible(a: &ConnectionConfig, b: &ConnectionConfig) -> bool {
|
||
fn normalize(mut config: ConnectionConfig) -> ConnectionConfig {
|
||
config.name.clear();
|
||
config.note.clear();
|
||
config.driver_label = None;
|
||
config.default_schema = None;
|
||
config.visible_databases = None;
|
||
config.visible_schemas = None;
|
||
config.show_system_schemas = false;
|
||
config.color = None;
|
||
config.docs_notes_path = None;
|
||
config.connect_timeout_secs = 0;
|
||
config.query_timeout_secs = 0;
|
||
config.idle_timeout_secs = 0;
|
||
config.keepalive_interval_secs = 0;
|
||
config.redis_key_separator.clear();
|
||
config.redis_scan_page_size = None;
|
||
config.redis_database_aliases.clear();
|
||
config.one_time = false;
|
||
config.read_only = false;
|
||
config.is_production = false;
|
||
config.production_databases.clear();
|
||
config.database_info = None;
|
||
|
||
if config.db_type == DatabaseType::Nacos {
|
||
if let Some(external) = config.external_config.as_mut().and_then(serde_json::Value::as_object_mut) {
|
||
external.remove("managedNamespaces");
|
||
external.remove("managed_namespaces");
|
||
if !config.save_password {
|
||
for key in ["auth", "rnacosConsoleAuth", "rnacos_console_auth"] {
|
||
if let Some(auth) = external.get_mut(key).and_then(serde_json::Value::as_object_mut) {
|
||
auth.insert("password".to_string(), serde_json::Value::String(String::new()));
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
if !config.save_password {
|
||
config.password.clear();
|
||
}
|
||
config
|
||
}
|
||
|
||
normalize(a.clone()) == normalize(b.clone())
|
||
}
|
||
|
||
fn pool_key_for_session_role(
|
||
config: Option<&ConnectionConfig>,
|
||
base_pool_key: String,
|
||
client_session_id: Option<&str>,
|
||
session_role: AgentSessionRole,
|
||
) -> 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))
|
||
{
|
||
format!("{pool_key}:role:metadata")
|
||
} else {
|
||
pool_key
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
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()),
|
||
PoolKind::Sqlite(p) => PoolKind::Sqlite(p.clone()),
|
||
PoolKind::Rqlite(client) => PoolKind::Rqlite(client.clone()),
|
||
PoolKind::Turso(client) => PoolKind::Turso(client.clone()),
|
||
PoolKind::CloudflareD1(client) => PoolKind::CloudflareD1(client.clone()),
|
||
#[cfg(feature = "duckdb-sidecar")]
|
||
PoolKind::DuckDbWorker(client) => PoolKind::DuckDbWorker(client.clone()),
|
||
#[cfg(not(feature = "duckdb-sidecar"))]
|
||
PoolKind::DuckDbWorker(_) => PoolKind::DuckDbWorker(()),
|
||
PoolKind::MongoDb(client) => PoolKind::MongoDb(client.clone()),
|
||
PoolKind::DynamoDb(client) => PoolKind::DynamoDb(client.clone()),
|
||
PoolKind::ClickHouse(client) => PoolKind::ClickHouse(client.clone()),
|
||
PoolKind::SqlServer(client) => PoolKind::SqlServer(client.clone()),
|
||
PoolKind::Elasticsearch(client) => PoolKind::Elasticsearch(client.clone()),
|
||
PoolKind::Easysearch(client) => PoolKind::Easysearch(client.clone()),
|
||
PoolKind::Meilisearch(client) => PoolKind::Meilisearch(client.clone()),
|
||
PoolKind::HBase(client) => PoolKind::HBase(client.clone()),
|
||
PoolKind::VectorDb(client) => PoolKind::VectorDb(client.clone()),
|
||
PoolKind::InfluxDb(client) => PoolKind::InfluxDb(client.clone()),
|
||
PoolKind::InfluxDb3(client) => PoolKind::InfluxDb3(client.clone()),
|
||
PoolKind::VictoriaMetrics(client) => PoolKind::VictoriaMetrics(client.clone()),
|
||
PoolKind::Agent(client) => PoolKind::Agent(client.clone()),
|
||
PoolKind::ExternalDriver { driver_id, config, session } => {
|
||
PoolKind::ExternalDriver { driver_id: driver_id.clone(), config: config.clone(), session: session.clone() }
|
||
}
|
||
PoolKind::PluginConnection(handle) => PoolKind::PluginConnection(handle.clone()),
|
||
PoolKind::MessageQueue => PoolKind::MessageQueue,
|
||
PoolKind::Nacos => PoolKind::Nacos,
|
||
PoolKind::Consul(client) => PoolKind::Consul(client.clone()),
|
||
#[cfg(feature = "mq-admin")]
|
||
PoolKind::Mqtt(client) => PoolKind::Mqtt(Arc::clone(client)),
|
||
PoolKind::Redis(_) => panic!("clone_pool_kind not supported for Redis — handled separately"),
|
||
}
|
||
}
|
||
|
||
async fn close_pool_kind(pool: PoolKind) -> Result<(), String> {
|
||
match pool {
|
||
PoolKind::Mysql(p, _) => {
|
||
let _ = p.disconnect().await;
|
||
}
|
||
PoolKind::Postgres(p) => p.close(),
|
||
PoolKind::Sqlite(pool) => {
|
||
pool.shutdown().await;
|
||
}
|
||
PoolKind::Rqlite(_) => {}
|
||
PoolKind::Turso(_) => {}
|
||
PoolKind::CloudflareD1(_) => {}
|
||
PoolKind::Redis(conn) => {
|
||
drop(conn);
|
||
}
|
||
#[cfg(feature = "duckdb-sidecar")]
|
||
PoolKind::DuckDbWorker(client) => {
|
||
client.shutdown().await;
|
||
}
|
||
#[cfg(not(feature = "duckdb-sidecar"))]
|
||
PoolKind::DuckDbWorker(_) => {}
|
||
PoolKind::MongoDb(client) => {
|
||
drop(client);
|
||
}
|
||
PoolKind::DynamoDb(client) => {
|
||
drop(client);
|
||
}
|
||
PoolKind::ClickHouse(client) => {
|
||
drop(client);
|
||
}
|
||
PoolKind::SqlServer(client) => {
|
||
drop(client);
|
||
}
|
||
PoolKind::Elasticsearch(client) => {
|
||
drop(client);
|
||
}
|
||
PoolKind::Easysearch(client) => {
|
||
drop(client);
|
||
}
|
||
PoolKind::Meilisearch(client) => {
|
||
drop(client);
|
||
}
|
||
PoolKind::HBase(client) => {
|
||
drop(client);
|
||
}
|
||
PoolKind::VectorDb(client) => {
|
||
drop(client);
|
||
}
|
||
PoolKind::InfluxDb(client) => {
|
||
drop(client);
|
||
}
|
||
PoolKind::InfluxDb3(client) => {
|
||
drop(client);
|
||
}
|
||
PoolKind::VictoriaMetrics(client) => {
|
||
drop(client);
|
||
}
|
||
PoolKind::Agent(client) => {
|
||
let mut client = client.lock().await;
|
||
client.disconnect().await?;
|
||
}
|
||
PoolKind::ExternalDriver { session, .. } => {
|
||
session.shutdown().await;
|
||
}
|
||
PoolKind::PluginConnection(handle) => {
|
||
handle.disconnect().await?;
|
||
}
|
||
PoolKind::MessageQueue => {}
|
||
PoolKind::Nacos => {}
|
||
PoolKind::Consul(_) => {}
|
||
#[cfg(feature = "mq-admin")]
|
||
PoolKind::Mqtt(client) => {
|
||
// 发送 DISCONNECT 并等待事件循环任务结束
|
||
client.disconnect().await;
|
||
}
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
async fn close_reclaimed_agent_pool(pool: PoolKind) -> Result<(), (PoolKind, String)> {
|
||
let client = match pool {
|
||
PoolKind::Agent(client) => client,
|
||
pool => return Err((pool, "Only Agent pools can be reclaimed after connection slot exhaustion".to_string())),
|
||
};
|
||
let result = {
|
||
let mut client = client.lock().await;
|
||
client.disconnect().await.map(|_| ())
|
||
};
|
||
match result {
|
||
Ok(()) => Ok(()),
|
||
Err(error) => Err((PoolKind::Agent(client), error)),
|
||
}
|
||
}
|
||
|
||
fn extract_auth_token_from_params(params: &str) -> Option<String> {
|
||
params
|
||
.trim()
|
||
.trim_start_matches('?')
|
||
.split('&')
|
||
.filter_map(|pair| pair.split_once('='))
|
||
.find(|(key, _)| {
|
||
let k = key.trim().to_ascii_lowercase();
|
||
k == "auth_token" || k == "authtoken" || k == "auth-token"
|
||
})
|
||
.map(|(_, value)| value.trim().to_string())
|
||
}
|
||
|
||
fn base_pool_key_for(
|
||
db_type: Option<DatabaseType>,
|
||
connection_id: &str,
|
||
database: Option<&str>,
|
||
include_elasticsearch_single_pool: bool,
|
||
) -> String {
|
||
base_pool_key_for_with_catalog(db_type, connection_id, database, None, include_elasticsearch_single_pool)
|
||
}
|
||
|
||
fn base_pool_key_for_with_catalog(
|
||
db_type: Option<DatabaseType>,
|
||
connection_id: &str,
|
||
database: Option<&str>,
|
||
catalog: 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
|
||
&& matches!(
|
||
db_type,
|
||
DatabaseType::Elasticsearch
|
||
| DatabaseType::Easysearch
|
||
| DatabaseType::Qdrant
|
||
| DatabaseType::Milvus
|
||
| DatabaseType::Weaviate
|
||
| DatabaseType::ChromaDb
|
||
));
|
||
is_single && (!database_capabilities::is_agent_type(db_type) || shares_database_pool_with_connection(db_type))
|
||
});
|
||
|
||
let key = if is_single_connection_pool {
|
||
connection_id.to_string()
|
||
} else {
|
||
match database.filter(|db| !db.trim().is_empty()) {
|
||
Some(db) => format!("{connection_id}:{db}"),
|
||
None => connection_id.to_string(),
|
||
}
|
||
};
|
||
match catalog.map(str::trim).filter(|value| !value.is_empty()) {
|
||
Some(catalog) => format!("{key}:catalog:{catalog}"),
|
||
None => key,
|
||
}
|
||
}
|
||
|
||
fn is_connection_slot_exhausted_error(error: &str) -> bool {
|
||
let lower = error.to_ascii_lowercase();
|
||
lower.contains("remaining connection slots are reserved")
|
||
|| lower.contains("too many connections")
|
||
|| lower.contains("maximum number of connections exceeded")
|
||
|| lower.contains("max client connections reached")
|
||
|| lower.contains("ora-00018")
|
||
}
|
||
|
||
fn shares_database_pool_with_connection(db_type: &DatabaseType) -> bool {
|
||
matches!(db_type, DatabaseType::Oracle)
|
||
}
|
||
|
||
#[cfg(test)]
|
||
fn uses_agent_connection_pool(db_type: &DatabaseType) -> bool {
|
||
matches!(*db_type, agent_connection_pool_database_type!())
|
||
}
|
||
|
||
fn should_validate_existing_pool_before_reuse(db_type: DatabaseType) -> bool {
|
||
// PostgreSQL and Agent-backed databases validate connections when they are
|
||
// checked out for actual work. An eager probe here would add a database
|
||
// round-trip before every request and can compete with active Agent leases.
|
||
db_type != DatabaseType::Postgres && !matches!(db_type, agent_connection_pool_database_type!())
|
||
}
|
||
|
||
#[cfg(test)]
|
||
fn uses_bare_mysql_pool(db_type: &DatabaseType) -> bool {
|
||
matches!(db_type, DatabaseType::Doris | DatabaseType::StarRocks | DatabaseType::ManticoreSearch)
|
||
}
|
||
|
||
fn default_plugin_dir() -> PathBuf {
|
||
default_dbx_dir().join("plugins")
|
||
}
|
||
|
||
pub fn default_agent_dir() -> PathBuf {
|
||
default_dbx_dir().join("agents")
|
||
}
|
||
|
||
fn default_dbx_dir() -> PathBuf {
|
||
let home = std::env::var(if cfg!(windows) { "USERPROFILE" } else { "HOME" }).unwrap_or_else(|_| ".".to_string());
|
||
PathBuf::from(home).join(".dbx")
|
||
}
|
||
|
||
pub fn connection_url_for_endpoint(config: &ConnectionConfig, host: &str, port: u16) -> String {
|
||
let normalized = native_postgres_url_config(config);
|
||
let config = normalized.as_ref().unwrap_or(config);
|
||
if host == config.host && port == config.port {
|
||
config.connection_url()
|
||
} else {
|
||
config.connection_url_with_host(host, port)
|
||
}
|
||
}
|
||
|
||
pub fn redacted_connection_url_for_endpoint(config: &ConnectionConfig, host: &str, port: u16) -> String {
|
||
let normalized = native_postgres_url_config(config);
|
||
let config = normalized.as_ref().unwrap_or(config);
|
||
if host == config.host && port == config.port {
|
||
config.redacted_connection_url()
|
||
} else {
|
||
config.redacted_connection_url_with_host(host, port)
|
||
}
|
||
}
|
||
|
||
pub fn agent_connect_timeout(config: &ConnectionConfig) -> std::time::Duration {
|
||
let min_timeout = if config.db_type == DatabaseType::Access {
|
||
ACCESS_AGENT_CONNECT_TIMEOUT_SECS
|
||
} else {
|
||
DEFAULT_AGENT_CONNECT_TIMEOUT_SECS
|
||
};
|
||
std::time::Duration::from_secs(config.effective_connect_timeout_secs().max(min_timeout))
|
||
}
|
||
|
||
fn external_driver_connect_timeout(config: &ConnectionConfig) -> std::time::Duration {
|
||
agent_connect_timeout(config)
|
||
}
|
||
|
||
fn native_postgres_url_config(config: &ConnectionConfig) -> Option<ConnectionConfig> {
|
||
match config.db_type {
|
||
DatabaseType::Gaussdb | DatabaseType::Kwdb | DatabaseType::OpenGauss | DatabaseType::Questdb => {
|
||
let mut normalized = config.clone();
|
||
normalized.database = normalized.effective_database().map(str::to_string);
|
||
if matches!(config.db_type, DatabaseType::Gaussdb | DatabaseType::Kwdb) {
|
||
let params = normalized.url_params.as_deref().unwrap_or("").trim().trim_start_matches('?');
|
||
if !params.to_lowercase().contains("sslmode=") {
|
||
normalized.url_params = Some(if params.is_empty() {
|
||
if config.ssl {
|
||
"sslmode=require".to_string()
|
||
} else {
|
||
"sslmode=prefer".to_string()
|
||
}
|
||
} else {
|
||
let sslmode = if config.ssl { "sslmode=require" } else { "sslmode=prefer" };
|
||
format!("{sslmode}&{params}")
|
||
});
|
||
}
|
||
}
|
||
normalized.db_type = DatabaseType::Postgres;
|
||
Some(normalized)
|
||
}
|
||
_ => None,
|
||
}
|
||
}
|
||
|
||
fn validate_connection_url_params(config: &ConnectionConfig) -> Result<(), String> {
|
||
let normalized = native_postgres_url_config(config);
|
||
normalized.as_ref().unwrap_or(config).validate_native_url_params()
|
||
}
|
||
|
||
pub async fn probe_connection_endpoint(config: &ConnectionConfig, host: &str, port: u16) -> Result<(), String> {
|
||
if !uses_tcp_probe(config, host, port) {
|
||
return Ok(());
|
||
}
|
||
let timeout = std::time::Duration::from_secs(config.effective_connect_timeout_secs());
|
||
|
||
let entries = connection_probe_endpoints(host, port);
|
||
|
||
if entries.is_empty() {
|
||
return Err("no host entries to probe".to_string());
|
||
}
|
||
|
||
// Probe each node sequentially; return success on the first reachable node.
|
||
// This matches the failover semantics of the real connection path.
|
||
let mut last_error = String::new();
|
||
for (entry_host, entry_port) in &entries {
|
||
match db::probe_tcp_endpoint(&format!("{:?}", config.db_type), entry_host, *entry_port, timeout).await {
|
||
Ok(()) => return Ok(()),
|
||
Err(e) => last_error = e,
|
||
}
|
||
}
|
||
Err(last_error)
|
||
}
|
||
|
||
fn connection_probe_endpoints(host: &str, default_port: u16) -> Vec<(String, u16)> {
|
||
host.split(',').filter_map(|part| parse_connection_probe_endpoint(part.trim(), default_port)).collect()
|
||
}
|
||
|
||
fn parse_connection_probe_endpoint(endpoint: &str, default_port: u16) -> Option<(String, u16)> {
|
||
if endpoint.is_empty() {
|
||
return None;
|
||
}
|
||
if let Some(rest) = endpoint.strip_prefix('[') {
|
||
let close = rest.find(']')?;
|
||
let host = rest[..close].to_string();
|
||
let port = rest
|
||
.get(close + 1..)
|
||
.and_then(|suffix| suffix.strip_prefix(':'))
|
||
.and_then(|value| value.parse::<u16>().ok())
|
||
.unwrap_or(default_port);
|
||
return Some((host, port));
|
||
}
|
||
if endpoint.matches(':').count() == 1 {
|
||
if let Some((host, raw_port)) = endpoint.rsplit_once(':') {
|
||
if let Ok(port) = raw_port.parse::<u16>() {
|
||
return Some((host.to_string(), port));
|
||
}
|
||
}
|
||
}
|
||
Some((endpoint.to_string(), default_port))
|
||
}
|
||
|
||
fn validate_h2_file_connection(config: &ConnectionConfig) -> Result<(), String> {
|
||
if !is_h2_file_connection(config) {
|
||
return Ok(());
|
||
}
|
||
let path = config
|
||
.connection_string
|
||
.as_deref()
|
||
.and_then(h2_file_path_from_jdbc_url)
|
||
.filter(|path| !path.trim().is_empty())
|
||
.unwrap_or_else(|| config.host.clone());
|
||
validate_h2_database_path(&path)
|
||
}
|
||
|
||
fn validate_h2_database_path(path: &str) -> Result<(), String> {
|
||
let first_err = match db::validate_file_path(path, |_| false) {
|
||
Ok(()) => return Ok(()),
|
||
Err(err) => err,
|
||
};
|
||
|
||
for suffix in [".mv.db", ".h2.db"] {
|
||
if path.ends_with(suffix) {
|
||
continue;
|
||
}
|
||
let candidate = format!("{path}{suffix}");
|
||
if db::validate_file_path(&candidate, |_| false).is_ok() {
|
||
return Ok(());
|
||
}
|
||
}
|
||
|
||
Err(first_err)
|
||
}
|
||
|
||
fn uses_tcp_probe(config: &ConnectionConfig, host: &str, port: u16) -> bool {
|
||
if config.db_type == DatabaseType::MongoDb
|
||
&& config.connection_string.as_deref().is_some_and(|value| !value.is_empty())
|
||
{
|
||
return false;
|
||
}
|
||
if database_capabilities::skips_tcp_probe(&config.db_type) {
|
||
return false;
|
||
}
|
||
if is_original_endpoint(config, host, port) {
|
||
return false;
|
||
}
|
||
true
|
||
}
|
||
|
||
fn is_original_endpoint(config: &ConnectionConfig, host: &str, port: u16) -> bool {
|
||
host == config.host && port == config.port
|
||
}
|
||
|
||
async fn detect_ob_oracle_mode(config: &ConnectionConfig, pool: &db::mysql::MySqlPool) -> MysqlMode {
|
||
let profile = config.driver_profile.as_deref().unwrap_or("").to_lowercase();
|
||
if !profile.contains("oceanbase") {
|
||
return MysqlMode::Normal;
|
||
}
|
||
let mut conn = match pool.get_conn().await {
|
||
Ok(c) => c,
|
||
Err(_) => return MysqlMode::Normal,
|
||
};
|
||
let result = conn.query_iter("SHOW VARIABLES LIKE 'ob_compatibility_mode'").await;
|
||
let rows: Vec<MysqlRow> = match result {
|
||
Ok(r) => match r.collect_and_drop().await {
|
||
Ok(rows) => rows,
|
||
Err(_) => return MysqlMode::Normal,
|
||
},
|
||
Err(_) => return MysqlMode::Normal,
|
||
};
|
||
match rows.first() {
|
||
Some(row) => {
|
||
let val: String = row.get(1).unwrap_or_default();
|
||
if val.to_lowercase() == "oracle" {
|
||
MysqlMode::OceanBaseOracle
|
||
} else {
|
||
MysqlMode::Normal
|
||
}
|
||
}
|
||
None => MysqlMode::Normal,
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::{
|
||
agent_connect_timeout, connection_configs_pool_equivalent, connection_configs_session_credentials_compatible,
|
||
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,
|
||
};
|
||
use crate::agent_connection::{
|
||
agent_connect_params, mongo_legacy_error_with_auth_hint, mongo_uses_legacy_driver,
|
||
oracle_alternate_connect_config, should_retry_mongo_with_legacy_driver,
|
||
};
|
||
use crate::agent_manager::{AgentState, JavaRuntimeConfig, JavaRuntimeMode, DEFAULT_JRE_KEY};
|
||
use crate::database_capabilities;
|
||
use crate::db;
|
||
use crate::models::connection::{
|
||
default_connect_timeout_secs, default_redis_key_separator, AttachedDatabaseConfig, ConnectionConfig,
|
||
DatabaseType, HttpTunnelConfig, ProxyTunnelConfig, ProxyType, SshTunnelConfig, TransportLayerConfig,
|
||
};
|
||
use crate::query;
|
||
use crate::schema;
|
||
use crate::storage::Storage;
|
||
use std::sync::Arc;
|
||
use std::time::{Duration, Instant};
|
||
|
||
fn mysql_config(database: Option<&str>) -> ConnectionConfig {
|
||
ConnectionConfig {
|
||
docs_notes_path: None,
|
||
id: "conn".to_string(),
|
||
name: "MySQL".to_string(),
|
||
note: String::new(),
|
||
db_type: DatabaseType::Mysql,
|
||
driver_profile: None,
|
||
driver_label: None,
|
||
url_params: None,
|
||
agent_java_options: Vec::new(),
|
||
host: "127.0.0.1".to_string(),
|
||
port: 3306,
|
||
username: "root".to_string(),
|
||
password: "secret".to_string(),
|
||
database: database.map(str::to_string),
|
||
default_schema: None,
|
||
visible_databases: None,
|
||
visible_database_patterns: None,
|
||
visible_schemas: None,
|
||
show_system_schemas: false,
|
||
attached_databases: Vec::new(),
|
||
init_script: None,
|
||
color: None,
|
||
transport_layers: Vec::new(),
|
||
connect_timeout_secs: default_connect_timeout_secs(),
|
||
query_timeout_secs: crate::models::connection::default_query_timeout_secs(),
|
||
idle_timeout_secs: crate::models::connection::default_idle_timeout_secs(),
|
||
keepalive_interval_secs: crate::models::connection::default_keepalive_interval_secs(),
|
||
ssl: false,
|
||
ca_cert_path: String::new(),
|
||
client_cert_path: String::new(),
|
||
client_key_path: String::new(),
|
||
sysdba: false,
|
||
oracle_connection_type: None,
|
||
connection_string: None,
|
||
redis_connection_mode: None,
|
||
redis_sentinel_master: String::new(),
|
||
redis_sentinel_nodes: String::new(),
|
||
redis_sentinel_username: String::new(),
|
||
redis_sentinel_password: String::new(),
|
||
redis_sentinel_tls: false,
|
||
redis_cluster_nodes: String::new(),
|
||
redis_key_separator: default_redis_key_separator(),
|
||
redis_scan_page_size: None,
|
||
redis_database_aliases: Default::default(),
|
||
redis_key_templates: Vec::new(),
|
||
redis_key_grouping: None,
|
||
etcd_endpoints: String::new(),
|
||
gbase_server: String::new(),
|
||
informix_server: String::new(),
|
||
external_config: None,
|
||
plugin_id: None,
|
||
plugin_connection_provider: None,
|
||
plugin_connection_type: None,
|
||
connection_secrets: Default::default(),
|
||
jdbc_driver_class: None,
|
||
jdbc_driver_paths: Vec::new(),
|
||
one_time: false,
|
||
save_password: true,
|
||
read_only: false,
|
||
is_production: false,
|
||
production_databases: vec![],
|
||
database_info: None,
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn upsert_connection_url_param_inserts_and_replaces_catalog() {
|
||
assert_eq!(upsert_connection_url_param(None, "catalog", "paimon"), "catalog=paimon");
|
||
assert_eq!(
|
||
upsert_connection_url_param(Some("charset=utf8mb4"), "catalog", "hive_catalog"),
|
||
"charset=utf8mb4&catalog=hive%5Fcatalog"
|
||
);
|
||
assert_eq!(
|
||
upsert_connection_url_param(Some("catalog=old&charset=utf8mb4"), "catalog", "new_cat"),
|
||
"charset=utf8mb4&catalog=new%5Fcat"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn database_connection_config_with_catalog_keeps_database_for_use_after_set_catalog() {
|
||
let mut config = mysql_config(None);
|
||
config.db_type = DatabaseType::StarRocks;
|
||
let db_config = database_connection_config_with_catalog(&config, Some("ads"), Some("paimon_catalog"));
|
||
assert_eq!(db_config.database.as_deref(), Some("ads"));
|
||
assert_eq!(db_config.url_params.as_deref(), Some("catalog=paimon%5Fcatalog"));
|
||
let url = db_config.connection_url();
|
||
assert!(url.contains("/ads"), "url should include database for USE setup: {url}");
|
||
assert!(url.contains("catalog=paimon%5Fcatalog"), "url should include catalog param: {url}");
|
||
}
|
||
|
||
#[test]
|
||
fn metadata_connection_config_clears_starrocks_default_database() {
|
||
let mut config = mysql_config(Some("ads"));
|
||
config.db_type = DatabaseType::StarRocks;
|
||
assert_eq!(metadata_connection_config(&config).database, None);
|
||
}
|
||
|
||
#[test]
|
||
fn connection_configs_pool_equivalent_ignores_password_differences() {
|
||
let mut a = mysql_config(None);
|
||
a.save_password = false;
|
||
a.password = "session-secret".to_string();
|
||
// 持久化同步后 password 为空(save_password=false):视为同一连接,不销毁池。
|
||
let mut b = a.clone();
|
||
b.password.clear();
|
||
assert!(connection_configs_pool_equivalent(&a, &b));
|
||
assert!(connection_configs_pool_equivalent(&b, &a));
|
||
|
||
// 两个非空但不同的密码:同样忽略(密码不应触发池重建)。
|
||
let mut c = a.clone();
|
||
c.password = "changed-secret".to_string();
|
||
assert!(connection_configs_pool_equivalent(&a, &c));
|
||
}
|
||
|
||
#[test]
|
||
fn connection_configs_pool_equivalent_detects_saved_password_change() {
|
||
let mut a = mysql_config(None);
|
||
a.save_password = true;
|
||
a.password = "old-secret".to_string();
|
||
let mut b = a.clone();
|
||
b.password = "new-secret".to_string();
|
||
// save_password=true:密码是真实连接参数,保存新密码后旧池不得继续用旧密码认证。
|
||
assert!(!connection_configs_pool_equivalent(&a, &b));
|
||
assert!(!connection_configs_pool_equivalent(&b, &a));
|
||
}
|
||
|
||
#[test]
|
||
fn connection_configs_pool_equivalent_detects_real_parameter_changes() {
|
||
let mut a = mysql_config(None);
|
||
a.password = "secret".to_string();
|
||
// host / port / username / database 等真实连接参数变化应视为不同连接(销毁池)。
|
||
let mut host = a.clone();
|
||
host.host = "other-host".to_string();
|
||
assert!(!connection_configs_pool_equivalent(&a, &host));
|
||
|
||
let mut port = a.clone();
|
||
port.port = 5433;
|
||
assert!(!connection_configs_pool_equivalent(&a, &port));
|
||
|
||
let mut user = a.clone();
|
||
user.username = "other-user".to_string();
|
||
assert!(!connection_configs_pool_equivalent(&a, &user));
|
||
|
||
let mut ssl = a.clone();
|
||
ssl.ssl = true;
|
||
assert!(!connection_configs_pool_equivalent(&a, &ssl));
|
||
|
||
// 完全相等(含相同密码)→ true。
|
||
assert!(connection_configs_pool_equivalent(&a, &a.clone()));
|
||
}
|
||
|
||
#[test]
|
||
fn session_credentials_survive_nacos_scope_and_display_updates() {
|
||
let mut initial = mysql_config(None);
|
||
initial.db_type = DatabaseType::Nacos;
|
||
initial.save_password = false;
|
||
initial.password = "old-session-secret".to_string();
|
||
initial.visible_databases = Some(vec!["namespace-a".to_string()]);
|
||
initial.external_config = Some(serde_json::json!({
|
||
"implementation": "nacos",
|
||
"versionMode": "v3",
|
||
"apiPlane": "admin",
|
||
"serverAddr": "http://127.0.0.1:8848",
|
||
"managedNamespaces": ["namespace-a"],
|
||
"auth": {
|
||
"kind": "usernamePassword",
|
||
"username": "ordinary-user",
|
||
"password": "old-session-secret"
|
||
}
|
||
}));
|
||
|
||
let mut updated = initial.clone();
|
||
updated.name = "Renamed Nacos".to_string();
|
||
updated.note = "Local note".to_string();
|
||
updated.visible_databases = Some(vec!["namespace-a".to_string(), "namespace-b".to_string()]);
|
||
updated.external_config.as_mut().unwrap()["managedNamespaces"] =
|
||
serde_json::json!(["namespace-a", "namespace-b"]);
|
||
updated.external_config.as_mut().unwrap()["auth"]["password"] = serde_json::json!("new-session-secret");
|
||
|
||
assert!(connection_configs_session_credentials_compatible(&initial, &updated));
|
||
|
||
let mut scrubbed = updated.clone();
|
||
scrubbed.external_config.as_mut().unwrap()["auth"].as_object_mut().unwrap().remove("password");
|
||
assert!(connection_configs_session_credentials_compatible(&initial, &scrubbed));
|
||
|
||
updated.external_config.as_mut().unwrap()["serverAddr"] = serde_json::json!("http://127.0.0.1:8080");
|
||
assert!(!connection_configs_session_credentials_compatible(&initial, &updated));
|
||
}
|
||
|
||
#[test]
|
||
fn session_credentials_do_not_survive_nacos_account_updates() {
|
||
let mut initial = mysql_config(None);
|
||
initial.db_type = DatabaseType::Nacos;
|
||
initial.save_password = false;
|
||
initial.external_config = Some(serde_json::json!({
|
||
"serverAddr": "http://127.0.0.1:8848",
|
||
"auth": { "kind": "usernamePassword", "username": "user-a", "password": "secret" }
|
||
}));
|
||
let mut updated = initial.clone();
|
||
updated.external_config.as_mut().unwrap()["auth"]["username"] = serde_json::json!("user-b");
|
||
|
||
assert!(!connection_configs_session_credentials_compatible(&initial, &updated));
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn apply_session_credential_injects_saved_password_only_for_no_save_connections() {
|
||
let dir = std::env::temp_dir().join(format!("dbx-core-session-cred-{}", std::process::id()));
|
||
std::fs::create_dir_all(&dir).unwrap();
|
||
let storage = Storage::open(&dir.join("storage.db")).await.unwrap();
|
||
let state = AppState::new_with_plugin_dir(storage, dir.join("plugins"));
|
||
let _ = state.session_credentials.set("", "conn-a", "s3cret");
|
||
|
||
let mut config = mysql_config(None);
|
||
config.id = "conn-a".to_string();
|
||
config.save_password = false;
|
||
config.password.clear();
|
||
|
||
// save_password=false + db_config 密码为空 → 从会话凭据仓库注入。
|
||
let mut db_config = metadata_connection_config(&config);
|
||
state.apply_session_credential(&config, &mut db_config, &config.id);
|
||
assert_eq!(db_config.password, "s3cret");
|
||
|
||
// save_password=true → 不注入(走持久化水合的密码,若为空则保持空)。
|
||
config.save_password = true;
|
||
let mut db_config = metadata_connection_config(&config);
|
||
state.apply_session_credential(&config, &mut db_config, &config.id);
|
||
assert_eq!(db_config.password, "");
|
||
|
||
// 无会话凭据 → 保持空密码("断开并忘记"后重新输入)。
|
||
state.session_credentials.remove("", "conn-a");
|
||
config.save_password = false;
|
||
let mut db_config = metadata_connection_config(&config);
|
||
state.apply_session_credential(&config, &mut db_config, &config.id);
|
||
assert_eq!(db_config.password, "");
|
||
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn apply_session_credential_reads_owner_scoped_credentials_only() {
|
||
let dir = std::env::temp_dir().join(format!("dbx-core-session-cred-owner-{}", std::process::id()));
|
||
std::fs::create_dir_all(&dir).unwrap();
|
||
let storage = Storage::open(&dir.join("storage.db")).await.unwrap();
|
||
let state = AppState::new_with_plugin_dir(storage, dir.join("plugins"));
|
||
|
||
let mut config = mysql_config(None);
|
||
config.id = "conn-a".to_string();
|
||
config.save_password = false;
|
||
config.password.clear();
|
||
|
||
// 会话 X 输入了密码,会话 Y 未输入:Y 的请求(owner=token-y)不得注入 X 的密码。
|
||
let _ = state.session_credentials.set("token-x", "conn-a", "x-secret");
|
||
crate::session_credentials::with_credential_owner(Some("token-y".to_string()), async {
|
||
let mut db_config = metadata_connection_config(&config);
|
||
state.apply_session_credential(&config, &mut db_config, &config.id);
|
||
assert_eq!(db_config.password, "");
|
||
})
|
||
.await;
|
||
|
||
// 会话 X 自身能读到自己的密码。
|
||
crate::session_credentials::with_credential_owner(Some("token-x".to_string()), async {
|
||
let mut db_config = metadata_connection_config(&config);
|
||
state.apply_session_credential(&config, &mut db_config, &config.id);
|
||
assert_eq!(db_config.password, "x-secret");
|
||
})
|
||
.await;
|
||
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn pool_credential_owner_mismatch_prevents_cross_session_pool_reuse() {
|
||
let dir = std::env::temp_dir().join(format!("dbx-core-pool-owner-{}", std::process::id()));
|
||
std::fs::create_dir_all(&dir).unwrap();
|
||
let storage = Storage::open(&dir.join("storage.db")).await.unwrap();
|
||
let state = AppState::new_with_plugin_dir(storage, dir.join("plugins"));
|
||
|
||
let mut config = mysql_config(None);
|
||
config.id = "conn-a".to_string();
|
||
config.save_password = false;
|
||
|
||
// 模拟会话 X 创建的 no-save 池。
|
||
state.session_credentials.record_pool_owner("conn-a", "token-x");
|
||
|
||
// 同一会话 X 复用 → 无冲突。
|
||
crate::session_credentials::with_credential_owner(Some("token-x".to_string()), async {
|
||
assert!(!state.pool_credential_owner_mismatch(&config, "conn-a").await);
|
||
})
|
||
.await;
|
||
|
||
// 另一会话 Y 请求同一 no-save 池 → 冲突,需以 Y 的凭据重建,避免复用 X 的密码。
|
||
crate::session_credentials::with_credential_owner(Some("token-y".to_string()), async {
|
||
assert!(state.pool_credential_owner_mismatch(&config, "conn-a").await);
|
||
})
|
||
.await;
|
||
|
||
// owner 标记缺失时按不可信处理,避免全局配置失效与异步移除旧池之间复用旧池。
|
||
state.session_credentials.clear_connection("conn-a");
|
||
crate::session_credentials::with_credential_owner(Some("token-x".to_string()), async {
|
||
assert!(state.pool_credential_owner_mismatch(&config, "conn-a").await);
|
||
})
|
||
.await;
|
||
|
||
// 桌面端(无 owner)单用户 → 不冲突,按既有逻辑复用。
|
||
assert!(!state.pool_credential_owner_mismatch(&config, "conn-a").await);
|
||
|
||
// 已保存密码连接始终共享池 → 不冲突。
|
||
config.save_password = true;
|
||
crate::session_credentials::with_credential_owner(Some("token-y".to_string()), async {
|
||
assert!(!state.pool_credential_owner_mismatch(&config, "conn-a").await);
|
||
})
|
||
.await;
|
||
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[test]
|
||
fn task_client_session_ids_are_stable_and_isolated() {
|
||
assert_eq!(task_client_session_id("table-export", "job-1"), "table-export:job-1");
|
||
assert_ne!(task_client_session_id("table-export", "job-1"), task_client_session_id("database-export", "job-1"));
|
||
assert_ne!(task_client_session_id("table-export", "job-1"), task_client_session_id("table-export", "job-2"));
|
||
}
|
||
|
||
#[test]
|
||
fn access_agent_connect_timeout_has_longer_default_floor() {
|
||
let mut config = mysql_config(None);
|
||
config.db_type = DatabaseType::Access;
|
||
config.connect_timeout_secs = 5;
|
||
assert_eq!(agent_connect_timeout(&config).as_secs(), 30);
|
||
|
||
config.connect_timeout_secs = 45;
|
||
assert_eq!(agent_connect_timeout(&config).as_secs(), 45);
|
||
}
|
||
|
||
#[test]
|
||
fn non_access_agent_connect_timeout_uses_standard_floor() {
|
||
let mut config = mysql_config(None);
|
||
config.db_type = DatabaseType::Oracle;
|
||
config.connect_timeout_secs = 5;
|
||
assert_eq!(agent_connect_timeout(&config).as_secs(), 30);
|
||
|
||
config.connect_timeout_secs = 45;
|
||
assert_eq!(agent_connect_timeout(&config).as_secs(), 45);
|
||
}
|
||
|
||
#[test]
|
||
fn prestosql_jdbc_config_sets_presto_url_and_driver_class() {
|
||
let mut config = mysql_config(Some("hive/default"));
|
||
config.db_type = DatabaseType::PrestoSql;
|
||
config.host = "presto.example.com".to_string();
|
||
config.port = 9090;
|
||
|
||
let jdbc_config = prestosql_jdbc_config_for_endpoint(&config, "127.0.0.1", 19090).unwrap();
|
||
|
||
assert_eq!(jdbc_config.connection_string.as_deref(), Some("jdbc:presto://127.0.0.1:19090/hive/default"));
|
||
assert_eq!(jdbc_config.jdbc_driver_class.as_deref(), Some(PRESTOSQL_JDBC_DRIVER_CLASS));
|
||
}
|
||
|
||
#[test]
|
||
fn prestosql_jdbc_config_preserves_custom_driver_class_and_paths() {
|
||
let mut config = mysql_config(Some("hive"));
|
||
config.db_type = DatabaseType::PrestoSql;
|
||
config.jdbc_driver_class = Some("custom.PrestoDriver".to_string());
|
||
config.jdbc_driver_paths = vec!["D:\\software\\jar\\presto-jdbc-350.jar".to_string()];
|
||
|
||
let jdbc_config = prestosql_jdbc_config_for_endpoint(&config, "presto.example.com", 9090).unwrap();
|
||
|
||
assert_eq!(jdbc_config.jdbc_driver_class.as_deref(), Some("custom.PrestoDriver"));
|
||
assert_eq!(jdbc_config.jdbc_driver_paths, vec!["D:\\software\\jar\\presto-jdbc-350.jar"]);
|
||
}
|
||
|
||
#[test]
|
||
fn prestosql_jdbc_config_consumes_url_params_after_building_url() {
|
||
let mut config = mysql_config(None);
|
||
config.db_type = DatabaseType::PrestoSql;
|
||
config.ssl = true;
|
||
config.url_params = Some(
|
||
"SSL=true&SSLKeyStorePassword=secret&SSLKeyStorePath=D:/keystore/presto/presto_keystore.jks".to_string(),
|
||
);
|
||
|
||
let jdbc_config = prestosql_jdbc_config_for_endpoint(&config, "presto.example.com", 8443).unwrap();
|
||
|
||
assert_eq!(
|
||
jdbc_config.connection_string.as_deref(),
|
||
Some(
|
||
"jdbc:presto://presto.example.com:8443?SSL=true&SSLKeyStorePassword=secret&SSLKeyStorePath=D:/keystore/presto/presto_keystore.jks"
|
||
)
|
||
);
|
||
assert_eq!(jdbc_config.url_params, None);
|
||
}
|
||
|
||
#[test]
|
||
fn gaussdb_m_profile_uses_vendor_jdbc_url_and_driver() {
|
||
let mut config = mysql_config(Some("业务库"));
|
||
config.db_type = DatabaseType::Gaussdb;
|
||
config.host = "db.internal".to_string();
|
||
config.port = 8000;
|
||
config.driver_profile = Some(GAUSSDB_M_JDBC_DRIVER_PROFILE.to_string());
|
||
config.url_params = Some("currentSchema=app".to_string());
|
||
|
||
assert!(gaussdb_uses_m_jdbc_driver(&config));
|
||
let jdbc = gaussdb_m_jdbc_config_for_endpoint(&config, "127.0.0.1", 18000);
|
||
assert_eq!(
|
||
jdbc.connection_string.as_deref(),
|
||
Some(
|
||
"jdbc:gaussdb://127.0.0.1:18000/%E4%B8%9A%E5%8A%A1%E5%BA%93?currentSchema=app&sslmode=prefer&ssl=true"
|
||
)
|
||
);
|
||
assert_eq!(jdbc.jdbc_driver_class.as_deref(), Some(GAUSSDB_M_JDBC_DRIVER_CLASS));
|
||
|
||
config.driver_profile = Some("gaussdb".to_string());
|
||
assert!(!gaussdb_uses_m_jdbc_driver(&config));
|
||
}
|
||
|
||
#[test]
|
||
fn gaussdb_m_jdbc_url_normalizes_tls_parameters() {
|
||
let mut config = mysql_config(Some("postgres"));
|
||
config.db_type = DatabaseType::Gaussdb;
|
||
config.driver_profile = Some(GAUSSDB_M_JDBC_DRIVER_PROFILE.to_string());
|
||
|
||
config.ssl = true;
|
||
config.url_params = None;
|
||
assert_eq!(
|
||
gaussdb_m_jdbc_config_for_endpoint(&config, "db.internal", 8000).connection_string.as_deref(),
|
||
Some("jdbc:gaussdb://db.internal:8000/postgres?sslmode=require&ssl=true")
|
||
);
|
||
|
||
config.ssl = false;
|
||
config.url_params = Some("?sslmode=disable¤tSchema=app".to_string());
|
||
assert_eq!(
|
||
gaussdb_m_jdbc_config_for_endpoint(&config, "db.internal", 8000).connection_string.as_deref(),
|
||
Some("jdbc:gaussdb://db.internal:8000/postgres?currentSchema=app&sslmode=disable&ssl=false")
|
||
);
|
||
|
||
config.url_params = Some("ssl=true¤tSchema=legacy".to_string());
|
||
assert_eq!(
|
||
gaussdb_m_jdbc_config_for_endpoint(&config, "db.internal", 8000).connection_string.as_deref(),
|
||
Some("jdbc:gaussdb://db.internal:8000/postgres?currentSchema=legacy&sslmode=prefer&ssl=true")
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn gaussdb_m_jdbc_target_server_type_preserves_url_params_and_allows_explicit_override() {
|
||
let mut config = mysql_config(Some("postgres"));
|
||
config.db_type = DatabaseType::Gaussdb;
|
||
config.driver_profile = Some(GAUSSDB_M_JDBC_DRIVER_PROFILE.to_string());
|
||
|
||
let default_url = gaussdb_m_jdbc_config_for_endpoint(&config, "db.internal", 8000).connection_string.unwrap();
|
||
assert!(!default_url.to_ascii_lowercase().contains("targetservertype="));
|
||
|
||
config.url_params = Some("targetServerType=slave¤tSchema=app".to_string());
|
||
assert_eq!(
|
||
gaussdb_m_jdbc_config_for_endpoint(&config, "db.internal", 8000).connection_string.as_deref(),
|
||
Some(
|
||
"jdbc:gaussdb://db.internal:8000/postgres?targetServerType=slave¤tSchema=app&sslmode=prefer&ssl=true"
|
||
)
|
||
);
|
||
|
||
config.external_config = Some(serde_json::json!({ "gaussdbTargetServerType": "master" }));
|
||
assert_eq!(
|
||
gaussdb_m_jdbc_config_for_endpoint(&config, "db.internal", 8000).connection_string.as_deref(),
|
||
Some(
|
||
"jdbc:gaussdb://db.internal:8000/postgres?currentSchema=app&sslmode=prefer&ssl=true&targetServerType=master"
|
||
)
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn gaussdb_jdbc_compatibility_query_result_selects_identifier_quote() {
|
||
let result = crate::types::QueryResult {
|
||
columns: vec!["datcompatibility".to_string()],
|
||
column_types: vec!["text".to_string()],
|
||
column_sortables: vec![true],
|
||
spatial_columns: vec![],
|
||
spatial_values: vec![],
|
||
rows: vec![vec![serde_json::json!("M")]],
|
||
affected_rows: 0,
|
||
execution_time_ms: 1,
|
||
truncated: false,
|
||
session_id: None,
|
||
has_more: false,
|
||
elasticsearch_raw_body: None,
|
||
messages: Vec::new(),
|
||
};
|
||
|
||
assert_eq!(gaussdb_identifier_quote_from_query_result(&result).as_deref(), Some("`"));
|
||
}
|
||
|
||
#[test]
|
||
fn sqlserver_legacy_agent_config_marks_hidden_profile() {
|
||
let mut config = mysql_config(Some("master"));
|
||
config.db_type = DatabaseType::SqlServer;
|
||
|
||
let legacy = sqlserver_legacy_agent_config(&config);
|
||
|
||
assert_eq!(legacy.db_type, DatabaseType::SqlServer);
|
||
assert_eq!(legacy.driver_profile.as_deref(), Some(crate::db::sqlserver::SQLSERVER_LEGACY_DRIVER_PROFILE));
|
||
assert_eq!(legacy.driver_label.as_deref(), Some(crate::db::sqlserver::SQLSERVER_LEGACY_DRIVER_LABEL));
|
||
assert!(sqlserver_uses_legacy_driver(&legacy));
|
||
}
|
||
|
||
#[test]
|
||
fn sqlserver_legacy_url_param_requires_canonicalization_before_agent_driver_selection() {
|
||
let mut config = mysql_config(Some("master"));
|
||
config.db_type = DatabaseType::SqlServer;
|
||
config.url_params = Some("applicationName=dbx;sqlserverEncryption=disabled".to_string());
|
||
|
||
assert!(!sqlserver_uses_legacy_driver(&config));
|
||
assert!(sqlserver_uses_legacy_driver(&config.canonicalized()));
|
||
}
|
||
|
||
#[test]
|
||
fn sqlserver_legacy_driver_error_mentions_driver_manager_when_missing() {
|
||
let message = sqlserver_legacy_driver_error(
|
||
"sqlserver-legacy driver is not installed. Please install it from the Driver Manager.",
|
||
);
|
||
|
||
assert!(message.contains("Driver Manager"));
|
||
assert!(message.contains("enable SQL Server legacy compatibility mode again"));
|
||
}
|
||
|
||
#[test]
|
||
fn sqlserver_legacy_driver_hint_preserves_structured_agent_error() {
|
||
let error = crate::db::agent_driver::AgentCallError::Structured {
|
||
rpc_code: -6007,
|
||
message: "driver is not installed".to_string(),
|
||
context: crate::db::agent_driver::AgentErrorContext {
|
||
contract_version: 1,
|
||
category: crate::db::agent_driver::AgentErrorCategory::Connection,
|
||
retryable: false,
|
||
session_disposition: crate::db::agent_driver::AgentSessionDisposition::Quarantine,
|
||
stage: crate::db::agent_driver::AgentErrorStage::Connect,
|
||
operation_outcome: crate::db::agent_driver::AgentOperationOutcome::NotStarted,
|
||
agent_session_id: None,
|
||
sql_state: None,
|
||
vendor_code: None,
|
||
exception_class: None,
|
||
},
|
||
}
|
||
.into_legacy_string();
|
||
|
||
let message = sqlserver_legacy_driver_error(&error);
|
||
|
||
assert!(matches!(
|
||
crate::db::agent_driver::agent_error_from_legacy(&message, None),
|
||
crate::db::agent_driver::AgentCallError::Structured { .. }
|
||
));
|
||
}
|
||
|
||
#[test]
|
||
fn agent_connect_params_include_url_params() {
|
||
let mut config = mysql_config(Some("testdb"));
|
||
config.username = "informix".to_string();
|
||
config.password = "in4mix".to_string();
|
||
config.url_params = Some("INFORMIXSERVER=informix;CLIENT_LOCALE=en_US.utf8".to_string());
|
||
|
||
let params = agent_connect_params(&config, "172.26.128.159", 20013, "testdb").unwrap();
|
||
|
||
assert_eq!(params["host"], "172.26.128.159");
|
||
assert_eq!(params["port"], 20013);
|
||
assert_eq!(params["database"], "testdb");
|
||
assert_eq!(params["username"], "informix");
|
||
assert_eq!(params["password"], "in4mix");
|
||
assert_eq!(params["url_params"], "INFORMIXSERVER=informix;CLIENT_LOCALE=en_US.utf8");
|
||
}
|
||
|
||
#[test]
|
||
fn agent_connect_params_include_sqlserver_explicit_port_state() {
|
||
let mut config = mysql_config(Some("master"));
|
||
config.db_type = DatabaseType::SqlServer;
|
||
config.external_config = Some(serde_json::json!({ "portExplicit": true }));
|
||
|
||
let params = agent_connect_params(&config, r"db.example.com\SQLEXPRESS", 1433, "master").unwrap();
|
||
|
||
assert_eq!(params["port_explicit"], true);
|
||
}
|
||
|
||
#[test]
|
||
fn databend_uses_agent_pool_not_bare_mysql_pool() {
|
||
assert!(uses_bare_mysql_pool(&DatabaseType::Doris));
|
||
assert!(uses_bare_mysql_pool(&DatabaseType::StarRocks));
|
||
assert!(uses_bare_mysql_pool(&DatabaseType::ManticoreSearch));
|
||
assert!(!uses_bare_mysql_pool(&DatabaseType::Databend));
|
||
assert!(database_capabilities::is_agent_type(&DatabaseType::Databend));
|
||
assert!(super::uses_agent_connection_pool(&DatabaseType::ZooKeeper));
|
||
}
|
||
|
||
#[test]
|
||
fn drivers_with_internal_recovery_skip_eager_pool_validation() {
|
||
assert!(!super::should_validate_existing_pool_before_reuse(DatabaseType::Postgres));
|
||
assert!(!super::should_validate_existing_pool_before_reuse(DatabaseType::Etcd));
|
||
assert!(!super::should_validate_existing_pool_before_reuse(DatabaseType::Dameng));
|
||
assert!(!super::should_validate_existing_pool_before_reuse(DatabaseType::Oracle));
|
||
assert!(super::should_validate_existing_pool_before_reuse(DatabaseType::Mysql));
|
||
}
|
||
|
||
#[test]
|
||
fn validates_h2_database_base_path_when_mv_db_file_exists() {
|
||
let dir = std::env::temp_dir().join(format!("dbx-h2-test-{}", uuid::Uuid::new_v4()));
|
||
std::fs::create_dir_all(&dir).unwrap();
|
||
let file_path = dir.join("app.mv.db");
|
||
std::fs::write(&file_path, b"h2").unwrap();
|
||
let base_path = dir.join("app");
|
||
|
||
validate_h2_database_path(base_path.to_str().unwrap()).unwrap();
|
||
std::fs::remove_dir_all(&dir).unwrap();
|
||
}
|
||
|
||
#[test]
|
||
fn rejects_missing_h2_database_path() {
|
||
let dir = std::env::temp_dir().join(format!("dbx-h2-test-{}", uuid::Uuid::new_v4()));
|
||
std::fs::create_dir_all(&dir).unwrap();
|
||
let missing_path = dir.join("missing");
|
||
|
||
let err = validate_h2_database_path(missing_path.to_str().unwrap()).unwrap_err();
|
||
|
||
assert!(err.contains("File does not exist"));
|
||
std::fs::remove_dir_all(&dir).unwrap();
|
||
}
|
||
|
||
#[test]
|
||
fn agent_connect_params_build_mongodb_connection_string_from_form_fields() {
|
||
let mut config = mysql_config(Some("RestCloud_V45PUB_Gateway"));
|
||
config.db_type = DatabaseType::MongoDb;
|
||
config.host = "172.22.4.42".to_string();
|
||
config.port = 27017;
|
||
config.username = "mongouser".to_string();
|
||
config.password = "secret".to_string();
|
||
config.url_params = Some("authSource=admin&authMechanism=SCRAM-SHA-1".to_string());
|
||
|
||
let params = agent_connect_params(&config, "172.22.4.42", 27017, "RestCloud_V45PUB_Gateway").unwrap();
|
||
|
||
assert_eq!(params["connection_string"], "mongodb://mongouser:secret@172.22.4.42:27017/RestCloud%5FV45PUB%5FGateway?authSource=admin&authMechanism=SCRAM-SHA-1");
|
||
}
|
||
|
||
#[test]
|
||
fn agent_connect_params_mongodb_uses_connection_string_database_when_database_is_empty() {
|
||
let mut config = mysql_config(None);
|
||
config.db_type = DatabaseType::MongoDb;
|
||
config.connection_string =
|
||
Some("mongodb://mongouser:secret@172.22.4.42:27017/RestCloud_V45PUB_Gateway?authSource=admin".to_string());
|
||
|
||
let params = agent_connect_params(&config, "172.22.4.42", 27017, "").unwrap();
|
||
|
||
assert_eq!(params["database"], "RestCloud_V45PUB_Gateway");
|
||
}
|
||
|
||
#[test]
|
||
fn mongo_legacy_auth_error_adds_auth_source_hint() {
|
||
let err = "Agent RPC error: Exception authenticating MongoCredential{mechanism=SCRAM-SHA-1, userName='rwuser', source='gray_lite_twin_fat'}";
|
||
|
||
assert_eq!(
|
||
mongo_legacy_error_with_auth_hint(err),
|
||
"Agent RPC error: Exception authenticating MongoCredential{mechanism=SCRAM-SHA-1, userName='rwuser', source='gray_lite_twin_fat'}\n\nCurrent authentication database: gray_lite_twin_fat. The server rejected these credentials. Verify the username and password, and confirm that the user was created in gray_lite_twin_fat. If the user was created in admin, set Authentication database to admin or add authSource=admin to URL params."
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn oceanbase_mysql_setup_queries_follow_query_timeout() {
|
||
let mut config = mysql_config(Some("dbx"));
|
||
config.driver_profile = Some("oceanbase".to_string());
|
||
config.query_timeout_secs = 30;
|
||
|
||
assert_eq!(oceanbase_mysql_setup_queries(&config), vec!["SET ob_query_timeout = 30000000"]);
|
||
}
|
||
|
||
#[test]
|
||
fn oceanbase_mysql_query_timeout_sql_accepts_large_timeout() {
|
||
let mut config = mysql_config(Some("dbx"));
|
||
config.driver_profile = Some("oceanbase".to_string());
|
||
|
||
assert_eq!(
|
||
db::oceanbase_mysql::query_timeout_sql(&config, 300_000),
|
||
Some("SET ob_query_timeout = 300000000000".to_string())
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn oceanbase_mysql_setup_queries_skip_disabled_timeout() {
|
||
let mut config = mysql_config(Some("dbx"));
|
||
config.driver_profile = Some("oceanbase".to_string());
|
||
config.query_timeout_secs = 0;
|
||
|
||
assert!(oceanbase_mysql_setup_queries(&config).is_empty());
|
||
}
|
||
|
||
#[test]
|
||
fn oceanbase_mysql_setup_queries_do_not_apply_to_plain_mysql() {
|
||
let mut config = mysql_config(Some("dbx"));
|
||
config.query_timeout_secs = 30;
|
||
|
||
assert!(oceanbase_mysql_setup_queries(&config).is_empty());
|
||
}
|
||
|
||
#[test]
|
||
fn doris_pool_setup_uses_switch_for_configured_catalog() {
|
||
let mut config = mysql_config(Some("bi"));
|
||
config.db_type = DatabaseType::Doris;
|
||
|
||
assert_eq!(
|
||
mysql_pool_setup_queries(&config, "mysql://root:secret@localhost:9030/bi?catalog=paimon%5Fcatalog"),
|
||
vec!["SWITCH `paimon_catalog`"]
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn starrocks_pool_setup_uses_set_catalog_for_configured_catalog() {
|
||
let mut config = mysql_config(Some("bi"));
|
||
config.db_type = DatabaseType::StarRocks;
|
||
|
||
assert_eq!(
|
||
mysql_pool_setup_queries(&config, "mysql://root:secret@localhost:9030/bi?catalog=paimon%5Fcatalog"),
|
||
vec!["SET CATALOG `paimon_catalog`"]
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn mongo_legacy_retry_covers_old_server_handshake_eof() {
|
||
let err = r#"MongoDB connection failed: Kind: Server selection timeout: No available servers. Topology: { Type: Unknown, Servers: [ { Address: db.example.com:27017, Type: Unknown, Error: Kind: I/O error: unexpected end of file } ] }"#;
|
||
|
||
assert!(mongo_uses_legacy_driver(&ConnectionConfig {
|
||
driver_profile: Some("mongodb-legacy".to_string()),
|
||
..mysql_config(None)
|
||
}));
|
||
assert!(should_retry_mongo_with_legacy_driver(err));
|
||
assert!(should_retry_mongo_with_legacy_driver("server reports wire version 5, but this driver requires 8"));
|
||
assert!(!should_retry_mongo_with_legacy_driver("Authentication failed."));
|
||
}
|
||
|
||
#[test]
|
||
fn agent_connect_params_build_oracle_service_connection_string() {
|
||
let mut config = mysql_config(Some("ORCLPDB1"));
|
||
config.db_type = DatabaseType::Oracle;
|
||
config.host = "oracle.example.com".to_string();
|
||
config.port = 1521;
|
||
config.username = "system".to_string();
|
||
config.password = "oracle".to_string();
|
||
config.sysdba = true;
|
||
config.oracle_connection_type = Some("service_name".to_string());
|
||
|
||
let params = agent_connect_params(&config, "oracle.example.com", 1521, "ORCLPDB1").unwrap();
|
||
|
||
assert_eq!(params["database"], "SYSDBA:ORCLPDB1");
|
||
assert_eq!(params["sysdba"], true);
|
||
assert_eq!(params["connection_string"], "jdbc:oracle:thin:@//oracle.example.com:1521/ORCLPDB1");
|
||
}
|
||
|
||
#[test]
|
||
fn agent_connect_params_build_postgres_like_agent_connection_string_for_selected_database() {
|
||
let cases = [
|
||
(
|
||
DatabaseType::Kingbase,
|
||
"kingbase.example.com",
|
||
54321,
|
||
"jdbc:kingbase8://kingbase.example.com:54321/platform_face_jgj",
|
||
"jdbc:kingbase8://kingbase.example.com:54321/platform_face_freezer_jgj?sslmode=disable",
|
||
),
|
||
(
|
||
DatabaseType::Highgo,
|
||
"highgo.example.com",
|
||
5866,
|
||
"jdbc:highgo://highgo.example.com:5866/highgo",
|
||
"jdbc:highgo://highgo.example.com:5866/platform_face_freezer_jgj?sslmode=disable",
|
||
),
|
||
(
|
||
DatabaseType::Vastbase,
|
||
"vastbase.example.com",
|
||
5432,
|
||
"jdbc:vastbase://vastbase.example.com:5432/postgres",
|
||
"jdbc:vastbase://vastbase.example.com:5432/platform_face_freezer_jgj?sslmode=disable",
|
||
),
|
||
];
|
||
|
||
for (db_type, host, port, stale_connection_string, expected_connection_string) in cases {
|
||
let mut config = mysql_config(Some("platform_face_jgj"));
|
||
config.db_type = db_type;
|
||
config.host = host.to_string();
|
||
config.port = port;
|
||
config.username = "system".to_string();
|
||
config.password = "secret".to_string();
|
||
config.url_params = Some("sslmode=disable".to_string());
|
||
config.connection_string = Some(stale_connection_string.to_string());
|
||
|
||
let params = agent_connect_params(&config, host, port, "platform_face_freezer_jgj").unwrap();
|
||
|
||
assert_eq!(params["database"], "platform_face_freezer_jgj");
|
||
assert_eq!(params["connection_string"], expected_connection_string);
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn agent_connect_params_build_oracle_sid_connection_string() {
|
||
let mut config = mysql_config(Some("ORCL"));
|
||
config.db_type = DatabaseType::Oracle;
|
||
config.oracle_connection_type = Some("sid".to_string());
|
||
|
||
let params = agent_connect_params(&config, "127.0.0.1", 11521, "ORCL").unwrap();
|
||
|
||
assert_eq!(params["connection_string"], "jdbc:oracle:thin:@127.0.0.1:11521:ORCL");
|
||
}
|
||
|
||
#[test]
|
||
fn agent_connect_params_preserve_legacy_oracle_configs_as_service_name() {
|
||
let mut config = mysql_config(Some("ORCL"));
|
||
config.db_type = DatabaseType::Oracle;
|
||
config.oracle_connection_type = None;
|
||
|
||
let params = agent_connect_params(&config, "127.0.0.1", 11521, "ORCL").unwrap();
|
||
|
||
assert_eq!(params["connection_string"], "jdbc:oracle:thin:@//127.0.0.1:11521/ORCL");
|
||
}
|
||
|
||
#[test]
|
||
fn oracle_listener_errors_can_retry_with_alternate_connect_descriptor() {
|
||
let mut config = mysql_config(Some("ORCL"));
|
||
config.db_type = DatabaseType::Oracle;
|
||
config.driver_profile = Some("oracle".to_string());
|
||
config.oracle_connection_type = Some("service_name".to_string());
|
||
|
||
let retry = oracle_alternate_connect_config(
|
||
&config,
|
||
"Agent RPC error (-1): ORA-12514: listener does not currently know of service requested",
|
||
)
|
||
.expect("listener errors should allow alternate descriptor retry");
|
||
assert_eq!(retry.driver_profile.as_deref(), Some("oracle"));
|
||
assert_eq!(retry.connection_string.as_deref(), Some("jdbc:oracle:thin:@127.0.0.1:3306:ORCL"));
|
||
|
||
let mut sid_config = config.clone();
|
||
sid_config.oracle_connection_type = Some("sid".to_string());
|
||
let service_retry = oracle_alternate_connect_config(
|
||
&sid_config,
|
||
"Agent RPC error (-1): ORA-12505: listener does not currently know of SID given",
|
||
)
|
||
.expect("SID listener errors should allow service-name retry");
|
||
assert_eq!(service_retry.connection_string.as_deref(), Some("jdbc:oracle:thin:@//127.0.0.1:3306/ORCL"));
|
||
|
||
assert!(oracle_alternate_connect_config(&config, "ORA-12541: TNS:no listener").is_some());
|
||
}
|
||
|
||
#[test]
|
||
fn oracle_alternate_descriptor_retry_skips_non_listener_errors() {
|
||
let mut config = mysql_config(Some("ORCL"));
|
||
config.db_type = DatabaseType::Oracle;
|
||
config.driver_profile = Some("oracle".to_string());
|
||
|
||
assert!(oracle_alternate_connect_config(&config, "ORA-01017: invalid username/password").is_none());
|
||
}
|
||
|
||
#[test]
|
||
fn oracle_alternate_descriptor_retry_skips_custom_connection_strings() {
|
||
let mut config = mysql_config(Some("ORCL"));
|
||
config.db_type = DatabaseType::Oracle;
|
||
config.driver_profile = Some("oracle".to_string());
|
||
config.connection_string = Some("jdbc:oracle:thin:@//oracle.example.com:1521/ORCL".to_string());
|
||
|
||
assert!(oracle_alternate_connect_config(&config, "ORA-12514: listener does not know service").is_none());
|
||
}
|
||
|
||
#[test]
|
||
fn agent_connect_params_build_saphana_connection_string_from_database_and_url_params() {
|
||
let mut config = mysql_config(Some("TENANT1"));
|
||
config.db_type = DatabaseType::SapHana;
|
||
config.host = "hana.example.com".to_string();
|
||
config.port = 30013;
|
||
config.username = "SYSTEM".to_string();
|
||
config.password = "secret".to_string();
|
||
config.url_params = Some("encrypt=true".to_string());
|
||
|
||
let params = agent_connect_params(&config, "hana.example.com", 30013, "TENANT1").unwrap();
|
||
|
||
assert_eq!(params["database"], "TENANT1");
|
||
assert_eq!(params["connection_string"], "jdbc:sap://hana.example.com:30013/?databaseName=TENANT1&encrypt=true");
|
||
}
|
||
|
||
async fn test_app_state() -> (AppState, std::path::PathBuf) {
|
||
let dir = std::env::temp_dir().join(format!("dbx-core-test-{}", uuid::Uuid::new_v4()));
|
||
std::fs::create_dir_all(&dir).unwrap();
|
||
let storage = Storage::open(&dir.join("storage.db")).await.unwrap();
|
||
(AppState::new(storage), dir)
|
||
}
|
||
|
||
fn agent_pool_stub() -> PoolKind {
|
||
PoolKind::agent(crate::db::agent_driver::AgentDriverClient::test_stub())
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn owned_pool_handle_does_not_block_registry_writes() {
|
||
let (state, dir) = test_app_state().await;
|
||
state
|
||
.update_connection_pools(|connections| {
|
||
connections.insert("slow-operation".to_string(), agent_pool_stub());
|
||
})
|
||
.await;
|
||
|
||
let handle = state.pool_handle("slow-operation").await.expect("pool handle");
|
||
tokio::time::timeout(
|
||
Duration::from_millis(100),
|
||
state.update_connection_pools(|connections| {
|
||
connections.insert("other-connection".to_string(), agent_pool_stub());
|
||
}),
|
||
)
|
||
.await
|
||
.expect("an owned handle must not retain the registry lock");
|
||
|
||
assert!(matches!(handle, PoolKind::Agent(_)));
|
||
assert!(state.pool_handle("other-connection").await.is_some());
|
||
drop(state);
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn redis_pool_handle_survives_registry_removal() {
|
||
let (state, dir) = test_app_state().await;
|
||
let redis = Arc::new(crate::db::redis_driver::redis_connection_test_stub());
|
||
state
|
||
.update_connection_pools(|connections| {
|
||
connections.insert("redis".to_string(), PoolKind::Redis(Arc::clone(&redis)));
|
||
})
|
||
.await;
|
||
|
||
let handle = state.pool_handle("redis").await.expect("Redis handle");
|
||
let removed = state.update_connection_pools(|connections| connections.remove("redis")).await;
|
||
|
||
assert!(matches!(handle, PoolKind::Redis(ref current) if Arc::ptr_eq(current, &redis)));
|
||
assert!(matches!(removed, Some(PoolKind::Redis(ref current)) if Arc::ptr_eq(current, &redis)));
|
||
assert!(state.pool_handle("redis").await.is_none());
|
||
drop(state);
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn shutdown_releases_connection_pools_and_agent_daemons() {
|
||
let (state, dir) = test_app_state().await;
|
||
state.connections.write().await.insert("conn".to_string(), agent_pool_stub());
|
||
state
|
||
.agent_manager
|
||
.daemons
|
||
.lock()
|
||
.await
|
||
.insert("dameng".to_string(), crate::db::agent_driver::AgentDriverClient::test_stub());
|
||
|
||
state.shutdown(Duration::from_secs(1)).await;
|
||
|
||
assert!(state.connections.read().await.is_empty());
|
||
assert!(state.agent_manager.active_daemon_keys().await.is_empty());
|
||
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn webview_reload_session_cancellation_marks_running_tasks_cancelled() {
|
||
let (state, dir) = test_app_state().await;
|
||
let registered = state.running_queries.register_task(
|
||
"exec-1".to_string(),
|
||
crate::query_cancel::RunningTaskMetadata::query("conn-1", "main", Some("tab-1".to_string())),
|
||
);
|
||
assert!(!registered.token().is_cancelled());
|
||
|
||
// The narrow renderer-reload boundary signals all running-query tasks
|
||
// (SQL execution and exports register here) without touching pools.
|
||
assert_eq!(state.cancel_webview_reload_session_tasks(), 1);
|
||
assert!(registered.token().is_cancelled());
|
||
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn agent_update_blockers_only_include_open_connections() {
|
||
let (state, dir) = test_app_state().await;
|
||
let mut config = mysql_config(None);
|
||
config.id = "dameng-conn".to_string();
|
||
config.db_type = DatabaseType::Dameng;
|
||
state.configs.write().await.insert(config.id.clone(), config);
|
||
state
|
||
.agent_manager
|
||
.daemons
|
||
.lock()
|
||
.await
|
||
.insert("oracle".to_string(), crate::db::agent_driver::AgentDriverClient::test_stub());
|
||
state.connections.write().await.insert("dameng-conn".to_string(), agent_pool_stub());
|
||
|
||
assert_eq!(
|
||
state.active_agent_connection_driver_keys().await,
|
||
std::collections::HashSet::from(["dameng".to_string()])
|
||
);
|
||
|
||
state.connections.write().await.remove("dameng-conn");
|
||
assert!(state.active_agent_connection_driver_keys().await.is_empty());
|
||
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn preparing_agent_update_stops_idle_runtime() {
|
||
let (state, dir) = test_app_state().await;
|
||
state
|
||
.agent_manager
|
||
.daemons
|
||
.lock()
|
||
.await
|
||
.insert("dameng".to_string(), crate::db::agent_driver::AgentDriverClient::test_stub());
|
||
|
||
let blockers = state.prepare_agent_driver_updates(&["dameng".to_string()]).await;
|
||
|
||
assert!(blockers.is_empty());
|
||
assert!(state.agent_manager.active_daemon_keys().await.is_empty());
|
||
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn preparing_agent_update_keeps_runtime_when_connection_is_open() {
|
||
let (state, dir) = test_app_state().await;
|
||
let mut config = mysql_config(None);
|
||
config.id = "dameng-conn".to_string();
|
||
config.db_type = DatabaseType::Dameng;
|
||
state.configs.write().await.insert(config.id.clone(), config);
|
||
state.connections.write().await.insert("dameng-conn".to_string(), agent_pool_stub());
|
||
state
|
||
.agent_manager
|
||
.daemons
|
||
.lock()
|
||
.await
|
||
.insert("dameng".to_string(), crate::db::agent_driver::AgentDriverClient::test_stub());
|
||
|
||
let blockers = state.prepare_agent_driver_updates(&["dameng".to_string()]).await;
|
||
|
||
assert_eq!(blockers, std::collections::HashSet::from(["dameng".to_string()]));
|
||
assert_eq!(state.agent_manager.active_daemon_keys().await, vec!["dameng".to_string()]);
|
||
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
fn touch_executable(path: &std::path::Path) {
|
||
if let Some(parent) = path.parent() {
|
||
std::fs::create_dir_all(parent).unwrap();
|
||
}
|
||
std::fs::write(path, b"test executable").unwrap();
|
||
#[cfg(unix)]
|
||
{
|
||
use std::os::unix::fs::PermissionsExt;
|
||
|
||
let mut permissions = std::fs::metadata(path).unwrap().permissions();
|
||
permissions.set_mode(0o755);
|
||
std::fs::set_permissions(path, permissions).unwrap();
|
||
}
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn app_state_uses_explicit_agent_dir() {
|
||
let dir = std::env::temp_dir().join(format!("dbx-core-agent-dir-test-{}", uuid::Uuid::new_v4()));
|
||
std::fs::create_dir_all(&dir).unwrap();
|
||
let storage = Storage::open(&dir.join("storage.db")).await.unwrap();
|
||
let agent_dir = dir.join("agents");
|
||
|
||
let state = AppState::new_with_plugin_and_agent_dir_and_app_version(
|
||
storage,
|
||
dir.join("plugins"),
|
||
agent_dir.clone(),
|
||
"0.0.0-test",
|
||
);
|
||
|
||
assert_eq!(state.agent_manager.base_dir(), &agent_dir);
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[test]
|
||
fn ssh_tunnel_test_includes_prerequisites_through_final_ssh_hop() {
|
||
let before = TransportLayerConfig::Proxy(ProxyTunnelConfig {
|
||
id: "before".to_string(),
|
||
name: String::new(),
|
||
enabled: true,
|
||
proxy_type: ProxyType::Socks5,
|
||
host: "proxy.internal".to_string(),
|
||
port: 1080,
|
||
username: String::new(),
|
||
password: String::new(),
|
||
test_target: None,
|
||
profile_id: String::new(),
|
||
});
|
||
let ssh = TransportLayerConfig::Ssh(ssh_layer("ssh", ""));
|
||
let after = TransportLayerConfig::Proxy(ProxyTunnelConfig {
|
||
id: "after".to_string(),
|
||
name: String::new(),
|
||
enabled: true,
|
||
proxy_type: ProxyType::Http,
|
||
host: "downstream.internal".to_string(),
|
||
port: 8080,
|
||
username: String::new(),
|
||
password: String::new(),
|
||
test_target: None,
|
||
profile_id: String::new(),
|
||
});
|
||
let layers = vec![before.clone(), ssh.clone(), after];
|
||
|
||
assert_eq!(transport_layers_through_last_ssh(&layers).unwrap(), &[before, ssh]);
|
||
}
|
||
|
||
#[test]
|
||
fn ssh_tunnel_test_rejects_chain_without_ssh() {
|
||
let layers = vec![TransportLayerConfig::Proxy(ProxyTunnelConfig {
|
||
id: "proxy".to_string(),
|
||
name: String::new(),
|
||
enabled: true,
|
||
proxy_type: ProxyType::Socks5,
|
||
host: "proxy.internal".to_string(),
|
||
port: 1080,
|
||
username: String::new(),
|
||
password: String::new(),
|
||
test_target: None,
|
||
profile_id: String::new(),
|
||
})];
|
||
|
||
assert_eq!(
|
||
transport_layers_through_last_ssh(&layers).unwrap_err(),
|
||
"Connection has no enabled SSH tunnel layer"
|
||
);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn test_tunnel_profile_rejects_non_ssh_and_missing_host() {
|
||
let (state, dir) = test_app_state().await;
|
||
|
||
// Proxy profiles now attempt a connection; with no proxy running at the
|
||
// test address the result is a connection error, not an SSH-only guard.
|
||
let test_port = portpicker::pick_unused_port().expect("no port available");
|
||
let proxy = TransportLayerConfig::Proxy(ProxyTunnelConfig {
|
||
id: "p1".to_string(),
|
||
name: String::new(),
|
||
enabled: true,
|
||
proxy_type: ProxyType::Socks5,
|
||
host: "127.0.0.1".to_string(),
|
||
port: test_port,
|
||
username: String::new(),
|
||
password: String::new(),
|
||
test_target: None,
|
||
profile_id: String::new(),
|
||
});
|
||
let err = state.test_tunnel_profile(&proxy).await.unwrap_err();
|
||
assert!(!err.contains("SSH"), "proxy test should not return SSH error, got: {err}");
|
||
|
||
// An SSH profile with no host fails fast rather than dialing an empty host.
|
||
let ssh = TransportLayerConfig::Ssh(SshTunnelConfig {
|
||
id: "s1".to_string(),
|
||
name: String::new(),
|
||
enabled: true,
|
||
host: String::new(),
|
||
port: 22,
|
||
user: "root".to_string(),
|
||
password: String::new(),
|
||
key_path: String::new(),
|
||
key_passphrase: String::new(),
|
||
connect_timeout_secs: 5,
|
||
expose_lan: false,
|
||
use_ssh_agent: false,
|
||
ssh_agent_sock_path: String::new(),
|
||
auth_method: "password".to_string(),
|
||
allow_exec_channel_proxy: false,
|
||
profile_id: String::new(),
|
||
});
|
||
assert!(state.test_tunnel_profile(&ssh).await.is_err());
|
||
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn stale_connection_attempt_cannot_replace_newer_pool() {
|
||
let (state, dir) = test_app_state().await;
|
||
let mut config = mysql_config(None);
|
||
config.name = "SQLite".to_string();
|
||
config.db_type = DatabaseType::Sqlite;
|
||
config.host = dir.join("current.db").to_string_lossy().to_string();
|
||
let old_attempt = state.begin_connection_attempt("conn").await;
|
||
let new_attempt = state.begin_connection_attempt("conn").await;
|
||
let current_pool =
|
||
db::sqlite::connect_path_create_if_missing(&dir.join("current.db").to_string_lossy()).await.unwrap();
|
||
let stale_pool =
|
||
db::sqlite::connect_path_create_if_missing(&dir.join("stale.db").to_string_lossy()).await.unwrap();
|
||
|
||
state
|
||
.insert_connection_pool_for_attempt(
|
||
"conn",
|
||
new_attempt,
|
||
"conn".to_string(),
|
||
PoolKind::Sqlite(current_pool),
|
||
&config,
|
||
)
|
||
.await
|
||
.unwrap();
|
||
|
||
let result = state
|
||
.insert_connection_pool_for_attempt(
|
||
"conn",
|
||
old_attempt,
|
||
"conn".to_string(),
|
||
PoolKind::Sqlite(stale_pool),
|
||
&config,
|
||
)
|
||
.await;
|
||
|
||
assert!(result.is_err());
|
||
let conns = state.connections.read().await;
|
||
assert!(matches!(conns.get("conn"), Some(PoolKind::Sqlite(_))));
|
||
assert_eq!(conns.len(), 1);
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
#[ignore = "requires DBX_TEST_MYSQL_URL"]
|
||
async fn live_mysql_health_check_keeps_saturated_pool() {
|
||
let url = std::env::var("DBX_TEST_MYSQL_URL").expect("DBX_TEST_MYSQL_URL is required");
|
||
let (state, dir) = test_app_state().await;
|
||
let config = mysql_config(Some("testdb"));
|
||
let pool = db::mysql::connect_bare_with_pool_limit(&url, Duration::from_secs(5), 1).await.unwrap();
|
||
state
|
||
.insert_connection_pool("conn".to_string(), PoolKind::Mysql(pool.clone(), MysqlMode::Normal), &config)
|
||
.await
|
||
.unwrap();
|
||
let held_connection = pool.get_conn().await.unwrap();
|
||
|
||
let started = Instant::now();
|
||
state.check_connection_health("conn").await.unwrap();
|
||
|
||
assert!(started.elapsed() < Duration::from_secs(2));
|
||
assert!(state.connections.read().await.contains_key("conn"));
|
||
drop(held_connection);
|
||
state.remove_connection_pools_detached("conn").await;
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn mysql_health_check_keeps_pool_when_connection_creation_exceeds_probe_budget() {
|
||
let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0)).await.unwrap();
|
||
let address = listener.local_addr().unwrap();
|
||
let server = tokio::spawn(async move {
|
||
let (_socket, _) = listener.accept().await.unwrap();
|
||
std::future::pending::<()>().await;
|
||
});
|
||
let pool_options =
|
||
mysql_async::PoolOpts::new().with_constraints(mysql_async::PoolConstraints::new(1, 2).unwrap());
|
||
let options = mysql_async::OptsBuilder::default()
|
||
.ip_or_hostname(address.ip().to_string())
|
||
.tcp_port(address.port())
|
||
.user(Some("fault-injection"))
|
||
.pass(Some("fault-injection"))
|
||
.pool_opts(Some(pool_options));
|
||
let pool = db::mysql::MySqlPool::new(options, 2);
|
||
let (state, dir) = test_app_state().await;
|
||
state.connections.write().await.insert("conn".to_string(), PoolKind::Mysql(pool.clone(), MysqlMode::Normal));
|
||
|
||
let started = Instant::now();
|
||
assert!(!state.remove_stale_connection_pool("conn").await);
|
||
|
||
assert!(started.elapsed() >= super::HEALTH_CHECK_POOL_ACQUIRE_TIMEOUT);
|
||
assert!(matches!(
|
||
state.connections.read().await.get("conn"),
|
||
Some(PoolKind::Mysql(current, _)) if pool.is_same_pool(current)
|
||
));
|
||
state.connections.write().await.remove("conn");
|
||
server.abort();
|
||
let _ = tokio::time::timeout(Duration::from_secs(1), pool.disconnect()).await;
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn jdbc_plugin_env_uses_managed_jre_when_installed() {
|
||
let dir = std::env::temp_dir().join(format!("dbx-core-jdbc-managed-jre-{}", uuid::Uuid::new_v4()));
|
||
std::fs::create_dir_all(&dir).unwrap();
|
||
let storage = Storage::open(&dir.join("storage.db")).await.unwrap();
|
||
let state = AppState::new_with_plugin_and_agent_dir_and_app_version(
|
||
storage,
|
||
dir.join("plugins"),
|
||
dir.join("agents"),
|
||
"0.0.0-test",
|
||
);
|
||
let java = state.agent_manager.jre_java_path(DEFAULT_JRE_KEY);
|
||
touch_executable(&java);
|
||
|
||
let env = state.external_driver_runtime_env("jdbc").unwrap();
|
||
|
||
assert_eq!(env.get("DBX_JAVA_BIN"), Some(java.to_string_lossy().as_ref()));
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn jdbc_plugin_env_keeps_wrapper_fallback_when_managed_jre_is_missing() {
|
||
let dir = std::env::temp_dir().join(format!("dbx-core-jdbc-missing-jre-{}", uuid::Uuid::new_v4()));
|
||
std::fs::create_dir_all(&dir).unwrap();
|
||
let storage = Storage::open(&dir.join("storage.db")).await.unwrap();
|
||
let state = AppState::new_with_plugin_and_agent_dir_and_app_version(
|
||
storage,
|
||
dir.join("plugins"),
|
||
dir.join("agents"),
|
||
"0.0.0-test",
|
||
);
|
||
|
||
let env = state.external_driver_runtime_env("jdbc").unwrap();
|
||
|
||
assert_eq!(env.get("DBX_JAVA_BIN"), None);
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn jdbc_plugin_env_uses_custom_java_runtime() {
|
||
let dir = std::env::temp_dir().join(format!("dbx-core-jdbc-custom-jre-{}", uuid::Uuid::new_v4()));
|
||
std::fs::create_dir_all(&dir).unwrap();
|
||
let storage = Storage::open(&dir.join("storage.db")).await.unwrap();
|
||
let state = AppState::new_with_plugin_and_agent_dir_and_app_version(
|
||
storage,
|
||
dir.join("plugins"),
|
||
dir.join("agents"),
|
||
"0.0.0-test",
|
||
);
|
||
let java = dir.join("custom").join("bin").join(if cfg!(windows) { "java.exe" } else { "java" });
|
||
touch_executable(&java);
|
||
state
|
||
.agent_manager
|
||
.save_state(&AgentState {
|
||
java_runtime: JavaRuntimeConfig {
|
||
mode: JavaRuntimeMode::Custom,
|
||
custom_java_path: Some(java.to_string_lossy().to_string()),
|
||
},
|
||
..AgentState::default()
|
||
})
|
||
.unwrap();
|
||
|
||
let env = state.external_driver_runtime_env("jdbc").unwrap();
|
||
|
||
assert_eq!(env.get("DBX_JAVA_BIN"), Some(java.to_string_lossy().as_ref()));
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
fn live_postgres_like_config(
|
||
db_type: DatabaseType,
|
||
host: &str,
|
||
port: u16,
|
||
username: &str,
|
||
password: &str,
|
||
url_params: Option<&str>,
|
||
) -> ConnectionConfig {
|
||
let mut config = mysql_config(Some("postgres"));
|
||
config.db_type = db_type;
|
||
config.host = host.to_string();
|
||
config.port = port;
|
||
config.username = username.to_string();
|
||
config.password = password.to_string();
|
||
config.url_params = url_params.map(str::to_string);
|
||
config
|
||
}
|
||
|
||
async fn assert_live_postgres_like_query(config: ConnectionConfig) {
|
||
let url = connection_url_for_endpoint(&config, &config.host, config.port);
|
||
let pool = db::postgres::connect(&url, std::time::Duration::from_secs(config.effective_connect_timeout_secs()))
|
||
.await
|
||
.unwrap_or_else(|err| {
|
||
panic!("failed to connect to {:?} at {}:{}: {}", config.db_type, config.host, config.port, err)
|
||
});
|
||
let result =
|
||
db::postgres::execute_query(&pool, "SELECT current_database(), current_schema()").await.unwrap_or_else(
|
||
|err| panic!("failed to query {:?} at {}:{}: {}", config.db_type, config.host, config.port, err),
|
||
);
|
||
assert_eq!(result.rows.len(), 1);
|
||
pool.close();
|
||
}
|
||
|
||
#[test]
|
||
fn mysql_metadata_connection_ignores_saved_default_database() {
|
||
let config = mysql_config(Some("app"));
|
||
|
||
let metadata = metadata_connection_config(&config);
|
||
|
||
assert_eq!(metadata.database, None);
|
||
assert_eq!(metadata.db_type, DatabaseType::Mysql);
|
||
}
|
||
|
||
#[test]
|
||
fn mysql_metadata_fallback_uses_saved_default_database() {
|
||
let config = mysql_config(Some("app"));
|
||
let metadata = metadata_connection_config(&config);
|
||
|
||
assert_eq!(
|
||
mysql_metadata_fallback_url(&config, &metadata, &config.host, config.port),
|
||
Some("mysql://root:secret@127.0.0.1:3306/app?ssl-mode=disabled&charset=utf8mb4".to_string())
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn mysql_metadata_fallback_is_unavailable_without_default_database() {
|
||
let config = mysql_config(None);
|
||
let metadata = metadata_connection_config(&config);
|
||
|
||
assert_eq!(mysql_metadata_fallback_url(&config, &metadata, &config.host, config.port), None);
|
||
}
|
||
|
||
#[test]
|
||
fn mysql_database_connection_keeps_requested_database() {
|
||
let config = mysql_config(Some("app"));
|
||
|
||
let scoped = database_connection_config(&config, Some("analytics"));
|
||
|
||
assert_eq!(scoped.database.as_deref(), Some("analytics"));
|
||
}
|
||
|
||
#[test]
|
||
fn gaussdb_database_connection_keeps_requested_database() {
|
||
let mut config = mysql_config(Some("postgres"));
|
||
config.db_type = DatabaseType::Gaussdb;
|
||
|
||
let scoped = database_connection_config(&config, Some("analytics"));
|
||
|
||
assert_eq!(scoped.database.as_deref(), Some("analytics"));
|
||
}
|
||
|
||
#[test]
|
||
fn gaussdb_endpoint_url_uses_postgres_scheme_for_native_driver() {
|
||
let mut config = mysql_config(Some("postgres"));
|
||
config.db_type = DatabaseType::Gaussdb;
|
||
config.username = "gaussdb".to_string();
|
||
config.password = "secret".to_string();
|
||
|
||
assert_eq!(
|
||
connection_url_for_endpoint(&config, &config.host, config.port),
|
||
"postgres://gaussdb:secret@127.0.0.1:3306/postgres?sslmode=prefer"
|
||
);
|
||
assert_eq!(
|
||
redacted_connection_url_for_endpoint(&config, &config.host, config.port),
|
||
"postgres://127.0.0.1:3306/postgres?sslmode=prefer"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn gaussdb_probe_endpoints_parse_legacy_and_ipv6_hosts() {
|
||
assert_eq!(connection_probe_endpoints("db.example.com:5433", 5432), vec![("db.example.com".to_string(), 5433)]);
|
||
assert_eq!(
|
||
connection_probe_endpoints("[2001:db8::1]:5433,[2001:db8::2]:5434", 5432),
|
||
vec![("2001:db8::1".to_string(), 5433), ("2001:db8::2".to_string(), 5434)]
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn kwdb_endpoint_url_uses_postgres_scheme_for_native_driver() {
|
||
let mut config = mysql_config(None);
|
||
config.db_type = DatabaseType::Kwdb;
|
||
config.username = "root".to_string();
|
||
config.password = "secret".to_string();
|
||
config.port = 26257;
|
||
|
||
assert_eq!(
|
||
connection_url_for_endpoint(&config, &config.host, config.port),
|
||
"postgres://root:secret@127.0.0.1:26257/defaultdb?sslmode=prefer"
|
||
);
|
||
assert_eq!(
|
||
redacted_connection_url_for_endpoint(&config, &config.host, config.port),
|
||
"postgres://127.0.0.1:26257/defaultdb?sslmode=prefer"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn kwdb_endpoint_url_keeps_explicit_sslmode() {
|
||
let mut config = mysql_config(None);
|
||
config.db_type = DatabaseType::Kwdb;
|
||
config.username = "root".to_string();
|
||
config.password = "secret".to_string();
|
||
config.port = 26257;
|
||
config.url_params = Some("sslmode=require&application_name=dbx".to_string());
|
||
|
||
assert_eq!(
|
||
connection_url_for_endpoint(&config, &config.host, config.port),
|
||
"postgres://root:secret@127.0.0.1:26257/defaultdb?sslmode=require&application_name=dbx"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn opengauss_endpoint_url_uses_postgres_scheme_for_native_driver() {
|
||
let mut config = mysql_config(Some("postgres"));
|
||
config.db_type = DatabaseType::OpenGauss;
|
||
config.username = "gaussdb".to_string();
|
||
config.password = "secret".to_string();
|
||
|
||
assert_eq!(
|
||
connection_url_for_endpoint(&config, &config.host, config.port),
|
||
"postgres://gaussdb:secret@127.0.0.1:3306/postgres?sslmode=prefer"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn gaussdb_endpoint_url_keeps_explicit_sslmode() {
|
||
let mut config = mysql_config(Some("postgres"));
|
||
config.db_type = DatabaseType::Gaussdb;
|
||
config.username = "gaussdb".to_string();
|
||
config.password = "secret".to_string();
|
||
config.url_params = Some("sslmode=require&application_name=dbx".to_string());
|
||
|
||
assert_eq!(
|
||
connection_url_for_endpoint(&config, &config.host, config.port),
|
||
"postgres://gaussdb:secret@127.0.0.1:3306/postgres?sslmode=require&application_name=dbx"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn gaussdb_endpoint_url_uses_require_sslmode_when_tls_enabled() {
|
||
let mut config = mysql_config(Some("postgres"));
|
||
config.db_type = DatabaseType::Gaussdb;
|
||
config.username = "gaussdb".to_string();
|
||
config.password = "secret".to_string();
|
||
config.ssl = true;
|
||
|
||
assert_eq!(
|
||
connection_url_for_endpoint(&config, &config.host, config.port),
|
||
"postgres://gaussdb:secret@127.0.0.1:3306/postgres?sslmode=require"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn gaussdb_endpoint_url_prepends_default_sslmode_to_custom_params() {
|
||
let mut config = mysql_config(Some("postgres"));
|
||
config.db_type = DatabaseType::Gaussdb;
|
||
config.username = "gaussdb".to_string();
|
||
config.password = "secret".to_string();
|
||
config.url_params = Some("application_name=dbx".to_string());
|
||
|
||
assert_eq!(
|
||
connection_url_for_endpoint(&config, &config.host, config.port),
|
||
"postgres://gaussdb:secret@127.0.0.1:3306/postgres?sslmode=prefer&application_name=dbx"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn postgres_url_validation_rejects_invalid_stringtype_before_connect() {
|
||
let mut config = mysql_config(Some("postgres"));
|
||
config.db_type = DatabaseType::Postgres;
|
||
config.url_params = Some("currentSchema=public&stringtype=text".to_string());
|
||
|
||
assert_eq!(
|
||
validate_connection_url_params(&config).unwrap_err(),
|
||
"Unsupported value for PostgreSQL stringtype parameter: text. Expected 'unspecified' or 'varchar'."
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn mongodb_database_connection_keeps_saved_database_for_auth() {
|
||
let mut config = mysql_config(Some("admin"));
|
||
config.db_type = DatabaseType::MongoDb;
|
||
|
||
let scoped = database_connection_config(&config, Some("shop"));
|
||
|
||
assert_eq!(scoped.database.as_deref(), Some("admin"));
|
||
}
|
||
|
||
#[test]
|
||
fn cloudflare_d1_query_namespace_does_not_replace_database_id() {
|
||
let mut config = mysql_config(Some("database-uuid"));
|
||
config.db_type = DatabaseType::CloudflareD1;
|
||
|
||
let scoped = database_connection_config(&config, Some("main"));
|
||
|
||
assert_eq!(scoped.database.as_deref(), Some("database-uuid"));
|
||
assert_eq!(
|
||
super::base_pool_key_for(Some(DatabaseType::CloudflareD1), "d1-conn", Some("main"), false),
|
||
"d1-conn"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn oracle_database_connection_ignores_requested_database() {
|
||
let mut config = mysql_config(Some("ORCL"));
|
||
config.db_type = DatabaseType::Oracle;
|
||
|
||
let scoped = database_connection_config(&config, Some("analytics"));
|
||
|
||
assert_eq!(scoped.database.as_deref(), Some("ORCL"));
|
||
}
|
||
|
||
#[test]
|
||
fn oracle_reuses_connection_scoped_pool_for_schema_database_keys() {
|
||
assert_eq!(
|
||
super::base_pool_key_for(Some(DatabaseType::Oracle), "oracle-conn", Some("ORCLPDB1"), false),
|
||
"oracle-conn"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn oracle_uses_isolated_pool_keys_for_tab_sessions() {
|
||
let mut config = mysql_config(Some("ORCLPDB1"));
|
||
config.db_type = DatabaseType::Oracle;
|
||
|
||
assert_eq!(
|
||
super::session_scoped_pool_key_for(Some(&config), "oracle-conn".to_string(), Some("table-tab-1")),
|
||
"oracle-conn:session:table-tab-1"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn other_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::MongoDb), "mongo-conn", Some("shop"), false),
|
||
"mongo-conn:shop"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn database_scoped_pool_keys_preserve_identifier_whitespace() {
|
||
assert_eq!(
|
||
super::base_pool_key_for(Some(DatabaseType::Mysql), "mysql-conn", Some(" analytics"), false),
|
||
"mysql-conn: analytics"
|
||
);
|
||
assert_eq!(
|
||
super::base_pool_key_for(Some(DatabaseType::Mysql), "mysql-conn", Some("analytics"), false),
|
||
"mysql-conn:analytics"
|
||
);
|
||
assert_eq!(
|
||
super::base_pool_key_for(Some(DatabaseType::Postgres), "pg-conn", Some("analytics "), false),
|
||
"pg-conn:analytics "
|
||
);
|
||
}
|
||
|
||
#[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 session_scoped_pool_keys_are_sanitized_and_detected() {
|
||
let mysql = mysql_config(Some("analytics"));
|
||
let key =
|
||
super::session_scoped_pool_key_for(Some(&mysql), "mysql-conn:analytics".to_string(), Some("tab-1:count"));
|
||
|
||
assert_eq!(key, "mysql-conn:analytics:session:tab-1_count");
|
||
assert!(super::is_session_scoped_pool_key(&key));
|
||
assert!(!super::is_session_scoped_pool_key("mysql-conn:analytics"));
|
||
|
||
let mut duckdb = mysql_config(None);
|
||
duckdb.db_type = DatabaseType::DuckDb;
|
||
assert_eq!(
|
||
super::session_scoped_pool_key_for(Some(&duckdb), "duckdb-conn".to_string(), Some("tab-1")),
|
||
"duckdb-conn"
|
||
);
|
||
let mut sqlite_memory = mysql_config(None);
|
||
sqlite_memory.db_type = DatabaseType::Sqlite;
|
||
sqlite_memory.host = " :MeMoRy: ".to_string();
|
||
assert_eq!(
|
||
super::session_scoped_pool_key_for(Some(&sqlite_memory), "sqlite-memory".to_string(), Some("tab-1")),
|
||
"sqlite-memory"
|
||
);
|
||
|
||
sqlite_memory.host = "/tmp/dbx-session-test.sqlite".to_string();
|
||
assert_eq!(
|
||
super::session_scoped_pool_key_for(Some(&sqlite_memory), "sqlite-file".to_string(), Some("tab-1")),
|
||
"sqlite-file:session:tab-1"
|
||
);
|
||
|
||
let mut cloudflare_d1 = mysql_config(None);
|
||
cloudflare_d1.db_type = DatabaseType::CloudflareD1;
|
||
assert_eq!(
|
||
super::session_scoped_pool_key_for(Some(&cloudflare_d1), "d1-conn".to_string(), Some("tab-1")),
|
||
"d1-conn"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn agent_metadata_pool_keys_are_isolated_from_workload_keys() {
|
||
let mut config = mysql_config(Some("analytics"));
|
||
config.db_type = DatabaseType::Dameng;
|
||
|
||
let workload = super::pool_key_for_session_role(
|
||
Some(&config),
|
||
"conn:analytics".to_string(),
|
||
Some("task:1"),
|
||
crate::agent_connection::AgentSessionRole::Workload,
|
||
);
|
||
let metadata = super::pool_key_for_session_role(
|
||
Some(&config),
|
||
"conn:analytics".to_string(),
|
||
Some("task:1"),
|
||
crate::agent_connection::AgentSessionRole::Metadata,
|
||
);
|
||
let base_metadata = super::pool_key_for_session_role(
|
||
Some(&config),
|
||
"conn:analytics".to_string(),
|
||
None,
|
||
crate::agent_connection::AgentSessionRole::Metadata,
|
||
);
|
||
|
||
assert_eq!(workload, "conn:analytics:session:task_1");
|
||
assert_eq!(metadata, "conn:analytics:session:task_1:role:metadata");
|
||
assert_eq!(base_metadata, "conn:analytics:role:metadata");
|
||
assert_ne!(metadata, workload);
|
||
}
|
||
|
||
#[test]
|
||
fn non_agent_metadata_role_preserves_existing_pool_keys() {
|
||
let config = mysql_config(Some("analytics"));
|
||
|
||
let workload = super::pool_key_for_session_role(
|
||
Some(&config),
|
||
"conn:analytics".to_string(),
|
||
Some("task:1"),
|
||
crate::agent_connection::AgentSessionRole::Workload,
|
||
);
|
||
let metadata = super::pool_key_for_session_role(
|
||
Some(&config),
|
||
"conn:analytics".to_string(),
|
||
Some("task:1"),
|
||
crate::agent_connection::AgentSessionRole::Metadata,
|
||
);
|
||
|
||
assert_eq!(metadata, workload);
|
||
assert_eq!(metadata, "conn:analytics: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 };
|
||
|
||
assert_eq!(redis_sentinel_transport_prefix("redis-prod"), "redis-prod:redis-sentinel:");
|
||
assert_eq!(
|
||
redis_sentinel_transport_id("redis-prod", "master", &endpoint),
|
||
"redis-prod:redis-sentinel:master:10.0.0.8:6379"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn mysql_pool_size_keeps_session_pools_single_connection() {
|
||
assert_eq!(super::mysql_pool_max_connections_for_session(None), 10);
|
||
assert_eq!(super::mysql_pool_max_connections_for_session(Some("")), 10);
|
||
assert_eq!(super::mysql_pool_max_connections_for_session(Some("tab-1")), 1);
|
||
}
|
||
|
||
#[test]
|
||
fn postgres_pool_size_keeps_session_pools_single_connection() {
|
||
assert_eq!(super::postgres_pool_max_connections_for_session(None), 10);
|
||
assert_eq!(super::postgres_pool_max_connections_for_session(Some("")), 10);
|
||
assert_eq!(super::postgres_pool_max_connections_for_session(Some("mcp:batch-1")), 1);
|
||
}
|
||
|
||
#[test]
|
||
fn stale_mysql_pool_observation_does_not_remove_replacement_generation() {
|
||
let checked = crate::db::mysql::MySqlPool::new("mysql://root@127.0.0.1:3306/app", 10);
|
||
let checked_clone = checked.clone();
|
||
let replacement = crate::db::mysql::MySqlPool::new("mysql://root@127.0.0.1:3306/app", 10);
|
||
assert!(checked.is_same_pool(&checked_clone));
|
||
assert!(!checked.is_same_pool(&replacement));
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn stale_postgres_cleanup_preserves_concurrent_replacement_publication() {
|
||
let (state, dir) = test_app_state().await;
|
||
let state = std::sync::Arc::new(state);
|
||
let pool_key = "conn:app";
|
||
let postgres_pool = || {
|
||
let manager = deadpool_postgres::Manager::new(tokio_postgres::Config::new(), tokio_postgres::NoTls);
|
||
deadpool_postgres::Pool::builder(manager)
|
||
.runtime(deadpool_postgres::Runtime::Tokio1)
|
||
.build()
|
||
.expect("build PostgreSQL test pool")
|
||
};
|
||
let checked = postgres_pool();
|
||
let replacement = postgres_pool();
|
||
state.connections.write().await.insert(pool_key.to_string(), PoolKind::Postgres(checked));
|
||
state.pool_activity.write().await.insert(pool_key.to_string(), super::PoolActivity::now());
|
||
let checked_publication = state.pool_publication_snapshot(pool_key).await.unwrap().publication;
|
||
|
||
let cleanup_ready = std::sync::Arc::new(tokio::sync::Barrier::new(2));
|
||
let continue_cleanup = std::sync::Arc::new(tokio::sync::Barrier::new(2));
|
||
let cleanup_state = state.clone();
|
||
let cleanup_publication = checked_publication.clone();
|
||
let cleanup_ready_for_task = cleanup_ready.clone();
|
||
let cleanup_continue = continue_cleanup.clone();
|
||
let cleanup = tokio::spawn(async move {
|
||
cleanup_state
|
||
.remove_stale_pool_if_current_inner(
|
||
pool_key,
|
||
&cleanup_publication,
|
||
Some(super::StalePoolCleanupBarriers {
|
||
before_removal: Some((cleanup_ready_for_task, cleanup_continue)),
|
||
after_removal: None,
|
||
}),
|
||
)
|
||
.await
|
||
});
|
||
cleanup_ready.wait().await;
|
||
|
||
state.connections.write().await.insert(pool_key.to_string(), PoolKind::Postgres(replacement));
|
||
let replacement_publication = state.pool_publication_snapshot(pool_key).await.unwrap().publication;
|
||
assert!(!checked_publication.is_same(&replacement_publication));
|
||
continue_cleanup.wait().await;
|
||
assert!(!cleanup.await.unwrap());
|
||
|
||
let current = state.pool_publication_snapshot(pool_key).await.expect("replacement must remain routable");
|
||
assert!(matches!(current.pool, PoolKind::Postgres(_)));
|
||
assert!(current.publication.is_same(&replacement_publication));
|
||
assert!(state.pool_activity.read().await.contains_key(pool_key));
|
||
|
||
state.shutdown(Duration::from_secs(1)).await;
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn stale_mysql_cleanup_preserves_replacement_support_state() {
|
||
let (state, dir) = test_app_state().await;
|
||
let state = std::sync::Arc::new(state);
|
||
let pool_key = "conn:app";
|
||
let checked = crate::db::mysql::MySqlPool::new("mysql://root@127.0.0.1:3306/app", 10);
|
||
let replacement = crate::db::mysql::MySqlPool::new("mysql://root@127.0.0.1:3306/app", 10);
|
||
let mut config = mysql_config(Some("app"));
|
||
config.keepalive_interval_secs = 60;
|
||
state
|
||
.connections
|
||
.write()
|
||
.await
|
||
.insert(pool_key.to_string(), PoolKind::Mysql(checked.clone(), MysqlMode::Normal));
|
||
state.pool_activity.write().await.insert(pool_key.to_string(), super::PoolActivity::now());
|
||
let checked_publication = state.pool_publication_snapshot(pool_key).await.unwrap().publication;
|
||
state.start_keepalive_task(
|
||
pool_key,
|
||
&PoolKind::Mysql(checked, MysqlMode::Normal),
|
||
&config,
|
||
#[cfg(feature = "mq-admin")]
|
||
None,
|
||
);
|
||
assert_eq!(state.supervised_task_count(), 1);
|
||
|
||
let route_removed = std::sync::Arc::new(tokio::sync::Barrier::new(2));
|
||
let continue_cleanup = std::sync::Arc::new(tokio::sync::Barrier::new(2));
|
||
let cleanup_state = state.clone();
|
||
let cleanup_publication = checked_publication.clone();
|
||
let cleanup_route_removed = route_removed.clone();
|
||
let cleanup_continue = continue_cleanup.clone();
|
||
let cleanup = tokio::spawn(async move {
|
||
cleanup_state
|
||
.remove_stale_pool_if_current_inner(
|
||
pool_key,
|
||
&cleanup_publication,
|
||
Some(super::StalePoolCleanupBarriers {
|
||
before_removal: None,
|
||
after_removal: Some((cleanup_route_removed, cleanup_continue)),
|
||
}),
|
||
)
|
||
.await
|
||
});
|
||
route_removed.wait().await;
|
||
|
||
let publish_state = state.clone();
|
||
let publish_config = config.clone();
|
||
let publish_replacement = replacement.clone();
|
||
let mut publish = tokio::spawn(async move {
|
||
publish_state
|
||
.insert_connection_pool(
|
||
pool_key.to_string(),
|
||
PoolKind::Mysql(publish_replacement, MysqlMode::Normal),
|
||
&publish_config,
|
||
)
|
||
.await
|
||
});
|
||
|
||
assert!(
|
||
tokio::time::timeout(Duration::from_millis(50), &mut publish).await.is_err(),
|
||
"replacement publication must wait while stale support state is cleaned"
|
||
);
|
||
continue_cleanup.wait().await;
|
||
assert!(cleanup.await.unwrap());
|
||
publish.await.unwrap().unwrap();
|
||
|
||
let current = state.pool_publication_snapshot(pool_key).await.expect("replacement must remain routable");
|
||
assert!(matches!(current.pool, PoolKind::Mysql(ref pool, _) if replacement.is_same_pool(pool)));
|
||
assert!(!current.publication.is_same(&checked_publication));
|
||
assert!(state.pool_activity.read().await.contains_key(pool_key));
|
||
assert_eq!(state.supervised_task_count(), 1);
|
||
|
||
state.shutdown(Duration::from_secs(1)).await;
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[test]
|
||
fn metadata_concurrency_reserves_query_capacity() {
|
||
assert_eq!(super::metadata_concurrency_limit(DatabaseType::Mysql, 10), 6);
|
||
assert_eq!(super::metadata_concurrency_limit(DatabaseType::Postgres, 10), 6);
|
||
assert_eq!(super::metadata_concurrency_limit(DatabaseType::Mysql, 3), 1);
|
||
assert_eq!(super::metadata_concurrency_limit(DatabaseType::Postgres, 1), 1);
|
||
assert_eq!(super::metadata_concurrency_limit(DatabaseType::SqlServer, 10), 1);
|
||
}
|
||
|
||
#[test]
|
||
fn session_metadata_gate_allowance_matches_session_pool_capacity() {
|
||
// PostgreSQL session pools are full-size (10 connections, see
|
||
// db/postgres.rs Pool::builder().max_size(10)), so the gate must not
|
||
// serialize the database-export prefetch down to one permit.
|
||
assert_eq!(super::metadata_gate_session_allowance(DatabaseType::Postgres), 10);
|
||
// MySQL and other session-scoped pools are single-connection.
|
||
assert_eq!(super::metadata_gate_session_allowance(DatabaseType::Mysql), 1);
|
||
// SQL Server gate is intentionally shared/serialized regardless.
|
||
assert_eq!(super::metadata_gate_session_allowance(DatabaseType::SqlServer), 1);
|
||
}
|
||
|
||
#[test]
|
||
fn metadata_gate_is_limited_to_supported_pool_drivers() {
|
||
assert!(super::uses_metadata_gate(DatabaseType::Mysql));
|
||
assert!(super::uses_metadata_gate(DatabaseType::Postgres));
|
||
assert!(super::uses_metadata_gate(DatabaseType::SqlServer));
|
||
assert!(!super::uses_metadata_gate(DatabaseType::Oracle));
|
||
assert!(!super::uses_metadata_gate(DatabaseType::Sqlite));
|
||
assert!(!super::uses_metadata_gate(DatabaseType::MongoDb));
|
||
}
|
||
|
||
#[test]
|
||
fn sqlserver_metadata_gate_is_shared_across_client_sessions() {
|
||
assert_eq!(
|
||
super::metadata_gate_key("conn", Some("app"), DatabaseType::SqlServer, Some("metadata-1")),
|
||
super::metadata_gate_key("conn", Some("app"), DatabaseType::SqlServer, Some("metadata-2"))
|
||
);
|
||
assert_ne!(
|
||
super::metadata_gate_key("conn", Some("app"), DatabaseType::Mysql, Some("tab-1")),
|
||
super::metadata_gate_key("conn", Some("app"), DatabaseType::Mysql, Some("tab-2"))
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn prestosql_uses_external_driver_pool_not_agent_pool() {
|
||
assert!(!super::uses_agent_connection_pool(&DatabaseType::PrestoSql));
|
||
assert!(super::uses_agent_connection_pool(&DatabaseType::Trino));
|
||
}
|
||
|
||
#[test]
|
||
fn mysql_hostname_connections_skip_tcp_probe() {
|
||
let mut config = mysql_config(Some("app"));
|
||
config.host = "mysql.example.com".to_string();
|
||
|
||
assert!(!uses_tcp_probe(&config, "mysql.example.com", 3306));
|
||
config.host = "192.0.2.10".to_string();
|
||
assert!(!uses_tcp_probe(&config, "192.0.2.10", 3306));
|
||
assert!(uses_tcp_probe(&config, "127.0.0.1", 53306));
|
||
}
|
||
|
||
#[test]
|
||
fn native_direct_connections_skip_tcp_probe() {
|
||
for db_type in [
|
||
DatabaseType::Postgres,
|
||
DatabaseType::Redshift,
|
||
DatabaseType::Redis,
|
||
DatabaseType::ClickHouse,
|
||
DatabaseType::SqlServer,
|
||
DatabaseType::Elasticsearch,
|
||
DatabaseType::Easysearch,
|
||
DatabaseType::Kwdb,
|
||
] {
|
||
let mut config = mysql_config(Some("app"));
|
||
config.db_type = db_type;
|
||
config.host = "db.example.com".to_string();
|
||
|
||
assert!(!uses_tcp_probe(&config, "db.example.com", config.port), "{db_type:?} hostname");
|
||
config.host = "192.0.2.10".to_string();
|
||
assert!(!uses_tcp_probe(&config, "192.0.2.10", config.port), "{db_type:?} ip");
|
||
assert!(uses_tcp_probe(&config, "127.0.0.1", 54000), "{db_type:?} forwarded");
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn h2_agent_connections_skip_tcp_probe_for_file_and_tcp_modes() {
|
||
let mut file_config = mysql_config(None);
|
||
file_config.db_type = DatabaseType::H2;
|
||
file_config.host = "/tmp/app.mv.db".to_string();
|
||
file_config.port = 0;
|
||
|
||
assert!(!uses_tcp_probe(&file_config, "/tmp/app.mv.db", 0));
|
||
|
||
let mut tcp_config = mysql_config(Some("test"));
|
||
tcp_config.db_type = DatabaseType::H2;
|
||
tcp_config.host = "127.0.0.1".to_string();
|
||
tcp_config.port = 9092;
|
||
|
||
assert!(!uses_tcp_probe(&tcp_config, "127.0.0.1", 9092));
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn sqlite_get_or_create_pool_initializes_connection_for_web_route() {
|
||
let (state, dir) = test_app_state().await;
|
||
let db_path = dir.join("app.db");
|
||
std::fs::File::create(&db_path).unwrap();
|
||
let mut config = mysql_config(None);
|
||
config.id = "sqlite-conn".to_string();
|
||
config.name = "SQLite".to_string();
|
||
config.db_type = DatabaseType::Sqlite;
|
||
config.host = db_path.to_string_lossy().to_string();
|
||
config.port = 0;
|
||
// This exercises a plain SQLite file. The shared MySQL fixture carries
|
||
// credentials, and a SQLite password intentionally opts into SQLCipher.
|
||
config.username.clear();
|
||
config.password.clear();
|
||
|
||
state.configs.write().await.insert(config.id.clone(), config);
|
||
|
||
let pool_key = state.get_or_create_pool("sqlite-conn", None).await.unwrap();
|
||
assert_eq!(pool_key, "sqlite-conn");
|
||
|
||
let databases = schema::list_databases_core(&state, "sqlite-conn").await.unwrap();
|
||
assert_eq!(databases.len(), 1);
|
||
assert_eq!(databases[0].name, "main");
|
||
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn sqlite_connection_restores_attached_databases() {
|
||
let (state, dir) = test_app_state().await;
|
||
let main_path = dir.join("main.sqlite");
|
||
let attached_path = dir.join("analytics.sqlite");
|
||
drop(db::sqlite::connect_path_create_if_missing(main_path.to_str().unwrap()).await.unwrap());
|
||
let attached = db::sqlite::connect_path_create_if_missing(attached_path.to_str().unwrap()).await.unwrap();
|
||
db::sqlite::execute_query(&attached, "CREATE TABLE events(id INTEGER PRIMARY KEY);").await.unwrap();
|
||
drop(attached);
|
||
|
||
let mut config = mysql_config(None);
|
||
config.id = "sqlite-conn".to_string();
|
||
config.name = "SQLite".to_string();
|
||
config.db_type = DatabaseType::Sqlite;
|
||
config.host = main_path.to_string_lossy().to_string();
|
||
config.port = 0;
|
||
config.password.clear();
|
||
config.attached_databases.push(AttachedDatabaseConfig {
|
||
name: "analytics".to_string(),
|
||
path: attached_path.to_string_lossy().to_string(),
|
||
});
|
||
state.configs.write().await.insert(config.id.clone(), config);
|
||
|
||
state.get_or_create_pool("sqlite-conn", None).await.unwrap();
|
||
let databases = schema::list_databases_core(&state, "sqlite-conn").await.unwrap();
|
||
assert!(databases.iter().any(|database| database.name == "analytics"));
|
||
let tables =
|
||
schema::list_tables_core(&state, "sqlite-conn", "analytics", "analytics", None, None, None, None, None)
|
||
.await
|
||
.unwrap();
|
||
assert!(tables.iter().any(|table| table.name == "events"));
|
||
|
||
state.connections.write().await.clear();
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn remove_connection_pools_clears_base_and_database_scoped_pools() {
|
||
let (state, dir) = test_app_state().await;
|
||
let pool = crate::db::sqlite::connect_path(":memory:").await.unwrap();
|
||
|
||
{
|
||
let mut conns = state.connections.write().await;
|
||
conns.insert("conn".to_string(), PoolKind::Sqlite(pool.clone()));
|
||
conns.insert("conn:analytics".to_string(), PoolKind::Sqlite(pool.clone()));
|
||
conns.insert("conn:session:tab-1".to_string(), PoolKind::Sqlite(pool.clone()));
|
||
conns.insert("conn:analytics:session:tab-1".to_string(), PoolKind::Sqlite(pool.clone()));
|
||
conns.insert("other".to_string(), PoolKind::Sqlite(pool));
|
||
}
|
||
|
||
state.remove_connection_pools("conn").await;
|
||
|
||
let conns = state.connections.read().await;
|
||
assert!(!conns.contains_key("conn"));
|
||
assert!(!conns.contains_key("conn:analytics"));
|
||
assert!(!conns.contains_key("conn:session:tab-1"));
|
||
assert!(!conns.contains_key("conn:analytics:session:tab-1"));
|
||
assert!(conns.contains_key("other"));
|
||
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn connection_slot_recovery_waits_for_agent_metadata_and_confirms_close() {
|
||
let (state, dir) = test_app_state().await;
|
||
let script_path = dir.join("agent.py");
|
||
let metadata_started_path = dir.join("metadata-started");
|
||
let session_closed_path = dir.join("session-closed");
|
||
let metadata_started = serde_json::to_string(&metadata_started_path.to_string_lossy()).unwrap();
|
||
let session_closed = serde_json::to_string(&session_closed_path.to_string_lossy()).unwrap();
|
||
std::fs::write(
|
||
&script_path,
|
||
format!(
|
||
r#"import json, pathlib, sys, time
|
||
metadata_started = pathlib.Path({metadata_started})
|
||
session_closed = pathlib.Path({session_closed})
|
||
print(json.dumps({{'ready': True}}), flush=True)
|
||
for line in sys.stdin:
|
||
req = json.loads(line)
|
||
method = req['method']
|
||
if method == 'handshake':
|
||
result = {{'protocolVersion': 2, 'agentProtocolVersion': 2, 'capabilities': ['multi_session']}}
|
||
elif method == 'list_databases':
|
||
metadata_started.write_text('started')
|
||
time.sleep(3.2)
|
||
result = []
|
||
elif method == 'close_session':
|
||
session_closed.write_text(req['params']['agentSessionId'])
|
||
result = {{}}
|
||
else:
|
||
result = {{}}
|
||
print(json.dumps({{'jsonrpc': '2.0', 'id': req['id'], 'result': result}}), flush=True)
|
||
"#
|
||
),
|
||
)
|
||
.unwrap();
|
||
|
||
let python = if cfg!(windows) { "python" } else { "python3" };
|
||
let runtime = crate::db::agent_driver::AgentRuntimeClient::spawn(
|
||
crate::db::agent_driver::AgentLaunchSpec::new(python)
|
||
.with_args([script_path.to_string_lossy().to_string()]),
|
||
"test",
|
||
)
|
||
.await
|
||
.unwrap();
|
||
runtime.increment_session_count();
|
||
let client = std::sync::Arc::new(crate::db::agent_driver::PooledAgentClient::new(
|
||
crate::db::agent_driver::AgentDriverClient::shared_session(runtime.clone(), "metadata-session".to_string()),
|
||
));
|
||
state.connections.write().await.insert("conn:analytics".to_string(), PoolKind::Agent(client.clone()));
|
||
|
||
let metadata_client = client.clone();
|
||
drop(client);
|
||
let metadata = tokio::spawn(async move {
|
||
metadata_client.lock().await.list_databases::<Vec<String>>(Some(Duration::from_secs(6))).await
|
||
});
|
||
for _ in 0..100 {
|
||
if metadata_started_path.exists() {
|
||
break;
|
||
}
|
||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||
}
|
||
assert!(metadata_started_path.exists());
|
||
|
||
let reclaim_started = Instant::now();
|
||
assert!(!state.reclaim_idle_base_pool_for_session("conn", "conn:analytics").await);
|
||
assert!(reclaim_started.elapsed() < Duration::from_secs(1));
|
||
assert!(state.connections.read().await.contains_key("conn:analytics"));
|
||
assert!(!session_closed_path.exists());
|
||
|
||
assert_eq!(metadata.await.unwrap().unwrap(), Vec::<String>::new());
|
||
assert!(state.reclaim_idle_base_pool_for_session("conn", "conn:analytics").await);
|
||
assert!(!state.connections.read().await.contains_key("conn:analytics"));
|
||
assert_eq!(std::fs::read_to_string(&session_closed_path).unwrap(), "metadata-session");
|
||
assert_eq!(runtime.active_session_count(), 0);
|
||
|
||
runtime.kill();
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn connection_slot_recovery_keeps_active_metadata_pool() {
|
||
let (state, dir) = test_app_state().await;
|
||
state.connections.write().await.insert("conn:analytics".to_string(), agent_pool_stub());
|
||
|
||
let _query = state.running_queries.register("active-metadata".to_string());
|
||
state.running_queries.set_pool_key("active-metadata", "conn:analytics");
|
||
|
||
assert!(!state.reclaim_idle_base_pool_for_session("conn", "conn:analytics").await);
|
||
assert!(state.connections.read().await.contains_key("conn:analytics"));
|
||
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn busy_agent_pool_skips_health_probe_without_waiting() {
|
||
let (state, dir) = test_app_state().await;
|
||
let client = std::sync::Arc::new(crate::db::agent_driver::PooledAgentClient::new(
|
||
crate::db::agent_driver::AgentDriverClient::test_stub(),
|
||
));
|
||
state.connections.write().await.insert("conn".to_string(), PoolKind::Agent(client.clone()));
|
||
let _busy = client.lock().await;
|
||
|
||
let started = Instant::now();
|
||
assert!(!state.remove_stale_connection_pool("conn").await);
|
||
assert!(started.elapsed() < Duration::from_millis(100));
|
||
assert!(state.connections.read().await.contains_key("conn"));
|
||
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[test]
|
||
fn detects_database_connection_slot_exhaustion_errors() {
|
||
assert!(super::is_connection_slot_exhausted_error(
|
||
"Agent RPC error (-1): FATAL: remaining connection slots are reserved for superuser manager connections"
|
||
));
|
||
assert!(super::is_connection_slot_exhausted_error("ORA-00018: maximum number of sessions exceeded"));
|
||
assert!(!super::is_connection_slot_exhausted_error("password authentication failed"));
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn pool_activity_touch_updates_existing_pool_only() {
|
||
let (state, dir) = test_app_state().await;
|
||
let pool = crate::db::sqlite::connect_path(":memory:").await.unwrap();
|
||
let pool_key = "conn:session:tab-1";
|
||
|
||
state.connections.write().await.insert(pool_key.to_string(), PoolKind::Sqlite(pool));
|
||
state
|
||
.pool_activity
|
||
.write()
|
||
.await
|
||
.insert(pool_key.to_string(), super::PoolActivity::idle_for(std::time::Duration::from_secs(10)));
|
||
|
||
{
|
||
let _touch = state.pool_activity_touch(pool_key);
|
||
}
|
||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||
let elapsed = state.pool_activity.read().await.get(pool_key).unwrap().elapsed();
|
||
assert!(elapsed < std::time::Duration::from_secs(10));
|
||
|
||
{
|
||
let _touch = state.pool_activity_touch(pool_key);
|
||
state.connections.write().await.remove(pool_key);
|
||
state.pool_activity.write().await.remove(pool_key);
|
||
}
|
||
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
|
||
assert!(!state.pool_activity.read().await.contains_key(pool_key));
|
||
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[test]
|
||
fn pool_activity_touch_never_moves_backwards() {
|
||
let activity = super::PoolActivity::idle_for(std::time::Duration::from_secs(0));
|
||
let newer = u64::MAX;
|
||
activity.last_used_at_ms.store(newer, std::sync::atomic::Ordering::Relaxed);
|
||
// touch 的当前时间早于已存值:不得倒退
|
||
activity.touch();
|
||
assert_eq!(activity.last_used_at_ms.load(std::sync::atomic::Ordering::Relaxed), newer);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn session_scoped_pool_is_not_closed_by_idle_timeout() {
|
||
let (state, dir) = test_app_state().await;
|
||
let pool_key = "conn:session:tab-1";
|
||
let pool = crate::db::sqlite::connect_path(":memory:").await.unwrap();
|
||
let mut config = mysql_config(None);
|
||
config.idle_timeout_secs = 1;
|
||
config.keepalive_interval_secs = 0;
|
||
|
||
state.connections.write().await.insert(pool_key.to_string(), PoolKind::Sqlite(pool));
|
||
state
|
||
.pool_activity
|
||
.write()
|
||
.await
|
||
.insert(pool_key.to_string(), super::PoolActivity::idle_for(std::time::Duration::from_secs(10)));
|
||
let pool = super::clone_pool_kind(state.connections.read().await.get(pool_key).unwrap());
|
||
state.start_keepalive_task(
|
||
pool_key,
|
||
&pool,
|
||
&config,
|
||
#[cfg(feature = "mq-admin")]
|
||
None,
|
||
);
|
||
|
||
tokio::time::sleep(std::time::Duration::from_millis(20)).await;
|
||
|
||
assert!(state.connections.read().await.contains_key(pool_key));
|
||
assert_eq!(state.supervised_task_count(), 0);
|
||
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn close_client_session_pool_releases_session_scoped_pool() {
|
||
let (state, dir) = test_app_state().await;
|
||
let mut config = mysql_config(None);
|
||
config.id = "conn".to_string();
|
||
config.db_type = DatabaseType::Sqlite;
|
||
state.configs.write().await.insert(config.id.clone(), config);
|
||
|
||
let pool_key = "conn:session:tab-1";
|
||
let pool = crate::db::sqlite::connect_path(":memory:").await.unwrap();
|
||
state.connections.write().await.insert(pool_key.to_string(), PoolKind::Sqlite(pool));
|
||
state.pool_activity.write().await.insert(pool_key.to_string(), super::PoolActivity::now());
|
||
|
||
assert!(state.close_client_session_pool("conn", None, "tab-1").await.unwrap());
|
||
assert!(!state.connections.read().await.contains_key(pool_key));
|
||
assert!(!state.pool_activity.read().await.contains_key(pool_key));
|
||
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn detach_client_session_pool_removes_pool_before_background_close() {
|
||
let (state, dir) = test_app_state().await;
|
||
let mut config = mysql_config(None);
|
||
config.id = "conn".to_string();
|
||
config.db_type = DatabaseType::Sqlite;
|
||
state.configs.write().await.insert(config.id.clone(), config);
|
||
|
||
let pool_key = "conn:session:import-1";
|
||
let pool = crate::db::sqlite::connect_path(":memory:").await.unwrap();
|
||
state.connections.write().await.insert(pool_key.to_string(), PoolKind::Sqlite(pool));
|
||
state.pool_activity.write().await.insert(pool_key.to_string(), super::PoolActivity::now());
|
||
|
||
assert!(state.detach_client_session_pool("conn", None, "import-1").await.unwrap());
|
||
assert!(!state.connections.read().await.contains_key(pool_key));
|
||
assert!(!state.pool_activity.read().await.contains_key(pool_key));
|
||
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn replace_runtime_for_base_metadata_pool_detaches_routing_and_kills_runtime() {
|
||
let (state, dir) = test_app_state().await;
|
||
let script_path = dir.join("replace-runtime-agent.py");
|
||
let request_started_path = dir.join("replace-runtime-request-started");
|
||
let request_started = serde_json::to_string(&request_started_path.to_string_lossy()).unwrap();
|
||
std::fs::write(
|
||
&script_path,
|
||
format!(
|
||
r#"import json, pathlib, sys, time
|
||
request_started = pathlib.Path({request_started})
|
||
print(json.dumps({{'ready': True}}), flush=True)
|
||
for line in sys.stdin:
|
||
req = json.loads(line)
|
||
if req['method'] == 'handshake':
|
||
result = {{'protocolVersion': 2, 'agentProtocolVersion': 2, 'capabilities': ['multi_session']}}
|
||
elif req['method'] == 'list_databases':
|
||
request_started.write_text('started')
|
||
time.sleep(30)
|
||
result = []
|
||
else:
|
||
result = {{}}
|
||
print(json.dumps({{'jsonrpc': '2.0', 'id': req['id'], 'result': result}}), flush=True)
|
||
"#
|
||
),
|
||
)
|
||
.unwrap();
|
||
let python = if cfg!(windows) { "python" } else { "python3" };
|
||
let runtime = crate::db::agent_driver::AgentRuntimeClient::spawn(
|
||
crate::db::agent_driver::AgentLaunchSpec::new(python)
|
||
.with_args([script_path.to_string_lossy().to_string()]),
|
||
"test",
|
||
)
|
||
.await
|
||
.unwrap();
|
||
runtime.increment_session_count();
|
||
let metadata_client = std::sync::Arc::new(crate::db::agent_driver::PooledAgentClient::new(
|
||
crate::db::agent_driver::AgentDriverClient::shared_session(runtime.clone(), "metadata-session".to_string()),
|
||
));
|
||
runtime.increment_session_count();
|
||
let workload_client = std::sync::Arc::new(crate::db::agent_driver::PooledAgentClient::new(
|
||
crate::db::agent_driver::AgentDriverClient::shared_session(runtime.clone(), "workload-session".to_string()),
|
||
));
|
||
|
||
let mut config = mysql_config(Some("analytics"));
|
||
config.id = "conn".to_string();
|
||
config.db_type = DatabaseType::Dameng;
|
||
state.configs.write().await.insert(config.id.clone(), config);
|
||
let pool_key = "conn:analytics:role:metadata";
|
||
let sibling_pool_key = "conn:analytics:session:workload";
|
||
{
|
||
let mut connections = state.connections.write().await;
|
||
connections.insert(pool_key.to_string(), PoolKind::Agent(metadata_client.clone()));
|
||
connections.insert(sibling_pool_key.to_string(), PoolKind::Agent(workload_client));
|
||
}
|
||
{
|
||
let mut activity = state.pool_activity.write().await;
|
||
activity.insert(pool_key.to_string(), super::PoolActivity::now());
|
||
activity.insert(sibling_pool_key.to_string(), super::PoolActivity::now());
|
||
}
|
||
|
||
let blocked_client = metadata_client.clone();
|
||
let blocked_request = tokio::spawn(async move {
|
||
blocked_client.lock().await.list_databases::<Vec<String>>(Some(Duration::from_secs(30))).await
|
||
});
|
||
for _ in 0..100 {
|
||
if request_started_path.exists() {
|
||
break;
|
||
}
|
||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||
}
|
||
assert!(request_started_path.exists());
|
||
|
||
let replacement = tokio::time::timeout(
|
||
Duration::from_secs(1),
|
||
state.replace_runtime_for_metadata_pool("conn", Some("analytics"), None),
|
||
)
|
||
.await;
|
||
if replacement.is_err() {
|
||
runtime.kill();
|
||
}
|
||
let _ = tokio::time::timeout(Duration::from_secs(2), blocked_request).await;
|
||
|
||
assert!(replacement.expect("runtime replacement must not wait for the Agent client lock"));
|
||
assert!(!state.connections.read().await.contains_key(pool_key));
|
||
assert!(!state.connections.read().await.contains_key(sibling_pool_key));
|
||
assert!(!state.pool_activity.read().await.contains_key(pool_key));
|
||
assert!(!state.pool_activity.read().await.contains_key(sibling_pool_key));
|
||
assert!(runtime.is_failed());
|
||
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
async fn replace_runtime_on_error_clients(
|
||
dir: &std::path::Path,
|
||
error_method: &str,
|
||
) -> (
|
||
std::sync::Arc<crate::db::agent_driver::AgentRuntimeClient>,
|
||
std::sync::Arc<crate::db::agent_driver::PooledAgentClient>,
|
||
std::sync::Arc<crate::db::agent_driver::PooledAgentClient>,
|
||
) {
|
||
let script_path = dir.join("close-replace-runtime-agent.py");
|
||
let script = r#"import json, sys
|
||
print(json.dumps({'ready': True}), flush=True)
|
||
for line in sys.stdin:
|
||
req = json.loads(line)
|
||
if req['method'] == 'handshake':
|
||
response = {
|
||
'jsonrpc': '2.0',
|
||
'id': req['id'],
|
||
'result': {'protocolVersion': 2, 'agentProtocolVersion': 2, 'capabilities': ['multi_session']}
|
||
}
|
||
elif req['method'] == '__ERROR_METHOD__':
|
||
response = {
|
||
'jsonrpc': '2.0',
|
||
'id': req['id'],
|
||
'error': {
|
||
'code': -1,
|
||
'message': 'Agent runtime resource limit reached',
|
||
'data': {
|
||
'category': 'resource',
|
||
'retryable': False,
|
||
'sessionDisposition': 'replace_runtime',
|
||
'stage': 'close'
|
||
}
|
||
}
|
||
}
|
||
else:
|
||
response = {'jsonrpc': '2.0', 'id': req['id'], 'result': {}}
|
||
print(json.dumps(response), flush=True)
|
||
"#
|
||
.replace("'__ERROR_METHOD__'", &serde_json::to_string(error_method).unwrap());
|
||
std::fs::write(&script_path, script).unwrap();
|
||
let python = if cfg!(windows) { "python" } else { "python3" };
|
||
let runtime = crate::db::agent_driver::AgentRuntimeClient::spawn(
|
||
crate::db::agent_driver::AgentLaunchSpec::new(python)
|
||
.with_args([script_path.to_string_lossy().to_string()]),
|
||
"test",
|
||
)
|
||
.await
|
||
.unwrap();
|
||
runtime.increment_session_count();
|
||
let metadata_client = std::sync::Arc::new(crate::db::agent_driver::PooledAgentClient::new(
|
||
crate::db::agent_driver::AgentDriverClient::shared_session(
|
||
runtime.clone(),
|
||
"metadata-agent-session".to_string(),
|
||
),
|
||
));
|
||
runtime.increment_session_count();
|
||
let workload_client = std::sync::Arc::new(crate::db::agent_driver::PooledAgentClient::new(
|
||
crate::db::agent_driver::AgentDriverClient::shared_session(
|
||
runtime.clone(),
|
||
"workload-agent-session".to_string(),
|
||
),
|
||
));
|
||
(runtime, metadata_client, workload_client)
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn metadata_close_replace_runtime_detaches_shared_runtime_siblings() {
|
||
let (state, dir) = test_app_state().await;
|
||
let (runtime, metadata_client, workload_client) = replace_runtime_on_error_clients(&dir, "close_session").await;
|
||
let mut config = mysql_config(Some("analytics"));
|
||
config.id = "conn".to_string();
|
||
config.db_type = DatabaseType::Dameng;
|
||
state.configs.write().await.insert(config.id.clone(), config);
|
||
let metadata_pool_key = "conn:analytics:session:metadata-session:role:metadata";
|
||
let workload_pool_key = "conn:analytics:session:workload-session";
|
||
{
|
||
let mut connections = state.connections.write().await;
|
||
connections.insert(metadata_pool_key.to_string(), PoolKind::Agent(metadata_client));
|
||
connections.insert(workload_pool_key.to_string(), PoolKind::Agent(workload_client));
|
||
}
|
||
{
|
||
let mut activity = state.pool_activity.write().await;
|
||
activity.insert(metadata_pool_key.to_string(), super::PoolActivity::now());
|
||
activity.insert(workload_pool_key.to_string(), super::PoolActivity::now());
|
||
}
|
||
|
||
assert!(state.close_metadata_session_pool("conn", Some("analytics"), "metadata-session").await.unwrap());
|
||
|
||
assert!(!state.connections.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));
|
||
assert!(runtime.is_failed());
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn metadata_close_replace_runtime_preserves_manual_transaction_sibling() {
|
||
let (state, dir) = test_app_state().await;
|
||
let (runtime, metadata_client, manual_txn_client) =
|
||
replace_runtime_on_error_clients(&dir, "close_session").await;
|
||
let mut config = mysql_config(Some("analytics"));
|
||
config.id = "conn".to_string();
|
||
config.db_type = DatabaseType::Dameng;
|
||
state.configs.write().await.insert(config.id.clone(), config);
|
||
let metadata_pool_key = "conn:analytics:session:metadata-session:role:metadata";
|
||
let manual_txn_pool_key = "conn:analytics:session:manual-txn-test";
|
||
{
|
||
let mut connections = state.connections.write().await;
|
||
connections.insert(metadata_pool_key.to_string(), PoolKind::Agent(metadata_client));
|
||
connections.insert(manual_txn_pool_key.to_string(), PoolKind::Agent(manual_txn_client));
|
||
}
|
||
{
|
||
let mut activity = state.pool_activity.write().await;
|
||
activity.insert(metadata_pool_key.to_string(), super::PoolActivity::now());
|
||
activity.insert(manual_txn_pool_key.to_string(), super::PoolActivity::now());
|
||
}
|
||
|
||
assert!(state.close_metadata_session_pool("conn", Some("analytics"), "metadata-session").await.unwrap());
|
||
|
||
assert!(!state.connections.read().await.contains_key(metadata_pool_key));
|
||
assert!(state.connections.read().await.contains_key(manual_txn_pool_key));
|
||
assert!(state.pool_activity.read().await.contains_key(manual_txn_pool_key));
|
||
assert!(!runtime.is_failed());
|
||
|
||
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;
|
||
let (runtime, reclaimed_client, sibling_client) = replace_runtime_on_error_clients(&dir, "close_session").await;
|
||
let mut config = mysql_config(Some("analytics"));
|
||
config.id = "conn".to_string();
|
||
config.db_type = DatabaseType::Dameng;
|
||
state.configs.write().await.insert(config.id.clone(), config);
|
||
let reclaimed_pool_key = "conn:analytics";
|
||
let sibling_pool_key = "conn:analytics:session:workload";
|
||
{
|
||
let mut connections = state.connections.write().await;
|
||
connections.insert(reclaimed_pool_key.to_string(), PoolKind::Agent(reclaimed_client));
|
||
connections.insert(sibling_pool_key.to_string(), PoolKind::Agent(sibling_client));
|
||
}
|
||
state.pool_activity.write().await.insert(reclaimed_pool_key.to_string(), super::PoolActivity::now());
|
||
|
||
let reclaimed = state.try_reclaim_idle_agent_pool(reclaimed_pool_key).await;
|
||
|
||
assert!(reclaimed);
|
||
assert!(!state.connections.read().await.contains_key(reclaimed_pool_key));
|
||
assert!(!state.connections.read().await.contains_key(sibling_pool_key));
|
||
assert!(runtime.is_failed());
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn inserting_replacement_returns_error_when_previous_close_replaces_shared_runtime() {
|
||
let (state, dir) = test_app_state().await;
|
||
let (runtime, previous_client, replacement_client) =
|
||
replace_runtime_on_error_clients(&dir, "close_session").await;
|
||
let mut config = mysql_config(Some("analytics"));
|
||
config.id = "conn".to_string();
|
||
config.db_type = DatabaseType::Dameng;
|
||
state.configs.write().await.insert(config.id.clone(), config.clone());
|
||
let pool_key = "conn:analytics";
|
||
state.connections.write().await.insert(pool_key.to_string(), PoolKind::Agent(previous_client));
|
||
|
||
let result =
|
||
state.insert_connection_pool(pool_key.to_string(), PoolKind::Agent(replacement_client), &config).await;
|
||
|
||
assert!(result.is_err());
|
||
assert!(!state.connections.read().await.contains_key(pool_key));
|
||
assert!(runtime.is_failed());
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn reconnect_waits_for_agent_runtime_replacement_before_publishing_new_pool() {
|
||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||
|
||
let (state, dir) = test_app_state().await;
|
||
let script_path = dir.join("delayed-close-replace-runtime-agent.py");
|
||
let close_finished_path = dir.join("close-finished");
|
||
let close_finished = serde_json::to_string(&close_finished_path.to_string_lossy()).unwrap();
|
||
std::fs::write(
|
||
&script_path,
|
||
format!(
|
||
r#"import json, pathlib, sys, time
|
||
close_finished = pathlib.Path({close_finished})
|
||
print(json.dumps({{'ready': True}}), flush=True)
|
||
for line in sys.stdin:
|
||
req = json.loads(line)
|
||
if req['method'] == 'handshake':
|
||
response = {{
|
||
'jsonrpc': '2.0',
|
||
'id': req['id'],
|
||
'result': {{'protocolVersion': 2, 'agentProtocolVersion': 2, 'capabilities': ['multi_session']}}
|
||
}}
|
||
elif req['method'] == 'close_session':
|
||
time.sleep(0.5)
|
||
close_finished.write_text('done')
|
||
response = {{
|
||
'jsonrpc': '2.0',
|
||
'id': req['id'],
|
||
'error': {{
|
||
'code': -1,
|
||
'message': 'Agent runtime resource limit reached',
|
||
'data': {{
|
||
'category': 'resource',
|
||
'retryable': False,
|
||
'sessionDisposition': 'replace_runtime',
|
||
'stage': 'close'
|
||
}}
|
||
}}
|
||
}}
|
||
else:
|
||
response = {{'jsonrpc': '2.0', 'id': req['id'], 'result': {{}}}}
|
||
print(json.dumps(response), flush=True)
|
||
"#
|
||
),
|
||
)
|
||
.unwrap();
|
||
|
||
let python = if cfg!(windows) { "python" } else { "python3" };
|
||
let runtime = crate::db::agent_driver::AgentRuntimeClient::spawn(
|
||
crate::db::agent_driver::AgentLaunchSpec::new(python)
|
||
.with_args([script_path.to_string_lossy().to_string()]),
|
||
"test",
|
||
)
|
||
.await
|
||
.unwrap();
|
||
runtime.increment_session_count();
|
||
let client = std::sync::Arc::new(crate::db::agent_driver::PooledAgentClient::new(
|
||
crate::db::agent_driver::AgentDriverClient::shared_session(
|
||
runtime.clone(),
|
||
"metadata-agent-session".to_string(),
|
||
),
|
||
));
|
||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||
let address = listener.local_addr().unwrap();
|
||
let server = tokio::spawn(async move {
|
||
let (mut socket, _) = listener.accept().await.unwrap();
|
||
let mut request = Vec::new();
|
||
let mut chunk = [0_u8; 1024];
|
||
while !request.windows(4).any(|window| window == b"\r\n\r\n") {
|
||
let read = socket.read(&mut chunk).await.unwrap();
|
||
assert!(read > 0, "rqlite probe ended before the request headers were complete");
|
||
request.extend_from_slice(&chunk[..read]);
|
||
}
|
||
assert!(String::from_utf8_lossy(&request).starts_with("POST /db/query HTTP/1.1"));
|
||
let body = r#"{"results":[{"columns":["1"],"values":[[1]]}]}"#;
|
||
let response = format!(
|
||
"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
|
||
body.len()
|
||
);
|
||
socket.write_all(response.as_bytes()).await.unwrap();
|
||
});
|
||
let mut config = mysql_config(None);
|
||
config.id = "conn".to_string();
|
||
config.db_type = DatabaseType::Rqlite;
|
||
config.host = address.ip().to_string();
|
||
config.port = address.port();
|
||
config.keepalive_interval_secs = 0;
|
||
state.configs.write().await.insert(config.id.clone(), config);
|
||
state.connections.write().await.insert("conn".to_string(), PoolKind::Agent(client));
|
||
|
||
let pool_key = state.reconnect_metadata_pool_for_session("conn", None, None).await.unwrap();
|
||
server.await.unwrap();
|
||
|
||
assert_eq!(pool_key, "conn");
|
||
assert!(close_finished_path.exists(), "reconnect returned before the old Agent close completed");
|
||
assert!(runtime.is_failed());
|
||
assert!(matches!(state.connections.read().await.get("conn"), Some(PoolKind::Rqlite(_))));
|
||
|
||
state.shutdown(Duration::from_secs(1)).await;
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn inserting_agent_pool_rejects_runtime_failed_before_publish() {
|
||
let (state, dir) = test_app_state().await;
|
||
let (runtime, client, sibling_client) = replace_runtime_on_error_clients(&dir, "unused").await;
|
||
let mut config = mysql_config(Some("analytics"));
|
||
config.id = "conn".to_string();
|
||
config.db_type = DatabaseType::Dameng;
|
||
state.connections.write().await.insert("conn:billing".to_string(), PoolKind::Agent(sibling_client));
|
||
runtime.kill();
|
||
|
||
let result = state.insert_connection_pool("conn:analytics".to_string(), PoolKind::Agent(client), &config).await;
|
||
|
||
assert!(result.is_err());
|
||
assert!(state.connections.read().await.is_empty());
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn inserting_failed_agent_pool_preserves_healthy_existing_route_state() {
|
||
let (state, dir) = test_app_state().await;
|
||
let (healthy_runtime, healthy_client, _healthy_sibling) =
|
||
replace_runtime_on_error_clients(&dir, "unused").await;
|
||
let (failed_runtime, failed_client, _failed_sibling) = replace_runtime_on_error_clients(&dir, "unused").await;
|
||
let mut config = mysql_config(Some("analytics"));
|
||
config.id = "conn".to_string();
|
||
config.db_type = DatabaseType::Dameng;
|
||
config.keepalive_interval_secs = 60;
|
||
let pool_key = "conn:analytics";
|
||
state
|
||
.insert_connection_pool(pool_key.to_string(), PoolKind::Agent(healthy_client.clone()), &config)
|
||
.await
|
||
.unwrap();
|
||
state
|
||
.pool_activity
|
||
.write()
|
||
.await
|
||
.insert(pool_key.to_string(), super::PoolActivity::idle_for(Duration::from_secs(600)));
|
||
assert_eq!(state.supervised_task_count(), 1);
|
||
failed_runtime.kill();
|
||
|
||
let result = state.insert_connection_pool(pool_key.to_string(), PoolKind::Agent(failed_client), &config).await;
|
||
|
||
assert!(result.is_err());
|
||
let connections = state.connections.read().await;
|
||
let PoolKind::Agent(routed_client) = connections.get(pool_key).expect("healthy route must remain") else {
|
||
panic!("existing route must remain an Agent pool");
|
||
};
|
||
assert!(healthy_client.shares_runtime_with(routed_client));
|
||
drop(connections);
|
||
assert!(
|
||
state.pool_activity.read().await.get(pool_key).expect("activity must remain").elapsed().as_secs() >= 300
|
||
);
|
||
assert_eq!(state.supervised_task_count(), 1);
|
||
assert!(!healthy_runtime.is_failed());
|
||
|
||
state.shutdown(Duration::from_secs(1)).await;
|
||
healthy_runtime.kill();
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn inserting_pool_does_not_wait_for_activity_while_holding_route_lock() {
|
||
let (state, dir) = test_app_state().await;
|
||
let state = std::sync::Arc::new(state);
|
||
let activity_guard = state.pool_activity.read().await;
|
||
let mut config = mysql_config(Some("analytics"));
|
||
config.keepalive_interval_secs = 0;
|
||
let publishing_state = state.clone();
|
||
let publish = tokio::spawn(async move {
|
||
publishing_state.insert_connection_pool("conn:analytics".to_string(), agent_pool_stub(), &config).await
|
||
});
|
||
tokio::time::sleep(Duration::from_millis(20)).await;
|
||
|
||
let route_read = tokio::time::timeout(Duration::from_millis(100), state.connections.read()).await;
|
||
|
||
assert!(route_read.is_ok(), "pool publication must not await activity while holding the route lock");
|
||
drop(route_read);
|
||
drop(activity_guard);
|
||
assert!(tokio::time::timeout(Duration::from_secs(1), publish).await.unwrap().unwrap().is_ok());
|
||
state.shutdown(Duration::from_secs(1)).await;
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn open_session_replace_runtime_error_detaches_existing_shared_runtime_pools() {
|
||
let (state, dir) = test_app_state().await;
|
||
let (runtime, first_client, second_client) = replace_runtime_on_error_clients(&dir, "unused").await;
|
||
{
|
||
let mut connections = state.connections.write().await;
|
||
connections.insert("conn:analytics".to_string(), PoolKind::Agent(first_client));
|
||
connections.insert("conn:billing".to_string(), PoolKind::Agent(second_client));
|
||
}
|
||
let error = crate::agent_runtime::SharedConnectionOpenError {
|
||
message: "open session requested runtime replacement".to_string(),
|
||
runtime: Some(runtime.clone()),
|
||
};
|
||
|
||
let message = state.handle_shared_connection_open_error(error).await;
|
||
|
||
assert_eq!(message, "open session requested runtime replacement");
|
||
assert!(state.connections.read().await.is_empty());
|
||
assert!(runtime.is_failed());
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn stale_probe_replace_runtime_detaches_shared_runtime_siblings() {
|
||
let (state, dir) = test_app_state().await;
|
||
let (runtime, target_client, sibling_client) =
|
||
replace_runtime_on_error_clients(&dir, "validate_connection").await;
|
||
let target_pool_key = "conn:analytics";
|
||
let sibling_pool_key = "conn:analytics:session:workload";
|
||
{
|
||
let mut connections = state.connections.write().await;
|
||
connections.insert(target_pool_key.to_string(), PoolKind::Agent(target_client));
|
||
connections.insert(sibling_pool_key.to_string(), PoolKind::Agent(sibling_client));
|
||
}
|
||
|
||
assert!(state.remove_stale_connection_pool(target_pool_key).await);
|
||
|
||
assert!(!state.connections.read().await.contains_key(target_pool_key));
|
||
assert!(!state.connections.read().await.contains_key(sibling_pool_key));
|
||
assert!(runtime.is_failed());
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn global_health_validates_current_session_and_fail_stops_shared_runtime() {
|
||
let (state, dir) = test_app_state().await;
|
||
let (runtime, target_client, sibling_client) =
|
||
replace_runtime_on_error_clients(&dir, "validate_connection").await;
|
||
let mut config = mysql_config(None);
|
||
config.transport_layers = vec![TransportLayerConfig::Proxy(ProxyTunnelConfig {
|
||
profile_id: String::new(),
|
||
id: "proxy".to_string(),
|
||
name: String::new(),
|
||
enabled: true,
|
||
proxy_type: ProxyType::Socks5,
|
||
host: "127.0.0.1".to_string(),
|
||
port: 65000,
|
||
username: String::new(),
|
||
password: String::new(),
|
||
test_target: None,
|
||
})];
|
||
state.configs.write().await.insert(config.id.clone(), config.clone());
|
||
state.connection_host_port(&config.id, &config).await.unwrap();
|
||
{
|
||
let mut connections = state.connections.write().await;
|
||
connections.insert("conn:analytics".to_string(), PoolKind::Agent(target_client));
|
||
connections.insert("conn:billing".to_string(), PoolKind::Agent(sibling_client));
|
||
}
|
||
|
||
state.refresh_connections().await;
|
||
|
||
assert!(state.connections.read().await.is_empty());
|
||
assert!(state.proxy_tunnels.local_port("conn:transport:0").await.is_none());
|
||
assert!(runtime.is_failed());
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn global_health_preserves_transport_for_connections_without_failed_pools() {
|
||
let (state, dir) = test_app_state().await;
|
||
let mut config = mysql_config(None);
|
||
config.id = "proxied-connection".to_string();
|
||
config.transport_layers = vec![TransportLayerConfig::Proxy(ProxyTunnelConfig {
|
||
profile_id: String::new(),
|
||
id: "proxy".to_string(),
|
||
name: String::new(),
|
||
enabled: true,
|
||
proxy_type: ProxyType::Socks5,
|
||
host: "127.0.0.1".to_string(),
|
||
port: 65000,
|
||
username: String::new(),
|
||
password: String::new(),
|
||
test_target: None,
|
||
})];
|
||
state.configs.write().await.insert(config.id.clone(), config.clone());
|
||
let (_, local_port) = state.connection_host_port(&config.id, &config).await.unwrap();
|
||
|
||
state.refresh_connections().await;
|
||
|
||
assert_eq!(state.proxy_tunnels.local_port("proxied-connection:transport:0").await, Some(local_port));
|
||
state.reset_connection_transport(&config.id).await;
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn stale_agent_failure_preserves_newer_pool_generation() {
|
||
let (state, dir) = test_app_state().await;
|
||
let (stale_runtime, stale_client, _stale_sibling) = replace_runtime_on_error_clients(&dir, "unused").await;
|
||
let (current_runtime, current_client, _current_sibling) =
|
||
replace_runtime_on_error_clients(&dir, "unused").await;
|
||
let pool_key = "conn:analytics";
|
||
state.connections.write().await.insert(pool_key.to_string(), PoolKind::Agent(current_client.clone()));
|
||
|
||
assert!(!state.detach_agent_pool_if_current(pool_key, &stale_client, true).await);
|
||
|
||
let connections = state.connections.read().await;
|
||
let PoolKind::Agent(routed_client) = connections.get(pool_key).expect("new generation must remain routed")
|
||
else {
|
||
panic!("current route must remain an Agent pool");
|
||
};
|
||
assert!(std::sync::Arc::ptr_eq(routed_client, ¤t_client));
|
||
drop(connections);
|
||
assert!(stale_runtime.is_failed());
|
||
assert!(!current_runtime.is_failed());
|
||
current_runtime.kill();
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn stale_agent_failure_detaches_current_routes_on_the_same_runtime() {
|
||
let (state, dir) = test_app_state().await;
|
||
let (runtime, stale_client, current_client) = replace_runtime_on_error_clients(&dir, "unused").await;
|
||
let pool_key = "conn:analytics";
|
||
let sibling_pool_key = "conn:billing";
|
||
{
|
||
let mut connections = state.connections.write().await;
|
||
connections.insert(pool_key.to_string(), PoolKind::Agent(current_client.clone()));
|
||
connections.insert(sibling_pool_key.to_string(), PoolKind::Agent(current_client));
|
||
}
|
||
|
||
assert!(state.detach_agent_pool_if_current(pool_key, &stale_client, true).await);
|
||
|
||
assert!(!state.connections.read().await.contains_key(pool_key));
|
||
assert!(!state.connections.read().await.contains_key(sibling_pool_key));
|
||
assert!(runtime.is_failed());
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn stale_metadata_error_preserves_newer_session_generation() {
|
||
let (state, dir) = test_app_state().await;
|
||
let (current_runtime, current_client, _current_sibling) =
|
||
replace_runtime_on_error_clients(&dir, "unused").await;
|
||
let mut config = mysql_config(Some("analytics"));
|
||
config.id = "conn".to_string();
|
||
config.db_type = DatabaseType::Dameng;
|
||
state.configs.write().await.insert(config.id.clone(), config);
|
||
let pool_key = "conn:analytics:role:metadata";
|
||
state.connections.write().await.insert(pool_key.to_string(), PoolKind::Agent(current_client.clone()));
|
||
assert!(
|
||
!state
|
||
.detach_metadata_pool_after_recovery("conn", Some("analytics"), None, Some("stale-session"), true,)
|
||
.await
|
||
);
|
||
|
||
let connections = state.connections.read().await;
|
||
let PoolKind::Agent(routed_client) = connections.get(pool_key).expect("new metadata generation must remain")
|
||
else {
|
||
panic!("metadata route must remain an Agent pool");
|
||
};
|
||
assert!(std::sync::Arc::ptr_eq(routed_client, ¤t_client));
|
||
drop(connections);
|
||
assert!(!current_runtime.is_failed());
|
||
current_runtime.kill();
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn keepalive_probe_reports_replacement_without_killing_runtime_directly() {
|
||
let (_state, dir) = test_app_state().await;
|
||
let (runtime, target_client, _sibling_client) =
|
||
replace_runtime_on_error_clients(&dir, "validate_connection").await;
|
||
let mut target = super::KeepaliveTarget::Agent(target_client);
|
||
|
||
let error = super::ping_keepalive_target(&mut target, Duration::from_secs(1)).await.unwrap_err();
|
||
|
||
assert!(matches!(error.recovery_decision(), Some(crate::agent_recovery::RecoveryDecision::ReplaceRuntime)));
|
||
assert!(!runtime.is_failed());
|
||
runtime.kill();
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn detach_pool_by_key_removes_routing_before_background_close() {
|
||
let (state, dir) = test_app_state().await;
|
||
let pool_key = "conn:session:timeout";
|
||
let pool = crate::db::sqlite::connect_path(":memory:").await.unwrap();
|
||
state.connections.write().await.insert(pool_key.to_string(), PoolKind::Sqlite(pool));
|
||
state.pool_activity.write().await.insert(pool_key.to_string(), super::PoolActivity::now());
|
||
|
||
assert!(state.detach_pool_by_key(pool_key, false).await);
|
||
assert!(!state.connections.read().await.contains_key(pool_key));
|
||
assert!(!state.pool_activity.read().await.contains_key(pool_key));
|
||
|
||
for _ in 0..100 {
|
||
if state.supervised_task_count() == 0 {
|
||
break;
|
||
}
|
||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||
}
|
||
assert_eq!(state.supervised_task_count(), 0);
|
||
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn client_session_cleanup_guard_detaches_pool_when_request_is_dropped() {
|
||
let (state, dir) = test_app_state().await;
|
||
let mut config = mysql_config(None);
|
||
config.id = "conn".to_string();
|
||
state.configs.write().await.insert(config.id.clone(), config);
|
||
|
||
let pool_key = "conn:session:completion-objects_request-1";
|
||
let pool = crate::db::sqlite::connect_path(":memory:").await.unwrap();
|
||
state.connections.write().await.insert(pool_key.to_string(), PoolKind::Sqlite(pool));
|
||
state.pool_activity.write().await.insert(pool_key.to_string(), super::PoolActivity::now());
|
||
|
||
let guard = state
|
||
.client_session_pool_cleanup_guard_for_role(
|
||
"conn",
|
||
None,
|
||
"completion-objects:request-1",
|
||
crate::agent_connection::AgentSessionRole::Workload,
|
||
)
|
||
.await
|
||
.unwrap();
|
||
drop(guard);
|
||
|
||
for _ in 0..100 {
|
||
if !state.connections.read().await.contains_key(pool_key) {
|
||
break;
|
||
}
|
||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||
}
|
||
assert!(!state.connections.read().await.contains_key(pool_key));
|
||
assert!(!state.pool_activity.read().await.contains_key(pool_key));
|
||
for _ in 0..100 {
|
||
if state.supervised_task_count() == 0 {
|
||
break;
|
||
}
|
||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||
}
|
||
assert_eq!(state.supervised_task_count(), 0);
|
||
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn manual_transaction_connection_drop_detaches_session_pool() {
|
||
let (state, dir) = test_app_state().await;
|
||
let (runtime, manual_txn_client, _sibling_client) = replace_runtime_on_error_clients(&dir, "unused").await;
|
||
let mut config = mysql_config(Some("analytics"));
|
||
config.id = "conn".to_string();
|
||
config.db_type = DatabaseType::Dameng;
|
||
state.configs.write().await.insert(config.id.clone(), config);
|
||
|
||
let client_session_id = "manual-txn-test";
|
||
let pool_key = "conn:analytics:session:manual-txn-test";
|
||
state.connections.write().await.insert(pool_key.to_string(), PoolKind::Agent(manual_txn_client.clone()));
|
||
state.pool_activity.write().await.insert(pool_key.to_string(), super::PoolActivity::now());
|
||
let cleanup_guard =
|
||
state.workload_session_pool_cleanup_guard("conn", Some("analytics"), client_session_id).await.unwrap();
|
||
|
||
drop(TxnConnection::Agent {
|
||
client: manual_txn_client,
|
||
client_session_id: client_session_id.to_string(),
|
||
database: Some("analytics".to_string()),
|
||
cleanup_guard,
|
||
});
|
||
|
||
for _ in 0..100 {
|
||
if !state.connections.read().await.contains_key(pool_key) {
|
||
break;
|
||
}
|
||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||
}
|
||
assert!(!state.connections.read().await.contains_key(pool_key));
|
||
assert!(!state.pool_activity.read().await.contains_key(pool_key));
|
||
assert!(!runtime.is_failed());
|
||
|
||
state.shutdown(Duration::from_secs(1)).await;
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn closing_oracle_table_tab_keeps_connection_scoped_pool() {
|
||
let (state, dir) = test_app_state().await;
|
||
let mut config = mysql_config(Some("ORCLPDB1"));
|
||
config.id = "oracle-conn".to_string();
|
||
config.db_type = DatabaseType::Oracle;
|
||
state.configs.write().await.insert(config.id.clone(), config);
|
||
|
||
let pool = crate::db::sqlite::connect_path(":memory:").await.unwrap();
|
||
state.connections.write().await.insert("oracle-conn".to_string(), PoolKind::Sqlite(pool));
|
||
state.pool_activity.write().await.insert("oracle-conn".to_string(), super::PoolActivity::now());
|
||
|
||
assert!(!state.close_client_session_pool("oracle-conn", Some("APP"), "table-tab-1").await.unwrap());
|
||
assert!(state.connections.read().await.contains_key("oracle-conn"));
|
||
assert!(state.pool_activity.read().await.contains_key("oracle-conn"));
|
||
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[test]
|
||
fn agent_validate_connection_unknown_method_is_not_stale() {
|
||
assert!(super::is_agent_validate_connection_unsupported(
|
||
"Agent RPC error (-1): Unknown method: validate_connection"
|
||
));
|
||
assert!(super::is_agent_validate_connection_unsupported(
|
||
"Agent RPC error (-32601): Method not found: validate_connection"
|
||
));
|
||
assert!(!super::is_agent_validate_connection_unsupported("Agent RPC error (-1): Connection timed out"));
|
||
assert!(!super::is_agent_validate_connection_unsupported("Agent RPC error (-1): Unknown method: kv_put"));
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn sqlite_memory_client_sessions_share_the_same_database() {
|
||
let (state, dir) = test_app_state().await;
|
||
let mut config = mysql_config(None);
|
||
config.id = "sqlite-memory".to_string();
|
||
config.name = "SQLite memory".to_string();
|
||
config.db_type = DatabaseType::Sqlite;
|
||
config.host = ":memory:".to_string();
|
||
config.password.clear();
|
||
config.port = 0;
|
||
|
||
state.configs.write().await.insert(config.id.clone(), config);
|
||
let base_pool_key = state.get_or_create_pool("sqlite-memory", None).await.unwrap();
|
||
let query_pool_key =
|
||
state.get_or_create_pool_for_session("sqlite-memory", Some("main"), Some("query-tab")).await.unwrap();
|
||
let data_pool_key =
|
||
state.get_or_create_pool_for_session("sqlite-memory", Some("main"), Some("data-tab")).await.unwrap();
|
||
|
||
assert_eq!(query_pool_key, base_pool_key);
|
||
assert_eq!(data_pool_key, base_pool_key);
|
||
|
||
let handle = {
|
||
let connections = state.connections.read().await;
|
||
match connections.get(&base_pool_key) {
|
||
Some(PoolKind::Sqlite(handle)) => handle.clone(),
|
||
_ => panic!("expected SQLite pool"),
|
||
}
|
||
};
|
||
handle
|
||
.with_connection(|connection| {
|
||
connection
|
||
.execute_batch("CREATE TABLE test (id INTEGER); INSERT INTO test VALUES (42);")
|
||
.map_err(|error| error.to_string())
|
||
})
|
||
.unwrap();
|
||
let value = handle
|
||
.with_connection(|connection| {
|
||
connection
|
||
.query_row("SELECT id FROM test", [], |row| row.get::<_, i64>(0))
|
||
.map_err(|error| error.to_string())
|
||
})
|
||
.unwrap();
|
||
assert_eq!(value, 42);
|
||
|
||
assert!(!state.close_client_session_pool("sqlite-memory", Some("main"), "query-tab").await.unwrap());
|
||
assert!(state.connections.read().await.contains_key(&base_pool_key));
|
||
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn close_database_pool_removes_database_and_session_scoped_pools_only() {
|
||
let (state, dir) = test_app_state().await;
|
||
let mut config = mysql_config(None);
|
||
config.id = "conn".to_string();
|
||
state.configs.write().await.insert(config.id.clone(), config);
|
||
let pool = crate::db::sqlite::connect_path(":memory:").await.unwrap();
|
||
|
||
{
|
||
let mut conns = state.connections.write().await;
|
||
conns.insert("conn".to_string(), PoolKind::Sqlite(pool.clone()));
|
||
conns.insert("conn:analytics".to_string(), PoolKind::Sqlite(pool.clone()));
|
||
conns.insert("conn:analytics:role:metadata".to_string(), PoolKind::Sqlite(pool.clone()));
|
||
conns.insert("conn:analytics:session:tab-1".to_string(), PoolKind::Sqlite(pool.clone()));
|
||
conns.insert("conn:billing".to_string(), PoolKind::Sqlite(pool));
|
||
}
|
||
|
||
assert!(state.close_database_pool("conn", Some("analytics")).await.unwrap());
|
||
|
||
let conns = state.connections.read().await;
|
||
assert!(conns.contains_key("conn"));
|
||
assert!(!conns.contains_key("conn:analytics"));
|
||
assert!(!conns.contains_key("conn:analytics:role:metadata"));
|
||
assert!(!conns.contains_key("conn:analytics:session:tab-1"));
|
||
assert!(conns.contains_key("conn:billing"));
|
||
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
fn ssh_layer(id: &str, profile_id: &str) -> SshTunnelConfig {
|
||
SshTunnelConfig {
|
||
id: id.to_string(),
|
||
name: String::new(),
|
||
enabled: true,
|
||
host: String::new(),
|
||
port: 22,
|
||
user: String::new(),
|
||
password: String::new(),
|
||
key_path: String::new(),
|
||
key_passphrase: String::new(),
|
||
connect_timeout_secs: 5,
|
||
expose_lan: false,
|
||
use_ssh_agent: false,
|
||
ssh_agent_sock_path: String::new(),
|
||
auth_method: String::new(),
|
||
allow_exec_channel_proxy: false,
|
||
profile_id: profile_id.to_string(),
|
||
}
|
||
}
|
||
|
||
fn proxy_layer(id: &str, profile_id: &str) -> ProxyTunnelConfig {
|
||
ProxyTunnelConfig {
|
||
id: id.to_string(),
|
||
name: String::new(),
|
||
enabled: true,
|
||
proxy_type: ProxyType::Socks5,
|
||
host: String::new(),
|
||
port: 1080,
|
||
username: String::new(),
|
||
password: String::new(),
|
||
test_target: None,
|
||
profile_id: profile_id.to_string(),
|
||
}
|
||
}
|
||
|
||
fn http_tunnel_layer(id: &str, profile_id: &str) -> HttpTunnelConfig {
|
||
HttpTunnelConfig {
|
||
id: id.to_string(),
|
||
name: String::new(),
|
||
enabled: true,
|
||
url: String::new(),
|
||
token: String::new(),
|
||
connect_timeout_secs: 10,
|
||
profile_id: profile_id.to_string(),
|
||
}
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn oracle_tns_connection_rejects_transport_layers() {
|
||
let (state, dir) = test_app_state().await;
|
||
let mut config = mysql_config(Some("DBX_FAILOVER"));
|
||
config.db_type = DatabaseType::Oracle;
|
||
config.oracle_connection_type = Some("tns".to_string());
|
||
config.transport_layers = vec![TransportLayerConfig::Ssh(ssh_layer("tns-tunnel", ""))];
|
||
|
||
let error = state.connection_host_port("oracle-tns", &config).await.unwrap_err();
|
||
assert!(error.contains("cannot be combined with SSH, proxy, or HTTP tunnel"));
|
||
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn hive_zookeeper_discovery_rejects_static_transport_layers() {
|
||
let (state, dir) = test_app_state().await;
|
||
let mut config = mysql_config(Some("default"));
|
||
config.db_type = DatabaseType::Hive;
|
||
config.connection_string = Some(
|
||
"jdbc:hive2://zk1.example.com:2181,zk2.example.com:2181/default;serviceDiscoveryMode=zooKeeper".to_string(),
|
||
);
|
||
config.url_params = Some("serviceDiscoveryMode=zooKeeper".to_string());
|
||
config.transport_layers = vec![TransportLayerConfig::Ssh(ssh_layer("hive-zk-tunnel", ""))];
|
||
|
||
let error = state.connection_host_port("hive-zookeeper", &config).await.unwrap_err();
|
||
assert!(error.contains("discovered HiveServer2 nodes would bypass the configured transport"));
|
||
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn resolved_transport_layers_substitutes_shared_profiles() {
|
||
let (state, dir) = test_app_state().await;
|
||
|
||
let mut profile = ssh_layer("shared-bastion", "");
|
||
profile.name = "Bastion".to_string();
|
||
profile.host = "bastion.example.com".to_string();
|
||
profile.user = "deploy".to_string();
|
||
profile.password = "s3cret".to_string();
|
||
profile.auth_method = "password".to_string();
|
||
state.storage.save_tunnel_profiles(&[TransportLayerConfig::Ssh(profile)]).await.unwrap();
|
||
|
||
let mut config = mysql_config(Some("app"));
|
||
config.transport_layers = vec![TransportLayerConfig::Ssh(ssh_layer("layer-1", "shared-bastion"))];
|
||
|
||
let resolved = state.resolved_transport_layers(&config).await.unwrap();
|
||
assert_eq!(resolved.len(), 1);
|
||
match &resolved[0] {
|
||
TransportLayerConfig::Ssh(ssh) => {
|
||
// Profile supplies the configuration; the layer keeps its identity.
|
||
assert_eq!(ssh.id, "layer-1");
|
||
assert_eq!(ssh.profile_id, "shared-bastion");
|
||
assert_eq!(ssh.host, "bastion.example.com");
|
||
assert_eq!(ssh.user, "deploy");
|
||
assert_eq!(ssh.password, "s3cret");
|
||
}
|
||
other => panic!("expected ssh layer, got {other:?}"),
|
||
}
|
||
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn resolved_transport_layers_fails_closed_on_missing_profile() {
|
||
let (state, dir) = test_app_state().await;
|
||
|
||
let mut config = mysql_config(Some("app"));
|
||
config.transport_layers = vec![TransportLayerConfig::Ssh(ssh_layer("layer-1", "deleted-profile"))];
|
||
|
||
let err = state.resolved_transport_layers(&config).await.unwrap_err();
|
||
assert!(err.contains("no longer exists"), "unexpected error: {err}");
|
||
|
||
// Disabled reference layers are filtered out before resolution, so a
|
||
// dangling reference on a disabled layer must not block connecting.
|
||
let mut disabled = ssh_layer("layer-1", "deleted-profile");
|
||
disabled.enabled = false;
|
||
config.transport_layers = vec![TransportLayerConfig::Ssh(disabled)];
|
||
assert!(state.resolved_transport_layers(&config).await.unwrap().is_empty());
|
||
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn resolved_transport_layers_rejects_mismatched_profile_types() {
|
||
let (state, dir) = test_app_state().await;
|
||
|
||
let mismatches = [
|
||
(
|
||
TransportLayerConfig::Ssh(ssh_layer("layer", "shared")),
|
||
TransportLayerConfig::Proxy(proxy_layer("shared", "")),
|
||
),
|
||
(
|
||
TransportLayerConfig::Proxy(proxy_layer("layer", "shared")),
|
||
TransportLayerConfig::HttpTunnel(http_tunnel_layer("shared", "")),
|
||
),
|
||
(
|
||
TransportLayerConfig::HttpTunnel(http_tunnel_layer("layer", "shared")),
|
||
TransportLayerConfig::Ssh(ssh_layer("shared", "")),
|
||
),
|
||
];
|
||
|
||
for (layer, profile) in mismatches {
|
||
state.storage.save_tunnel_profiles(&[profile]).await.unwrap();
|
||
let mut config = mysql_config(Some("app"));
|
||
config.transport_layers = vec![layer];
|
||
|
||
let error = state.resolved_transport_layers(&config).await.unwrap_err();
|
||
assert!(error.contains("different type"));
|
||
}
|
||
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn proxy_connection_uses_local_forward_endpoint() {
|
||
let (state, dir) = test_app_state().await;
|
||
let mut config = mysql_config(Some("app"));
|
||
config.transport_layers = vec![TransportLayerConfig::Proxy(ProxyTunnelConfig {
|
||
profile_id: String::new(),
|
||
id: "proxy".to_string(),
|
||
name: String::new(),
|
||
enabled: true,
|
||
proxy_type: ProxyType::Socks5,
|
||
host: "127.0.0.1".to_string(),
|
||
port: 65000,
|
||
username: String::new(),
|
||
password: String::new(),
|
||
test_target: None,
|
||
})];
|
||
|
||
let (host, _port) = state.connection_host_port("proxied", &config).await.unwrap();
|
||
|
||
assert_eq!(host, "127.0.0.1");
|
||
state.proxy_tunnels.stop_tunnel("proxied:transport:0").await;
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[test]
|
||
fn mq_remote_endpoint_comes_from_admin_url_when_host_fields_are_empty() {
|
||
let mut config = mysql_config(None);
|
||
config.db_type = DatabaseType::MessageQueue;
|
||
config.host = String::new();
|
||
config.port = 0;
|
||
config.external_config = Some(serde_json::json!({
|
||
"systemKind": "pulsar",
|
||
"adminUrl": "https://broker.internal:8443/pulsar-admin?tenant=public",
|
||
"auth": { "kind": "none" }
|
||
}));
|
||
|
||
assert_eq!(connection_remote_endpoint(&config), ("broker.internal".to_string(), 8443));
|
||
}
|
||
|
||
#[test]
|
||
fn kafka_loopback_bootstrap_endpoint_requires_one_ipv4_loopback_server() {
|
||
assert_eq!(
|
||
kafka_single_loopback_bootstrap_endpoint(&serde_json::json!({
|
||
"bootstrapServers": "127.0.0.1:9093"
|
||
})),
|
||
Some(("127.0.0.1".to_string(), 9093))
|
||
);
|
||
assert_eq!(
|
||
kafka_single_loopback_bootstrap_endpoint(&serde_json::json!({
|
||
"bootstrap_servers": "PLAINTEXT://127.0.0.1:19093"
|
||
})),
|
||
Some(("127.0.0.1".to_string(), 19093))
|
||
);
|
||
assert_eq!(
|
||
kafka_single_loopback_bootstrap_endpoint(&serde_json::json!({
|
||
"bootstrapServers": "127.0.0.1:9093,127.0.0.1:9094"
|
||
})),
|
||
None
|
||
);
|
||
assert_eq!(
|
||
kafka_single_loopback_bootstrap_endpoint(&serde_json::json!({
|
||
"bootstrapServers": "broker.internal:9093"
|
||
})),
|
||
None
|
||
);
|
||
}
|
||
|
||
#[cfg(feature = "mq-admin")]
|
||
#[tokio::test]
|
||
async fn mq_admin_config_preserves_admin_url_and_uses_forwarded_connect_override() {
|
||
let (state, dir) = test_app_state().await;
|
||
let mut config = mysql_config(None);
|
||
config.id = "proxied-mq".to_string();
|
||
config.db_type = DatabaseType::MessageQueue;
|
||
config.host = String::new();
|
||
config.port = 0;
|
||
config.external_config = Some(serde_json::json!({
|
||
"systemKind": "pulsar",
|
||
"adminUrl": "https://broker.internal:8443/pulsar-admin?tenant=public",
|
||
"auth": { "kind": "none" }
|
||
}));
|
||
config.transport_layers = vec![TransportLayerConfig::Proxy(ProxyTunnelConfig {
|
||
profile_id: String::new(),
|
||
id: "proxy".to_string(),
|
||
name: String::new(),
|
||
enabled: true,
|
||
proxy_type: ProxyType::Socks5,
|
||
host: "127.0.0.1".to_string(),
|
||
port: 65000,
|
||
username: String::new(),
|
||
password: String::new(),
|
||
test_target: None,
|
||
})];
|
||
|
||
let mqc = state.mq_admin_config_for_connection("proxied-mq", &config).await.unwrap();
|
||
|
||
assert_eq!(mqc.admin_url, "https://broker.internal:8443/pulsar-admin?tenant=public");
|
||
let connect_override = mqc.connect_override.expect("MQ transport should set a connect override");
|
||
assert_eq!(connect_override.host, "127.0.0.1");
|
||
assert_ne!(connect_override.port, 8443);
|
||
state.proxy_tunnels.stop_tunnel("proxied-mq:transport:0").await;
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[cfg(feature = "mq-admin")]
|
||
#[tokio::test]
|
||
async fn rocketmq_transport_passes_socks_proxy_to_multi_endpoint_client() {
|
||
let (state, dir) = test_app_state().await;
|
||
let mut config = mysql_config(None);
|
||
config.id = "proxied-rocketmq".to_string();
|
||
config.db_type = DatabaseType::MessageQueue;
|
||
config.host = "172.19.191.166".to_string();
|
||
config.port = 9876;
|
||
config.external_config = Some(serde_json::json!({
|
||
"systemKind": "rocketmq",
|
||
"adminUrl": "",
|
||
"auth": { "kind": "none" },
|
||
"extra": { "namesrvAddr": "172.19.191.166:9876" }
|
||
}));
|
||
config.transport_layers = vec![TransportLayerConfig::Proxy(ProxyTunnelConfig {
|
||
profile_id: String::new(),
|
||
id: "proxy".to_string(),
|
||
name: String::new(),
|
||
enabled: true,
|
||
proxy_type: ProxyType::Socks5,
|
||
host: "proxy.internal".to_string(),
|
||
port: 1080,
|
||
username: "proxy-user".to_string(),
|
||
password: "proxy-secret".to_string(),
|
||
test_target: None,
|
||
})];
|
||
|
||
let endpoint = state.connection_host_port("proxied-rocketmq", &config).await.unwrap();
|
||
assert_eq!(endpoint, ("172.19.191.166".to_string(), 9876));
|
||
assert!(state.proxy_tunnels.local_port("proxied-rocketmq:transport:0").await.is_none());
|
||
|
||
let mqc = state.mq_admin_config_for_connection("proxied-rocketmq", &config).await.unwrap();
|
||
|
||
let socks_proxy = mqc.socks_proxy.expect("RocketMQ transport should configure SOCKS5 routing");
|
||
assert_eq!(socks_proxy.host, "proxy.internal");
|
||
assert_eq!(socks_proxy.port, 1080);
|
||
assert_eq!(socks_proxy.username, "proxy-user");
|
||
assert_eq!(socks_proxy.password, "proxy-secret");
|
||
assert!(mqc.connect_override.is_none());
|
||
assert!(state.proxy_tunnels.local_port("proxied-rocketmq:transport:0").await.is_none());
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[cfg(feature = "mq-admin")]
|
||
#[tokio::test]
|
||
async fn rabbitmq_transport_uses_separate_amqp_and_management_tunnels() {
|
||
let (state, dir) = test_app_state().await;
|
||
let mut config = mysql_config(None);
|
||
config.id = "proxied-rabbitmq".to_string();
|
||
config.db_type = DatabaseType::MessageQueue;
|
||
config.host = "rabbit.internal".to_string();
|
||
config.port = 5672;
|
||
config.external_config = Some(serde_json::json!({
|
||
"systemKind": "rabbitmq",
|
||
"adminUrl": "http://management.internal:15672/rmq",
|
||
"auth": { "kind": "none" },
|
||
"extra": {
|
||
"addresses": "rabbit.internal:5672",
|
||
"virtualHost": "/"
|
||
}
|
||
}));
|
||
config.transport_layers = vec![TransportLayerConfig::Proxy(ProxyTunnelConfig {
|
||
profile_id: String::new(),
|
||
id: "proxy".to_string(),
|
||
name: String::new(),
|
||
enabled: true,
|
||
proxy_type: ProxyType::Socks5,
|
||
host: "127.0.0.1".to_string(),
|
||
port: 65000,
|
||
username: String::new(),
|
||
password: String::new(),
|
||
test_target: None,
|
||
})];
|
||
|
||
let mqc = state.mq_admin_config_for_connection("proxied-rabbitmq", &config).await.unwrap();
|
||
let amqp_override = mqc.connect_override.expect("RabbitMQ AMQP tunnel override");
|
||
let management_override = mqc.management_connect_override.expect("RabbitMQ Management tunnel override");
|
||
assert_eq!(amqp_override.host, "127.0.0.1");
|
||
assert_eq!(management_override.host, "127.0.0.1");
|
||
assert_ne!(amqp_override.port, management_override.port);
|
||
assert_eq!(state.proxy_tunnels.local_port("proxied-rabbitmq:transport:0").await, Some(amqp_override.port));
|
||
assert_eq!(
|
||
state.proxy_tunnels.local_port("proxied-rabbitmq:rabbitmq-management:transport:0").await,
|
||
Some(management_override.port)
|
||
);
|
||
|
||
state.reset_connection_transport_for_config("proxied-rabbitmq", &config).await;
|
||
assert!(state.proxy_tunnels.local_port("proxied-rabbitmq:transport:0").await.is_none());
|
||
assert!(state.proxy_tunnels.local_port("proxied-rabbitmq:rabbitmq-management:transport:0").await.is_none());
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[cfg(feature = "mq-admin")]
|
||
#[tokio::test]
|
||
async fn rabbitmq_transport_does_not_reuse_prestarted_management_tunnel_for_amqp() {
|
||
let (state, dir) = test_app_state().await;
|
||
let mut config = mysql_config(None);
|
||
config.id = "proxied-rabbitmq-custom-ports".to_string();
|
||
config.db_type = DatabaseType::MessageQueue;
|
||
config.host = "127.0.0.1".to_string();
|
||
config.port = 35672;
|
||
config.external_config = Some(serde_json::json!({
|
||
"systemKind": "rabbitmq",
|
||
"adminUrl": "http://127.0.0.1:35673",
|
||
"auth": { "kind": "none" },
|
||
"extra": {
|
||
"addresses": "127.0.0.1:35672",
|
||
"virtualHost": "/"
|
||
}
|
||
}));
|
||
config.transport_layers = vec![TransportLayerConfig::Proxy(ProxyTunnelConfig {
|
||
profile_id: String::new(),
|
||
id: "proxy".to_string(),
|
||
name: String::new(),
|
||
enabled: true,
|
||
proxy_type: ProxyType::Socks5,
|
||
host: "127.0.0.1".to_string(),
|
||
port: 65000,
|
||
username: String::new(),
|
||
password: String::new(),
|
||
test_target: None,
|
||
})];
|
||
|
||
assert_eq!(connection_remote_endpoint(&config), ("127.0.0.1".to_string(), 35672));
|
||
|
||
let prestarted_amqp_port =
|
||
state.connection_host_port("proxied-rabbitmq-custom-ports", &config).await.unwrap().1;
|
||
let mqc = state.mq_admin_config_for_connection("proxied-rabbitmq-custom-ports", &config).await.unwrap();
|
||
let amqp_override = mqc.connect_override.expect("RabbitMQ AMQP tunnel override");
|
||
let management_override = mqc.management_connect_override.expect("RabbitMQ Management tunnel override");
|
||
|
||
assert_eq!(amqp_override.port, prestarted_amqp_port);
|
||
assert_ne!(amqp_override.port, management_override.port);
|
||
|
||
state.reset_connection_transport_for_config("proxied-rabbitmq-custom-ports", &config).await;
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn nacos_admin_config_allows_domain_server_addr_without_transport_override() {
|
||
let (state, dir) = test_app_state().await;
|
||
let mut config = mysql_config(None);
|
||
config.id = "aliyun-nacos".to_string();
|
||
config.db_type = DatabaseType::Nacos;
|
||
config.host = "example.com".to_string();
|
||
config.port = 8848;
|
||
config.external_config = Some(serde_json::json!({
|
||
"serverAddr": "https://nacos.aliyuncs.com:8848",
|
||
"namespace": "public",
|
||
"contextPath": "/nacos",
|
||
"auth": { "kind": "none" }
|
||
}));
|
||
|
||
let nacos_config = state.nacos_admin_config_for_connection("aliyun-nacos", &config).await.unwrap();
|
||
|
||
assert_eq!(nacos_config.server_addr, "https://nacos.aliyuncs.com:8848");
|
||
assert!(nacos_config.connect_override.is_none());
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn nacos_admin_config_rewrites_server_addr_to_forwarded_endpoint() {
|
||
let (state, dir) = test_app_state().await;
|
||
let mut config = mysql_config(None);
|
||
config.id = "proxied-nacos".to_string();
|
||
config.db_type = DatabaseType::Nacos;
|
||
config.host = "192.168.2.51".to_string();
|
||
config.port = 10840;
|
||
config.external_config = Some(serde_json::json!({
|
||
"serverAddr": "http://192.168.2.51:10840",
|
||
"rnacosConsoleAddr": "http://192.168.2.51:10848",
|
||
"namespace": "public",
|
||
"contextPath": "",
|
||
"auth": { "kind": "none" }
|
||
}));
|
||
config.transport_layers = vec![TransportLayerConfig::Proxy(ProxyTunnelConfig {
|
||
profile_id: String::new(),
|
||
id: "proxy".to_string(),
|
||
name: String::new(),
|
||
enabled: true,
|
||
proxy_type: ProxyType::Socks5,
|
||
host: "127.0.0.1".to_string(),
|
||
port: 65000,
|
||
username: String::new(),
|
||
password: String::new(),
|
||
test_target: None,
|
||
})];
|
||
|
||
let nacos_config = state.nacos_admin_config_for_connection("proxied-nacos", &config).await.unwrap();
|
||
|
||
assert!(nacos_config.server_addr.starts_with("http://127.0.0.1:"));
|
||
assert_ne!(nacos_config.server_addr, "http://192.168.2.51:10840");
|
||
assert!(nacos_config.rnacos_console_addr.starts_with("http://127.0.0.1:"));
|
||
assert_ne!(nacos_config.rnacos_console_addr, "http://192.168.2.51:10848");
|
||
assert!(nacos_config.connect_override.is_none());
|
||
state.proxy_tunnels.stop_tunnel("proxied-nacos:transport:0").await;
|
||
state.proxy_tunnels.stop_tunnel("proxied-nacos:rnacos-console:transport:0").await;
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
|
||
#[tokio::test]
|
||
#[ignore = "requires a reachable GaussDB instance via environment variables"]
|
||
async fn live_gaussdb_native_connection_succeeds() {
|
||
let host = std::env::var("DBX_TEST_GAUSSDB_HOST").expect("DBX_TEST_GAUSSDB_HOST not set");
|
||
let port = std::env::var("DBX_TEST_GAUSSDB_PORT")
|
||
.expect("DBX_TEST_GAUSSDB_PORT not set")
|
||
.parse::<u16>()
|
||
.expect("DBX_TEST_GAUSSDB_PORT should be a u16");
|
||
let username = std::env::var("DBX_TEST_GAUSSDB_USER").expect("DBX_TEST_GAUSSDB_USER not set");
|
||
let password = std::env::var("DBX_TEST_GAUSSDB_PASSWORD").expect("DBX_TEST_GAUSSDB_PASSWORD not set");
|
||
let url_params = std::env::var("DBX_TEST_GAUSSDB_URL_PARAMS").ok();
|
||
|
||
assert_live_postgres_like_query(live_postgres_like_config(
|
||
DatabaseType::Gaussdb,
|
||
&host,
|
||
port,
|
||
&username,
|
||
&password,
|
||
url_params.as_deref(),
|
||
))
|
||
.await;
|
||
}
|
||
|
||
#[tokio::test]
|
||
#[ignore = "requires a reachable openGauss instance via environment variables"]
|
||
async fn live_opengauss_native_connection_succeeds() {
|
||
let host = std::env::var("DBX_TEST_OPENGAUSS_HOST").expect("DBX_TEST_OPENGAUSS_HOST not set");
|
||
let port = std::env::var("DBX_TEST_OPENGAUSS_PORT")
|
||
.expect("DBX_TEST_OPENGAUSS_PORT not set")
|
||
.parse::<u16>()
|
||
.expect("DBX_TEST_OPENGAUSS_PORT should be a u16");
|
||
let username = std::env::var("DBX_TEST_OPENGAUSS_USER").expect("DBX_TEST_OPENGAUSS_USER not set");
|
||
let password = std::env::var("DBX_TEST_OPENGAUSS_PASSWORD").expect("DBX_TEST_OPENGAUSS_PASSWORD not set");
|
||
let url_params = std::env::var("DBX_TEST_OPENGAUSS_URL_PARAMS").ok();
|
||
|
||
assert_live_postgres_like_query(live_postgres_like_config(
|
||
DatabaseType::OpenGauss,
|
||
&host,
|
||
port,
|
||
&username,
|
||
&password,
|
||
url_params.as_deref(),
|
||
))
|
||
.await;
|
||
}
|
||
|
||
#[tokio::test]
|
||
#[ignore = "requires a reachable KWDB instance via environment variables"]
|
||
async fn live_kwdb_native_connection_succeeds() {
|
||
let host = std::env::var("DBX_TEST_KWDB_HOST").expect("DBX_TEST_KWDB_HOST not set");
|
||
let port = std::env::var("DBX_TEST_KWDB_PORT")
|
||
.expect("DBX_TEST_KWDB_PORT not set")
|
||
.parse::<u16>()
|
||
.expect("DBX_TEST_KWDB_PORT should be a u16");
|
||
let username = std::env::var("DBX_TEST_KWDB_USER").unwrap_or_else(|_| "root".to_string());
|
||
let password = std::env::var("DBX_TEST_KWDB_PASSWORD").unwrap_or_default();
|
||
let database = std::env::var("DBX_TEST_KWDB_DATABASE").unwrap_or_else(|_| "defaultdb".to_string());
|
||
let url_params = std::env::var("DBX_TEST_KWDB_URL_PARAMS").unwrap_or_else(|_| "sslmode=disable".to_string());
|
||
|
||
let mut config = mysql_config(Some(&database));
|
||
config.id = "kwdb-live".to_string();
|
||
config.db_type = DatabaseType::Kwdb;
|
||
config.host = host;
|
||
config.port = port;
|
||
config.username = username;
|
||
config.password = password;
|
||
config.url_params = Some(url_params);
|
||
|
||
let (state, dir) = test_app_state().await;
|
||
state.configs.write().await.insert(config.id.clone(), config);
|
||
let pool_key = state.get_or_create_pool("kwdb-live", None).await.unwrap();
|
||
let pool = {
|
||
let connections = state.connections.read().await;
|
||
match connections.get(&pool_key).expect("KWDB pool should be created") {
|
||
PoolKind::Postgres(pool) => pool.clone(),
|
||
_ => panic!("KWDB should use the PostgreSQL pool path"),
|
||
}
|
||
};
|
||
let result = query::execute_sql_statement(
|
||
&state,
|
||
"kwdb-live",
|
||
&database,
|
||
"SELECT current_database(), current_schema()",
|
||
None,
|
||
None,
|
||
)
|
||
.await
|
||
.unwrap_or_else(|err| panic!("failed to query live KWDB: {err}"));
|
||
assert_eq!(result.rows.len(), 1);
|
||
let database_column_index = result
|
||
.columns
|
||
.iter()
|
||
.position(|column| column == "current_database")
|
||
.expect("current_database column should be present");
|
||
assert_eq!(result.rows[0].get(database_column_index).and_then(|value| value.as_str()), Some(database.as_str()));
|
||
let databases = schema::list_databases_core(&state, "kwdb-live").await.unwrap();
|
||
assert!(databases.iter().any(|database| database.name == "defaultdb"));
|
||
let test_schema = "dbx_kwdb_live";
|
||
db::postgres::execute_query(&pool, &format!("DROP SCHEMA IF EXISTS {test_schema} CASCADE"))
|
||
.await
|
||
.unwrap_or_else(|err| panic!("failed to clean KWDB test schema: {err}"));
|
||
query::execute_sql_statement(
|
||
&state,
|
||
"kwdb-live",
|
||
&database,
|
||
&format!("CREATE SCHEMA {test_schema}"),
|
||
None,
|
||
None,
|
||
)
|
||
.await
|
||
.unwrap_or_else(|err| panic!("failed to create KWDB test schema: {err}"));
|
||
query::execute_sql_statement(
|
||
&state,
|
||
"kwdb-live",
|
||
&database,
|
||
"CREATE TABLE devices (id INT PRIMARY KEY, name STRING, active BOOL)",
|
||
Some(test_schema),
|
||
None,
|
||
)
|
||
.await
|
||
.unwrap_or_else(|err| panic!("failed to create KWDB test table: {err}"));
|
||
query::execute_sql_statement(
|
||
&state,
|
||
"kwdb-live",
|
||
&database,
|
||
"INSERT INTO devices (id, name, active) VALUES (1, 'meter-a', true)",
|
||
Some(test_schema),
|
||
None,
|
||
)
|
||
.await
|
||
.unwrap_or_else(|err| panic!("failed to insert KWDB test row: {err}"));
|
||
let query_result = query::execute_sql_statement(
|
||
&state,
|
||
"kwdb-live",
|
||
&database,
|
||
"SELECT name, active FROM devices WHERE id = 1",
|
||
Some(test_schema),
|
||
None,
|
||
)
|
||
.await
|
||
.unwrap_or_else(|err| panic!("failed to query KWDB test row: {err}"));
|
||
assert_eq!(query_result.rows.len(), 1);
|
||
assert_eq!(query_result.rows[0].first().and_then(|value| value.as_str()), Some("meter-a"));
|
||
|
||
let schemas = schema::list_schemas_core(&state, "kwdb-live", &database).await.unwrap();
|
||
assert!(schemas.iter().any(|schema| schema == test_schema));
|
||
let tables =
|
||
schema::list_tables_core(&state, "kwdb-live", &database, test_schema, None, None, None, None, None)
|
||
.await
|
||
.unwrap();
|
||
assert!(tables.iter().any(|table| table.name == "devices" && table.table_type == "BASE TABLE"));
|
||
let columns = schema::get_columns_core(&state, "kwdb-live", &database, test_schema, "devices").await.unwrap();
|
||
let id_column = columns.iter().find(|column| column.name == "id").expect("id column should be listed");
|
||
assert!(id_column.data_type.to_lowercase().contains("int"));
|
||
let name_column = columns.iter().find(|column| column.name == "name").expect("name column should be listed");
|
||
assert!(name_column.data_type.to_lowercase().contains("text"));
|
||
let active_column =
|
||
columns.iter().find(|column| column.name == "active").expect("active column should be listed");
|
||
assert!(active_column.data_type.to_lowercase().contains("bool"));
|
||
db::postgres::execute_query(&pool, &format!("DROP SCHEMA {test_schema} CASCADE"))
|
||
.await
|
||
.unwrap_or_else(|err| panic!("failed to drop KWDB test schema: {err}"));
|
||
pool.close();
|
||
let _ = std::fs::remove_dir_all(dir);
|
||
}
|
||
}
|