github/dbx
4
0 Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-11-14 22:50:26 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(jdbc): support manual transactions and auto-commit control (#6997)

Browse Source 
* 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:
zipgandzipg authored and GitHub committed 2026-08-24 14:50:18 +08:00
1 parent c8dbeeaf7d
commit c021fe009b
6 files changed
+662 -55
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

crates/dbx-core/src/connection.rs
+8
View File
@@ -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 {
crates/dbx-core/src/query.rs
+224 -8
View File
@@ -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() {