fix(postgres): recover schema session after aborted transaction

This commit is contained in:
t8y2
2026-06-05 14:46:00 +08:00
parent a2c9503a3d
commit 249a1ed649
2 changed files with 86 additions and 0 deletions
+28
View File
@@ -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]
@@ -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");
}