diff --git a/crates/dbx-core/src/connection/mod.rs b/crates/dbx-core/src/connection/mod.rs index 13d805600f..1e1329f93e 100644 --- a/crates/dbx-core/src/connection/mod.rs +++ b/crates/dbx-core/src/connection/mod.rs @@ -1478,14 +1478,13 @@ impl AppState { }?; let pool_database = metadata_pool_database(Some(&config), database); let mut base_pool_keys = - vec![base_pool_key_for_with_catalog(Some(config.db_type), connection_id, pool_database, None, false)]; + vec![base_pool_key_for_config(Some(&config), connection_id, pool_database, None, false)]; // MongoDB document operations use a connection-level pool and send the // requested database in the command. A host connection can therefore // legitimately be registered either under the selected database or // under the connection-level key; probe both without creating either. if config.db_type == DatabaseType::MongoDb { - let connection_pool_key = - base_pool_key_for_with_catalog(Some(config.db_type), connection_id, None, None, false); + let connection_pool_key = base_pool_key_for_config(Some(&config), connection_id, None, None, false); if !base_pool_keys.contains(&connection_pool_key) { base_pool_keys.push(connection_pool_key); } @@ -2657,11 +2656,10 @@ impl AppState { 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 base_pool_key = base_pool_key_for_config(Some(&config), 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 { @@ -4678,14 +4676,13 @@ impl AppState { 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 pool_database = if session_role == AgentSessionRole::Metadata { metadata_pool_database(config.as_ref(), database) } else { database }; - let base_pool_key = base_pool_key_for_with_catalog(db_type, connection_id, pool_database, catalog, true); + let base_pool_key = base_pool_key_for_config(config.as_ref(), connection_id, pool_database, catalog, true); let pool_key = pool_key_for_session_role(config.as_ref(), base_pool_key, client_session_id, session_role); if self.uses_forwarded_transport(connection_id).await { self.remove_connection_pools(connection_id).await; @@ -4788,13 +4785,12 @@ impl AppState { let configs = self.configs.read().await; configs.get(connection_id).cloned() }; - let db_type = config.as_ref().map(|config| config.db_type); let pool_database = if session_role == AgentSessionRole::Metadata { metadata_pool_database(config.as_ref(), database) } else { database }; - let base_pool_key = base_pool_key_for(db_type, connection_id, pool_database, false); + let base_pool_key = base_pool_key_for_config(config.as_ref(), connection_id, pool_database, None, 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 { @@ -4832,9 +4828,8 @@ impl AppState { let configs = self.configs.read().await; configs.get(connection_id).cloned() }; - let db_type = config.as_ref().map(|config| config.db_type); let pool_database = metadata_pool_database(config.as_ref(), database); - let base_pool_key = base_pool_key_for(db_type, connection_id, pool_database, false); + let base_pool_key = base_pool_key_for_config(config.as_ref(), connection_id, pool_database, None, 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 @@ -4852,9 +4847,8 @@ impl AppState { let configs = self.configs.read().await; configs.get(connection_id).cloned() }; - let db_type = config.as_ref().map(|config| config.db_type); let pool_database = metadata_pool_database(config.as_ref(), database); - let base_pool_key = base_pool_key_for(db_type, connection_id, pool_database, false); + let base_pool_key = base_pool_key_for_config(config.as_ref(), connection_id, pool_database, None, 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 { @@ -4903,8 +4897,7 @@ impl AppState { 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 base_pool_key = base_pool_key_for_config(config.as_ref(), connection_id, database, None, 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); @@ -5059,19 +5052,24 @@ impl AppState { } pub async fn close_database_pool(&self, connection_id: &str, database: Option<&str>) -> Result { - let (db_type, default_database) = { + let config = { let configs = self.configs.read().await; - configs - .get(connection_id) - .map_or((None, None), |config| (Some(config.db_type), config.effective_database().map(str::to_string))) + configs.get(connection_id).cloned() }; + let db_type = config.as_ref().map(|config| config.db_type); + let default_database = config.as_ref().and_then(|config| config.effective_database()); if database.is_some() && db_type.is_some_and(|db_type| shares_database_pool_with_connection(&db_type)) { return Ok(false); } - let target_database = database.map(str::trim).filter(|database| !database.is_empty()); - let mut base_pool_keys = vec![base_pool_key_for(db_type, connection_id, target_database, false)]; - if target_database.is_some() && target_database == default_database.as_deref() { - let connection_pool_key = base_pool_key_for(db_type, connection_id, None, false); + let target_database = if config.as_ref().is_some_and(is_postgres_jdbc) { + database.filter(|database| !database.trim().is_empty()) + } else { + database.map(str::trim).filter(|database| !database.is_empty()) + }; + let mut base_pool_keys = + vec![base_pool_key_for_config(config.as_ref(), connection_id, target_database, None, false)]; + if target_database.is_some() && target_database == default_database { + let connection_pool_key = base_pool_key_for_config(config.as_ref(), connection_id, None, None, false); if !base_pool_keys.contains(&connection_pool_key) { base_pool_keys.push(connection_pool_key); } @@ -6473,9 +6471,17 @@ fn is_embedded_h2_jdbc(config: &ConnectionConfig) -> bool { && !location.is_empty() } +fn is_postgres_jdbc(config: &ConnectionConfig) -> bool { + config.db_type == DatabaseType::Jdbc + && config.connection_string.as_deref().is_some_and(|url| url.trim().starts_with("jdbc:postgresql:")) +} + #[cfg(all(test, unix))] mod h2_jdbc_tests; +#[cfg(all(test, unix))] +mod postgres_jdbc_tests; + pub fn task_client_session_id(task_kind: &str, task_id: &str) -> String { format!("{task_kind}:{task_id}") } @@ -6668,7 +6674,9 @@ fn pool_key_for_session_role( 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) || sqlserver_uses_legacy_driver(config) + database_capabilities::is_agent_type(&config.db_type) + || sqlserver_uses_legacy_driver(config) + || is_postgres_jdbc(config) }) { // The legacy SQL Server Agent borrows one connection-level pool for metadata across @@ -6841,6 +6849,35 @@ fn extract_auth_token_from_params(params: &str) -> Option { .map(|(_, value)| value.trim().to_string()) } +fn base_pool_key_for_config( + config: Option<&ConnectionConfig>, + connection_id: &str, + database: Option<&str>, + catalog: Option<&str>, + include_elasticsearch_single_pool: bool, +) -> String { + if config.is_some_and(is_postgres_jdbc) { + let key = match database.filter(|database| !database.trim().is_empty()) { + Some(database) => format!( + "{connection_id}:database:{}", + percent_encoding::utf8_percent_encode(database, percent_encoding::NON_ALPHANUMERIC) + ), + None => connection_id.to_string(), + }; + return match catalog.map(str::trim).filter(|catalog| !catalog.is_empty()) { + Some(catalog) => format!("{key}:catalog:{catalog}"), + None => key, + }; + } + base_pool_key_for_with_catalog( + config.map(|config| config.db_type), + connection_id, + database, + catalog, + include_elasticsearch_single_pool, + ) +} + fn base_pool_key_for( db_type: Option, connection_id: &str, diff --git a/crates/dbx-core/src/connection/postgres_jdbc_tests.rs b/crates/dbx-core/src/connection/postgres_jdbc_tests.rs new file mode 100644 index 0000000000..a9095ad660 --- /dev/null +++ b/crates/dbx-core/src/connection/postgres_jdbc_tests.rs @@ -0,0 +1,158 @@ +use super::*; + +fn config(url: &str) -> ConnectionConfig { + let mut config = super::tests::mysql_config(None); + config.db_type = DatabaseType::Jdbc; + config.connection_string = Some(url.into()); + config.jdbc_driver_class = Some("org.postgresql.Driver".into()); + config +} + +#[test] +fn postgres_jdbc_pool_keys_preserve_database_role_and_session_identity() { + let config = config("jdbc:postgresql://localhost/a?sslmode=require"); + let key = |database, session, role| { + pool_key_for_session_role( + Some(&config), + base_pool_key_for_config(Some(&config), "conn", database, None, false), + session, + role, + ) + }; + let mut keys = std::collections::HashSet::new(); + for database in [None, Some("a"), Some("b"), Some(" a"), Some("a "), Some("a:session:x")] { + for session in [None, Some("tab-a"), Some("tab-b")] { + for role in [AgentSessionRole::Metadata, AgentSessionRole::Workload] { + assert!(keys.insert(key(database, session, role))); + } + } + } + assert_eq!(key(Some("a:session:x"), None, AgentSessionRole::Workload), "conn:database:a%3Asession%3Ax"); + assert_eq!(key(Some(""), None, AgentSessionRole::Workload), "conn"); +} + +#[test] +fn postgres_jdbc_routing_scope_keeps_other_vendor_pool_defaults() { + for url in [ + "jdbc:h2:file:./sample", + "jdbc:h2:tcp://localhost/sample", + "jdbc:h2:mem:sample", + "jdbc:hive2://localhost/default", + "jdbc:kingbase8://localhost/a", + "jdbc:postgresqlx://localhost/a", + ] { + let config = config(url); + assert!(!is_postgres_jdbc(&config)); + assert_eq!(base_pool_key_for_config(Some(&config), "conn", Some("b"), None, false), "conn"); + } + let config = config("jdbc:postgresql://localhost/a"); + let target = database_connection_config(&config, Some(" b ")); + assert_eq!(target.database.as_deref(), Some(" b ")); + assert_eq!(target.connection_string, config.connection_string); + assert_eq!(database_connection_config(&config, None).database, None); +} + +async fn fixture() -> (AppState, std::path::PathBuf) { + use std::os::unix::fs::PermissionsExt; + let (mut state, directory) = super::tests::test_app_state().await; + let plugin = directory.join("plugins/jdbc"); + std::fs::create_dir_all(&plugin).unwrap(); + let script = plugin.join("driver.py"); + std::fs::write( + &script, + r#"#!/usr/bin/python3 +import json, os, sys +database = None +for line in sys.stdin: + request = json.loads(line) + method, params = request['method'], request.get('params', {}) + response = {'id': request['id']} + result = {'pid': os.getpid(), 'database': database} + if method == 'connect': + database = params['connection'].get('database') or 'url_default' + if database in ('missing', 'denied'): + response['error'] = {'message': 'target database unavailable'} + elif method == 'listSchemas': + result = [database + '_only', 'public'] + response.setdefault('result', result) + print(json.dumps(response), flush=True) + if method == 'close': + break +"#, + ) + .unwrap(); + std::fs::set_permissions(&script, std::fs::Permissions::from_mode(0o755)).unwrap(); + std::fs::write( + plugin.join("manifest.json"), + serde_json::json!({ + "id": "jdbc", "name": "JDBC", "version": "0.1.41", "protocol_version": 1, + "executable": "driver.py", "drivers": [{"id": "jdbc", "label": "JDBC", "kind": "external", "database_type": "jdbc"}] + }).to_string(), + ).unwrap(); + state.plugins = PluginRegistry::new_with_app_version(directory.join("plugins"), "0.6.30"); + state.configs.write().await.insert("conn".into(), config("jdbc:postgresql://localhost/a")); + (state, directory) +} + +async fn identity(state: &AppState, key: &str) -> serde_json::Value { + let PoolKind::ExternalDriver { session, .. } = state.pool_handle(key).await.unwrap() else { panic!() }; + session.invoke("identity", serde_json::json!({})).await.unwrap() +} + +#[tokio::test] +async fn postgres_jdbc_owning_entries_route_metadata_cache_hits_and_failures() { + let (state, directory) = fixture().await; + let mut owners = HashMap::new(); + for database in ["a", "b", "a"] { + assert_eq!( + crate::schema::list_schemas_core(&state, "conn", database).await.unwrap(), + vec![format!("{database}_only"), "public".into()] + ); + let key = state.get_or_create_metadata_pool_for_session("conn", Some(database), None).await.unwrap(); + let current = identity(&state, &key).await; + assert_eq!(current["database"], database); + if let Some(original) = owners.insert(database, current.clone()) { + assert_eq!(original, current); + } + assert_eq!(state.existing_metadata_pool_key_for_session("conn", Some(database), None).await, Some(key)); + } + assert_ne!(owners["a"]["pid"], owners["b"]["pid"]); + for database in ["missing", "denied"] { + assert!(state.get_or_create_metadata_pool_for_session("conn", Some(database), None).await.is_err()); + assert_eq!(state.existing_metadata_pool_key_for_session("conn", Some(database), None).await, None); + } + let database_a_pool = state.get_or_create_metadata_pool_for_session("conn", Some("a"), None).await.unwrap(); + assert_eq!(identity(&state, &database_a_pool).await, owners["a"]); + assert!(state.close_database_pool("conn", Some("b")).await.unwrap()); + assert_eq!(identity(&state, &database_a_pool).await, owners["a"]); + assert_eq!(state.existing_metadata_pool_key_for_session("conn", Some("b"), None).await, None); + state.shutdown(Duration::from_secs(5)).await; + std::fs::remove_dir_all(directory).unwrap(); +} + +#[tokio::test] +async fn postgres_jdbc_owning_entries_separate_concurrent_roles_sessions_and_cleanup() { + let (state, directory) = fixture().await; + let (metadata, workload, other_database, other_tab) = tokio::join!( + state.get_or_create_metadata_pool_for_session("conn", Some("a"), Some("tab")), + state.get_or_create_pool_for_session("conn", Some("a"), Some("tab")), + state.get_or_create_pool_for_session("conn", Some("b"), Some("tab")), + state.get_or_create_pool_for_session("conn", Some("a"), Some("other")), + ); + let keys = [metadata.unwrap(), workload.unwrap(), other_database.unwrap(), other_tab.unwrap()]; + let mut pids = std::collections::HashSet::new(); + for key in &keys { + assert!(pids.insert(identity(&state, key).await["pid"].as_i64().unwrap())); + } + assert!(state.close_metadata_session_pool("conn", Some("a"), "tab").await.unwrap()); + assert!(state.pool_handle(&keys[0]).await.is_none()); + assert!(state.detach_client_session_pool("conn", Some("a"), "tab").await.unwrap()); + assert!(state.pool_handle(&keys[1]).await.is_none()); + assert_eq!(identity(&state, &keys[2]).await["database"], "b"); + assert_eq!(identity(&state, &keys[3]).await["database"], "a"); + assert!(!state.detach_metadata_pool_after_recovery("conn", Some("b"), None, None, false).await); + let default = state.get_or_create_pool("conn", None).await.unwrap(); + assert_eq!(identity(&state, &default).await["database"], "url_default"); + state.shutdown(Duration::from_secs(5)).await; + std::fs::remove_dir_all(directory).unwrap(); +} diff --git a/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java b/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java index 48fa2e9a47..0cd2ca1b93 100644 --- a/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java +++ b/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java @@ -393,7 +393,8 @@ public final class DbxJdbcPlugin { registerDrivers(connection); return handle(method, params, connection); } catch (Exception | LinkageError error) { - if ("connect".equals(method)) { + if ("connect".equals(method) + && (state.sharedConnection == null || !isPostgresJdbcUrl(optionalText(connection, "connection_string")))) { closeSharedConnection(); state.closing = true; LOGICAL_SESSIONS.remove(id, state); @@ -527,6 +528,15 @@ public final class DbxJdbcPlugin { } private static JsonNode handle(String method, JsonNode params, JsonNode connection) throws Exception { + String database = optionalText(params, "database"); + if (isPostgresJdbcUrl(optionalText(connection, "connection_string"))) { + database = postgresDatabase(params); + if (database != null) { + ObjectNode targeted = connection.deepCopy(); + targeted.put("database", database); + connection = targeted; + } + } return switch (method) { case "testConnection" -> connectionTestResult(openConnection(connection)); case "connect" -> { @@ -539,7 +549,7 @@ public final class DbxJdbcPlugin { case "executeQuery" -> executeQuery( connection, requireText(params, "sql"), - optionalText(params, "database"), + database, optionalText(params, "schema"), positiveInt(params, "maxRows", MAX_ROWS), nonNegativeInt(params, "fetchSize", 0), @@ -548,13 +558,13 @@ public final class DbxJdbcPlugin { ); case "beginManualTransaction", "begin_manual_transaction" -> beginManualTransaction( connection, - optionalText(params, "database"), + database, optionalText(params, "schema") ); case "executeInManualTransaction", "execute_in_manual_transaction" -> executeInManualTransaction( connection, requireText(params, "sql"), - optionalText(params, "database"), + database, optionalText(params, "schema"), positiveInt(params, "maxRows", MAX_ROWS), nonNegativeInt(params, "fetchSize", 0), @@ -566,7 +576,7 @@ public final class DbxJdbcPlugin { case "executeQueryPage", "execute_query_page" -> executeQueryPage( connection, requireText(params, "sql"), - optionalText(params, "database"), + database, optionalText(params, "schema"), positiveInt(params, "pageSize", 100), positiveInt(params, "maxRows", MAX_ROWS), @@ -579,10 +589,10 @@ public final class DbxJdbcPlugin { ); case "closeQuerySession", "close_query_session" -> closeQuerySessionResult(requireText(params, "sessionId")); case "listDatabases" -> listDatabases(connection); - case "listSchemas" -> listSchemas(connection, optionalText(params, "database")); + case "listSchemas" -> listSchemas(connection, database); case "listTables" -> listTables( connection, - optionalText(params, "database"), + database, optionalText(params, "schema"), optionalText(params, "filter"), nonNegativeInt(params, "limit", 0), @@ -591,7 +601,7 @@ public final class DbxJdbcPlugin { ); case "listObjects", "list_objects" -> listObjects( connection, - optionalText(params, "database"), + database, optionalText(params, "schema"), optionalText(params, "filter"), nonNegativeInt(params, "limit", 0), @@ -600,28 +610,28 @@ public final class DbxJdbcPlugin { ); case "listIndexes", "list_indexes" -> listIndexes( connection, - optionalText(params, "database"), + database, optionalText(params, "schema"), requireText(params, "table") ); - case "listDataTypes", "list_data_types" -> listDataTypes(connection, optionalText(params, "database")); + case "listDataTypes", "list_data_types" -> listDataTypes(connection, database); case "getObjectSource", "get_object_source" -> getObjectSource( connection, - optionalText(params, "database"), + database, optionalText(params, "schema"), requireText(params, "name"), requireText(params, "object_type") ); case "getColumns" -> getColumns( connection, - optionalText(params, "database"), + database, optionalText(params, "schema"), requireText(params, "table") ); case "getExplainInfo" -> getExplainInfo( connection, requireText(params, "sql"), - optionalText(params, "database"), + database, optionalText(params, "schema"), nonNegativeInt(params, "timeoutSecs", -1), optionalText(params, "mode") @@ -807,7 +817,12 @@ public final class DbxJdbcPlugin { configureOrdinaryAutoCommit(state.sharedConnection); return state.sharedConnection; } - closeSharedConnection(); + boolean postgres = isPostgresJdbcUrl(url); + if (postgres && state.sharedConnection != null + && (state.manualTransactionActive || hasActiveQuerySession(state.sharedConnection))) { + throw new SQLException("Cannot change PostgreSQL database while a transaction or query session is active"); + } + if (!postgres) closeSharedConnection(); JdbcUrlCredentials urlCredentials = extractJdbcUrlCredentials(url); url = urlCredentials.url; @@ -835,7 +850,21 @@ public final class DbxJdbcPlugin { // Prefer the explicitly registered driver. DriverManager.getConnection only catches // SQLException; Hive/Inceptor drivers may throw UnsupportedOperationException for optional // methods, which aborts connect before the intended driver is reached. - state.sharedConnection = connectWithRegisteredDriver(url, properties); + Connection opened = connectWithRegisteredDriver(url, properties); + if (postgres) { + try { + configureOrdinaryAutoCommit(opened); + } catch (SQLException | RuntimeException | Error error) { + try { + opened.close(); + } catch (SQLException closeError) { + error.addSuppressed(closeError); + } + throw error; + } + closeSharedConnection(); + } + state.sharedConnection = opened; state.sharedConnectionKey = key; configureOrdinaryAutoCommit(state.sharedConnection); return state.sharedConnection; @@ -3692,6 +3721,7 @@ public final class DbxJdbcPlugin { private static String connectionKey(JsonNode connection) { String connectionString = optionalText(connection, "connection_string"); + if (isPostgresJdbcUrl(connectionString)) connectionString = jdbcUrl(connection); String jdbcxSecurityKey = isJdbcxUrl(connectionString) ? "|jdbcxHighPrivilegeExtensions=" + jdbcxHighPrivilegeExtensionsEnabled(connection) : ""; @@ -3739,9 +3769,53 @@ public final class DbxJdbcPlugin { static String jdbcUrl(JsonNode connection) { String url = appendJdbcUrlParams(optionalText(connection, "connection_string"), optionalText(connection, "url_params")); + url = postgresDatabaseUrl(url, postgresDatabase(connection)); return jdbcUrlWithPasswordKey(url, optionalText(connection, "password")); } + private static boolean isPostgresJdbcUrl(String url) { + return url != null && url.startsWith("jdbc:postgresql:"); + } + + private static String postgresDatabase(JsonNode node) { + String database = node.path("database").asText(null); + return database == null || database.isBlank() ? null : database; + } + + private static String postgresDatabaseUrl(String url, String database) { + if (!isPostgresJdbcUrl(url) || database == null) return url; + String prefix = "jdbc:postgresql:"; + int queryStart = url.indexOf('?'); + String location = url.substring(prefix.length(), queryStart < 0 ? url.length() : queryStart); + String authority = ""; + String originalDatabase = location; + if (location.equals("//") || location.equals("///")) { + originalDatabase = ""; + } else if (location.startsWith("//")) { + int slash = location.indexOf('/', 2); + if (slash < 0 || location.indexOf('/', slash + 1) >= 0) { + throw new IllegalArgumentException("Invalid PostgreSQL JDBC URL database path"); + } + authority = location.substring(0, slash + 1); + originalDatabase = location.substring(slash + 1); + } else if (location.startsWith("/")) { + throw new IllegalArgumentException("Invalid PostgreSQL JDBC URL database path"); + } + try { + URLDecoder.decode(originalDatabase, StandardCharsets.UTF_8); + } catch (IllegalArgumentException error) { + throw new IllegalArgumentException("Invalid PostgreSQL JDBC URL database encoding"); + } + String targeted = prefix + authority + URLEncoder.encode(database, StandardCharsets.UTF_8); + if (queryStart < 0) return targeted; + List parameters = new ArrayList<>(); + for (String part : url.substring(queryStart + 1).split("&", -1)) { + String name = partName(part); + if (!name.equals("PGDBNAME") && !name.equals("dbname")) parameters.add(part); + } + return targeted + "?" + String.join("&", parameters); + } + private record JdbcUrlCredentials(String url, String username, String password) {} static JdbcUrlCredentials extractJdbcUrlCredentials(String url) { diff --git a/plugins/jdbc/src/test/java/app/dbx/jdbc/PostgresDatabaseRoutingTest.java b/plugins/jdbc/src/test/java/app/dbx/jdbc/PostgresDatabaseRoutingTest.java new file mode 100644 index 0000000000..8d5d0a2c5a --- /dev/null +++ b/plugins/jdbc/src/test/java/app/dbx/jdbc/PostgresDatabaseRoutingTest.java @@ -0,0 +1,269 @@ +package app.dbx.jdbc; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; + +import java.lang.reflect.Proxy; +import java.net.URLDecoder; +import java.nio.charset.StandardCharsets; +import java.sql.*; +import java.util.*; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.logging.Logger; + +import static org.junit.jupiter.api.Assertions.*; + +final class PostgresDatabaseRoutingTest { + private static final ObjectMapper MAPPER = new ObjectMapper(); + private static final List PHYSICAL = new CopyOnWriteArrayList<>(); + private final List sessions = new ArrayList<>(); + + @AfterEach + void cleanup() throws Exception { + for (String session : sessions) call("closeJdbcSession", params("a", null).put("jdbcSessionId", session)); + call("close", MAPPER.createObjectNode()); + var field = DbxJdbcPlugin.class.getDeclaredField("logicalDriverKey"); + field.setAccessible(true); + field.set(null, null); + PHYSICAL.clear(); + } + + @Test + void metadataBindsTheRequestedPhysicalDatabaseAcrossAToBToA() throws Exception { + for (String database : List.of("a", "b", "a")) { + JsonNode result = ok(call("listSchemas", params("a", database))); + assertEquals(List.of(database + "_only", "public"), MAPPER.convertValue(result, List.class)); + } + assertEquals(List.of("a", "b", "a"), PHYSICAL.stream().map(physical -> physical.database).toList()); + } + + @Test + void urlBindingPreservesEndpointsOptionsAndEscapesDatabaseNames() { + String target = " sales/季度?&+#% "; + ObjectNode config = config("jdbc:postgresql://[::1]:5433,standby:5434/a%20b?sslmode=require&user=u%40x&password=p%26x&dbname=old"); + config.put("database", target).put("url_params", "ApplicationName=dbx&PGDBNAME=also_old&options=-c%20search_path%3Dpublic"); + String url = DbxJdbcPlugin.jdbcUrl(config); + assertEquals(target, database(url)); + assertTrue(url.startsWith("jdbc:postgresql://[::1]:5433,standby:5434/")); + assertTrue(url.contains("sslmode=require&user=u%40x&password=p%26x")); + assertTrue(url.contains("ApplicationName=dbx")); + assertTrue(url.contains("options=-c%20search_path%3Dpublic")); + assertFalse(url.contains("old")); + assertFalse(url.contains("季度")); + } + + @Test + void targetConnectionKeepsRoutedEndpointCredentialsAndOptions() throws Exception { + ObjectNode request = MAPPER.createObjectNode().put("database", "b"); + ObjectNode connection = config("jdbc:postgresql://127.0.0.1:25432/a?sslmode=require&user=url%40user&password=p%26%2Bword"); + connection.put("url_params", "ApplicationName=dbx&options=-c%20search_path%3Dpublic"); + request.set("connection", connection); + ok(call("connect", request)); + Physical opened = PHYSICAL.getFirst(); + assertEquals("jdbc:postgresql://127.0.0.1:25432/b?sslmode=require&ApplicationName=dbx&options=-c%20search_path%3Dpublic", opened.url); + assertEquals("url@user", opened.properties.getProperty("user")); + assertEquals("p&+word", opened.properties.getProperty("password")); + connection.put("username", "form@user").put("password", "form-password"); + ok(call("connect", request)); + assertEquals("form@user", PHYSICAL.getLast().properties.getProperty("user")); + assertEquals("form-password", PHYSICAL.getLast().properties.getProperty("password")); + } + + @Test + void defaultAndUnspecifiedDatabasePreserveTheOriginalUrl() throws Exception { + for (String url : List.of("jdbc:postgresql:a", "jdbc:postgresql:", "jdbc:postgresql://host/", "jdbc:postgresql://host/a?ApplicationName=dbx")) { + assertEquals(url, DbxJdbcPlugin.jdbcUrl(config(url))); + ObjectNode selected = config(url).put("database", "b"); + assertEquals("b", database(DbxJdbcPlugin.jdbcUrl(selected))); + } + assertEquals("a_only", ok(call("listSchemas", params("a", null))).get(0).asText()); + assertEquals("b_only", ok(call("listSchemas", params("b", null))).get(0).asText()); + } + + @Test + void otherVendorsAndMalformedPostgresUrlsKeepTheirBoundaries() { + for (String url : List.of("jdbc:h2:file:./sample", "jdbc:h2:tcp://host/sample", "jdbc:hive2://host/default", "jdbc:kingbase8://host/a")) { + assertEquals(url, DbxJdbcPlugin.jdbcUrl(config(url).put("database", "b"))); + } + for (String url : List.of("jdbc:postgresql://host", "jdbc:postgresql:/a", "jdbc:postgresql://host/a/b", "jdbc:postgresql://host/%xx")) { + assertThrows(IllegalArgumentException.class, () -> DbxJdbcPlugin.jdbcUrl(config(url).put("database", "b"))); + } + } + + @Test + void missingAndDeniedTargetsNeverReturnTheOriginalDatabaseOrCloseIt() throws Exception { + ok(call("listSchemas", params("a", "a"))); + Physical original = PHYSICAL.getFirst(); + for (String target : List.of("missing", "denied")) { + JsonNode response = call("listSchemas", params("a", target)); + assertTrue(response.has("error"), response.toString()); + assertTrue(response.path("error").path("message").asText().contains(target)); + assertFalse(original.closed); + } + assertEquals("a_only", ok(call("listSchemas", params("a", "a"))).get(0).asText()); + assertEquals(1, PHYSICAL.size()); + } + + @Test + void activeTransactionsRejectDatabaseChangesWithoutLosingTheirOwner() throws Exception { + ok(call("beginManualTransaction", params("a", "a"))); + Physical owner = PHYSICAL.getFirst(); + assertFalse(owner.autoCommit); + assertTrue(call("listSchemas", params("a", "b")).has("error")); + assertTrue(call("executeInManualTransaction", params("a", "b").put("sql", "SELECT 1")).has("error")); + assertFalse(owner.closed); + assertFalse(owner.autoCommit); + ok(call("rollbackManualTransaction", params("a", "a"))); + assertEquals(1, owner.rollbacks); + assertEquals("b_only", ok(call("listSchemas", params("a", "b"))).get(0).asText()); + } + + @Test + void concurrentLogicalSessionsKeepMetadataAndTransactionOwnersSeparate() throws Exception { + List> pending = new ArrayList<>(); + for (int i = 0; i < 8; i++) { + String session = "postgres-" + i; + sessions.add(session); + ObjectNode request = params("a", i % 2 == 0 ? "a" : "b").put("jdbcSessionId", session); + ok(call("openJdbcSession", request)); + pending.add(CompletableFuture.supplyAsync(() -> { + try { return call("listSchemas", request); } + catch (Exception error) { throw new RuntimeException(error); } + })); + } + for (int i = 0; i < pending.size(); i++) { + assertEquals((i % 2 == 0 ? "a" : "b") + "_only", ok(pending.get(i).get()).get(0).asText()); + } + assertEquals(8, PHYSICAL.size()); + ok(call("beginManualTransaction", params("a", "a").put("jdbcSessionId", sessions.getFirst()))); + assertEquals(1, PHYSICAL.stream().filter(physical -> !physical.autoCommit).count()); + ok(call("closeJdbcSession", params("a", null).put("jdbcSessionId", sessions.getFirst()))); + assertEquals(1, PHYSICAL.stream().filter(physical -> physical.closed).count()); + } + + @Test + void failedReconnectKeepsTheExistingLogicalSessionAndTransaction() throws Exception { + String session = "reconnect"; + sessions.add(session); + ObjectNode original = params("a", "a").put("jdbcSessionId", session); + ok(call("openJdbcSession", original)); + ok(call("connect", original)); + Physical owner = PHYSICAL.getFirst(); + assertTrue(call("connect", params("a", "missing").put("jdbcSessionId", session)).has("error")); + assertFalse(owner.closed); + ok(call("beginManualTransaction", original)); + assertTrue(call("connect", params("a", "b").put("jdbcSessionId", session)).has("error")); + assertFalse(owner.closed); + ok(call("rollbackManualTransaction", original)); + assertEquals(1, owner.rollbacks); + assertEquals("a_only", ok(call("listSchemas", original)).get(0).asText()); + } + + private static ObjectNode config(String url) { + ObjectNode config = MAPPER.createObjectNode().put("connection_string", url) + .put("jdbc_driver_class", FakeDriver.class.getName()); + config.putArray("jdbc_driver_paths"); + return config; + } + + private static ObjectNode params(String savedDatabase, String target) { + ObjectNode params = MAPPER.createObjectNode(); + params.set("connection", config("jdbc:postgresql://localhost/" + savedDatabase)); + if (target != null) params.put("database", target); + return params; + } + + private static JsonNode call(String method, ObjectNode params) throws Exception { + ObjectNode request = MAPPER.createObjectNode().put("id", 1).put("method", method); + request.set("params", params); + var handler = DbxJdbcPlugin.class.getDeclaredMethod("handleLine", String.class); + handler.setAccessible(true); + return (JsonNode) handler.invoke(null, request.toString()); + } + + private static JsonNode ok(JsonNode response) { + assertFalse(response.has("error"), response.toString()); + return response.path("result"); + } + + private static String database(String url) { + String location = url.substring("jdbc:postgresql:".length()).split("\\?", 2)[0]; + String value = location.startsWith("//") ? location.substring(location.indexOf('/', 2) + 1) : location; + for (String part : url.substring(url.indexOf('?') + 1).split("&")) { + if (part.startsWith("PGDBNAME=") || part.startsWith("dbname=")) value = part.substring(part.indexOf('=') + 1); + } + return URLDecoder.decode(value, StandardCharsets.UTF_8); + } + + public static final class FakeDriver implements Driver { + public Connection connect(String url, Properties properties) throws SQLException { + if (!acceptsURL(url)) return null; + String database = database(url); + if (Set.of("missing", "denied").contains(database)) throw new SQLException(database + " database unavailable"); + Physical physical = new Physical(database, url, properties); + PHYSICAL.add(physical); + return physical.connection; + } + public boolean acceptsURL(String url) { return url.startsWith("jdbc:postgresql:"); } + public DriverPropertyInfo[] getPropertyInfo(String url, Properties info) { return new DriverPropertyInfo[0]; } + public int getMajorVersion() { return 1; } + public int getMinorVersion() { return 0; } + public boolean jdbcCompliant() { return false; } + public Logger getParentLogger() { return Logger.getGlobal(); } + } + + private static final class Physical { + final String database; + final String url; + final Properties properties; + final Connection connection; + boolean closed; + boolean autoCommit = true; + int rollbacks; + + Physical(String database, String url, Properties properties) { + this.database = database; + this.url = url; + this.properties = properties; + DatabaseMetaData metadata = proxy(DatabaseMetaData.class, (object, method, args) -> { + if (method.getName().equals("supportsTransactions")) return true; + if (method.getName().equals("getSchemas")) { + List schemas = args == null || args[0] == null || database.equals(args[0]) + ? List.of(database + "_only", "public") : List.of(); + int[] row = {-1}; + return proxy(ResultSet.class, (rs, operation, values) -> switch (operation.getName()) { + case "next" -> ++row[0] < schemas.size(); + case "getString" -> "TABLE_CATALOG".equals(values[0]) ? database : schemas.get(row[0]); + default -> defaultValue(operation.getReturnType()); + }); + } + return defaultValue(method.getReturnType()); + }); + connection = proxy(Connection.class, (object, method, args) -> switch (method.getName()) { + case "getMetaData" -> metadata; + case "getCatalog" -> database; + case "getAutoCommit" -> autoCommit; + case "setAutoCommit" -> { autoCommit = (boolean) args[0]; yield null; } + case "isClosed" -> closed; + case "close" -> { closed = true; yield null; } + case "rollback" -> { rollbacks++; yield null; } + default -> defaultValue(method.getReturnType()); + }); + } + } + + private static T proxy(Class type, java.lang.reflect.InvocationHandler handler) { + return type.cast(Proxy.newProxyInstance(type.getClassLoader(), new Class[]{type}, handler)); + } + + private static Object defaultValue(Class type) { + if (type == boolean.class) return false; + if (type == int.class) return 0; + if (type == long.class) return 0L; + return null; + } +}