From 249a1ed6497eb968bc0f605c8d4af3945b1baa58 Mon Sep 17 00:00:00 2001
From: t8y2 <1156263951@qq.com>
Date: Fri, 5 Jun 2026 14:46:00 +0800
Subject: [PATCH] fix(postgres): recover schema session after aborted
 transaction

---
 crates/dbx-core/src/db/postgres.rs            | 28 +++++++++
 .../live_postgres_transaction_recovery.rs     | 58 +++++++++++++++++++
 2 files changed, 86 insertions(+)
 create mode 100644 crates/dbx-core/tests/live_postgres_transaction_recovery.rs

diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs
index 2a7a816822..5023540d3e 100644
--- a/crates/dbx-core/src/db/postgres.rs
+++ b/crates/dbx-core/src/db/postgres.rs
@@ -1435,6 +1435,14 @@ pub async fn execute_query_with_schema_and_max_rows(
         start.elapsed().as_millis(),
         schema
     );
+    if is_transaction_recovery_statement(sql) {
+        log::info!(
+            "[postgres][execute_with_schema:skip-search-path] total_ms={} reason=transaction-recovery",
+            start.elapsed().as_millis()
+        );
+        return execute_query_with_max_rows_inner(&client, sql, max_rows).await;
+    }
+
     let set_schema_start = Instant::now();
     client.execute(&format!("SET search_path TO {}", pg_quote_ident(schema)), &[]).await.map_err(pg_error_to_string)?;
     log::info!(
@@ -1464,6 +1472,10 @@ pub async fn execute_query_with_schema_and_max_rows(
     result
 }
 
+fn is_transaction_recovery_statement(sql: &str) -> bool {
+    starts_with_executable_sql_keyword(sql, &["ROLLBACK", "ABORT", "COMMIT", "END"])
+}
+
 async fn execute_query_with_max_rows_inner(
     client: &deadpool_postgres::Client,
     sql: &str,
@@ -2051,6 +2063,22 @@ mod tests {
         assert!(list_objects_sql(false).contains("pg_catalog.pg_proc"));
     }
 
+    #[test]
+    fn transaction_recovery_statement_detection_matches_common_postgres_commands() {
+        assert!(is_transaction_recovery_statement("ROLLBACK"));
+        assert!(is_transaction_recovery_statement("rollback work"));
+        assert!(is_transaction_recovery_statement("ABORT TRANSACTION"));
+        assert!(is_transaction_recovery_statement("commit"));
+        assert!(is_transaction_recovery_statement("END"));
+    }
+
+    #[test]
+    fn transaction_recovery_statement_detection_ignores_regular_queries() {
+        assert!(!is_transaction_recovery_statement("SELECT 1"));
+        assert!(!is_transaction_recovery_statement("BEGIN"));
+        assert!(!is_transaction_recovery_statement("UPDATE users SET name = 'dbx'"));
+    }
+
     // --- execute_batch ---
 
     #[tokio::test]
diff --git a/crates/dbx-core/tests/live_postgres_transaction_recovery.rs b/crates/dbx-core/tests/live_postgres_transaction_recovery.rs
new file mode 100644
index 0000000000..540f2c78ac
--- /dev/null
+++ b/crates/dbx-core/tests/live_postgres_transaction_recovery.rs
@@ -0,0 +1,58 @@
+use std::time::Duration;
+
+fn scalar_string(value: &serde_json::Value) -> &str {
+    value.as_str().expect("query result should be a string")
+}
+
+#[tokio::test]
+#[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a writable PostgreSQL database"]
+async fn live_postgres_schema_queries_can_recover_after_transaction_abort() {
+    let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL");
+    let pool = dbx_core::db::postgres::connect(&url, Duration::from_secs(5)).await.expect("connect");
+    let schema = format!("dbx_tx_recovery_{}", uuid::Uuid::new_v4().simple());
+
+    dbx_core::db::postgres::execute_batch(
+        &pool,
+        &[
+            format!("DROP SCHEMA IF EXISTS \"{schema}\" CASCADE"),
+            format!("CREATE SCHEMA \"{schema}\""),
+            format!("CREATE TABLE \"{schema}\".\"tx_probe\" (value integer NOT NULL)"),
+            format!("INSERT INTO \"{schema}\".\"tx_probe\" (value) VALUES (1)"),
+        ],
+    )
+    .await
+    .expect("prepare schema");
+
+    dbx_core::db::postgres::execute_query_with_schema(&pool, &schema, "BEGIN").await.expect("begin transaction");
+
+    let failed =
+        dbx_core::db::postgres::execute_query_with_schema(&pool, &schema, "UPDATE missing_tx_probe SET value = 2")
+            .await
+            .expect_err("invalid statement should fail");
+    assert!(failed.contains("missing_tx_probe"));
+
+    let aborted = dbx_core::db::postgres::execute_query_with_schema(&pool, &schema, "SELECT value FROM tx_probe")
+        .await
+        .expect_err("session should remain aborted before rollback");
+    assert!(aborted.contains("current transaction is aborted"));
+
+    dbx_core::db::postgres::execute_query_with_schema(&pool, &schema, "ROLLBACK")
+        .await
+        .expect("rollback should recover the session");
+
+    let recovered = dbx_core::db::postgres::execute_query_with_schema(
+        &pool,
+        &schema,
+        "SELECT current_schema(), value::text FROM tx_probe",
+    )
+    .await
+    .expect("query should succeed after rollback");
+
+    assert_eq!(recovered.rows.len(), 1);
+    assert_eq!(scalar_string(&recovered.rows[0][0]), schema);
+    assert_eq!(scalar_string(&recovered.rows[0][1]), "1");
+
+    dbx_core::db::postgres::execute_batch(&pool, &[format!("DROP SCHEMA IF EXISTS \"{schema}\" CASCADE")])
+        .await
+        .expect("cleanup schema");
+}