From c021fe009b96e37185cf5caeefed61d2331693da Mon Sep 17 00:00:00 2001 From: zipg Date: Mon, 24 Aug 2026 14:50:18 +0800 Subject: [PATCH] fix(jdbc): support manual transactions and auto-commit control (#6997) * fix(jdbc): support manual transactions * fix(jdbc): preserve paged query transaction state --------- Co-authored-by: zipg <4047349+zipg@users.noreply.github.com> --- .../database/databaseFeatureSupport.spec.ts | 1 + .../lib/database/databaseFeatureSupport.ts | 2 +- crates/dbx-core/src/connection.rs | 8 + crates/dbx-core/src/query.rs | 232 +++++++++++- .../main/java/app/dbx/jdbc/DbxJdbcPlugin.java | 120 +++++- .../java/app/dbx/jdbc/DbxJdbcPluginTest.java | 354 ++++++++++++++++-- 6 files changed, 662 insertions(+), 55 deletions(-) diff --git a/apps/desktop/src/lib/__tests__/database/databaseFeatureSupport.spec.ts b/apps/desktop/src/lib/__tests__/database/databaseFeatureSupport.spec.ts index d15a56b5f1..5847a89c69 100644 --- a/apps/desktop/src/lib/__tests__/database/databaseFeatureSupport.spec.ts +++ b/apps/desktop/src/lib/__tests__/database/databaseFeatureSupport.spec.ts @@ -91,6 +91,7 @@ describe("supportsTransaction", () => { expect(supportsTransaction("postgres")).toBe(true); expect(supportsTransaction("mysql")).toBe(true); expect(supportsTransaction("oracle")).toBe(true); + expect(supportsTransaction("jdbc")).toBe(true); }); it("returns false for unsupported database types", () => { diff --git a/apps/desktop/src/lib/database/databaseFeatureSupport.ts b/apps/desktop/src/lib/database/databaseFeatureSupport.ts index 0923d600c9..3a005e1c31 100644 --- a/apps/desktop/src/lib/database/databaseFeatureSupport.ts +++ b/apps/desktop/src/lib/database/databaseFeatureSupport.ts @@ -238,7 +238,7 @@ export function usesPostgresLikeStructureCopy(dbType?: DatabaseType): boolean { return !!dbType && PG_LIKE_STRUCTURE_TYPES.has(dbType); } -const TRANSACTION_SUPPORTED_TYPES: readonly string[] = ["postgres", "mysql", "oracle"]; +const TRANSACTION_SUPPORTED_TYPES: readonly string[] = ["postgres", "mysql", "oracle", "jdbc"]; /** * Returns true if the given database type supports explicit transaction control diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index b2e1d7eb27..49a40fbf88 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -190,6 +190,14 @@ pub enum TxnConnection { database: Option, cleanup_guard: ClientSessionPoolCleanupGuard, }, + /// Dedicated external-driver process whose shared JDBC connection owns the TX. + ExternalDriver { + session: Arc, + config: Arc, + client_session_id: String, + database: Option, + cleanup_guard: ClientSessionPoolCleanupGuard, + }, } pub struct TransactionSession { diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index 604c880a0f..147b5b8b9e 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -4348,6 +4348,7 @@ async fn begin_transaction_session( Postgres(deadpool_postgres::Pool), Mysql(db::mysql::MySqlPool), Agent, + ExternalDriver, } let pool_handle = { let connections = state.connections.read().await; @@ -4355,6 +4356,7 @@ async fn begin_transaction_session( PoolKind::Postgres(pg) => TxnPoolHandle::Postgres(pg.clone()), PoolKind::Mysql(mp, _) => TxnPoolHandle::Mysql(mp.clone()), PoolKind::Agent(_) if !consistent_snapshot => TxnPoolHandle::Agent, + PoolKind::ExternalDriver { .. } if !consistent_snapshot => TxnPoolHandle::ExternalDriver, _ => return Err("Manual transaction is not supported for this database type".to_string()), } }; // connections lock released here @@ -4438,6 +4440,53 @@ async fn begin_transaction_session( agent_pool_key, ) } + TxnPoolHandle::ExternalDriver => { + let client_session_id = format!("manual-txn-{}", uuid::Uuid::new_v4()); + let external_pool_key = + state.get_or_create_pool_for_session(connection_id, pool_database, Some(&client_session_id)).await?; + let (config, session) = { + let connections = state.connections.read().await; + match connections.get(&external_pool_key) { + Some(PoolKind::ExternalDriver { config, session, .. }) => (config.clone(), session.clone()), + _ => { + let _ = state.close_client_session_pool(connection_id, pool_database, &client_session_id).await; + return Err("External driver connection not found for manual transaction".to_string()); + } + } + }; + let begin_params = serde_json::json!({ + "connection": config.as_ref(), + "database": database, + "schema": schema, + }); + if let Err(error) = session + .invoke_with_timeout::( + "beginManualTransaction", + begin_params, + query_timeout_duration(Some(config.effective_query_timeout_secs())), + ) + .await + { + let _ = state.close_client_session_pool(connection_id, pool_database, &client_session_id).await; + return Err(format!("BEGIN manual transaction failed: {error}")); + } + let Some(cleanup_guard) = + state.workload_session_pool_cleanup_guard(connection_id, pool_database, &client_session_id).await + else { + let _ = state.close_client_session_pool(connection_id, pool_database, &client_session_id).await; + return Err("Manual transaction requires a dedicated external driver session".to_string()); + }; + ( + TxnConnection::ExternalDriver { + session, + config, + client_session_id, + database: pool_database.map(str::to_string), + cleanup_guard, + }, + external_pool_key, + ) + } }; let txn_session_id = uuid::Uuid::new_v4().to_string(); @@ -4589,7 +4638,7 @@ pub async fn execute_in_manual_transaction_with_options( if let Some(session) = connection { let mut conn = session.connection.lock().await; let _ = rollback_manual_txn_connection(&mut conn).await; - release_manual_txn_agent_pool(state, &session.connection_id, &mut conn).await; + release_manual_txn_session_pool(state, &session.connection_id, &mut conn).await; return Err(MANUAL_TRANSACTION_IDLE_ROLLBACK_ERROR.to_string()); } @@ -4622,6 +4671,10 @@ pub async fn execute_in_manual_transaction_with_options( ) .await } + TxnConnection::ExternalDriver { session, config, .. } => { + execute_manual_txn_external_driver_statement(session, config, statement, database, schema, row_limit) + .await + } }; match result { Ok(query_result) => { @@ -4639,7 +4692,7 @@ pub async fn execute_in_manual_transaction_with_options( }; if should_rollback { let _ = rollback_manual_txn_connection(&mut conn).await; - release_manual_txn_agent_pool(state, &connection_id, &mut conn).await; + release_manual_txn_session_pool(state, &connection_id, &mut conn).await; } return Err(format!("Statement {} failed: {}. The manual transaction was rolled back.", i + 1, e)); } @@ -4774,6 +4827,9 @@ where TxnConnection::Agent { .. } => { Err("Streaming rows inside an agent manual transaction is not supported".to_string()) } + TxnConnection::ExternalDriver { .. } => { + Err("Streaming rows inside an external driver manual transaction is not supported".to_string()) + } }; if let Err(err) = &stream_result { @@ -4783,7 +4839,7 @@ where }; if let Some(session) = removed { let _ = rollback_manual_txn_connection(&mut conn).await; - release_manual_txn_agent_pool(state, &session.connection_id, &mut conn).await; + release_manual_txn_session_pool(state, &session.connection_id, &mut conn).await; } return Err(format!("{err}. Transaction was auto-rolled back.")); } @@ -4824,13 +4880,31 @@ async fn rollback_manual_txn_connection(conn: &mut TxnConnection) -> Result<(), // the pool map entry is cleaned up later. let _ = locked.disconnect().await; } + TxnConnection::ExternalDriver { session, config, .. } => { + match session + .invoke_with_timeout::( + "rollbackManualTransaction", + serde_json::json!({ "connection": config.as_ref() }), + query_timeout_duration(Some(config.effective_query_timeout_secs())), + ) + .await + { + Ok(_) => {} + Err(error) if error.to_ascii_lowercase().contains("no manual transaction") => {} + Err(error) => return Err(format!("ROLLBACK failed: {error}")), + } + } } Ok(()) } -async fn release_manual_txn_agent_pool(state: &AppState, connection_id: &str, conn: &mut TxnConnection) { - let TxnConnection::Agent { client_session_id, database, cleanup_guard, .. } = conn else { - return; +async fn release_manual_txn_session_pool(state: &AppState, connection_id: &str, conn: &mut TxnConnection) { + let (client_session_id, database, cleanup_guard) = match conn { + TxnConnection::Agent { client_session_id, database, cleanup_guard, .. } + | TxnConnection::ExternalDriver { client_session_id, database, cleanup_guard, .. } => { + (client_session_id, database, cleanup_guard) + } + _ => return, }; if state.detach_client_session_pool(connection_id, database.as_deref(), client_session_id).await.unwrap_or(false) { cleanup_guard.disarm(); @@ -4863,6 +4937,31 @@ async fn execute_manual_txn_agent_statement( .map_err(|error| error.into_legacy_string()) } +async fn execute_manual_txn_external_driver_statement( + session: &Arc, + config: &Arc, + statement: &str, + database: &str, + schema: Option<&str>, + row_limit: usize, +) -> Result { + let timeout_secs = config.effective_query_timeout_secs(); + let options = QueryExecutionOptions { + max_rows: Some(row_limit.max(1)), + timeout_secs: Some(timeout_secs), + ..QueryExecutionOptions::default() + }; + let params = external_driver_query_params(config.as_ref(), statement, database, schema, &options); + session + .invoke_with_timeout::( + "executeInManualTransaction", + params, + query_timeout_duration(Some(timeout_secs)), + ) + .await + .map(|result| truncate_result_with_max_rows(result, Some(row_limit.max(1)))) +} + /// Spawn a background task that removes and rolls back a transaction session /// after 5 minutes of inactivity. The task does not hold the global lock across /// I/O: it briefly checks the map, and if the session exists and is expired, @@ -5019,8 +5118,18 @@ pub async fn commit_manual_transaction(state: &AppState, txn_session_id: &str) - .map_err(|error| format!("COMMIT failed: {error}"))?; let _ = locked.disconnect().await; } + TxnConnection::ExternalDriver { session, config, .. } => { + session + .invoke_with_timeout::( + "commitManualTransaction", + serde_json::json!({ "connection": config.as_ref() }), + query_timeout_duration(Some(config.effective_query_timeout_secs())), + ) + .await + .map_err(|error| format!("COMMIT failed: {error}"))?; + } } - release_manual_txn_agent_pool(state, &session.connection_id, &mut conn).await; + release_manual_txn_session_pool(state, &session.connection_id, &mut conn).await; log::info!("[query][manual_txn:commit] session_id={}", txn_session_id); Ok(db::QueryResult { @@ -5049,7 +5158,7 @@ pub async fn rollback_manual_transaction(state: &AppState, txn_session_id: &str) let mut conn = session.connection.lock().await; rollback_manual_txn_connection(&mut conn).await?; - release_manual_txn_agent_pool(state, &session.connection_id, &mut conn).await; + release_manual_txn_session_pool(state, &session.connection_id, &mut conn).await; log::info!("[query][manual_txn:rollback] session_id={}", txn_session_id); Ok(db::QueryResult { @@ -6757,6 +6866,113 @@ for line in sys.stdin: )); } + #[cfg(unix)] + #[tokio::test] + async fn external_driver_manual_transaction_executes_commits_and_releases_session_pool() { + use std::os::unix::fs::PermissionsExt; + + let dir = std::env::temp_dir().join(format!("dbx-jdbc-manual-txn-test-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&dir).unwrap(); + let executable = dir.join("plugin.sh"); + let calls = dir.join("calls.log"); + std::fs::write( + &executable, + format!( + "#!/bin/sh\nwhile IFS= read -r line; do\n id=$(printf '%s' \"$line\" | sed -E 's/.*\"id\":([0-9]+).*/\\1/')\n case \"$line\" in\n *'\"method\":\"executeInManualTransaction\"'*)\n echo executeInManualTransaction >> '{}'\n printf '{{\"id\":%s,\"result\":{{\"columns\":[\"value\"],\"column_types\":[],\"column_sortables\":[],\"rows\":[[42]],\"affected_rows\":0,\"execution_time_ms\":1,\"truncated\":false,\"session_id\":null,\"has_more\":false}}}}\\n' \"$id\"\n ;;\n *'\"method\":\"beginManualTransaction\"'*)\n echo beginManualTransaction >> '{}'\n printf '{{\"id\":%s,\"result\":{{\"ok\":true}}}}\\n' \"$id\"\n ;;\n *'\"method\":\"commitManualTransaction\"'*)\n echo commitManualTransaction >> '{}'\n printf '{{\"id\":%s,\"result\":{{\"ok\":true}}}}\\n' \"$id\"\n ;;\n esac\ndone\n", + calls.display(), + calls.display(), + calls.display() + ), + ) + .unwrap(); + let mut permissions = std::fs::metadata(&executable).unwrap().permissions(); + permissions.set_mode(0o755); + std::fs::set_permissions(&executable, permissions).unwrap(); + + let plugin = InstalledPlugin { + manifest: PluginManifest { + id: "jdbc".to_string(), + name: "JDBC".to_string(), + version: "test".to_string(), + protocol_version: 1, + description: String::new(), + executable: Some("plugin.sh".to_string()), + drivers: vec![PluginDriverManifest { + id: "jdbc".to_string(), + label: "JDBC".to_string(), + kind: "external".to_string(), + database_type: Some("jdbc".to_string()), + }], + }, + path: dir.clone(), + }; + let session = Arc::new( + PluginDriverSession::start_for_test(plugin, "jdbc".to_string(), PluginRuntimeEnv::default()) + .await + .expect("plugin should start"), + ); + let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); + let state = AppState::new(storage); + let mut config = test_connection_config(DatabaseType::Jdbc); + config.id = "jdbc-conn".to_string(); + config.database = Some("dbx_test".to_string()); + config.connection_string = Some("jdbc:test".to_string()); + state.configs.write().await.insert(config.id.clone(), config.clone()); + + session + .invoke::( + "beginManualTransaction", + serde_json::json!({ "connection": &config, "database": "dbx_test" }), + ) + .await + .unwrap(); + + let client_session_id = "manual-txn-test"; + let pool_key = "jdbc-conn:session:manual-txn-test"; + let config = Arc::new(config); + state.connections.write().await.insert( + pool_key.to_string(), + PoolKind::ExternalDriver { + driver_id: "jdbc".to_string(), + config: config.clone(), + session: session.clone(), + }, + ); + let cleanup_guard = + state.workload_session_pool_cleanup_guard("jdbc-conn", Some("dbx_test"), client_session_id).await.unwrap(); + state.transaction_sessions.write().await.insert( + "txn-test".to_string(), + TransactionSession { + connection: Arc::new(tokio::sync::Mutex::new(TxnConnection::ExternalDriver { + session, + config, + client_session_id: client_session_id.to_string(), + database: Some("dbx_test".to_string()), + cleanup_guard, + })), + pool_key: pool_key.to_string(), + last_activity: std::time::Instant::now(), + busy: false, + connection_id: "jdbc-conn".to_string(), + database: "dbx_test".to_string(), + schema: None, + }, + ); + + let results = + execute_in_manual_transaction(&state, "txn-test", "SELECT 42", "dbx_test", None, Some(10)).await.unwrap(); + assert_eq!(results[0].rows, vec![vec![serde_json::json!(42)]]); + commit_manual_transaction(&state, "txn-test").await.unwrap(); + assert!(!state.connections.read().await.contains_key(pool_key)); + assert_eq!( + std::fs::read_to_string(&calls).unwrap(), + "beginManualTransaction\nexecuteInManualTransaction\ncommitManualTransaction\n" + ); + + state.shutdown(Duration::from_secs(1)).await; + let _ = std::fs::remove_dir_all(dir); + } + #[cfg(unix)] #[tokio::test] async fn external_driver_query_page_falls_back_to_legacy_execute_query() { 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 e1e3b729a9..13cf36a779 100644 --- a/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java +++ b/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java @@ -144,6 +144,7 @@ public final class DbxJdbcPlugin { private static String registeredDriverKey = ""; private static String sharedConnectionKey = ""; private static Connection sharedConnection; + private static boolean manualTransactionActive; private static final Map QUERY_SESSIONS = new HashMap<>(); record JdbcDriverQuirks( @@ -353,6 +354,23 @@ public final class DbxJdbcPlugin { nonNegativeInt(params, "rowOffset", 0), nonNegativeInt(params, "timeoutSecs", -1) ); + case "beginManualTransaction", "begin_manual_transaction" -> beginManualTransaction( + connection, + optionalText(params, "database"), + optionalText(params, "schema") + ); + case "executeInManualTransaction", "execute_in_manual_transaction" -> executeInManualTransaction( + connection, + requireText(params, "sql"), + optionalText(params, "database"), + optionalText(params, "schema"), + positiveInt(params, "maxRows", MAX_ROWS), + nonNegativeInt(params, "fetchSize", 0), + nonNegativeInt(params, "rowOffset", 0), + nonNegativeInt(params, "timeoutSecs", -1) + ); + case "commitManualTransaction", "commit_manual_transaction" -> commitManualTransaction(); + case "rollbackManualTransaction", "rollback_manual_transaction" -> rollbackManualTransaction(); case "executeQueryPage", "execute_query_page" -> executeQueryPage( connection, requireText(params, "sql"), @@ -462,6 +480,10 @@ public final class DbxJdbcPlugin { ); putMetadataText(info, "driverName", metadata::getDriverName); putMetadataText(info, "driverVersion", metadata::getDriverVersion); + Boolean supportsTransactions = readMetadata(metadata::supportsTransactions); + if (supportsTransactions != null) { + info.put("supportsTransactions", supportsTransactions); + } Integer jdbcMajor = readMetadata(metadata::getJDBCMajorVersion); Integer jdbcMinor = readMetadata(metadata::getJDBCMinorVersion); @@ -554,6 +576,7 @@ public final class DbxJdbcPlugin { } String key = connectionKey(connection); if (sharedConnection != null && key.equals(sharedConnectionKey) && !sharedConnection.isClosed()) { + configureOrdinaryAutoCommit(sharedConnection); return sharedConnection; } closeSharedConnection(); @@ -582,19 +605,22 @@ public final class DbxJdbcPlugin { applyOracleProperties(connection, properties); } sharedConnection = DriverManager.getConnection(url, properties); - configurePhoenixAutoCommit(connection, url, sharedConnection); sharedConnectionKey = key; + configureOrdinaryAutoCommit(sharedConnection); return sharedConnection; } - private static void configurePhoenixAutoCommit(JsonNode connection, String url, Connection jdbcConnection) - throws SQLException { - if (!isPhoenixConnection(connection, url) || jdbcConnection.getAutoCommit()) { + private static void configureOrdinaryAutoCommit(Connection jdbcConnection) throws SQLException { + if (manualTransactionActive || hasActiveQuerySession(jdbcConnection) || jdbcConnection.getAutoCommit()) { return; } jdbcConnection.setAutoCommit(true); } + private static boolean hasActiveQuerySession(Connection jdbcConnection) { + return QUERY_SESSIONS.values().stream().anyMatch(session -> session.connection == jdbcConnection); + } + private static boolean isPhoenixConnection(JsonNode connection, String url) { if (isPhoenixUrl(url)) { return true; @@ -775,8 +801,22 @@ public final class DbxJdbcPlugin { int rowOffset, int timeoutSecs ) throws Exception { - long start = System.nanoTime(); Connection conn = openConnection(connection); + return executeQueryOnConnection(connection, conn, sql, database, schema, maxRows, fetchSize, rowOffset, timeoutSecs); + } + + private static JsonNode executeQueryOnConnection( + JsonNode connection, + Connection conn, + String sql, + String database, + String schema, + int maxRows, + int fetchSize, + int rowOffset, + int timeoutSecs + ) throws Exception { + long start = System.nanoTime(); applyExecutionContext(connection, conn, database, schema); JdbcDriverQuirks quirks = driverQuirks(connection); boolean preserveOracleDateTime = isOracleUrl(jdbcUrl(connection)); @@ -826,6 +866,69 @@ public final class DbxJdbcPlugin { } } + private static ObjectNode beginManualTransaction(JsonNode connection, String database, String schema) + throws SQLException { + if (manualTransactionActive) { + throw new SQLException("A manual transaction is already active"); + } + Connection conn = openConnection(connection); + DatabaseMetaData metadata = readMetadata(conn::getMetaData); + Boolean supportsTransactions = metadata == null ? null : readMetadata(metadata::supportsTransactions); + if (Boolean.FALSE.equals(supportsTransactions)) { + throw new SQLFeatureNotSupportedException("This JDBC driver does not support transactions"); + } + applyExecutionContext(connection, conn, database, schema); + conn.setAutoCommit(false); + manualTransactionActive = true; + return okResult(); + } + + private static JsonNode executeInManualTransaction( + JsonNode connection, + String sql, + String database, + String schema, + int maxRows, + int fetchSize, + int rowOffset, + int timeoutSecs + ) throws Exception { + Connection conn = activeManualTransactionConnection(connection); + return executeQueryOnConnection(connection, conn, sql, database, schema, maxRows, fetchSize, rowOffset, timeoutSecs); + } + + private static ObjectNode commitManualTransaction() throws SQLException { + Connection conn = activeManualTransactionConnection(null); + conn.commit(); + conn.setAutoCommit(true); + manualTransactionActive = false; + return okResult(); + } + + private static ObjectNode rollbackManualTransaction() throws SQLException { + Connection conn = activeManualTransactionConnection(null); + conn.rollback(); + conn.setAutoCommit(true); + manualTransactionActive = false; + return okResult(); + } + + private static Connection activeManualTransactionConnection(JsonNode connection) throws SQLException { + if (!manualTransactionActive || sharedConnection == null || sharedConnection.isClosed()) { + throw new SQLException("No manual transaction is active"); + } + if (connection != null && !connectionKey(connection).equals(sharedConnectionKey)) { + throw new SQLException("The manual transaction belongs to a different JDBC connection"); + } + return sharedConnection; + } + + private static ObjectNode okResult() { + ObjectNode result = MAPPER.createObjectNode(); + result.put("ok", true); + return result; + } + private record ExecutedStatement(ResultSet resultSet, int updateCount) { } @@ -2836,6 +2939,12 @@ public final class DbxJdbcPlugin { private static void closeSharedConnection() { closeAllQuerySessions(); if (sharedConnection != null) { + if (manualTransactionActive) { + try { + sharedConnection.rollback(); + } catch (SQLException ignored) { + } + } try { sharedConnection.close(); } catch (SQLException ignored) { @@ -2843,6 +2952,7 @@ public final class DbxJdbcPlugin { sharedConnection = null; sharedConnectionKey = ""; } + manualTransactionActive = false; } private static String driverKey(JsonNode connection) { diff --git a/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java b/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java index 80c35cbef7..50d5a6cbcb 100644 --- a/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java +++ b/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java @@ -214,6 +214,120 @@ final class DbxJdbcPluginTest { assertFalse(info.path("driverName").asText().isEmpty()); assertFalse(info.path("driverVersion").asText().isEmpty()); assertFalse(info.path("jdbcVersion").asText().isEmpty()); + assertEquals(true, info.path("supportsTransactions").asBoolean()); + } + + @Test + void manualTransactionsCommitAndRollbackOnTheDedicatedConnection() throws Exception { + request("executeQuery", """ + { + "connection": %s, + "sql": "CREATE TABLE IF NOT EXISTS manual_tx_rows(id INT PRIMARY KEY)" + } + """.formatted(CONNECTION)); + request("executeQuery", """ + { + "connection": %s, + "sql": "DELETE FROM manual_tx_rows" + } + """.formatted(CONNECTION)); + + JsonNode begun = request("beginManualTransaction", """ + { "connection": %s } + """.formatted(CONNECTION)); + assertFalse(begun.has("error"), begun.toString()); + JsonNode insertedThenRolledBack = request("executeInManualTransaction", """ + { + "connection": %s, + "sql": "INSERT INTO manual_tx_rows VALUES (1)" + } + """.formatted(CONNECTION)); + assertFalse(insertedThenRolledBack.has("error"), insertedThenRolledBack.toString()); + JsonNode rolledBack = request("rollbackManualTransaction", """ + { "connection": %s } + """.formatted(CONNECTION)); + assertFalse(rolledBack.has("error"), rolledBack.toString()); + assertEquals(0, manualTransactionRowCount()); + + request("begin_manual_transaction", """ + { "connection": %s } + """.formatted(CONNECTION)); + request("execute_in_manual_transaction", """ + { + "connection": %s, + "sql": "INSERT INTO manual_tx_rows VALUES (2)" + } + """.formatted(CONNECTION)); + JsonNode committed = request("commit_manual_transaction", """ + { "connection": %s } + """.formatted(CONNECTION)); + assertFalse(committed.has("error"), committed.toString()); + assertEquals(1, manualTransactionRowCount()); + } + + @Test + void manualTransactionFallsBackWhenTransactionMetadataIsUnavailable() throws Exception { + for (Throwable failure : List.of( + new UnsupportedOperationException("unsupported"), + new AbstractMethodError("unsupported") + )) { + String connection = """ + { "connection_string": "jdbc:dbx-transaction-metadata:%s" } + """.formatted(failure.getClass().getSimpleName()); + Driver driver = testDriver( + "jdbc:dbx-transaction-metadata:", + transactionMetadataConnection(failure, true) + ); + DriverManager.registerDriver(driver); + try { + JsonNode begun = request("beginManualTransaction", """ + { "connection": %s } + """.formatted(connection)); + assertFalse(begun.has("error"), failure.getClass().getSimpleName() + ": " + begun); + + JsonNode rolledBack = request("rollbackManualTransaction", """ + { "connection": %s } + """.formatted(connection)); + assertFalse(rolledBack.has("error"), failure.getClass().getSimpleName() + ": " + rolledBack); + } finally { + closeAndDeregister(connection, driver); + } + } + } + + @Test + void manualTransactionRejectsExplicitUnsupportedMetadata() throws Exception { + String connection = """ + { "connection_string": "jdbc:dbx-transaction-metadata:false" } + """; + Driver driver = testDriver( + "jdbc:dbx-transaction-metadata:", + transactionMetadataConnection(null, false) + ); + DriverManager.registerDriver(driver); + try { + JsonNode response = request("beginManualTransaction", """ + { "connection": %s } + """.formatted(connection)); + + assertEquals( + "This JDBC driver does not support transactions", + response.path("error").path("message").asText() + ); + } finally { + closeAndDeregister(connection, driver); + } + } + + private static int manualTransactionRowCount() throws Exception { + JsonNode result = request("executeQuery", """ + { + "connection": %s, + "sql": "SELECT COUNT(*) FROM manual_tx_rows" + } + """.formatted(CONNECTION)); + assertFalse(result.has("error"), result.toString()); + return result.path("result").path("rows").path(0).path(0).asInt(); } @Test @@ -365,6 +479,51 @@ final class DbxJdbcPluginTest { assertEquals(true, third.path("result").path("session_id").isNull()); } + @Test + void ordinaryRequestDoesNotInvalidateActivePostgresPagedQuery() throws Exception { + List calls = new ArrayList<>(); + Driver driver = testDriver("jdbc:postgresql:dbx-interleaved", interleavedPostgresConnection(calls)); + DriverManager.registerDriver(driver); + String connection = """ + { "connection_string": "jdbc:postgresql:dbx-interleaved" } + """; + try { + JsonNode first = request("executeQueryPage", """ + { + "connection": %s, + "sql": "SELECT value FROM paged_values", + "pageSize": 1, + "maxRows": 10 + } + """.formatted(connection)); + assertFalse(first.has("error"), first.toString()); + assertEquals(1, first.path("result").path("rows").path(0).path(0).asInt()); + String sessionId = first.path("result").path("session_id").asText(); + + JsonNode ordinary = request("executeQuery", """ + { + "connection": %s, + "sql": "SELECT 99 AS ordinary_value" + } + """.formatted(connection)); + assertFalse(ordinary.has("error"), ordinary.toString()); + assertEquals(99, ordinary.path("result").path("rows").path(0).path(0).asInt()); + + JsonNode second = request("fetchQueryPage", """ + { + "connection": %s, + "sessionId": "%s", + "pageSize": 1 + } + """.formatted(connection, sessionId)); + assertFalse(second.has("error"), second.toString()); + assertEquals(2, second.path("result").path("rows").path(0).path(0).asInt()); + assertFalse(calls.contains("setAutoCommit:true"), calls.toString()); + } finally { + closeAndDeregister(connection, driver); + } + } + @Test void readValueFormatsDateColumnsWithoutMidnightTime() throws Exception { Method method = DbxJdbcPlugin.class.getDeclaredMethod( @@ -1017,22 +1176,15 @@ final class DbxJdbcPluginTest { } @Test - void phoenixConnectionsEnableAutoCommitWhenDriverDefaultsToManualTransactions() throws Exception { + void ordinaryConnectionsEnableAutoCommitWhenDriverDefaultsToManualTransactions() throws Exception { Method method = DbxJdbcPlugin.class.getDeclaredMethod( - "configurePhoenixAutoCommit", - JsonNode.class, - String.class, + "configureOrdinaryAutoCommit", Connection.class ); method.setAccessible(true); List calls = new ArrayList<>(); - JsonNode connection = MAPPER.readTree(""" - { - "connection_string": "jdbc:phoenix:localhost" - } - """); - method.invoke(null, connection, "jdbc:phoenix:localhost", pagedQueryConnection(calls, false)); + method.invoke(null, pagedQueryConnection(calls, false)); assertEquals(List.of("getAutoCommit", "setAutoCommit:true"), calls); } @@ -1149,43 +1301,15 @@ final class DbxJdbcPluginTest { } @Test - void phoenixAutoCommitConfigurationSkipsNonPhoenixConnections() throws Exception { + void ordinaryAutoCommitConfigurationDoesNotResetExistingAutoCommit() throws Exception { Method method = DbxJdbcPlugin.class.getDeclaredMethod( - "configurePhoenixAutoCommit", - JsonNode.class, - String.class, + "configureOrdinaryAutoCommit", Connection.class ); method.setAccessible(true); List calls = new ArrayList<>(); - JsonNode connection = MAPPER.readTree(""" - { - "connection_string": "jdbc:h2:mem:dbx" - } - """); - method.invoke(null, connection, "jdbc:h2:mem:dbx", pagedQueryConnection(calls, false)); - - assertEquals(List.of(), calls); - } - - @Test - void phoenixAutoCommitConfigurationDoesNotResetExistingAutoCommit() throws Exception { - Method method = DbxJdbcPlugin.class.getDeclaredMethod( - "configurePhoenixAutoCommit", - JsonNode.class, - String.class, - Connection.class - ); - method.setAccessible(true); - List calls = new ArrayList<>(); - JsonNode connection = MAPPER.readTree(""" - { - "jdbc_driver_class": "org.apache.phoenix.jdbc.PhoenixDriver" - } - """); - - method.invoke(null, connection, "jdbc:custom:phoenix", pagedQueryConnection(calls, true)); + method.invoke(null, pagedQueryConnection(calls, true)); assertEquals(List.of("getAutoCommit"), calls); } @@ -3150,6 +3274,154 @@ final class DbxJdbcPluginTest { ); } + private static Driver testDriver(String urlPrefix, Connection connection) { + return (Driver) Proxy.newProxyInstance( + DbxJdbcPluginTest.class.getClassLoader(), + new Class[] { Driver.class }, + (proxy, method, args) -> switch (method.getName()) { + case "connect" -> ((String) args[0]).startsWith(urlPrefix) ? connection : null; + case "acceptsURL" -> ((String) args[0]).startsWith(urlPrefix); + case "getPropertyInfo" -> new DriverPropertyInfo[0]; + case "getMajorVersion" -> 1; + case "getMinorVersion" -> 0; + case "jdbcCompliant" -> false; + case "getParentLogger" -> java.util.logging.Logger.getGlobal(); + default -> defaultValue(method.getReturnType()); + } + ); + } + + private static Connection transactionMetadataConnection(Throwable metadataFailure, boolean supportsTransactions) { + boolean[] autoCommit = { true }; + DatabaseMetaData metadata = (DatabaseMetaData) Proxy.newProxyInstance( + DbxJdbcPluginTest.class.getClassLoader(), + new Class[] { DatabaseMetaData.class }, + (proxy, method, args) -> { + if ("supportsTransactions".equals(method.getName())) { + if (metadataFailure != null) { + throw metadataFailure; + } + return supportsTransactions; + } + return defaultValue(method.getReturnType()); + } + ); + return (Connection) Proxy.newProxyInstance( + DbxJdbcPluginTest.class.getClassLoader(), + new Class[] { Connection.class }, + (proxy, method, args) -> switch (method.getName()) { + case "getAutoCommit" -> autoCommit[0]; + case "setAutoCommit" -> { + autoCommit[0] = (boolean) args[0]; + yield null; + } + case "getMetaData" -> metadata; + case "isClosed" -> false; + case "rollback", "close", "setCatalog", "setSchema" -> null; + default -> defaultValue(method.getReturnType()); + } + ); + } + + private static Connection interleavedPostgresConnection(List calls) { + boolean[] autoCommit = { true }; + boolean[] activePagedCursor = { false }; + boolean[] cursorInvalidated = { false }; + return (Connection) Proxy.newProxyInstance( + DbxJdbcPluginTest.class.getClassLoader(), + new Class[] { Connection.class }, + (proxy, method, args) -> switch (method.getName()) { + case "getAutoCommit" -> { + calls.add("getAutoCommit"); + yield autoCommit[0]; + } + case "setAutoCommit" -> { + boolean next = (boolean) args[0]; + calls.add("setAutoCommit:" + next); + if (next && !autoCommit[0] && activePagedCursor[0]) { + cursorInvalidated[0] = true; + } + autoCommit[0] = next; + yield null; + } + case "createStatement" -> interleavedStatement(activePagedCursor, cursorInvalidated); + case "isClosed" -> false; + case "rollback", "close", "setCatalog", "setSchema" -> null; + default -> defaultValue(method.getReturnType()); + } + ); + } + + private static Statement interleavedStatement(boolean[] activePagedCursor, boolean[] cursorInvalidated) { + return (Statement) Proxy.newProxyInstance( + DbxJdbcPluginTest.class.getClassLoader(), + new Class[] { Statement.class }, + new java.lang.reflect.InvocationHandler() { + private ResultSet resultSet; + + @Override + public Object invoke(Object proxy, Method method, Object[] args) { + return switch (method.getName()) { + case "execute" -> { + String sql = (String) args[0]; + boolean paged = sql.contains("paged_values"); + if (paged) { + activePagedCursor[0] = true; + } + resultSet = integerResultSet( + paged ? new int[] { 1, 2, 3 } : new int[] { 99 }, + paged, + activePagedCursor, + cursorInvalidated + ); + yield true; + } + case "getResultSet" -> resultSet; + case "getUpdateCount" -> -1; + case "setMaxRows", "setFetchSize", "setQueryTimeout", "close" -> null; + default -> defaultValue(method.getReturnType()); + }; + } + } + ); + } + + private static ResultSet integerResultSet( + int[] values, + boolean paged, + boolean[] activePagedCursor, + boolean[] cursorInvalidated + ) { + return (ResultSet) Proxy.newProxyInstance( + DbxJdbcPluginTest.class.getClassLoader(), + new Class[] { ResultSet.class }, + new java.lang.reflect.InvocationHandler() { + private int index = -1; + + @Override + public Object invoke(Object proxy, Method method, Object[] args) throws SQLException { + return switch (method.getName()) { + case "next" -> { + if (paged && cursorInvalidated[0]) { + throw new SQLException("PostgreSQL cursor was invalidated by auto-commit"); + } + yield ++index < values.length; + } + case "getMetaData" -> columnMeta(Types.INTEGER); + case "getObject" -> values[index]; + case "close" -> { + if (paged) { + activePagedCursor[0] = false; + } + yield null; + } + default -> defaultValue(method.getReturnType()); + }; + } + } + ); + } + private static final class RecordingConnectDriver implements Driver { private final String urlPrefix; private final List urls = new ArrayList<>();