diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index d12802cdc3..19a721e238 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -97,10 +97,11 @@ pub async fn connect_mysql_metadata_pool( 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); if db_config.needs_bare_mysql() { - return match db::mysql::connect_bare(&url, connect_timeout).await { + return match db::mysql::connect_bare_with_pool_limit(&url, connect_timeout, max_connections).await { Ok(pool) => Ok((pool, MysqlMode::Bare)), Err(err) => { let fallback_url = mysql_metadata_fallback_url(config, db_config, host, port); @@ -108,7 +109,9 @@ pub async fn connect_mysql_metadata_pool( log::info!( "MySQL metadata connection without a default database failed ({err}); retrying with configured default database." ); - db::mysql::connect_bare(&fallback_url, connect_timeout).await.map(|pool| (pool, MysqlMode::Bare)) + db::mysql::connect_bare_with_pool_limit(&fallback_url, connect_timeout, max_connections) + .await + .map(|pool| (pool, MysqlMode::Bare)) } else { Err(err) } @@ -116,7 +119,14 @@ pub async fn connect_mysql_metadata_pool( }; } - match db::mysql::connect_with_ca_cert(&url, Some(&db_config.ca_cert_path), connect_timeout).await { + match db::mysql::connect_with_ca_cert_and_pool_limit( + &url, + Some(&db_config.ca_cert_path), + connect_timeout, + max_connections, + ) + .await + { Ok(pool) => { let mode = detect_ob_oracle_mode(config, &pool).await; Ok((pool, mode)) @@ -127,8 +137,13 @@ pub async fn connect_mysql_metadata_pool( log::info!( "MySQL metadata connection without a default database failed ({err}); retrying with configured default database." ); - let pool = - db::mysql::connect_with_ca_cert(&fallback_url, Some(&config.ca_cert_path), connect_timeout).await?; + let pool = db::mysql::connect_with_ca_cert_and_pool_limit( + &fallback_url, + Some(&config.ca_cert_path), + connect_timeout, + max_connections, + ) + .await?; let mode = detect_ob_oracle_mode(config, &pool).await; Ok((pool, mode)) } else { @@ -143,21 +158,22 @@ pub async fn connect_bare_metadata_pool( host: &str, port: u16, connect_timeout: std::time::Duration, + max_connections: usize, ) -> Result { let url = connection_url_for_endpoint(db_config, host, port); if db_config.effective_database().is_none() { - return db::mysql::connect_bare(&url, connect_timeout).await; + return db::mysql::connect_bare_with_pool_limit(&url, connect_timeout, max_connections).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 db::mysql::connect_bare(&url, connect_timeout).await; + return db::mysql::connect_bare_with_pool_limit(&url, connect_timeout, max_connections).await; } - let preferred = db::mysql::connect_bare(&url, connect_timeout); - let unscoped = db::mysql::connect_bare(&unscoped_url, connect_timeout); + let preferred = db::mysql::connect_bare_with_pool_limit(&url, connect_timeout, max_connections); + let unscoped = db::mysql::connect_bare_with_pool_limit(&unscoped_url, connect_timeout, max_connections); tokio::pin!(preferred); tokio::pin!(unscoped); @@ -315,17 +331,26 @@ impl AppState { probe_connection_endpoint(&db_config, &host, port).await?; 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 mysql_pool_max_connections = if normalize_client_session_id(client_session_id).is_some() { 1 } else { 3 }; let pool = match db_config.db_type { DatabaseType::Mysql => { - let (pool, mode) = - connect_mysql_metadata_pool(&config, &db_config, &host, port, connect_timeout).await?; + 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::Databend => { let pool = if database.is_none() { - connect_bare_metadata_pool(&db_config, &host, port, connect_timeout).await? + connect_bare_metadata_pool(&db_config, &host, port, connect_timeout, mysql_pool_max_connections) + .await? } else { - db::mysql::connect_bare(&url, connect_timeout).await? + db::mysql::connect_bare_with_pool_limit(&url, connect_timeout, mysql_pool_max_connections).await? }; PoolKind::Mysql(pool, MysqlMode::Bare) } diff --git a/crates/dbx-core/src/db/mysql.rs b/crates/dbx-core/src/db/mysql.rs index b376ad5957..15713d2846 100644 --- a/crates/dbx-core/src/db/mysql.rs +++ b/crates/dbx-core/src/db/mysql.rs @@ -332,16 +332,25 @@ pub async fn connect_with_ca_cert( url: &str, ca_cert_path: Option<&str>, fallback_timeout: Duration, +) -> Result { + connect_with_ca_cert_and_pool_limit(url, ca_cert_path, fallback_timeout, 3).await +} + +pub async fn connect_with_ca_cert_and_pool_limit( + url: &str, + ca_cert_path: Option<&str>, + fallback_timeout: Duration, + max_connections: usize, ) -> Result { let timeout = super::parse_connect_timeout_with_fallback(url, fallback_timeout); - let pool = create_pool(url, ca_cert_path)?; + let pool = create_pool(url, ca_cert_path, max_connections)?; let result = verify_pool_connection(&pool, timeout).await; if let Err(ref e) = result { if mysql_error_should_retry_without_ssl(e) { if let Some(fallback_url) = ssl_fallback_url(url) { log::info!("SSL handshake failed, retrying with ssl-mode=disabled"); - let fallback_pool = create_pool(&fallback_url, None)?; + let fallback_pool = create_pool(&fallback_url, None, max_connections)?; return match verify_pool_connection(&fallback_pool, timeout).await { Ok(()) => Ok(fallback_pool), Err(e) => Err(e), @@ -359,13 +368,14 @@ struct MySqlTlsFiles { sslkey: Option, } -fn create_pool(url: &str, ca_cert_path: Option<&str>) -> Result { +fn create_pool(url: &str, ca_cert_path: Option<&str>, max_connections: usize) -> Result { let tls_url = mysql_tls_url(url)?; let opts = mysql_async::Opts::from_url(&mysql_async_url(&tls_url.url)).map_err(|e| format!("Invalid MySQL URL: {e}"))?; let base_ssl_opts = opts.ssl_opts().cloned(); + let max_connections = max_connections.max(1); let pool_opts = mysql_async::PoolOpts::new() - .with_constraints(mysql_async::PoolConstraints::new(1, 3).unwrap()) + .with_constraints(mysql_async::PoolConstraints::new(1, max_connections).unwrap()) .with_inactive_connection_ttl(Duration::from_secs(300)); let mut builder = mysql_async::OptsBuilder::from_opts(opts) .stmt_cache_size(0) @@ -747,8 +757,16 @@ fn mysql_async_url(url: &str) -> Cow<'_, str> { } pub async fn connect_bare(url: &str, fallback_timeout: Duration) -> Result { + connect_bare_with_pool_limit(url, fallback_timeout, 3).await +} + +pub async fn connect_bare_with_pool_limit( + url: &str, + fallback_timeout: Duration, + max_connections: usize, +) -> Result { let timeout = super::parse_connect_timeout_with_fallback(url, fallback_timeout); - let pool = create_pool(url, None)?; + let pool = create_pool(url, None, max_connections)?; verify_pool_connection(&pool, timeout).await.map(|_| pool) } diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index 91f0d189e7..bc619e39ce 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -741,12 +741,11 @@ pub async fn execute_sql_statement_with_options( cancel_token: Option, options: QueryExecutionOptions, ) -> Result { - // When database is not set, fall back to the shared (non-session-scoped) pool - // to avoid creating a connection without a default database context. - // This is particularly important for Doris/StarRocks, where metadata connections - // omit the database and would cause "Current database is not selected" errors. + // When a query tab has a client session, keep even database-less execution + // on that tab-scoped pool so connection-level state (for example MySQL @vars) + // survives across runs. let pool_key = if database.is_empty() { - state.get_or_create_pool(connection_id, None).await? + state.get_or_create_pool_for_session(connection_id, None, options.client_session_id.as_deref()).await? } else { state .get_or_create_pool_for_session(connection_id, Some(database), options.client_session_id.as_deref()) @@ -766,7 +765,7 @@ pub async fn execute_sql_statement_with_options( Err(e) if is_connection_error(e) && !is_canceled(&cancel_token) => { let db_opt = if database.is_empty() { None } else { Some(database) }; let new_key = if database.is_empty() { - state.reconnect_pool(connection_id, db_opt).await? + state.reconnect_pool_for_session(connection_id, db_opt, options.client_session_id.as_deref()).await? } else { state.reconnect_pool_for_session(connection_id, db_opt, options.client_session_id.as_deref()).await? }; @@ -784,7 +783,7 @@ pub async fn close_query_session( client_session_id: Option<&str>, ) -> Result { let pool_key = if database.is_empty() { - state.get_or_create_pool(connection_id, None).await? + state.get_or_create_pool_for_session(connection_id, None, client_session_id).await? } else { state.get_or_create_pool_for_session(connection_id, Some(database), client_session_id).await? }; @@ -832,7 +831,7 @@ pub async fn execute_multi_core_with_options( options: QueryExecutionOptions, ) -> Result, String> { let pool_key = if database.is_empty() { - state.get_or_create_pool(connection_id, None).await? + state.get_or_create_pool_for_session(connection_id, None, options.client_session_id.as_deref()).await? } else { state .get_or_create_pool_for_session(connection_id, Some(database), options.client_session_id.as_deref()) diff --git a/src-tauri/src/commands/connection.rs b/src-tauri/src/commands/connection.rs index 3939ba7600..1cee2c7652 100644 --- a/src-tauri/src/commands/connection.rs +++ b/src-tauri/src/commands/connection.rs @@ -437,11 +437,12 @@ pub async fn connect_db(state: State<'_, Arc>, config: ConnectionConfi let pool = match db_config.db_type { DatabaseType::Mysql => { - let (pool, mode) = connect_mysql_metadata_pool(&config, &db_config, &host, port, connect_timeout).await?; + let (pool, mode) = + connect_mysql_metadata_pool(&config, &db_config, &host, port, connect_timeout, 3).await?; PoolKind::Mysql(pool, mode) } DatabaseType::Doris | DatabaseType::StarRocks => PoolKind::Mysql( - connect_bare_metadata_pool(&db_config, &host, port, connect_timeout).await?, + connect_bare_metadata_pool(&db_config, &host, port, connect_timeout, 3).await?, MysqlMode::Bare, ), DatabaseType::Postgres