diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index 36c52b9ab..d45a1956f 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -1,3 +1,4 @@ +use percent_encoding::percent_decode_str; use std::collections::HashMap; use std::path::PathBuf; use std::sync::Arc; @@ -499,6 +500,11 @@ pub fn redacted_connection_url_for_endpoint(config: &ConnectionConfig, host: &st } pub fn agent_connect_params(config: &ConnectionConfig, host: &str, port: u16, database: &str) -> serde_json::Value { + let agent_database = if config.db_type == DatabaseType::MongoDb { + mongo_agent_database(config, database) + } else { + database.to_string() + }; let connection_string = if config.db_type == DatabaseType::MongoDb { config.connection_url_with_host(host, port) } else if config.db_type == DatabaseType::Oracle { @@ -510,7 +516,7 @@ pub fn agent_connect_params(config: &ConnectionConfig, host: &str, port: u16, da serde_json::json!({ "host": host, "port": port, - "database": database, + "database": agent_database, "username": config.username, "password": config.password, "url_params": config.url_params.as_deref().unwrap_or(""), @@ -518,6 +524,34 @@ pub fn agent_connect_params(config: &ConnectionConfig, host: &str, port: u16, da }) } +fn mongo_agent_database(config: &ConnectionConfig, database: &str) -> String { + if let Some(database) = non_empty_database(database) { + return database.to_string(); + } + if let Some(database) = config.database.as_deref().and_then(non_empty_database) { + return database.to_string(); + } + if let Some(database) = config.connection_string.as_deref().and_then(mongo_uri_database) { + return database; + } + "admin".to_string() +} + +fn non_empty_database(database: &str) -> Option<&str> { + let database = database.trim(); + (!database.is_empty()).then_some(database) +} + +fn mongo_uri_database(uri: &str) -> Option { + let rest = uri.strip_prefix("mongodb://").or_else(|| uri.strip_prefix("mongodb+srv://"))?; + let (_, after_hosts) = rest.split_once('/')?; + let database = after_hosts.split(['?', '#']).next()?.trim(); + if database.is_empty() { + return None; + } + Some(percent_decode_str(database).decode_utf8_lossy().into_owned()) +} + pub fn mongo_legacy_error_with_auth_hint(err: &str) -> String { let Some(source_start) = err.find("source='") else { return err.to_string(); @@ -690,6 +724,18 @@ mod tests { 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, ""); + + 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'}"; diff --git a/crates/dbx-core/src/models/connection.rs b/crates/dbx-core/src/models/connection.rs index a803646f7..db97d8cca 100644 --- a/crates/dbx-core/src/models/connection.rs +++ b/crates/dbx-core/src/models/connection.rs @@ -187,11 +187,15 @@ impl ConnectionConfig { _ => Some("postgres"), }, DatabaseType::Redshift => Some("dev"), + DatabaseType::ClickHouse => Some("default"), DatabaseType::Gaussdb | DatabaseType::OpenGauss => Some("postgres"), DatabaseType::Kingbase | DatabaseType::Vastbase => Some("postgres"), DatabaseType::Highgo => Some("highgo"), DatabaseType::Yashandb => Some("yasdb"), DatabaseType::Firebird => Some("employee"), + DatabaseType::H2 => Some("test"), + DatabaseType::Informix => Some("sysmaster"), + DatabaseType::Neo4j => Some("neo4j"), _ => None, } } @@ -803,6 +807,45 @@ mod tests { assert!(!canonical.needs_bare_mysql()); } + #[test] + fn informix_empty_database_uses_sysmaster_for_connection() { + let mut config = mysql_config("informix", "in4mix", None); + config.db_type = DatabaseType::Informix; + config.port = 9088; + + assert_eq!(config.effective_database(), Some("sysmaster")); + assert_eq!(config.connection_url(), "informix://informix:in4mix@10.1.2.3:9088/sysmaster"); + } + + #[test] + fn h2_empty_database_uses_test_for_connection() { + let mut config = mysql_config("sa", "", None); + config.db_type = DatabaseType::H2; + config.port = 9092; + + assert_eq!(config.effective_database(), Some("test")); + assert_eq!(config.connection_url(), "h2://sa:@10.1.2.3:9092/test"); + } + + #[test] + fn neo4j_empty_database_uses_neo4j_for_connection() { + let mut config = mysql_config("neo4j", "secret", None); + config.db_type = DatabaseType::Neo4j; + config.port = 7687; + + assert_eq!(config.effective_database(), Some("neo4j")); + assert_eq!(config.connection_url(), "neo4j://neo4j:secret@10.1.2.3:7687/neo4j"); + } + + #[test] + fn clickhouse_empty_database_uses_default() { + let mut config = mysql_config("default", "", None); + config.db_type = DatabaseType::ClickHouse; + config.port = 8123; + + assert_eq!(config.effective_database(), Some("default")); + } + #[test] fn mysql_url_encodes_password_and_database() { let config = mysql_config("root", "p@ss:word#1", Some("db/name")); diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index be613b337..b66a74d39 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -1016,7 +1016,7 @@ async fn exec_tx_explicit_inner( } drop(conns); - do_execute(state, pool_key, None, "BEGIN", schema, None, QueryExecutionOptions::default()) + do_execute(state, pool_key, None, "BEGIN TRANSACTION", schema, None, QueryExecutionOptions::default()) .await .map_err(|e| format!("Failed to begin transaction: {}", e))?;