diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index cf1d99479..2268220e8 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -289,6 +289,21 @@ pub async fn connect_mysql_metadata_pool( ) .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) } @@ -327,6 +342,23 @@ pub async fn connect_mysql_metadata_pool( .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) } @@ -426,6 +458,40 @@ async fn connect_bare_mysql_pool_with_setup( } } +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 { + // 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_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, diff --git a/crates/dbx-core/src/db/mysql.rs b/crates/dbx-core/src/db/mysql.rs index 3dacbb16b..9f501cab6 100644 --- a/crates/dbx-core/src/db/mysql.rs +++ b/crates/dbx-core/src/db/mysql.rs @@ -425,17 +425,44 @@ pub async fn connect_with_ca_cert_pool_limit_idle_and_setup( max_connections: usize, idle_timeout_secs: Option, extra_setup_queries: &[String], +) -> Result { + connect_with_ca_cert_pool_limit_idle_and_setup_database( + url, + ca_cert_path, + fallback_timeout, + max_connections, + idle_timeout_secs, + None, + extra_setup_queries, + ) + .await +} + +pub async fn connect_with_ca_cert_pool_limit_idle_and_setup_database( + url: &str, + ca_cert_path: Option<&str>, + fallback_timeout: Duration, + max_connections: usize, + idle_timeout_secs: Option, + setup_database: Option<&str>, + extra_setup_queries: &[String], ) -> Result { let timeout = super::parse_connect_timeout_with_fallback(url, fallback_timeout); - let pool = create_pool(url, ca_cert_path, max_connections, idle_timeout_secs, extra_setup_queries)?; + let pool = create_pool(url, ca_cert_path, max_connections, idle_timeout_secs, setup_database, extra_setup_queries)?; 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, max_connections, idle_timeout_secs, extra_setup_queries)?; + let fallback_pool = create_pool( + &fallback_url, + None, + max_connections, + idle_timeout_secs, + setup_database, + extra_setup_queries, + )?; return match verify_pool_connection(&fallback_pool, timeout).await { Ok(()) => Ok(fallback_pool), Err(e) => Err(e), @@ -458,6 +485,7 @@ fn create_pool( ca_cert_path: Option<&str>, max_connections: usize, idle_timeout_secs: Option, + setup_database: Option<&str>, extra_setup_queries: &[String], ) -> Result { let tls_url = mysql_tls_url(url)?; @@ -474,12 +502,16 @@ fn create_pool( .with_constraints(mysql_async::PoolConstraints::new(1, max_connections).unwrap()) .with_inactive_connection_ttl(inactive_ttl) .with_reset_connection(max_connections > 1); + let setup_queries = match setup_database { + Some(database) => mysql_setup_queries_for_database(url, Some(database), extra_setup_queries), + None => mysql_setup_queries(url, extra_setup_queries), + }; let mut builder = mysql_async::OptsBuilder::from_opts(opts) .stmt_cache_size(0) .prefer_socket(false) .pool_opts(Some(pool_opts)) .tcp_keepalive(Some(MYSQL_TCP_KEEPALIVE_MS)) - .setup(mysql_setup_queries(url, extra_setup_queries)); + .setup(setup_queries); if let Some(ssl_opts) = mysql_ssl_opts(base_ssl_opts, url, ca_cert_path, &tls_url.files)? { builder = builder.ssl_opts(ssl_opts); } @@ -584,9 +616,17 @@ fn mysql_ssl_opts( } fn mysql_setup_queries(url: &str, extra_setup_queries: &[String]) -> Vec { + mysql_setup_queries_for_database(url, None, extra_setup_queries) +} + +fn mysql_setup_queries_for_database( + url: &str, + setup_database: Option<&str>, + extra_setup_queries: &[String], +) -> Vec { let charset = mysql_connection_charset(url).unwrap_or("utf8mb4"); let catalog = mysql_connection_catalog(url); - let database = mysql_connection_database(url); + let database = setup_database.map(ToOwned::to_owned).or_else(|| mysql_connection_database(url)); let mut queries = Vec::new(); if let Some(database) = database.as_deref() { queries.push(format!("USE {}", quote_identifier(database))); @@ -1144,9 +1184,20 @@ pub async fn connect_bare_with_pool_limit_and_setup( fallback_timeout: Duration, max_connections: usize, extra_setup_queries: &[String], +) -> Result { + connect_bare_with_pool_limit_and_setup_database(url, fallback_timeout, max_connections, None, extra_setup_queries) + .await +} + +pub async fn connect_bare_with_pool_limit_and_setup_database( + url: &str, + fallback_timeout: Duration, + max_connections: usize, + setup_database: Option<&str>, + extra_setup_queries: &[String], ) -> Result { let timeout = super::parse_connect_timeout_with_fallback(url, fallback_timeout); - let pool = create_pool(url, None, max_connections, None, extra_setup_queries)?; + let pool = create_pool(url, None, max_connections, None, setup_database, extra_setup_queries)?; verify_pool_connection(&pool, timeout).await.map(|_| pool) } @@ -3729,6 +3780,17 @@ UNIQUE KEY(`tenant_id`, `name``part`) assert_eq!(queries, vec!["USE `db/name`", "SET NAMES utf8mb4"]); } + #[test] + fn mysql_setup_queries_can_select_database_without_url_path() { + let queries = mysql_setup_queries_for_database( + "mysql://root:secret@localhost:3306?charset=utf8mb4", + Some("app`proxy"), + &[], + ); + + assert_eq!(queries, vec!["USE `app``proxy`", "SET NAMES utf8mb4"]); + } + #[test] fn mysql_datetime_utc_values_display_without_rfc3339_offset() { let value = NaiveDateTime::new(