mirror of
https://github.com/t8y2/dbx.git
synced 2026-10-19 16:16:45 +00:00
fix(postgres): recover schema session after aborted transaction
This commit is contained in:
@@ -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");
|
||||
}
|
||||
Reference in New Issue
Block a user