diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index 64a9c800ac..55791d0f3f 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -274,6 +274,24 @@ pub struct TransactionSession { pub schema: Option, } +impl TransactionSession { + /// Whether this session holds a natively connected MySQL/PostgreSQL + /// connection whose transaction keeps a read snapshot that a rollback + + /// fresh BEGIN can rotate in place. Agent and external-driver sessions are + /// excluded: their rollback path closes the dedicated session, so they + /// cannot reopen in place. + pub fn connection_is_native_snapshot_dialect(&self) -> bool { + // The connection is inspected under its own lock elsewhere; here only + // the variant is needed, and reading the enum discriminant without the + // lock would race. Callers hold the sessions lock and pass the same + // decision point, so mirror them by locking briefly. + match self.connection.try_lock() { + Ok(conn) => matches!(&*conn, TxnConnection::Mysql(Some(_)) | TxnConnection::Postgres(_)), + Err(_) => false, + } + } +} + macro_rules! agent_connection_pool_database_type { () => { DatabaseType::Dameng @@ -5018,17 +5036,45 @@ impl AppState { } pub async fn remove_connection_pools(&self, connection_id: &str) { + self.rollback_manual_transaction_sessions(connection_id).await; let removed = self.drain_connection_pools(connection_id).await; self.clear_metadata_gates_for_connection(connection_id).await; self.pool_routing_control().close_removed(removed).await; } pub async fn remove_connection_pools_detached(&self, connection_id: &str) { + self.rollback_manual_transaction_sessions(connection_id).await; let removed = self.drain_connection_pools(connection_id).await; self.clear_metadata_gates_for_connection(connection_id).await; self.pool_routing_control().close_removed_in_background(removed); } + /// Close and roll back every manual-transaction session of a connection + /// before its pools are drained. The transaction sessions hold dedicated + /// connections outside the pools being removed; without this the session + /// map survives a user disconnect with open transactions, and a later + /// reconnect/execution could observe or reuse the stale snapshot state. + /// Errors are logged and the session is dropped regardless: pool removal + /// is already the caller's decision. + async fn rollback_manual_transaction_sessions(&self, connection_id: &str) { + let sessions: Vec<(String, TransactionSession)> = { + let mut map = self.transaction_sessions.write().await; + let keys: Vec = map + .iter() + .filter(|(_, session)| session.connection_id == connection_id) + .map(|(id, _)| id.clone()) + .collect(); + keys.into_iter().filter_map(|id| map.remove(&id).map(|session| (id, session))).collect() + }; + for (session_id, session) in sessions { + let mut conn = session.connection.lock().await; + let outcome = crate::query::rollback_manual_txn_connection(&mut conn).await; + if let Err(error) = outcome { + log::warn!("[connection:manual-txn:rollback-on-disconnect] session={} error={}", session_id, error); + } + } + } + pub async fn invalidate_agent_pool_if_current( &self, pool_key: &str, diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index 8f5b15fcc2..73c73cb8da 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -3995,6 +3995,38 @@ async fn execute_multi_mysql( let statements_ms = statements_started_at.elapsed().as_millis(); drop(executor); + // Tab-scoped single-connection pools disable COM_RESET_CONNECTION on return + // (to preserve session state like temporary tables), so an open transaction + // left on the connection — a user-typed BEGIN/START TRANSACTION without + // COMMIT, or a canceled/aborted batch that skipped its cleanup — would pin + // the REPEATABLE READ snapshot for every later auto-commit query on that + // tab, making the tab read stale rows until disconnect. Closing any open + // transaction before returning the connection restores the auto-commit + // contract; ROLLBACK on an already-committed/implicit transaction is a + // server no-op, and a failure here only discards this connection. + { + let rollback_started_at = std::time::Instant::now(); + match conn.query_drop("ROLLBACK").await { + Ok(()) => { + if rollback_started_at.elapsed() > std::time::Duration::from_millis(5) { + log::info!( + "[query][mysql-batch] trace_id={} open_txn_rollback_ms={}", + trace_id, + rollback_started_at.elapsed().as_millis() + ); + } + } + Err(error) => { + // A failed ROLLBACK leaves the transaction state unknown: drop + // the connection instead of returning it to the session pool. + log::warn!("[query][mysql-batch] trace_id={} open_txn_rollback_failed error={}", trace_id, error); + let _ = tokio::time::timeout(Duration::from_secs(5), conn.disconnect()).await; + state.remove_pool_by_key(pool_key).await; + return Ok(results); + } + } + } + log::info!( "[query][mysql-batch] trace_id={} checkout_ms={} catalog_ms={} statements_ms={} total_ms={} result_count={} row_counts={:?}", trace_id, @@ -5856,6 +5888,29 @@ pub async fn execute_in_manual_transaction_with_options( } } } + + // Snapshot rotation for fully proven read-only batches (native MySQL/PG + // connections only): see rotate_clean_read_only_snapshot. Runs while the + // session is still marked busy so no concurrent execution can observe the + // half-rotated transaction. A rotation failure tears the session down and + // leans on the frontend rolled-back-session recovery (next execution opens + // a fresh session) instead of failing the already-successful batch. + if manual_txn_batch_fully_proven_read_only(&results) && { + let sessions = state.transaction_sessions.read().await; + sessions.get(txn_session_id).map(|session| session.connection_is_native_snapshot_dialect()).unwrap_or(false) + } { + if let Err(rotation_error) = rotate_clean_read_only_snapshot(&mut conn, schema).await { + let removed = { + let mut sessions = state.transaction_sessions.write().await; + sessions.remove(txn_session_id).is_some() + }; + if removed { + let _ = rollback_manual_txn_connection(&mut conn).await; + release_manual_txn_session_pool(state, &connection_id, &mut conn).await; + } + let _ = rotation_error; + } + } drop(conn); let should_watch = { @@ -6068,7 +6123,55 @@ where stream_result } -async fn rollback_manual_txn_connection(conn: &mut TxnConnection) -> Result<(), String> { +/// Whether a finished manual-transaction batch was fully proven read-only, so +/// the transaction still holds only a read view and no user data state. Factored +/// out so the snapshot-rotation decision is unit-testable without a database. +fn manual_txn_batch_fully_proven_read_only(results: &[ExecuteMultiResult]) -> bool { + !results.is_empty() && results.iter().all(|result| result.manual_transaction_proven_read_only) +} + +/// Rotate the transaction on a natively-connected MySQL/PostgreSQL session after +/// a fully proven read-only batch. MySQL REPEATABLE READ (and PostgreSQL when +/// the server default was changed to a snapshot isolation) pins the read view of +/// the first SELECT for the whole transaction: a user who keeps polling a clean +/// read-only session would otherwise never see rows committed by others, and the +/// clean-state toolbar hides Commit/Rollback, leaving disconnect/reconnect as +/// the only visible way out. A rollback of a read-only transaction and a fresh +/// BEGIN are both cheap metadata operations, and the batch gate keeps this off +/// every write path. Failures here never fail the (already successful) batch: +/// the session is torn down instead and the frontend's existing +/// rolled-back-session recovery transparently opens a new one on the next run. +async fn rotate_clean_read_only_snapshot(conn: &mut TxnConnection, schema: Option<&str>) -> Result<(), String> { + match conn { + TxnConnection::Mysql(Some(conn)) => { + conn.query_drop("ROLLBACK").await.map_err(|e| format!("ROLLBACK failed: {e}"))?; + conn.query_drop("START TRANSACTION").await.map_err(|e| format!("START TRANSACTION failed: {e}"))?; + Ok(()) + } + TxnConnection::Postgres(conn) => { + conn.execute_typed("ROLLBACK", &[]).await.map_err(|e| format!("ROLLBACK failed: {e}"))?; + conn.execute_typed("BEGIN", &[]).await.map_err(|e| format!("BEGIN failed: {e}"))?; + if let Some(schema) = schema { + db::postgres::set_postgres_search_path( + conn, + schema, + db::postgres::PostgresSearchPathContext::LocalTransaction, + db::connection_timeout(), + ) + .await + .map_err(|e| format!("SET search_path failed: {e}"))?; + } + Ok(()) + } + // Agent and external-driver sessions cannot reopen in place (their + // rollback path closes the dedicated session), so they keep the + // snapshot semantics; proven-read-only markers for those dialects do + // not reach this helper's native variants in practice. + _ => Ok(()), + } +} + +pub(crate) async fn rollback_manual_txn_connection(conn: &mut TxnConnection) -> Result<(), String> { rollback_manual_txn_connection_with_postgres_timeout(conn, None).await } @@ -6476,6 +6579,32 @@ mod tests { use super::*; use crate::query_cancel::RunningTaskMetadata; + #[test] + fn manual_txn_batch_proven_read_only_requires_non_empty_all_proven_results() { + // Empty batch (e.g. comments-only script) must not rotate: there is no + // snapshot to renew and rotating would churn a BEGIN for nothing. + assert!(!manual_txn_batch_fully_proven_read_only(&[])); + + let unproven = + vec![ExecuteMultiResult::success_with_optional_server_large_values(empty_query_result(1), false)]; + assert!(!manual_txn_batch_fully_proven_read_only(&unproven)); + + let proven = unproven + .clone() + .into_iter() + .map(|result| result.with_manual_transaction_proven_read_only()) + .collect::>(); + assert!(manual_txn_batch_fully_proven_read_only(&proven)); + + // One write statement anywhere in the batch keeps the transaction. + let mixed = vec![ + ExecuteMultiResult::success_with_optional_server_large_values(empty_query_result(1), false) + .with_manual_transaction_proven_read_only(), + ExecuteMultiResult::success_with_optional_server_large_values(empty_query_result(1), false), + ]; + assert!(!manual_txn_batch_fully_proven_read_only(&mixed)); + } + #[test] fn redshift_queries_prefer_text_protocol() { assert!(postgres_prefers_text_protocol(Some(DatabaseType::Redshift)));