diff --git a/CHANGELOG.md b/CHANGELOG.md index d970119..eb99765 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,25 @@ # Changelog +## Unreleased + +### Added + +- `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 + +- 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..cac9177 100644 --- a/README.md +++ b/README.md @@ -178,7 +178,8 @@ 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 | | `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..2d3ded0 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)); @@ -18,6 +19,37 @@ 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) { + // 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( + 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), @@ -42,41 +74,87 @@ 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); + + 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), + } +} - // 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), +/// 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> { + // 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 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 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(); + + // Acquire ONE connection for the entire batch (session state must survive) + let pg_client = match pinned { + Some(client) => client, + None => { + let pool = client::build_pool_pub(conn_params).await?; + pool.get() + .await + .map_err(|e| 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( - 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; + in_transaction = + transaction_effect(query).in_transaction_after(outcome.is_ok(), in_transaction); + match outcome { Ok(result) => results.push(json!({ "result": result, @@ -91,7 +169,28 @@ pub async fn execute_query_batch(id: Value, params: &Value) -> Value { } } - ok_response(id, json!(results)) + match slot.as_mut() { + Some(slot) if in_transaction => slot.pin(pg_client), + _ => { + if in_transaction { + session::rollback_and_release(pg_client).await; + } + } + } + + Ok((results, in_transaction)) +} + +/// 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 +396,130 @@ 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, + /// 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 | 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, + } + } +} + +/// 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 { + // 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 first.as_str() { + // `BEGIN` alone, `BEGIN TRANSACTION`, `BEGIN ISOLATION LEVEL …`. + "BEGIN" => TransactionEffect::Opens, + "START" if second == Some("TRANSACTION") => TransactionEffect::Opens, + // Two-phase commit acts on a prepared transaction, not this session's. + "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. + "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 79cb3d5..53ee778 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,172 @@ 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_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!(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)); + assert!(None.in_transaction_after(false, true)); + 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 [ + "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/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/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..f14b5ff --- /dev/null +++ b/src/session.rs @@ -0,0 +1,173 @@ +//! 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. +//! +//! 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::{Arc, OnceLock}; +use std::time::{Duration, Instant}; + +use deadpool_postgres::Client; +use tokio::sync::{Mutex, OwnedMutexGuard}; + +/// 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 session's pinned connection, if any. Holding the guard serializes runs. +#[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 { + 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>>; + +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}"); + } +} + +/// 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`], and forget +/// slots nothing holds. A slot in use by a run is skipped. +pub async fn sweep_idle() { + let mut expired = Vec::new(); + { + let now = Instant::now(); + 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; + } + 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; + } +} + +/// 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 = lock(session_id).await.take(); + 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. A session a run still holds is skipped +/// rather than waited for; its connection closes with the process. +pub async fn release_all() { + 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; + } + log::info!("Releasing {} pinned session(s)", clients.len()); + for client in clients { + 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..f05d940 --- /dev/null +++ b/src/session_tests.rs @@ -0,0 +1,65 @@ +//! 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(); +} + +#[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")); +} + +#[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); +} diff --git a/tests/live_db.rs b/tests/live_db.rs index 281dad9..9568b17 100644 --- a/tests/live_db.rs +++ b/tests/live_db.rs @@ -1688,3 +1688,53 @@ 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_ok( + "execute_query", + json!({ "params": params, "session_id": session, "query": "COMMIT" }), + ); + assert!( + 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( + "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" + ); +}