fix(mysql): clear leftover transactions that pin stale read snapshots

crates/dbx-core/src/connection.rs

@@ -274,6 +274,24 @@ pub struct TransactionSession {
     pub schema: Option<String>,
 }
 
+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<String> = 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,

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::<Vec<_>>();
+        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)));