From f31a08a45bf165c3e57b4098ae0914d70b67350f Mon Sep 17 00:00:00 2001 From: Egert Aia Date: Mon, 21 Sep 2026 22:30:13 +0300 Subject: [PATCH 01/10] feat: keep a session's transaction alive across batches execute_query_batch already runs its statements on one pooled connection, so BEGIN ... COMMIT inside a script works. The workflow the transaction exists for did not: run BEGIN, inspect, run the changes, verify, then COMMIT, each as its own run. Between runs the connection went back to the pool, so the next run could land on a different one - the verify step saw pre-transaction data and COMMIT reported no transaction in progress while the real one stayed open elsewhere. The pool recycles with RecyclingMethod::Fast, which resets nothing, so that stranded connection also went back still inside the transaction, holding its locks, and the next borrower ran inside it. - transaction_effect() classifies a statement's effect on the surrounding transaction from its leading keywords only, so a BEGIN in a string literal or a PL/pgSQL block body is not mistaken for transaction control, and ROLLBACK TO SAVEPOINT correctly leaves it open. - session: registry of connections pinned per host session, a ROLLBACK before any pinned connection returns to the pool, and an idle sweep so an abandoned session cannot hold locks indefinitely. - execute_query_batch honours session_id and answers with { results, in_transaction } when one is given, keeping the bare array for hosts that do not send one. - release_session RPC method; shutdown releases every pinned session. - search_path is not re-applied to a reused connection: it would run inside the open transaction and change what the rest of it sees. --- CHANGELOG.md | 19 +++++ README.md | 1 + src/handlers/connection.rs | 4 ++ src/handlers/query.rs | 138 +++++++++++++++++++++++++++++++---- src/handlers/query_tests.rs | 96 ++++++++++++++++++++++++- src/lib.rs | 1 + src/rpc.rs | 1 + src/session.rs | 139 ++++++++++++++++++++++++++++++++++++ 8 files changed, 384 insertions(+), 15 deletions(-) create mode 100644 src/session.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index d970119..a366986 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,24 @@ # Changelog +## Unreleased + +### Added + +- `execute_query_batch` accepts a `session_id` and keeps that session's + connection when the batch leaves an explicit transaction open, so `BEGIN`, + the changes, a verifying `SELECT` and `COMMIT` can each be their own run + instead of having to be one script. When `session_id` is given the reply + carries `{ results, in_transaction }`; without one it stays the bare array, + so older hosts are unaffected. +- `release_session` RPC method, called when the owning editor tab closes. + +### Fixed + +- A batch that left a transaction open returned its connection to the pool + as-is. The pool recycles with `RecyclingMethod::Fast`, which resets + nothing, so the next borrower inherited the open transaction and its + locks. A connection is now always rolled back before it goes back. + ## [1.0.0-rc.4] - 2026-09-17 ### Added diff --git a/README.md b/README.md index 21fb5c3..0fbceda 100644 --- a/README.md +++ b/README.md @@ -179,6 +179,7 @@ instead of reconnecting. | `get_routines` / `get_routine_parameters` / `get_routine_definition` | Function/procedure metadata | | `get_triggers` / `get_trigger_definition` / `create_trigger` / `drop_trigger` | Trigger lifecycle | | `execute_query` / `execute_query_batch` / `explain_query` | Query execution, multi-statement batches, and query plans | +| `release_session` | Rolls back and releases the connection a session pinned for an open transaction | | `insert_record` / `update_record` / `delete_record` | Row-level CRUD with type-aware value binding | | `get_create_table_sql` / `get_add_column_sql` / `get_alter_column_sql` / `get_create_index_sql` / `get_create_foreign_key_sql` / `drop_index` / `drop_foreign_key` | DDL generation and execution | | `save_blob_to_file` / `fetch_blob_as_data_url` | BLOB (`bytea`) export and preview | diff --git a/src/handlers/connection.rs b/src/handlers/connection.rs index ec7d61b..823c3b4 100644 --- a/src/handlers/connection.rs +++ b/src/handlers/connection.rs @@ -40,5 +40,9 @@ pub async fn test_connection(id: Value, params: &Value) -> Value { /// Graceful shutdown — drain pools and exit. pub async fn shutdown(id: Value, _params: &Value) -> Value { + // Any session still holding a connection has an open transaction. The + // pool recycles without resetting, so roll them back rather than let + // the server keep them open until it times the backend out. + crate::session::release_all().await; ok_response(id, Value::Null) } diff --git a/src/handlers/query.rs b/src/handlers/query.rs index f459e3f..365d970 100644 --- a/src/handlers/query.rs +++ b/src/handlers/query.rs @@ -8,6 +8,7 @@ use crate::client; use crate::extract::extract_value; use crate::models::{inner_params, ConnectionParams}; use crate::rpc::{error_response, ok_response}; +use crate::session; pub async fn execute_query(id: Value, params: &Value) -> Value { let conn_params = ConnectionParams::from_value(inner_params(params)); @@ -42,24 +43,44 @@ pub async fn execute_query_batch(id: Value, params: &Value) -> Value { .map(|v| v as u32); let page = params.get("page").and_then(Value::as_u64).unwrap_or(1) as u32; let schema = params.get("schema").and_then(Value::as_str); + // The host sends the editor tab's id. A batch that leaves a transaction + // open keeps its connection under this key so the next batch from the + // same tab continues that transaction. + let session_id = params.get("session_id").and_then(Value::as_str); + + // Reuse the session's pinned connection when it left a transaction open. + let pinned = match session_id { + Some(sid) => session::take(sid).await, + None => None, + }; + // A pinned connection only exists because its transaction is still open. + let mut in_transaction = pinned.is_some(); + let reused = pinned.is_some(); // Acquire ONE connection for the entire batch (session state must survive) - let pool = match client::build_pool_pub(&conn_params).await { - Ok(p) => p, - Err(e) => return error_response(id, -32603, &e), - }; - let pg_client = match pool.get().await { - Ok(c) => c, - Err(e) => { - return error_response( - id, - -32603, - &format!("Connection failed: {}", client::format_pool_error(&e)), - ) + let pg_client = match pinned { + Some(client) => client, + None => { + let pool = match client::build_pool_pub(&conn_params).await { + Ok(p) => p, + Err(e) => return error_response(id, -32603, &e), + }; + match pool.get().await { + Ok(c) => c, + Err(e) => { + return error_response( + id, + -32603, + &format!("Connection failed: {}", client::format_pool_error(&e)), + ) + } + } } }; - if let Some(s) = schema { + // Applying it to a reused connection would run inside the open + // transaction and change what the rest of it sees. + if let Some(s) = schema.filter(|_| !reused) { let set_path = format!("SET search_path TO \"{}\"", s.replace('"', "\"\"")); if let Err(e) = pg_client.batch_execute(&set_path).await { return error_response( @@ -77,6 +98,17 @@ pub async fn execute_query_batch(id: Value, params: &Value) -> Value { let outcome = exec_query_on_client(&pg_client, query, limit, page).await; let elapsed_ms = start.elapsed().as_secs_f64() * 1000.0; + // A failed statement changes nothing: a failed COMMIT leaves the + // transaction open (aborted), and the session stays pinned so the + // user can still ROLLBACK from the same tab. + if outcome.is_ok() { + match transaction_effect(query) { + TransactionEffect::Opens => in_transaction = true, + TransactionEffect::Closes => in_transaction = false, + TransactionEffect::None => {} + } + } + match outcome { Ok(result) => results.push(json!({ "result": result, @@ -91,7 +123,38 @@ pub async fn execute_query_batch(id: Value, params: &Value) -> Value { } } - ok_response(id, json!(results)) + match session_id { + Some(sid) if in_transaction => session::store(sid, pg_client).await, + _ => { + if in_transaction { + // Nothing to pin the connection to, and the pool recycles + // without resetting, so it must not go back mid-transaction. + session::rollback_and_release(pg_client).await; + } + } + } + + match session_id { + // Only a session-aware call gets the richer shape; a host that did + // not send a session_id still receives the bare array it expects. + Some(_) => ok_response( + id, + json!({ "results": results, "in_transaction": in_transaction }), + ), + None => ok_response(id, json!(results)), + } +} + +/// Roll back and release the connection pinned to a session. Called when +/// the owning editor tab closes. +pub async fn release_session(id: Value, params: &Value) -> Value { + match params.get("session_id").and_then(Value::as_str) { + Some(session_id) => { + session::release(session_id).await; + ok_response(id, json!({ "released": true })) + } + None => error_response(id, -32602, "session_id is required"), + } } pub async fn explain_query(id: Value, params: &Value) -> Value { @@ -297,6 +360,53 @@ async fn exec_query_on_client( })) } +/// What a statement does to the surrounding transaction. +/// +/// Decides whether the connection a batch ran on must be kept for the next +/// batch of the same session (an explicit transaction is still open) or may +/// go back to the pool. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TransactionEffect { + /// Opens an explicit transaction (`BEGIN`, `START TRANSACTION`). + Opens, + /// Closes the current transaction (`COMMIT`, `ROLLBACK`, `END`). + /// + /// `ROLLBACK TO SAVEPOINT` does not close it and is classified as + /// [`TransactionEffect::None`]. + Closes, + /// Leaves the transaction state as it was. + None, +} + +/// Classify a statement's effect on the transaction state. +/// +/// Only the leading keywords are inspected, so a `BEGIN` inside a string +/// literal or a later clause cannot be mistaken for transaction control. A +/// PL/pgSQL `BEGIN … END` body is not a concern: it arrives inside a `DO` +/// or `CREATE FUNCTION` statement, whose leading keyword is neither. +pub fn transaction_effect(query: &str) -> TransactionEffect { + let normalized = strip_leading_sql_comments(query); + let mut words = normalized + .split(|c: char| !c.is_ascii_alphanumeric() && c != '_') + .filter(|w| !w.is_empty()) + .map(|w| w.to_uppercase()); + + let Some(first) = words.next() else { + return TransactionEffect::None; + }; + let second = words.next(); + + match first.as_str() { + "BEGIN" => TransactionEffect::Opens, + "START" if second.as_deref() == Some("TRANSACTION") => TransactionEffect::Opens, + "COMMIT" => TransactionEffect::Closes, + "ROLLBACK" if second.as_deref() == Some("TO") => TransactionEffect::None, + "ROLLBACK" => TransactionEffect::Closes, + "END" => TransactionEffect::Closes, + _ => TransactionEffect::None, + } +} + /// Strip leading SQL comments (`-- …` line comments and `/* … */` block /// comments) and whitespace so the first statement keyword is at position 0. /// Matches the builtin driver's `drivers/common/query.rs::strip_leading_sql_comments` diff --git a/src/handlers/query_tests.rs b/src/handlers/query_tests.rs index 79cb3d5..a6a3ade 100644 --- a/src/handlers/query_tests.rs +++ b/src/handlers/query_tests.rs @@ -12,7 +12,7 @@ use super::{ raw_explain_output, returns_result_set, strip_leading_sql_comments, - supports_trailing_limit_clause, + supports_trailing_limit_clause, transaction_effect, TransactionEffect, }; #[test] @@ -190,3 +190,97 @@ fn raw_explain_output_payload_is_not_the_live_json_value() { "payload must be Value::String, got {payload_value:?}" ); } + +#[test] +fn transaction_effect_detects_opening_statements() { + for query in [ + "BEGIN", + "begin;", + "BEGIN TRANSACTION", + "BEGIN ISOLATION LEVEL SERIALIZABLE", + "START TRANSACTION", + "start transaction read write", + ] { + assert_eq!( + transaction_effect(query), + TransactionEffect::Opens, + "{query} should open a transaction" + ); + } +} + +#[test] +fn transaction_effect_detects_closing_statements() { + for query in ["COMMIT", "commit;", "ROLLBACK", "END", "END TRANSACTION"] { + assert_eq!( + transaction_effect(query), + TransactionEffect::Closes, + "{query} should close the transaction" + ); + } +} + +#[test] +fn transaction_effect_ignores_savepoint_rollback() { + // Unwinding to a savepoint leaves the transaction open, so the + // connection must stay pinned to the session. + assert_eq!( + transaction_effect("ROLLBACK TO SAVEPOINT before_update"), + TransactionEffect::None + ); + assert_eq!( + transaction_effect("ROLLBACK TO before_update"), + TransactionEffect::None + ); +} + +#[test] +fn transaction_effect_ignores_ordinary_statements() { + for query in [ + "SELECT 1", + "UPDATE t SET a = 1", + "SAVEPOINT before_update", + // `BEGIN` appearing as data, not as the leading keyword. + "SELECT 'BEGIN' AS word", + "INSERT INTO log (msg) VALUES ('COMMIT')", + ] { + assert_eq!( + transaction_effect(query), + TransactionEffect::None, + "{query} should not change the transaction state" + ); + } +} + +#[test] +fn transaction_effect_sees_through_leading_comments() { + assert_eq!( + transaction_effect("-- start the transaction\nBEGIN"), + TransactionEffect::Opens + ); + assert_eq!( + transaction_effect("/* done */ COMMIT"), + TransactionEffect::Closes + ); +} + +#[test] +fn transaction_effect_handles_empty_input() { + assert_eq!(transaction_effect(""), TransactionEffect::None); + assert_eq!(transaction_effect(" \n"), TransactionEffect::None); + assert_eq!( + transaction_effect("-- only a comment"), + TransactionEffect::None + ); +} + +#[test] +fn transaction_effect_does_not_match_plpgsql_block_bodies() { + // A PL/pgSQL `BEGIN … END` arrives inside a DO or CREATE FUNCTION + // statement, whose leading keyword is neither, so the block body cannot + // be mistaken for transaction control. + assert_eq!( + transaction_effect("DO $$ BEGIN RAISE NOTICE 'hi'; END $$"), + TransactionEffect::None + ); +} diff --git a/src/lib.rs b/src/lib.rs index e5ca155..f576fcc 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -18,5 +18,6 @@ mod extract_tests; pub mod handlers; pub mod models; pub mod rpc; +pub mod session; pub mod settings; pub mod utils; diff --git a/src/rpc.rs b/src/rpc.rs index 4b6e8df..331108f 100644 --- a/src/rpc.rs +++ b/src/rpc.rs @@ -71,6 +71,7 @@ pub async fn handle_line(line: &str) -> Value { // Query execution "execute_query" => handlers::query::execute_query(id, ¶ms).await, "execute_query_batch" => handlers::query::execute_query_batch(id, ¶ms).await, + "release_session" => handlers::query::release_session(id, ¶ms).await, "explain_query" => handlers::query::explain_query(id, ¶ms).await, // CRUD diff --git a/src/session.rs b/src/session.rs new file mode 100644 index 0000000..26e84c0 --- /dev/null +++ b/src/session.rs @@ -0,0 +1,139 @@ +//! Connections pinned to a host session (an editor tab). +//! +//! `execute_query_batch` already runs every statement of one batch on a +//! single pooled connection, so `BEGIN … COMMIT` inside a script works. +//! What did not work is the workflow the transaction exists for: run +//! `BEGIN`, inspect, run the changes, verify them, and only then `COMMIT`, +//! each as its own run. Between runs the connection went back to the pool, +//! so the next run could land on a different one and the transaction was +//! stranded. +//! +//! When the host sends a `session_id`, a batch that leaves a transaction +//! open keeps its connection here until the session commits, rolls back, +//! closes, or goes idle. The pool recycles with `RecyclingMethod::Fast`, +//! which resets nothing, so a pinned connection is always rolled back +//! before it goes back to the pool — otherwise the next borrower would +//! inherit the transaction and its locks. + +use std::collections::HashMap; +use std::sync::OnceLock; +use std::time::{Duration, Instant}; + +use deadpool_postgres::Client; +use tokio::sync::Mutex; + +/// A pooled client held between batches because the session that owns it +/// left an explicit transaction open. +struct PinnedSession { + client: Client, + last_used: Instant, +} + +/// A pinned connection holds its transaction's locks until the session ends +/// it. An abandoned session would hold them indefinitely, so one untouched +/// for this long is rolled back and released the next time any session is +/// looked up. +const MAX_IDLE: Duration = Duration::from_secs(30 * 60); + +type SessionMap = HashMap; + +fn sessions() -> &'static Mutex { + static SESSIONS: OnceLock> = OnceLock::new(); + SESSIONS.get_or_init(|| Mutex::new(HashMap::new())) +} + +/// End the transaction before the client goes back to the pool. +/// +/// A failing `ROLLBACK` is only logged: the connection is already unusable +/// and the pool will discard it. +pub async fn rollback_and_release(client: Client) { + if let Err(e) = client.batch_execute("ROLLBACK").await { + log::warn!("ROLLBACK while releasing a pinned session failed: {e}"); + } +} + +/// Take the connection pinned to `session_id`, if any, and sweep sessions +/// idle past [`MAX_IDLE`]. +/// +/// The caller owns the returned client and must either hand it back via +/// [`store`] or end the transaction itself. +pub async fn take(session_id: &str) -> Option { + let (taken, expired) = { + let mut map = sessions().lock().await; + let taken = map.remove(session_id).map(|s| s.client); + let now = Instant::now(); + let stale: Vec = map + .iter() + .filter(|(_, s)| now.duration_since(s.last_used) > MAX_IDLE) + .map(|(id, _)| id.clone()) + .collect(); + let expired: Vec = stale + .iter() + .filter_map(|id| map.remove(id).map(|s| s.client)) + .collect(); + if !stale.is_empty() { + log::info!( + "Releasing {} pinned session(s) idle for over {} minutes", + stale.len(), + MAX_IDLE.as_secs() / 60 + ); + } + (taken, expired) + }; + + for client in expired { + rollback_and_release(client).await; + } + + taken +} + +/// Pin `client` to `session_id` until the session ends its transaction. +pub async fn store(session_id: &str, client: Client) { + let previous = { + let mut map = sessions().lock().await; + map.insert( + session_id.to_string(), + PinnedSession { + client, + last_used: Instant::now(), + }, + ) + }; + + // Only reachable if two batches for one session overlapped; the older + // connection is no longer referenced by anything. + if let Some(stale) = previous { + rollback_and_release(stale.client).await; + } +} + +/// Roll back and release the connection pinned to `session_id`, if any. +pub async fn release(session_id: &str) { + let client = { + let mut map = sessions().lock().await; + map.remove(session_id).map(|s| s.client) + }; + + if let Some(client) = client { + log::info!("Releasing pinned session {session_id}"); + rollback_and_release(client).await; + } +} + +/// Roll back and release every pinned connection, so shutdown leaves no +/// transaction open on the server. +pub async fn release_all() { + let clients: Vec = { + let mut map = sessions().lock().await; + map.drain().map(|(_, s)| s.client).collect() + }; + + if clients.is_empty() { + return; + } + log::info!("Releasing {} pinned session(s)", clients.len()); + for client in clients { + rollback_and_release(client).await; + } +} From 44364d75bcddfce61084e108728a8815935b615d Mon Sep 17 00:00:00 2001 From: Egert Aia Date: Mon, 21 Sep 2026 22:56:43 +0300 Subject: [PATCH 02/10] fix: pin the connection for single statements too Running statements one at a time is exactly how a transaction is driven - BEGIN, the changes, a verifying SELECT, COMMIT - and the host sends each of those through execute_query, not execute_query_batch. Without session handling there, every run took a fresh pooled connection: the UPDATE auto-committed on its own connection and the later COMMIT/ROLLBACK hit a connection with no transaction, where PostgreSQL only warns and reports success. - Extract run_batch_in_session() from execute_query_batch so a single statement and a batch take the same session path. - execute_query honours session_id and answers with { result, in_transaction } when one is given, keeping the bare QueryResult for hosts that do not send one. --- CHANGELOG.md | 13 +++--- README.md | 2 +- src/handlers/query.rs | 102 +++++++++++++++++++++++++++++------------- 3 files changed, 78 insertions(+), 39 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a366986..eb99765 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,12 +4,13 @@ ### Added -- `execute_query_batch` accepts a `session_id` and keeps that session's - connection when the batch leaves an explicit transaction open, so `BEGIN`, - the changes, a verifying `SELECT` and `COMMIT` can each be their own run - instead of having to be one script. When `session_id` is given the reply - carries `{ results, in_transaction }`; without one it stays the bare array, - so older hosts are unaffected. +- `execute_query` and `execute_query_batch` accept a `session_id` and keep + that session's connection when the statement or batch leaves an explicit + transaction open, so `BEGIN`, the changes, a verifying `SELECT` and + `COMMIT` can each be their own run instead of having to be one script. + With a `session_id` the replies carry `{ result, in_transaction }` and + `{ results, in_transaction }`; without one they stay the bare + `QueryResult` and array, so older hosts are unaffected. - `release_session` RPC method, called when the owning editor tab closes. ### Fixed diff --git a/README.md b/README.md index 0fbceda..cac9177 100644 --- a/README.md +++ b/README.md @@ -178,7 +178,7 @@ instead of reconnecting. | `get_materialized_views` / `refresh_materialized_view` | Materialized view lifecycle | | `get_routines` / `get_routine_parameters` / `get_routine_definition` | Function/procedure metadata | | `get_triggers` / `get_trigger_definition` / `create_trigger` / `drop_trigger` | Trigger lifecycle | -| `execute_query` / `execute_query_batch` / `explain_query` | Query execution, multi-statement batches, and query plans | +| `execute_query` / `execute_query_batch` / `explain_query` | Query execution, multi-statement batches, and query plans. Both query methods accept an optional `session_id` that keeps the connection while an explicit transaction is open | | `release_session` | Rolls back and releases the connection a session pinned for an open transaction | | `insert_record` / `update_record` / `delete_record` | Row-level CRUD with type-aware value binding | | `get_create_table_sql` / `get_add_column_sql` / `get_alter_column_sql` / `get_create_index_sql` / `get_create_foreign_key_sql` / `drop_index` / `drop_foreign_key` | DDL generation and execution | diff --git a/src/handlers/query.rs b/src/handlers/query.rs index 365d970..1b9ddec 100644 --- a/src/handlers/query.rs +++ b/src/handlers/query.rs @@ -19,6 +19,33 @@ pub async fn execute_query(id: Value, params: &Value) -> Value { .map(|v| v as u32); let page = params.get("page").and_then(Value::as_u64).unwrap_or(1) as u32; let schema = params.get("schema").and_then(Value::as_str); + let session_id = params.get("session_id").and_then(Value::as_str); + + // One statement at a time is how a transaction is actually driven - + // BEGIN, the changes, a verifying SELECT, COMMIT - so this path needs + // the same pinning as a batch. Implemented as a one-statement batch so + // there is a single session code path. + if session_id.is_some() { + let queries = [query.to_string()]; + return match run_batch_in_session(&conn_params, &queries, limit, page, schema, session_id) + .await + { + Ok((mut results, in_transaction)) => { + let statement = results.pop().unwrap_or(Value::Null); + match statement.get("error").and_then(Value::as_str) { + Some(error) => error_response(id, -32603, error), + None => { + let result = statement.get("result").cloned().unwrap_or(Value::Null); + ok_response( + id, + json!({ "result": result, "in_transaction": in_transaction }), + ) + } + } + } + Err(e) => error_response(id, -32603, &e), + }; + } match exec_query(&conn_params, query, limit, page, schema).await { Ok(result) => ok_response(id, result), @@ -48,7 +75,39 @@ pub async fn execute_query_batch(id: Value, params: &Value) -> Value { // same tab continues that transaction. let session_id = params.get("session_id").and_then(Value::as_str); - // Reuse the session's pinned connection when it left a transaction open. + match run_batch_in_session(&conn_params, &queries, limit, page, schema, session_id).await { + // Only a session-aware call gets the richer shape; a host that did + // not send a session_id still receives the bare array it expects. + Ok((results, in_transaction)) => match session_id { + Some(_) => ok_response( + id, + json!({ "results": results, "in_transaction": in_transaction }), + ), + None => ok_response(id, json!(results)), + }, + Err(e) => error_response(id, -32603, &e), + } +} + +/// Run `queries` on one connection, reusing the session's pinned one when +/// it left a transaction open and re-pinning it if one is still open after. +/// +/// Returns one per-statement result object per query — the same shape +/// `execute_query_batch` replies with — plus whether the session is still +/// inside a transaction. Shared with `execute_query` so a single statement +/// and a batch take the same session path. +/// +/// A batch that leaves a transaction open with no session to pin it to is +/// rolled back: the pool recycles with `RecyclingMethod::Fast`, so the next +/// borrower would otherwise inherit the transaction and its locks. +async fn run_batch_in_session( + conn_params: &ConnectionParams, + queries: &[String], + limit: Option, + page: u32, + schema: Option<&str>, + session_id: Option<&str>, +) -> Result<(Vec, bool), String> { let pinned = match session_id { Some(sid) => session::take(sid).await, None => None, @@ -61,20 +120,10 @@ pub async fn execute_query_batch(id: Value, params: &Value) -> Value { let pg_client = match pinned { Some(client) => client, None => { - let pool = match client::build_pool_pub(&conn_params).await { - Ok(p) => p, - Err(e) => return error_response(id, -32603, &e), - }; - match pool.get().await { - Ok(c) => c, - Err(e) => { - return error_response( - id, - -32603, - &format!("Connection failed: {}", client::format_pool_error(&e)), - ) - } - } + let pool = client::build_pool_pub(conn_params).await?; + pool.get() + .await + .map_err(|e| format!("Connection failed: {}", client::format_pool_error(&e)))? } }; @@ -83,17 +132,16 @@ pub async fn execute_query_batch(id: Value, params: &Value) -> Value { if let Some(s) = schema.filter(|_| !reused) { let set_path = format!("SET search_path TO \"{}\"", s.replace('"', "\"\"")); if let Err(e) = pg_client.batch_execute(&set_path).await { - return error_response( - id, - -32603, - &format!("Failed to set search_path: {}", client::format_pg_error(&e)), - ); + return Err(format!( + "Failed to set search_path: {}", + client::format_pg_error(&e) + )); } } let mut results: Vec = Vec::new(); - for query in &queries { + for query in queries { let start = Instant::now(); let outcome = exec_query_on_client(&pg_client, query, limit, page).await; let elapsed_ms = start.elapsed().as_secs_f64() * 1000.0; @@ -127,22 +175,12 @@ pub async fn execute_query_batch(id: Value, params: &Value) -> Value { Some(sid) if in_transaction => session::store(sid, pg_client).await, _ => { if in_transaction { - // Nothing to pin the connection to, and the pool recycles - // without resetting, so it must not go back mid-transaction. session::rollback_and_release(pg_client).await; } } } - match session_id { - // Only a session-aware call gets the richer shape; a host that did - // not send a session_id still receives the bare array it expects. - Some(_) => ok_response( - id, - json!({ "results": results, "in_transaction": in_transaction }), - ), - None => ok_response(id, json!(results)), - } + Ok((results, in_transaction)) } /// Roll back and release the connection pinned to a session. Called when From 09849ce504855347418b88eae6683777b2be3c3d Mon Sep 17 00:00:00 2001 From: Egert Aia Date: Sat, 26 Sep 2026 12:18:04 +0300 Subject: [PATCH 03/10] fix: end the session's transaction when COMMIT itself fails PostgreSQL rolls the transaction back when COMMIT fails on its own (e.g. a deferred constraint), so keeping the session pinned left a connection held with no transaction behind it. Also classify ABORT, PREPARE TRANSACTION, COMMIT/ROLLBACK PREPARED and ... AND CHAIN, the last of which previously returned a connection to the pool with a fresh transaction still open. --- src/handlers/query.rs | 61 ++++++++++++++++++++++--------------- src/handlers/query_tests.rs | 45 +++++++++++++++++++++++++++ tests/live_db.rs | 45 +++++++++++++++++++++++++++ 3 files changed, 126 insertions(+), 25 deletions(-) diff --git a/src/handlers/query.rs b/src/handlers/query.rs index 1b9ddec..9cdd8b2 100644 --- a/src/handlers/query.rs +++ b/src/handlers/query.rs @@ -146,16 +146,8 @@ async fn run_batch_in_session( let outcome = exec_query_on_client(&pg_client, query, limit, page).await; let elapsed_ms = start.elapsed().as_secs_f64() * 1000.0; - // A failed statement changes nothing: a failed COMMIT leaves the - // transaction open (aborted), and the session stays pinned so the - // user can still ROLLBACK from the same tab. - if outcome.is_ok() { - match transaction_effect(query) { - TransactionEffect::Opens => in_transaction = true, - TransactionEffect::Closes => in_transaction = false, - TransactionEffect::None => {} - } - } + in_transaction = + transaction_effect(query).in_transaction_after(outcome.is_ok(), in_transaction); match outcome { Ok(result) => results.push(json!({ @@ -412,10 +404,24 @@ pub enum TransactionEffect { /// `ROLLBACK TO SAVEPOINT` does not close it and is classified as /// [`TransactionEffect::None`]. Closes, + /// Ends the current transaction and opens a new one (`COMMIT AND CHAIN`). + Chains, /// Leaves the transaction state as it was. None, } +impl TransactionEffect { + /// Whether a transaction is open after a statement with this effect ran. + pub fn in_transaction_after(self, succeeded: bool, before: bool) -> bool { + match (self, succeeded) { + (TransactionEffect::Opens, true) => true, + // PostgreSQL ends the transaction even when COMMIT itself fails, e.g. on a deferred constraint. + (TransactionEffect::Closes, _) | (TransactionEffect::Chains, false) => false, + _ => before, + } + } +} + /// Classify a statement's effect on the transaction state. /// /// Only the leading keywords are inspected, so a `BEGIN` inside a string @@ -424,23 +430,28 @@ pub enum TransactionEffect { /// or `CREATE FUNCTION` statement, whose leading keyword is neither. pub fn transaction_effect(query: &str) -> TransactionEffect { let normalized = strip_leading_sql_comments(query); - let mut words = normalized + // Five words cover the longest form, `COMMIT TRANSACTION AND NO CHAIN`. + let words: Vec = normalized .split(|c: char| !c.is_ascii_alphanumeric() && c != '_') .filter(|w| !w.is_empty()) - .map(|w| w.to_uppercase()); - - let Some(first) = words.next() else { - return TransactionEffect::None; - }; - let second = words.next(); - - match first.as_str() { - "BEGIN" => TransactionEffect::Opens, - "START" if second.as_deref() == Some("TRANSACTION") => TransactionEffect::Opens, - "COMMIT" => TransactionEffect::Closes, - "ROLLBACK" if second.as_deref() == Some("TO") => TransactionEffect::None, - "ROLLBACK" => TransactionEffect::Closes, - "END" => TransactionEffect::Closes, + .take(5) + .map(|w| w.to_uppercase()) + .collect(); + let second = words.get(1).map(String::as_str); + let chains = words.windows(2).any(|p| p[0] == "AND" && p[1] == "CHAIN"); + + match words.first().map(String::as_str) { + // `BEGIN` alone, `BEGIN TRANSACTION`, `BEGIN ISOLATION LEVEL …`. + Some("BEGIN") => TransactionEffect::Opens, + Some("START") if second == Some("TRANSACTION") => TransactionEffect::Opens, + // Two-phase commit acts on a prepared transaction, not this session's. + Some("COMMIT" | "ROLLBACK") if second == Some("PREPARED") => TransactionEffect::None, + // `ROLLBACK TO [SAVEPOINT] x` unwinds to a savepoint and leaves the transaction open. + Some("ROLLBACK") if second == Some("TO") => TransactionEffect::None, + Some("COMMIT" | "END" | "ROLLBACK" | "ABORT") if chains => TransactionEffect::Chains, + Some("COMMIT" | "END" | "ROLLBACK" | "ABORT") => TransactionEffect::Closes, + // `PREPARE TRANSACTION` dissociates the transaction from the session. + Some("PREPARE") if second == Some("TRANSACTION") => TransactionEffect::Closes, _ => TransactionEffect::None, } } diff --git a/src/handlers/query_tests.rs b/src/handlers/query_tests.rs index a6a3ade..94f1a74 100644 --- a/src/handlers/query_tests.rs +++ b/src/handlers/query_tests.rs @@ -234,6 +234,51 @@ fn transaction_effect_ignores_savepoint_rollback() { ); } +#[test] +fn transaction_effect_detects_chained_and_two_phase_statements() { + for query in [ + "COMMIT AND CHAIN", + "commit work and chain", + "ROLLBACK AND CHAIN", + "ABORT AND CHAIN", + ] { + assert_eq!( + transaction_effect(query), + TransactionEffect::Chains, + "{query}" + ); + } + for query in ["COMMIT AND NO CHAIN", "ABORT", "PREPARE TRANSACTION 'tx1'"] { + assert_eq!( + transaction_effect(query), + TransactionEffect::Closes, + "{query}" + ); + } + // Two-phase commit runs outside a transaction block and never touches the session's. + for query in ["COMMIT PREPARED 'tx1'", "ROLLBACK PREPARED 'tx1'"] { + assert_eq!( + transaction_effect(query), + TransactionEffect::None, + "{query}" + ); + } +} + +#[test] +fn in_transaction_after_treats_a_failed_commit_as_closing() { + use TransactionEffect::*; + // A COMMIT failing on a deferred constraint has already ended the transaction server-side. + assert!(!Closes.in_transaction_after(false, true)); + assert!(!Chains.in_transaction_after(false, true)); + assert!(Chains.in_transaction_after(true, true)); + assert!(Opens.in_transaction_after(true, false)); + // A failed BEGIN or ordinary statement leaves the state as it was. + assert!(!Opens.in_transaction_after(false, false)); + assert!(None.in_transaction_after(false, true)); + assert!(!None.in_transaction_after(true, false)); +} + #[test] fn transaction_effect_ignores_ordinary_statements() { for query in [ diff --git a/tests/live_db.rs b/tests/live_db.rs index 281dad9..6a53ad4 100644 --- a/tests/live_db.rs +++ b/tests/live_db.rs @@ -1688,3 +1688,48 @@ fn execute_query_decodes_pgvector_types_correctly() { assert_eq!(all_null[col_idx], Value::Null, "{col_name} must be null"); } } + +#[test] +fn commit_that_fails_on_a_deferred_constraint_releases_the_session() { + let mut plugin = Plugin::spawn(); + let params = conn_params(); + for query in [ + "DROP TABLE IF EXISTS live_db_deferred_fk_scratch", + "CREATE TABLE live_db_deferred_fk_scratch (id INT PRIMARY KEY, other_id INT \ + REFERENCES live_db_deferred_fk_scratch(id) DEFERRABLE INITIALLY DEFERRED)", + ] { + plugin.call_ok("execute_query", json!({ "params": params, "query": query })); + } + + let session = json!("live-db-session-commit-fails"); + plugin.call_ok( + "execute_query", + json!({ "params": params, "session_id": session, "query": "BEGIN" }), + ); + // The FK check is deferred, so the INSERT succeeds and only COMMIT fails. + plugin.call_ok( + "execute_query", + json!({ + "params": params, "session_id": session, + "query": "INSERT INTO live_db_deferred_fk_scratch VALUES (1, 999)" + }), + ); + let commit = plugin.call( + "execute_query", + json!({ "params": params, "session_id": session, "query": "COMMIT" }), + ); + assert!( + commit.get("error").is_some(), + "COMMIT violating a deferred FK must error" + ); + + let after = plugin.call_ok( + "execute_query", + json!({ "params": params, "session_id": session, "query": "SELECT 1" }), + ); + assert_eq!( + after["in_transaction"], + json!(false), + "a failed COMMIT already ended the transaction server-side" + ); +} From 14548f49660f5600f26c309913e8d785927f6552 Mon Sep 17 00:00:00 2001 From: Egert Aia Date: Sat, 26 Sep 2026 12:19:48 +0300 Subject: [PATCH 04/10] fix: sweep idle pinned sessions on the cleanup timer The idle check only ran inside take(), so an abandoned session kept its transaction's locks until some other session happened to run a query. Hang sweep_idle() off the existing 10-minute pool cleanup timer instead. --- src/main.rs | 5 ++++- src/session.rs | 37 +++++++++++++++++++------------------ 2 files changed, 23 insertions(+), 19 deletions(-) diff --git a/src/main.rs b/src/main.rs index c9727c1..f1a10bc 100644 --- a/src/main.rs +++ b/src/main.rs @@ -62,7 +62,10 @@ async fn run_pool_cleanup(mut shutdown_rx: watch::Receiver) { let mut timer = interval(POOL_CLEANUP_INTERVAL); loop { tokio::select! { - _ = timer.tick() => postgresql_plugin::client::cleanup_idle_pools(), + _ = timer.tick() => { + postgresql_plugin::client::cleanup_idle_pools(); + postgresql_plugin::session::sweep_idle().await; + } _ = shutdown_rx.changed() => break, } } diff --git a/src/session.rs b/src/session.rs index 26e84c0..060cea7 100644 --- a/src/session.rs +++ b/src/session.rs @@ -31,8 +31,8 @@ struct PinnedSession { /// A pinned connection holds its transaction's locks until the session ends /// it. An abandoned session would hold them indefinitely, so one untouched -/// for this long is rolled back and released the next time any session is -/// looked up. +/// for this long is rolled back and released by the periodic +/// [`sweep_idle`]. const MAX_IDLE: Duration = Duration::from_secs(30 * 60); type SessionMap = HashMap; @@ -52,40 +52,41 @@ pub async fn rollback_and_release(client: Client) { } } -/// Take the connection pinned to `session_id`, if any, and sweep sessions -/// idle past [`MAX_IDLE`]. +/// Take the connection pinned to `session_id`, if any. /// /// The caller owns the returned client and must either hand it back via /// [`store`] or end the transaction itself. pub async fn take(session_id: &str) -> Option { - let (taken, expired) = { + sessions().lock().await.remove(session_id).map(|s| s.client) +} + +/// Roll back and release every session idle past [`MAX_IDLE`]. +pub async fn sweep_idle() { + let expired: Vec = { let mut map = sessions().lock().await; - let taken = map.remove(session_id).map(|s| s.client); let now = Instant::now(); let stale: Vec = map .iter() .filter(|(_, s)| now.duration_since(s.last_used) > MAX_IDLE) .map(|(id, _)| id.clone()) .collect(); - let expired: Vec = stale + stale .iter() .filter_map(|id| map.remove(id).map(|s| s.client)) - .collect(); - if !stale.is_empty() { - log::info!( - "Releasing {} pinned session(s) idle for over {} minutes", - stale.len(), - MAX_IDLE.as_secs() / 60 - ); - } - (taken, expired) + .collect() }; + if expired.is_empty() { + return; + } + log::info!( + "Releasing {} pinned session(s) idle for over {} minutes", + expired.len(), + MAX_IDLE.as_secs() / 60 + ); for client in expired { rollback_and_release(client).await; } - - taken } /// Pin `client` to `session_id` until the session ends its transaction. From 7005c15df1777c00b010c6a3ffa5548439749e4a Mon Sep 17 00:00:00 2001 From: Egert Aia Date: Sun, 27 Sep 2026 13:16:04 +0300 Subject: [PATCH 05/10] fix: serialize runs for the same session Two overlapping calls for one session both found nothing pinned, each took a fresh connection, and the later store() rolled back the other's transaction. Each session now has its own lock held for the whole take, run, pin sequence, so an overlapping call waits its turn on the same connection, and release_session waits for a run still in flight. --- src/handlers/query.rs | 10 ++-- src/session.rs | 119 ++++++++++++++++++++++++------------------ src/session_tests.rs | 43 +++++++++++++++ 3 files changed, 116 insertions(+), 56 deletions(-) create mode 100644 src/session_tests.rs diff --git a/src/handlers/query.rs b/src/handlers/query.rs index 9cdd8b2..3517814 100644 --- a/src/handlers/query.rs +++ b/src/handlers/query.rs @@ -108,10 +108,12 @@ async fn run_batch_in_session( schema: Option<&str>, session_id: Option<&str>, ) -> Result<(Vec, bool), String> { - let pinned = match session_id { - Some(sid) => session::take(sid).await, + // Held for the whole run, so an overlapping call for the same session waits its turn. + let mut slot = match session_id { + Some(sid) => Some(session::lock(sid).await), None => None, }; + let pinned = slot.as_mut().and_then(|s| s.take()); // A pinned connection only exists because its transaction is still open. let mut in_transaction = pinned.is_some(); let reused = pinned.is_some(); @@ -163,8 +165,8 @@ async fn run_batch_in_session( } } - match session_id { - Some(sid) if in_transaction => session::store(sid, pg_client).await, + match slot.as_mut() { + Some(slot) if in_transaction => slot.pin(pg_client), _ => { if in_transaction { session::rollback_and_release(pg_client).await; diff --git a/src/session.rs b/src/session.rs index 060cea7..e706a36 100644 --- a/src/session.rs +++ b/src/session.rs @@ -14,13 +14,17 @@ //! which resets nothing, so a pinned connection is always rolled back //! before it goes back to the pool — otherwise the next borrower would //! inherit the transaction and its locks. +//! +//! Each session has its own lock, held for a whole run, so two overlapping +//! calls for one session run one after the other on the same connection +//! instead of each taking a fresh one. use std::collections::HashMap; -use std::sync::OnceLock; +use std::sync::{Arc, OnceLock}; use std::time::{Duration, Instant}; use deadpool_postgres::Client; -use tokio::sync::Mutex; +use tokio::sync::{Mutex, OwnedMutexGuard}; /// A pooled client held between batches because the session that owns it /// left an explicit transaction open. @@ -29,13 +33,32 @@ struct PinnedSession { last_used: Instant, } +/// A session's pinned connection, if any. Holding the guard serializes runs. +#[derive(Default)] +pub struct Slot(Option); + +impl Slot { + /// Take the pinned connection; the caller must [`Slot::pin`] it again or end its transaction. + pub fn take(&mut self) -> Option { + self.0.take().map(|s| s.client) + } + + /// Pin `client` until the session ends its transaction. + pub fn pin(&mut self, client: Client) { + self.0 = Some(PinnedSession { + client, + last_used: Instant::now(), + }); + } +} + /// A pinned connection holds its transaction's locks until the session ends /// it. An abandoned session would hold them indefinitely, so one untouched /// for this long is rolled back and released by the periodic /// [`sweep_idle`]. const MAX_IDLE: Duration = Duration::from_secs(30 * 60); -type SessionMap = HashMap; +type SessionMap = HashMap>>; fn sessions() -> &'static Mutex { static SESSIONS: OnceLock> = OnceLock::new(); @@ -52,29 +75,39 @@ pub async fn rollback_and_release(client: Client) { } } -/// Take the connection pinned to `session_id`, if any. -/// -/// The caller owns the returned client and must either hand it back via -/// [`store`] or end the transaction itself. -pub async fn take(session_id: &str) -> Option { - sessions().lock().await.remove(session_id).map(|s| s.client) +/// Lock `session_id`'s slot, waiting for any run already holding it. +pub async fn lock(session_id: &str) -> OwnedMutexGuard { + let slot = sessions() + .lock() + .await + .entry(session_id.to_string()) + .or_default() + .clone(); + slot.lock_owned().await } -/// Roll back and release every session idle past [`MAX_IDLE`]. +/// Roll back and release every session idle past [`MAX_IDLE`], and forget +/// slots nothing holds. A slot in use by a run is skipped. pub async fn sweep_idle() { - let expired: Vec = { - let mut map = sessions().lock().await; + let mut expired = Vec::new(); + { let now = Instant::now(); - let stale: Vec = map - .iter() - .filter(|(_, s)| now.duration_since(s.last_used) > MAX_IDLE) - .map(|(id, _)| id.clone()) - .collect(); - stale - .iter() - .filter_map(|id| map.remove(id).map(|s| s.client)) - .collect() - }; + let mut map = sessions().lock().await; + map.retain(|_, slot| { + let Ok(mut guard) = slot.try_lock() else { + return true; + }; + if guard + .0 + .as_ref() + .is_some_and(|s| now.duration_since(s.last_used) > MAX_IDLE) + { + expired.extend(guard.take()); + } + // Only the map holds it and nothing is pinned, so nobody can be waiting on it. + guard.0.is_some() || Arc::strong_count(slot) > 1 + }); + } if expired.is_empty() { return; @@ -89,33 +122,10 @@ pub async fn sweep_idle() { } } -/// Pin `client` to `session_id` until the session ends its transaction. -pub async fn store(session_id: &str, client: Client) { - let previous = { - let mut map = sessions().lock().await; - map.insert( - session_id.to_string(), - PinnedSession { - client, - last_used: Instant::now(), - }, - ) - }; - - // Only reachable if two batches for one session overlapped; the older - // connection is no longer referenced by anything. - if let Some(stale) = previous { - rollback_and_release(stale.client).await; - } -} - -/// Roll back and release the connection pinned to `session_id`, if any. +/// Roll back and release the connection pinned to `session_id`, if any, +/// after any run still in flight for it finishes. pub async fn release(session_id: &str) { - let client = { - let mut map = sessions().lock().await; - map.remove(session_id).map(|s| s.client) - }; - + let client = lock(session_id).await.take(); if let Some(client) = client { log::info!("Releasing pinned session {session_id}"); rollback_and_release(client).await; @@ -125,10 +135,11 @@ pub async fn release(session_id: &str) { /// Roll back and release every pinned connection, so shutdown leaves no /// transaction open on the server. pub async fn release_all() { - let clients: Vec = { - let mut map = sessions().lock().await; - map.drain().map(|(_, s)| s.client).collect() - }; + let slots: Vec>> = sessions().lock().await.drain().map(|(_, s)| s).collect(); + let mut clients = Vec::new(); + for slot in slots { + clients.extend(slot.lock().await.take()); + } if clients.is_empty() { return; @@ -138,3 +149,7 @@ pub async fn release_all() { rollback_and_release(client).await; } } + +#[cfg(test)] +#[path = "session_tests.rs"] +mod session_tests; diff --git a/src/session_tests.rs b/src/session_tests.rs new file mode 100644 index 0000000..de406b9 --- /dev/null +++ b/src/session_tests.rs @@ -0,0 +1,43 @@ +//! Unit tests for `session.rs`'s per-session lock. Pinning itself needs a +//! live connection and is covered by `tests/live_db.rs`. + +use super::{lock, sweep_idle}; +use std::time::Duration; + +#[tokio::test] +async fn a_second_run_for_the_same_session_waits_for_the_first() { + let first = lock("session-tests-same").await; + let second = tokio::spawn(async { drop(lock("session-tests-same").await) }); + + tokio::time::sleep(Duration::from_millis(50)).await; + assert!( + !second.is_finished(), + "overlapping run must wait for the lock" + ); + + drop(first); + tokio::time::timeout(Duration::from_secs(1), second) + .await + .expect("second run proceeds once the first releases") + .unwrap(); +} + +#[tokio::test] +async fn different_sessions_do_not_block_each_other() { + let _a = lock("session-tests-a").await; + tokio::time::timeout(Duration::from_secs(1), lock("session-tests-b")) + .await + .expect("another session must not wait"); +} + +#[tokio::test] +async fn sweep_keeps_a_slot_that_a_run_holds() { + let held = lock("session-tests-held").await; + sweep_idle().await; + let waiter = tokio::spawn(async { drop(lock("session-tests-held").await) }); + tokio::time::sleep(Duration::from_millis(50)).await; + // Had the sweep dropped the slot, the waiter would get a fresh one and not block. + assert!(!waiter.is_finished()); + drop(held); + waiter.await.unwrap(); +} From 6e9857cad596f1c254187378b052519dd0f2f957 Mon Sep 17 00:00:00 2001 From: Egert Aia Date: Sun, 27 Sep 2026 13:16:06 +0300 Subject: [PATCH 06/10] fix: treat a successful AND CHAIN as in a transaction It opens a new transaction whatever the prior state was. --- src/handlers/query.rs | 2 +- src/handlers/query_tests.rs | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/src/handlers/query.rs b/src/handlers/query.rs index 3517814..de8fe3d 100644 --- a/src/handlers/query.rs +++ b/src/handlers/query.rs @@ -416,7 +416,7 @@ impl TransactionEffect { /// Whether a transaction is open after a statement with this effect ran. pub fn in_transaction_after(self, succeeded: bool, before: bool) -> bool { match (self, succeeded) { - (TransactionEffect::Opens, true) => true, + (TransactionEffect::Opens | TransactionEffect::Chains, true) => true, // PostgreSQL ends the transaction even when COMMIT itself fails, e.g. on a deferred constraint. (TransactionEffect::Closes, _) | (TransactionEffect::Chains, false) => false, _ => before, diff --git a/src/handlers/query_tests.rs b/src/handlers/query_tests.rs index 94f1a74..3932017 100644 --- a/src/handlers/query_tests.rs +++ b/src/handlers/query_tests.rs @@ -272,6 +272,7 @@ fn in_transaction_after_treats_a_failed_commit_as_closing() { assert!(!Closes.in_transaction_after(false, true)); assert!(!Chains.in_transaction_after(false, true)); assert!(Chains.in_transaction_after(true, true)); + assert!(Chains.in_transaction_after(true, false)); assert!(Opens.in_transaction_after(true, false)); // A failed BEGIN or ordinary statement leaves the state as it was. assert!(!Opens.in_transaction_after(false, false)); From 3173317238ff09fb50fd48e815d2bba8ca42d8bd Mon Sep 17 00:00:00 2001 From: Egert Aia Date: Sun, 27 Sep 2026 13:53:09 +0300 Subject: [PATCH 07/10] fix: skip a busy session on shutdown instead of waiting for it Also covers the sweep forgetting a slot nothing uses. --- src/session.rs | 11 ++++++----- src/session_tests.rs | 10 ++++++++++ 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/src/session.rs b/src/session.rs index e706a36..06e58ac 100644 --- a/src/session.rs +++ b/src/session.rs @@ -133,13 +133,14 @@ pub async fn release(session_id: &str) { } /// Roll back and release every pinned connection, so shutdown leaves no -/// transaction open on the server. +/// transaction open on the server. A session a run still holds is skipped +/// rather than waited for; its connection closes with the process. pub async fn release_all() { let slots: Vec>> = sessions().lock().await.drain().map(|(_, s)| s).collect(); - let mut clients = Vec::new(); - for slot in slots { - clients.extend(slot.lock().await.take()); - } + let clients: Vec = slots + .iter() + .filter_map(|slot| slot.try_lock().ok().and_then(|mut s| s.take())) + .collect(); if clients.is_empty() { return; diff --git a/src/session_tests.rs b/src/session_tests.rs index de406b9..8b83feb 100644 --- a/src/session_tests.rs +++ b/src/session_tests.rs @@ -41,3 +41,13 @@ async fn sweep_keeps_a_slot_that_a_run_holds() { drop(held); waiter.await.unwrap(); } + +#[tokio::test] +async fn sweep_forgets_a_slot_nothing_uses() { + drop(lock("session-tests-unused").await); + sweep_idle().await; + assert!(!super::sessions() + .lock() + .await + .contains_key("session-tests-unused")); +} From 54b5b13d2fabea609e96487a5ccdbbc775770659 Mon Sep 17 00:00:00 2001 From: Egert Aia Date: Mon, 28 Sep 2026 10:17:40 +0300 Subject: [PATCH 08/10] fix: read transaction keywords past comments and optional WORK/TRANSACTION ROLLBACK WORK/TRANSACTION TO SAVEPOINT was classified as closing, returning a connection to the pool mid-transaction, and only leading comments were skipped, so COMMIT -- and chain read as chaining and START /* x */ TRANSACTION was missed. transaction_effect now reads keywords with every comment skipped and drops the optional noise word. --- src/handlers/query.rs | 96 +++++++++++++++++++++++++++++-------- src/handlers/query_tests.rs | 29 +++++++++++ 2 files changed, 106 insertions(+), 19 deletions(-) diff --git a/src/handlers/query.rs b/src/handlers/query.rs index de8fe3d..f5376b8 100644 --- a/src/handlers/query.rs +++ b/src/handlers/query.rs @@ -431,33 +431,91 @@ impl TransactionEffect { /// PL/pgSQL `BEGIN … END` body is not a concern: it arrives inside a `DO` /// or `CREATE FUNCTION` statement, whose leading keyword is neither. pub fn transaction_effect(query: &str) -> TransactionEffect { - let normalized = strip_leading_sql_comments(query); - // Five words cover the longest form, `COMMIT TRANSACTION AND NO CHAIN`. - let words: Vec = normalized - .split(|c: char| !c.is_ascii_alphanumeric() && c != '_') - .filter(|w| !w.is_empty()) - .take(5) - .map(|w| w.to_uppercase()) - .collect(); - let second = words.get(1).map(String::as_str); - let chains = words.windows(2).any(|p| p[0] == "AND" && p[1] == "CHAIN"); + // Six cover the longest form, `ROLLBACK TRANSACTION AND NO CHAIN`. + let mut words = leading_keywords(query, 6).into_iter(); + let Some(first) = words.next() else { + return TransactionEffect::None; + }; + let mut rest: Vec = words.collect(); + // `WORK` / `TRANSACTION` after a transaction-control verb is optional noise. + if matches!(first.as_str(), "COMMIT" | "END" | "ROLLBACK" | "ABORT") + && matches!( + rest.first().map(String::as_str), + Some("WORK" | "TRANSACTION") + ) + { + rest.remove(0); + } + let second = rest.first().map(String::as_str); + let chains = rest.windows(2).any(|p| p[0] == "AND" && p[1] == "CHAIN"); - match words.first().map(String::as_str) { + match first.as_str() { // `BEGIN` alone, `BEGIN TRANSACTION`, `BEGIN ISOLATION LEVEL …`. - Some("BEGIN") => TransactionEffect::Opens, - Some("START") if second == Some("TRANSACTION") => TransactionEffect::Opens, + "BEGIN" => TransactionEffect::Opens, + "START" if second == Some("TRANSACTION") => TransactionEffect::Opens, // Two-phase commit acts on a prepared transaction, not this session's. - Some("COMMIT" | "ROLLBACK") if second == Some("PREPARED") => TransactionEffect::None, - // `ROLLBACK TO [SAVEPOINT] x` unwinds to a savepoint and leaves the transaction open. - Some("ROLLBACK") if second == Some("TO") => TransactionEffect::None, - Some("COMMIT" | "END" | "ROLLBACK" | "ABORT") if chains => TransactionEffect::Chains, - Some("COMMIT" | "END" | "ROLLBACK" | "ABORT") => TransactionEffect::Closes, + "COMMIT" | "ROLLBACK" if second == Some("PREPARED") => TransactionEffect::None, + // `ROLLBACK [WORK] TO [SAVEPOINT] x` unwinds to a savepoint and leaves the transaction open. + "ROLLBACK" if second == Some("TO") => TransactionEffect::None, + "COMMIT" | "END" | "ROLLBACK" | "ABORT" if chains => TransactionEffect::Chains, + "COMMIT" | "END" | "ROLLBACK" | "ABORT" => TransactionEffect::Closes, // `PREPARE TRANSACTION` dissociates the transaction from the session. - Some("PREPARE") if second == Some("TRANSACTION") => TransactionEffect::Closes, + "PREPARE" if second == Some("TRANSACTION") => TransactionEffect::Closes, _ => TransactionEffect::None, } } +/// The first `n` keywords of `query`, uppercased. Every `--` and (nested) +/// `/* */` comment is skipped, wherever it appears, and reading stops at a +/// string literal, dollar quote or `;`, so no data is read as a keyword. +fn leading_keywords(query: &str, n: usize) -> Vec { + let mut words = Vec::new(); + let mut word = String::new(); + let mut chars = query.chars().peekable(); + let mut comment_depth = 0usize; + while let Some(c) = chars.next() { + if comment_depth > 0 { + if c == '*' && chars.peek() == Some(&'/') { + chars.next(); + comment_depth -= 1; + } else if c == '/' && chars.peek() == Some(&'*') { + chars.next(); + comment_depth += 1; + } + continue; + } + if c.is_ascii_alphanumeric() || c == '_' { + word.push(c.to_ascii_uppercase()); + continue; + } + if !word.is_empty() { + words.push(std::mem::take(&mut word)); + if words.len() == n { + return words; + } + } + match (c, chars.peek()) { + ('-', Some('-')) => { + for c in chars.by_ref() { + if c == '\n' { + break; + } + } + } + ('/', Some('*')) => { + chars.next(); + comment_depth = 1; + } + ('\'' | '"' | '$' | ';', _) => return words, + _ => {} + } + } + if !word.is_empty() && words.len() < n { + words.push(word); + } + words +} + /// Strip leading SQL comments (`-- …` line comments and `/* … */` block /// comments) and whitespace so the first statement keyword is at position 0. /// Matches the builtin driver's `drivers/common/query.rs::strip_leading_sql_comments` diff --git a/src/handlers/query_tests.rs b/src/handlers/query_tests.rs index 3932017..53ee778 100644 --- a/src/handlers/query_tests.rs +++ b/src/handlers/query_tests.rs @@ -280,6 +280,35 @@ fn in_transaction_after_treats_a_failed_commit_as_closing() { assert!(!None.in_transaction_after(true, false)); } +#[test] +fn transaction_effect_reads_optional_noise_words_and_inner_comments() { + for (query, effect) in [ + ("ROLLBACK WORK TO SAVEPOINT sp1", TransactionEffect::None), + ( + "ROLLBACK TRANSACTION TO SAVEPOINT sp1", + TransactionEffect::None, + ), + ("COMMIT WORK", TransactionEffect::Closes), + ( + "COMMIT -- and chain later if needed", + TransactionEffect::Closes, + ), + ("COMMIT /* and chain */", TransactionEffect::Closes), + ("START /* explicit */ TRANSACTION", TransactionEffect::Opens), + ( + "COMMIT /* outer /* nested */ still comment */ AND CHAIN", + TransactionEffect::Chains, + ), + ( + "ROLLBACK TRANSACTION AND NO CHAIN", + TransactionEffect::Closes, + ), + ("PREPARE TRANSACTION 'and chain'", TransactionEffect::Closes), + ] { + assert_eq!(transaction_effect(query), effect, "{query}"); + } +} + #[test] fn transaction_effect_ignores_ordinary_statements() { for query in [ From 66ee2f75052e3fb5e6d1e4edf4ec80e9d115c676 Mon Sep 17 00:00:00 2001 From: Egert Aia Date: Mon, 28 Sep 2026 10:18:16 +0300 Subject: [PATCH 09/10] fix: keep a busy session tracked on shutdown and close any dropped pin release_all drained every slot before skipping busy ones, so a run still in flight pinned its connection into a slot nothing tracked, and it went back to the pool without a ROLLBACK. Busy slots now stay in the map, and a slot dropped while holding a connection closes it instead of returning it. --- src/session.rs | 27 ++++++++++++++++++++++----- src/session_tests.rs | 12 ++++++++++++ 2 files changed, 34 insertions(+), 5 deletions(-) diff --git a/src/session.rs b/src/session.rs index 06e58ac..f14b5ff 100644 --- a/src/session.rs +++ b/src/session.rs @@ -37,6 +37,16 @@ struct PinnedSession { #[derive(Default)] pub struct Slot(Option); +/// A slot dropped while still holding a connection closes it rather than +/// returning it to the pool, where the next borrower would inherit its transaction. +impl Drop for Slot { + fn drop(&mut self) { + if let Some(pinned) = self.0.take() { + drop(Client::take(pinned.client)); + } + } +} + impl Slot { /// Take the pinned connection; the caller must [`Slot::pin`] it again or end its transaction. pub fn take(&mut self) -> Option { @@ -136,11 +146,18 @@ pub async fn release(session_id: &str) { /// transaction open on the server. A session a run still holds is skipped /// rather than waited for; its connection closes with the process. pub async fn release_all() { - let slots: Vec>> = sessions().lock().await.drain().map(|(_, s)| s).collect(); - let clients: Vec = slots - .iter() - .filter_map(|slot| slot.try_lock().ok().and_then(|mut s| s.take())) - .collect(); + let mut clients: Vec = Vec::new(); + // A busy slot stays in the map, so the run holding it can still pin into a tracked slot. + sessions() + .lock() + .await + .retain(|_, slot| match slot.try_lock() { + Ok(mut guard) => { + clients.extend(guard.take()); + false + } + Err(_) => true, + }); if clients.is_empty() { return; diff --git a/src/session_tests.rs b/src/session_tests.rs index 8b83feb..f05d940 100644 --- a/src/session_tests.rs +++ b/src/session_tests.rs @@ -51,3 +51,15 @@ async fn sweep_forgets_a_slot_nothing_uses() { .await .contains_key("session-tests-unused")); } + +#[tokio::test] +async fn release_all_keeps_tracking_a_slot_a_run_holds() { + let held = lock("session-tests-busy-shutdown").await; + super::release_all().await; + // The run would otherwise pin into a slot nothing tracks, and it would never be rolled back. + assert!(super::sessions() + .lock() + .await + .contains_key("session-tests-busy-shutdown")); + drop(held); +} From df5e44823d0431bf812b4d70f0e43f8fa1ef064c Mon Sep 17 00:00:00 2001 From: Egert Aia Date: Mon, 28 Sep 2026 10:19:08 +0300 Subject: [PATCH 10/10] fix: report in_transaction on a failed session execute_query A failing statement came back as an RPC error, so the host never learned that a failed COMMIT had ended the transaction. With a session_id it is now a reply carrying the error and in_transaction. --- src/handlers/query.rs | 6 +++++- tests/live_db.rs | 11 ++++++++--- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/src/handlers/query.rs b/src/handlers/query.rs index f5376b8..2d3ded0 100644 --- a/src/handlers/query.rs +++ b/src/handlers/query.rs @@ -33,7 +33,11 @@ pub async fn execute_query(id: Value, params: &Value) -> Value { Ok((mut results, in_transaction)) => { let statement = results.pop().unwrap_or(Value::Null); match statement.get("error").and_then(Value::as_str) { - Some(error) => error_response(id, -32603, error), + // Still a reply, not an RPC error, so the host learns the state a failed COMMIT left. + Some(error) => ok_response( + id, + json!({ "error": error, "in_transaction": in_transaction }), + ), None => { let result = statement.get("result").cloned().unwrap_or(Value::Null); ok_response( diff --git a/tests/live_db.rs b/tests/live_db.rs index 6a53ad4..9568b17 100644 --- a/tests/live_db.rs +++ b/tests/live_db.rs @@ -1714,13 +1714,18 @@ fn commit_that_fails_on_a_deferred_constraint_releases_the_session() { "query": "INSERT INTO live_db_deferred_fk_scratch VALUES (1, 999)" }), ); - let commit = plugin.call( + let commit = plugin.call_ok( "execute_query", json!({ "params": params, "session_id": session, "query": "COMMIT" }), ); assert!( - commit.get("error").is_some(), - "COMMIT violating a deferred FK must error" + commit["error"].is_string(), + "COMMIT violating a deferred FK must report its error" + ); + assert_eq!( + commit["in_transaction"], + json!(false), + "the failing COMMIT's own reply must say the transaction ended" ); let after = plugin.call_ok(