mirror of
https://github.com/t8y2/dbx.git
synced 2026-10-20 08:18:52 +00:00
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>
This commit is contained in:
zipg
@@ -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", () => {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -190,6 +190,14 @@ pub enum TxnConnection {
|
||||
database: Option<String>,
|
||||
cleanup_guard: ClientSessionPoolCleanupGuard,
|
||||
},
|
||||
/// Dedicated external-driver process whose shared JDBC connection owns the TX.
|
||||
ExternalDriver {
|
||||
session: Arc<PluginDriverSession>,
|
||||
config: Arc<ConnectionConfig>,
|
||||
client_session_id: String,
|
||||
database: Option<String>,
|
||||
cleanup_guard: ClientSessionPoolCleanupGuard,
|
||||
},
|
||||
}
|
||||
|
||||
pub struct TransactionSession {
|
||||
|
||||
@@ -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::<serde_json::Value>(
|
||||
"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::<serde_json::Value>(
|
||||
"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<crate::plugins::PluginDriverSession>,
|
||||
config: &Arc<crate::models::connection::ConnectionConfig>,
|
||||
statement: &str,
|
||||
database: &str,
|
||||
schema: Option<&str>,
|
||||
row_limit: usize,
|
||||
) -> Result<db::QueryResult, String> {
|
||||
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::<db::QueryResult>(
|
||||
"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::<serde_json::Value>(
|
||||
"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::<serde_json::Value>(
|
||||
"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() {
|
||||
|
||||
@@ -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<String, QuerySession> 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) {
|
||||
|
||||
@@ -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<String> 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<String> 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<String> 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<String> 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<String> 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<String> urls = new ArrayList<>();
|
||||
|
||||
Reference in New Issue
Block a user