From 6a7aea44bc8a40ae19061bd88ca8995659685d65 Mon Sep 17 00:00:00 2001 From: Adam J Esslinger Date: Fri, 23 Oct 2026 12:03:41 -0400 Subject: [PATCH 1/2] fix: keep bare [SET/CREATE TEMP, SELECT] batch on the shared client on a transaction pooler (#173) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PR #172 (#169 fix) routed a non-`in_transaction` result-set query in a batch to the one-shot pinned path (`exec_query_pooled_txn`) on a separate connection. A bare batch (no `session_id`, no `BEGIN`) that sets session-local state in an earlier non-result-set statement — a session `SET` (not `SET LOCAL`) or a `CREATE TEMP TABLE` outside a transaction — then runs a result-set `SELECT` silently loses that state, because the SELECT runs on a fresh one-shot connection that never saw the `SET`/temp table (the one-shot applies only the host's `schema` via `SET LOCAL` and does not replay prior session state). Result: the SELECT returns the wrong precision (GUC case) or `ERROR: relation "t" does not exist` (temp-table case). No error surfaces about the lost state. On `main` (rc.11, before #172), `run_batch_in_session` ran every query via `exec_query_on_client` on one shared `pg_client`, so a `SET` then `SELECT` in a single bare batch reliably shared state (same connection, same backend). #172 broke that for the pooler `!in_transaction` case. Fix: detect session-local state set earlier in the bare batch and keep the result-set SELECT on the shared batch client with a column-type pre-resolve, instead of the one-shot. A new pure classifier `sets_session_local_state` matches a session `SET` (without `LOCAL`) or a `CREATE TEMP[ORARY] TABLE`; transaction-scoped forms (`SET LOCAL`) and permanent `CREATE TABLE` are excluded — the `in_transaction` branch already keeps the former on the shared client, and the latter is visible to any connection. A sticky `session_state_set` flag in `run_batch_in_session` flips after such a statement, and the routing condition `in_transaction || session_state_set` reuses the existing `in_transaction` pre-resolve path verbatim. No regression to #166/#168/#160/#163/#169: the fix only adds a third routing condition that routes to the same `exec_query_on_client( pre_resolve_types=true)` path the `in_transaction` branch already uses. It never alters the direct-connection #166 path (`!is_transaction_pooler()`), the #164 one-shot path, or the #169 `in_transaction` path. Warming `CachedTypeInfo.types` via `prepare` means `query_typed`'s `get_type` is a cache hit and `typeinfo_statement` never fires mid-drain on any connection — the identical mechanism #166/#168 and the #169 `in_transaction` branch rely on, so no new hang/deadlock is introduced. Scope: restores the direct-PostgreSQL + `pooler_mode=transaction` config (one TCP connection = one backend, held) and the temp-table case — the clearest *new* regressions #173 identifies. On a real PgBouncer transaction pooler, a session `SET` between autocommit statements is already lost to backend reassignment on `main`, so the pure-GUC case is pre-existing there (not a new regression) and stays out of scope; the existing session-pinned `BEGIN`/`COMMIT` path (covered by `wide_enum_table_returns_rows_via_session_ pinned_batch_on_pooler`) remains the PgBouncer-canonical workaround. Adds 9 unit tests for `sets_session_local_state` and two live regression tests: `temp_table_visible_in_bare_batch_on_pooler` (the temp-table case) and `session_set_visible_in_bare_batch_on_pooler` (probed via `SELECT current_setting('extra_float_digits')` returning "3" vs the lost-SET default "1"). Verified: 434 unit + 13 live `client_live_tests` (incl. all #166/#168/#160/ #163/#169 guards) pass against a live PostgreSQL; 38/39 `live_db` end-to-end (the one failure is `extension "vector" is not available`, a pre-existing fixture gap needing the pgvector-test container). clippy clean, fmt applied. --- src/client_live_tests.rs | 248 ++++++++++++++++++++++++++++++++++++ src/handlers/query.rs | 144 +++++++++++++++++---- src/handlers/query_tests.rs | 156 ++++++++++++++++++++++- 3 files changed, 524 insertions(+), 24 deletions(-) diff --git a/src/client_live_tests.rs b/src/client_live_tests.rs index bf583ea..87986cb 100644 --- a/src/client_live_tests.rs +++ b/src/client_live_tests.rs @@ -1086,3 +1086,251 @@ async fn empty_result_returns_columns_via_session_pinned_batch_on_pooler() { let key = connection_key(¶ms); POOLS.lock().unwrap().remove(&key); } + +/// Regression guard for #173 (temp-table case): a bare non-session batch on +/// a `pooler_mode: "transaction"` connection that creates a temp table in an +/// earlier statement then selects from it must keep the `SELECT` on the +/// shared batch connection — not the #172 one-shot path, which runs the +/// `SELECT` on a fresh connection that never saw the temp table. +/// +/// On rc.11+#172 this errors `ERROR: relation "t" does not exist` (the +/// `SELECT` runs on a one-shot connection where `CREATE TEMP TABLE` never +/// ran). The #173 fix detects the `CREATE TEMP TABLE` as session-local state +/// and routes the `SELECT` to the shared client with a column-type +/// pre-resolve, where the temp table is visible. +/// +/// Run against direct Postgres with `pooler_mode: "transaction"` (the +/// clearest regression config — one TCP connection = one backend, held — +/// where `main` preserved within-batch state and #172 broke it). A real +/// PgBouncer reassigns the backend per autocommit transaction, so the +/// temp-table-in-a-bare-batch case needs `BEGIN`/`COMMIT` there regardless; +/// that path is covered by `wide_enum_table_returns_rows_via_session_pinned_ +/// batch_on_pooler`. +#[tokio::test] +#[ignore = "needs a live PostgreSQL (PG_PLUGIN_TEST_HOST=127.0.0.1 ... podman 54320)"] +async fn temp_table_visible_in_bare_batch_on_pooler() { + use crate::client; + use crate::handlers::query::execute_query_batch; + use serde_json::json; + + let _guard = POOLS_TEST_LOCK.lock().await; + + let params = live_params(&correct_password()) + .expect("PG_PLUGIN_TEST_HOST not set — set PG_PLUGIN_TEST_HOST=127.0.0.1 (and optionally PG_PLUGIN_TEST_PORT/USER/DB) to run live tests"); + + // A unique temp-table name so a leftover from a prior run can't mask the + // assertion. Temp tables are connection-scoped and dropped at session end, + // but the bare batch uses a cached pool client that may survive the test. + let table_name = "plugin_test_173_temp_tbl"; + + // Ensure no permanent leftover (defensive — temp tables don't persist, but + // a prior crashed run could have left a permanent table of the same name). + let _ = + client::execute_typed(¶ms, &format!("DROP TABLE IF EXISTS {table_name}"), &[]).await; + + // Bare non-session batch (no `session_id`): `CREATE TEMP TABLE` (a + // non-result-set statement → `else` branch on the shared client, sets + // `session_state_set`) then `SELECT * FROM t` (result-set → with the fix, + // `!in_transaction` + `session_state_set` routes to the shared client, + // not the one-shot). On rc.11+#172 the SELECT takes the one-shot path and + // errors `relation "t" does not exist`. + let request_params = json!({ + "driver": "postgresql", + "host": params.host, + "port": params.port, + "database": params.database, + "username": params.username, + "password": params.password, + "pooler_mode": "transaction", + "queries": [ + format!("CREATE TEMP TABLE {table_name} AS SELECT 1 AS id"), + format!("SELECT * FROM {table_name}"), + ], + "limit": 100, + "page": 1, + }); + + let response = tokio::time::timeout( + std::time::Duration::from_secs(30), + execute_query_batch(json!(1), &request_params), + ) + .await + .expect( + "bare [CREATE TEMP TABLE, SELECT] batch hung past 30s on a pooler-mode connection — \ + the #173 regression is present (the temp-table SELECT misroutes to the one-shot path)", + ); + + assert!( + response.get("error").is_none(), + "bare batch should succeed; got: {response}" + ); + // Non-session batch returns a bare `result` array. + let results = response + .get("result") + .and_then(|r| r.as_array()) + .expect("non-session batch returns a bare `result` array"); + assert_eq!(results.len(), 2, "two queries in the batch; got {response}"); + + // The CREATE TEMP TABLE statement (non-result-set) returns affected_rows, + // no error. + let create_result = &results[0]; + assert!( + create_result + .get("error") + .and_then(|e| e.as_str()) + .is_none(), + "CREATE TEMP TABLE should not error; got: {create_result}" + ); + + // The SELECT must succeed and see the temp table on the shared + // connection. On rc.11+#172 this is the failing statement (one-shot path + // → `relation "t" does not exist`). + let select_result = &results[1]; + assert!( + select_result + .get("error") + .and_then(|e| e.as_str()) + .is_none(), + "SELECT after CREATE TEMP TABLE should see the temp table (not error \ + `relation \"t\" does not exist`); got: {select_result}" + ); + let rows = select_result + .get("result") + .and_then(|r| r.get("rows")) + .and_then(|r| r.as_array()) + .expect("the SELECT's result.rows is an array"); + assert_eq!( + rows.len(), + 1, + "the temp table has one row; got empty: {select_result}" + ); + assert_eq!( + rows[0].as_array().unwrap()[0], + json!(1), + "the temp table's id column is 1; got: {select_result}" + ); + + // Teardown. The temp table is connection-scoped; the cached pool client + // holds it, so drop it explicitly before clearing the pool entry (a + // recycled connection would otherwise carry the temp table to the next + // borrower on direct PG). Then evict the pool so no state leaks to other + // tests. + let _ = + client::execute_typed(¶ms, &format!("DROP TABLE IF EXISTS {table_name}"), &[]).await; + let key = connection_key(¶ms); + POOLS.lock().unwrap().remove(&key); +} + +/// Regression guard for #173 (session-SET case): a bare non-session batch on +/// a `pooler_mode: "transaction"` connection that sets a session GUC in an +/// earlier statement then reads it back must keep the `SELECT` on the shared +/// batch connection so it sees the `SET`. +/// +/// `SET extra_float_digits = 3` (default is `1`) is read back via +/// `SELECT current_setting('extra_float_digits')`, which returns the GUC's +/// *current* value as `text`. On rc.11+#172 the `SELECT` runs on a one-shot +/// connection that never saw the `SET`, so it returns the default `"1"` +/// (the SET is lost). The #173 fix detects the `SET` as session-local state +/// and routes the `SELECT` to the shared client where the GUC took effect, +/// returning `"3"`. The probe is a direct `current_setting` read (not a +/// value whose rendering the GUC influences), so the assertion is independent +/// of how the plugin decodes/formats `float8`/`numeric` — it proves the +/// SELECT ran on the connection that saw the SET, unambiguously. +/// +/// Run against direct Postgres with `pooler_mode: "transaction"` (the +/// clearest regression config). A real PgBouncer loses the GUC to backend +/// reassignment between autocommit statements regardless (pre-existing on +/// `main`), so the GUC case is not asserted on a real pooler — only the +/// direct+flag config where `main` preserved it. +#[tokio::test] +#[ignore = "needs a live PostgreSQL (PG_PLUGIN_TEST_HOST=127.0.0.1 ... podman 54320)"] +async fn session_set_visible_in_bare_batch_on_pooler() { + use crate::handlers::query::execute_query_batch; + use serde_json::json; + + let _guard = POOLS_TEST_LOCK.lock().await; + + let params = live_params(&correct_password()) + .expect("PG_PLUGIN_TEST_HOST not set — set PG_PLUGIN_TEST_HOST=127.0.0.1 (and optionally PG_PLUGIN_TEST_PORT/USER/DB) to run live tests"); + + // Bare non-session batch: `SET extra_float_digits = 3` (non-result-set → + // shared client, sets `session_state_set`) then + // `SELECT current_setting('extra_float_digits')` (result-set → with the + // fix, routes to the shared client, not the one-shot). On rc.11+#172 the + // SELECT takes the one-shot path, loses the SET, and returns the default + // "1" instead of "3". + let request_params = json!({ + "driver": "postgresql", + "host": params.host, + "port": params.port, + "database": params.database, + "username": params.username, + "password": params.password, + "pooler_mode": "transaction", + "queries": [ + "SET extra_float_digits = 3", + "SELECT current_setting('extra_float_digits') AS digits", + ], + "limit": 100, + "page": 1, + }); + + let response = tokio::time::timeout( + std::time::Duration::from_secs(30), + execute_query_batch(json!(1), &request_params), + ) + .await + .expect( + "bare [SET, SELECT] batch hung past 30s on a pooler-mode connection — \ + the #173 regression is present", + ); + + assert!( + response.get("error").is_none(), + "bare batch should succeed; got: {response}" + ); + let results = response + .get("result") + .and_then(|r| r.as_array()) + .expect("non-session batch returns a bare `result` array"); + assert_eq!(results.len(), 2, "two queries in the batch; got {response}"); + + let set_result = &results[0]; + assert!( + set_result.get("error").and_then(|e| e.as_str()).is_none(), + "SET should not error; got: {set_result}" + ); + + let select_result = &results[1]; + assert!( + select_result + .get("error") + .and_then(|e| e.as_str()) + .is_none(), + "SELECT after SET should not error; got: {select_result}" + ); + let rows = select_result + .get("result") + .and_then(|r| r.get("rows")) + .and_then(|r| r.as_array()) + .expect("the SELECT's result.rows is an array"); + assert_eq!(rows.len(), 1, "one row; got: {select_result}"); + + // `current_setting('extra_float_digits')` returns the GUC's current value + // as text: "3" if the SELECT ran on the connection that saw the SET, "1" + // (the default) if it ran on a one-shot that lost the SET. Decoded as a + // JSON string (text column). + let val = &rows[0].as_array().unwrap()[0]; + assert_eq!( + val, + &json!("3"), + "current_setting('extra_float_digits') should be \"3\" (the SET took \ + effect on the SELECT's connection); got: {val} — if this is \"1\", the \ + SELECT ran on a one-shot connection that lost the SET (#173 regression)" + ); + + // Teardown: evict the pool entry so the SET GUC doesn't leak to other + // tests on the cached direct-PG connection. + let key = connection_key(¶ms); + POOLS.lock().unwrap().remove(&key); +} diff --git a/src/handlers/query.rs b/src/handlers/query.rs index 13fac97..aec8631 100644 --- a/src/handlers/query.rs +++ b/src/handlers/query.rs @@ -133,6 +133,13 @@ async fn run_batch_in_session( // A pinned connection only exists because its transaction is still open. let mut in_transaction = pinned.is_some(); let reused = pinned.is_some(); + // Whether any earlier non-result-set statement in this batch set + // session-local state (a session `SET`, a `CREATE TEMP TABLE`) that a + // later result-set `SELECT` would see only on this shared connection — + // not on a fresh one-shot connection (#173). Sticky for the rest of the + // batch: once set, a later `SELECT` that needs to see it stays on the + // shared client. See the routing comment below. + let mut session_state_set = false; // Acquire ONE connection for the entire batch (session state must survive) let pg_client = match pinned { @@ -173,7 +180,8 @@ async fn run_batch_in_session( // #166/#168). `exec_query` routes single queries to the one-shot // pinned path (`exec_query_pooled_txn`) for this; the batch path // bypasses `exec_query`, so it makes the same routing decision here. - // Two sub-cases, keyed on `in_transaction` (not `reused`): + // The routing keys on whether the SELECT must run on this shared + // connection, with two cases (plus the cancel-guard note): // // * `in_transaction` (the user's `BEGIN` has pinned the backend on // this shared connection — either a session-pinned batch reusing a @@ -183,28 +191,49 @@ async fn run_batch_in_session( // SELECT on the same connection where transaction-local state // (temp tables, SET variables) is visible. Run // `exec_query_on_client` with `pre_resolve_types = true`. - // * Not `in_transaction` (no open transaction pins the backend): a - // plain `prepare` could collide across reassignment (the `s0`→`s1` - // shift that killed the removed pre-warm, ca33b98). Route to the - // one-shot pinned path (`exec_query_pooled_txn`) — the proven - // #164 mechanism. The one-shot client dies with its stale + // * `session_state_set` (a prior non-result-set statement in this + // bare batch set session-local state — a session `SET` or a + // `CREATE TEMP TABLE` outside a transaction — #173): the SELECT + // must stay on this shared connection to see that state. A fresh + // one-shot connection never saw the `SET`/temp table (the one-shot + // applies only the host's `schema` via `SET LOCAL`), so routing it + // there would silently lose the state — wrong float precision, or + // `ERROR: relation "t" does not exist`. Keep it on the shared + // client with `pre_resolve_types = true` (the same path as the + // `in_transaction` case). On a real PgBouncer a session `SET` + // between autocommit statements is already lost to backend + // reassignment (pre-existing on `main`), so this primarily + // restores the direct-PostgreSQL + `pooler_mode=transaction` + // config (one TCP connection = one backend, held) and the + // temp-table case; see `sets_session_local_state` for the + // classifier and why transaction-scoped forms are excluded. + // * Otherwise (no open transaction, no prior session state): route + // to the one-shot pinned path (`exec_query_pooled_txn`) — the + // proven #164 mechanism. The one-shot client dies with its stale // client-side type cache; the shared batch client sits idle for - // that one query (acceptable: correctness over connection reuse - // for one statement; when there's no transaction there's no - // transaction-local state to lose by using a different connection). + // that one query (acceptable: when there's no transaction and no + // session state there's nothing to lose by using a different + // connection). A plain `prepare` here could collide across + // reassignment (the `s0`→`s1` shift that killed the removed + // pre-warm, ca33b98), so the one-shot is the safe choice. // // Non-result-set statements (DDL, writes) and direct connections // stay on the shared cached client — the unnamed protocol is // pooler-safe and needs no pinning (#149, #168). + let shared_client = in_transaction || session_state_set; let outcome = if conn_params.is_transaction_pooler() && returns_result_set(query) { - if in_transaction { - // The user's open transaction pins the backend on this - // connection. Pre-resolve column types via `prepare` inside - // it (safe — no reassignment can orphan the named statement) - // so `query_typed`'s `get_type` is a cache hit (no mid-stream - // `typeinfo_statement`, no deadlock). Keeps the SELECT on the - // shared connection where transaction-local state is visible. - // `cancel_guard` stays registered for this shared client. + if shared_client { + // Run on this shared connection (the user's open transaction + // pins it, or a prior statement set session-local state the + // SELECT must see — #173). Pre-resolve column types via + // `prepare` (safe — the backend is pinned by an open + // transaction, or there's no pooler reassignment to collide + // with on the direct+`pooler_mode=transaction` config) so + // `query_typed`'s `get_type` is a cache hit (no mid-stream + // `typeinfo_statement`, no deadlock). Keeps the SELECT on + // the shared connection where the session/transaction-local + // state is visible. `cancel_guard` stays registered for + // this shared client. exec_query_on_client( &pg_client, query, @@ -215,11 +244,11 @@ async fn run_batch_in_session( ) .await } else { - // No open transaction pins the backend. Route to the one-shot - // pinned path. The one-shot client acquired inside - // `exec_query_pooled_txn` is its own connection (never - // returned to a pool, so no stale client-side type cache - // hazard — #164's core property). + // No open transaction pins the backend and no prior statement + // set session state. Route to the one-shot pinned path. The + // one-shot client acquired inside `exec_query_pooled_txn` is + // its own connection (never returned to a pool, so no stale + // client-side type cache hazard — #164's core property). // // Cancel-guard coexistence: `exec_query_pooled_txn` registers // its own `CancelGuard` under the SAME `request_id` we hold @@ -261,6 +290,15 @@ async fn run_batch_in_session( in_transaction = transaction_effect(query).in_transaction_after(outcome.is_ok(), in_transaction); + // Track session-local state set by a non-result-set statement so a + // later result-set SELECT in this bare batch routes to the shared + // client above instead of the one-shot (#173). A `SET`/`CREATE TEMP + // TABLE` here flips the flag; result-set statements and transaction + // verbs don't, and `SET LOCAL` deliberately doesn't (transaction- + // scoped, not session-local). Sticky once set. + if sets_session_local_state(query) { + session_state_set = true; + } match outcome { Ok(result) => results.push(json!({ @@ -1144,6 +1182,68 @@ fn returns_result_set(query: &str) -> bool { || upper.starts_with("CALL") } +/// Whether `query` sets session-local state — a session `SET` (without +/// `LOCAL`) or a `CREATE TEMP[ORARY] TABLE` outside a transaction — that a +/// later statement in the same bare batch would see only on the *same* +/// connection (#173). +/// +/// `run_batch_in_session` uses this to keep a result-set `SELECT` on the +/// shared batch client (with a column-type pre-resolve) instead of the #172 +/// one-shot path when routing to a fresh connection would lose the state: +/// the one-shot applies only the host's `schema` via `SET LOCAL` and never +/// replays the user's prior `SET`/temp table. On `main` (rc.11, before +/// #172) every statement of a batch ran on one shared `pg_client`, so +/// `["SET extra_float_digits = 3", "SELECT 1.0/3.0"]` (or `["CREATE TEMP +/// TABLE t ...", "SELECT * FROM t"]`) reliably shared state; #172's +/// `!in_transaction` routing to the one-shot broke that for transaction +/// poolers. +/// +/// Transaction-scoped forms are deliberately **excluded**: +/// - `SET LOCAL ...` dies at transaction end. With no open transaction the +/// one-shot path wraps its own `BEGIN`/`COMMIT`, so a `SET LOCAL` is +/// correctly scoped there; with an open transaction the `in_transaction` +/// branch already keeps the SELECT on the shared client. Either way a +/// bare-batch SELECT on a fresh one-shot connection would never have seen +/// it, so excluding it keeps the cheaper, pooler-safe one-shot path. +/// - `CREATE TEMP TABLE` *inside* a `BEGIN` is likewise covered by the +/// `in_transaction` branch; this classifier only decides the +/// `!in_transaction` case. +/// +/// A permanent `CREATE TABLE`/`CREATE TYPE`/etc. is visible to any +/// connection, so the one-shot sees it — not session-local, excluded. +/// +/// Soundness over completeness: uncommon session-state forms +/// (`SELECT set_config(..., false)`, `RESET`, `SET SESSION AUTHORIZATION`) +/// may be missed and degrade gracefully to the one-shot (the current, pre-fix +/// behavior) — a false negative is never a *new* regression. A false positive +/// routes to the shared client with a pre-resolve, the same safe path the +/// `in_transaction` branch uses, so it is not dangerous, only occasionally +/// suboptimal. +/// +/// Shares [`leading_keywords`] with [`transaction_effect`], so comment +/// stripping and string-literal handling match (a comment-headed `SET` still +/// counts; a `SET` inside a string literal does not). +pub fn sets_session_local_state(query: &str) -> bool { + let mut words = leading_keywords(query, 3).into_iter(); + let Some(first) = words.next() else { + return false; + }; + let second = words.next(); + let third = words.next(); + match first.as_str() { + // Any `SET` that is not `SET LOCAL` is session-scoped (`SET `, + // `SET ROLE`, `SET SESSION ...`). + "SET" => second.as_deref() != Some("LOCAL"), + // `CREATE TEMP[ORARY] TABLE` lives for the session and is + // connection-scoped; permanent `CREATE TABLE` is visible everywhere. + "CREATE" => { + matches!(second.as_deref(), Some("TEMP") | Some("TEMPORARY")) + && third.as_deref() == Some("TABLE") + } + _ => false, + } +} + /// Check if a result-set-bearing statement's syntax actually accepts a /// trailing `LIMIT`/`OFFSET` clause — narrower than `returns_result_set`. /// Verified directly against a live PostgreSQL instance (`SELECT`, `WITH`, diff --git a/src/handlers/query_tests.rs b/src/handlers/query_tests.rs index 1f86ff8..b9dcc54 100644 --- a/src/handlers/query_tests.rs +++ b/src/handlers/query_tests.rs @@ -11,8 +11,9 @@ //! pagination is applied to a statement. use super::{ - build_search_path_sql, raw_explain_output, returns_result_set, strip_leading_sql_comments, - supports_trailing_limit_clause, transaction_effect, TransactionEffect, + build_search_path_sql, raw_explain_output, returns_result_set, sets_session_local_state, + strip_leading_sql_comments, supports_trailing_limit_clause, transaction_effect, + TransactionEffect, }; #[test] @@ -454,3 +455,154 @@ fn build_search_path_sql_local_quotes_schema_with_special_characters() { "SET LOCAL search_path TO \"a\"\"b\", public" ); } + +// --- sets_session_local_state (#173) ----------------------------------------- +// +// A bare batch (no `session_id`, no `BEGIN`) on a `pooler_mode = "transaction"` +// connection runs its result-set `SELECT` on the shared batch client when an +// earlier statement set session-local state, instead of the #172 one-shot +// path (a fresh connection that never saw the `SET` / temp table). This +// classifier decides which earlier statements flip that routing. See +// `run_batch_in_session`'s routing comment and the `sets_session_local_state` +// doc comment for why transaction-scoped forms are excluded (the +// `in_transaction` branch already keeps those on the shared client). + +#[test] +fn sets_session_local_state_detects_session_set() { + // A `SET` without `LOCAL` persists for the session on the held connection, + // so a later `SELECT` on the same connection sees it. The GUC name and + // `=`/`TO` forms both count. + for query in [ + "SET extra_float_digits = 3", + "SET search_path TO analytics", + "SET statement_timeout = 5000", + "SET timezone TO 'UTC'", + ] { + assert!( + sets_session_local_state(query), + "{query} should count as session-local state" + ); + } +} + +#[test] +fn sets_session_local_state_detects_set_role_and_session_authorization() { + // `SET ROLE` / `SET SESSION AUTHORIZATION` change session identity — + // session-scoped, visible to a later statement on the same connection. + for query in [ + "SET ROLE analyst", + "SET SESSION AUTHORIZATION postgres", + "SET SESSION CHARACTERISTICS AS TRANSACTION ISOLATION LEVEL READ COMMITTED", + ] { + assert!( + sets_session_local_state(query), + "{query} should count as session-local state" + ); + } +} + +#[test] +fn sets_session_local_state_detects_create_temp_table() { + // Both `TEMP` and `TEMPORARY`; a temp table created outside a transaction + // lives for the session and is visible only on the same connection. + for query in [ + "CREATE TEMP TABLE t (id int)", + "CREATE TEMPORARY TABLE t AS SELECT 1 AS id", + "CREATE TEMP TABLE t (id int) ON COMMIT DROP", + ] { + assert!( + sets_session_local_state(query), + "{query} should count as session-local state" + ); + } +} + +#[test] +fn sets_session_local_state_excludes_set_local() { + // `SET LOCAL` is transaction-scoped: it dies at transaction end, so on the + // one-shot path (which wraps its own `BEGIN`/`COMMIT`) it is correctly + // scoped and the `in_transaction` branch handles it when there is an open + // transaction. Either way it is not state a *bare-batch* SELECT on a fresh + // one-shot connection would have seen, so excluding it keeps the routing + // on the (cheaper, pooler-safe) one-shot path. + for query in [ + "SET LOCAL statement_timeout = 5000", + "SET LOCAL search_path TO analytics", + ] { + assert!( + !sets_session_local_state(query), + "{query} is transaction-scoped, not session-local" + ); + } +} + +#[test] +fn sets_session_local_state_excludes_permanent_create_and_other_ddl() { + // A permanent `CREATE TABLE` (or `CREATE TYPE`/`CREATE INDEX`) is visible to + // any connection, so the one-shot path sees it — routing to the shared + // client is unnecessary. Only *temp* tables are connection-scoped. + for query in [ + "CREATE TABLE t (id int)", + "CREATE TYPE mood AS ENUM ('happy', 'sad')", + "CREATE INDEX ON t (id)", + "CREATE SCHEMA s", + "CREATE OR REPLACE VIEW v AS SELECT 1", + ] { + assert!( + !sets_session_local_state(query), + "{query} is not session-local state" + ); + } +} + +#[test] +fn sets_session_local_state_excludes_result_set_and_other_statements() { + // A `SELECT`/`UPDATE`/transaction-control verb sets no session state a + // later statement depends on; routing them to shared is unnecessary. + for query in [ + "SELECT 1.0/3.0", + "SELECT * FROM t", + "UPDATE t SET x = 1", + "BEGIN", + "COMMIT", + "ROLLBACK", + "INSERT INTO t VALUES (1)", + "DELETE FROM t", + ] { + assert!( + !sets_session_local_state(query), + "{query} sets no session-local state" + ); + } +} + +#[test] +fn sets_session_local_state_sees_through_leading_comments() { + // `leading_keywords` strips `--`/`/* */` comments, so a comment-headed + // session `SET` is still classified — matching `transaction_effect`'s + // comment handling (a comment-headed `BEGIN` still opens a transaction). + assert!(sets_session_local_state( + "-- tune precision\nSET extra_float_digits = 3" + )); + assert!(sets_session_local_state( + "/* temp scratch */ CREATE TEMP TABLE t (id int)" + )); +} + +#[test] +fn sets_session_local_state_handles_empty_and_comment_only_input() { + assert!(!sets_session_local_state("")); + assert!(!sets_session_local_state(" \n")); + assert!(!sets_session_local_state("-- only a comment")); + assert!(!sets_session_local_state("/* block only */")); +} + +#[test] +fn sets_session_local_state_ignores_set_inside_string_literal() { + // `leading_keywords` stops at a string/quoted identifier/dollar quote, so + // a `SET` that is data (e.g. inside a `SELECT 'SET x = 1'`) is not read as + // a keyword — matching `transaction_effect`'s string-handling. + assert!(!sets_session_local_state( + "SELECT 'SET extra_float_digits = 3'" + )); +} From d2ffd38494f7cb5f65b06378a366b38d51587fd7 Mon Sep 17 00:00:00 2001 From: Adam J Esslinger Date: Fri, 23 Oct 2026 12:19:09 -0400 Subject: [PATCH 2/2] fix: address pre-handoff-review findings on #173 classifier and routing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pre-handoff ensemble review of #174 surfaced four verified findings: 1. `CREATE [GLOBAL|LOCAL] TEMP[ORARY] TABLE` was a false negative. PostgreSQL accepts an optional deprecated `GLOBAL`/`LOCAL` qualifier before `TEMP`/`TEMPORARY` (both still create a session-scoped temp table — verified live), but the 3-keyword lookahead classified `CREATE LOCAL TEMP TABLE ...` as non-session-local, so a later `SELECT * FROM t` in the bare batch still hit the exact #173 `relation does not exist` regression. Extend the `CREATE` arm to 4 keywords and skip the `GLOBAL`/`LOCAL` noise word (mirrors how `transaction_effect` skips optional `WORK`/`TRANSACTION`). 2. `SET TRANSACTION` / `SET CONSTRAINTS` were false positives — transaction-control verbs, not GUCs, that set no session state a later SELECT depends on. Exclude them alongside `SET LOCAL`. 3. `session_state_set` flipped on the query text alone, ignoring `outcome.is_ok()` — a `SET`/`CREATE TEMP TABLE` the server rejected established no session state but still steered a later SELECT onto the shared client. Gate on success, mirroring the sibling `in_transaction` update two lines above. 4. Docs drift: the `exec_query_on_client` pre-resolve comment listed "two cases where a plain `prepare` is safe" and said the no-open-transaction pooler path "can't safely `prepare`" — but #173 now routes that path to the shared-client pre-resolve. Rewrote the comment to list three cases and document the #173 case's real-PgBouncer tradeoff honestly (the prepare is non-fatal; safe on the direct+`pooler_mode=transaction` config the PR targets). Also fixed "two cases"→"three cases" in the routing comment and the `exec_query` batch-path description, and corrected the `sets_session_local_state` doc (it listed `SET SESSION AUTHORIZATION` as "may be missed" but the classifier catches it). Adds a unit test for the `GLOBAL`/`LOCAL` qualifier and folds `SET TRANSACTION`/`SET CONSTRAINTS` into the existing exclusion test. Verified: 435 unit + 13 live `client_live_tests` (incl. all #166/#168/#160/ #163/#169 guards) pass against a live PostgreSQL; clippy clean; fmt applied. --- src/handlers/query.rs | 90 +++++++++++++++++++++++++------------ src/handlers/query_tests.rs | 27 ++++++++++- 2 files changed, 88 insertions(+), 29 deletions(-) diff --git a/src/handlers/query.rs b/src/handlers/query.rs index aec8631..e95e39b 100644 --- a/src/handlers/query.rs +++ b/src/handlers/query.rs @@ -181,7 +181,7 @@ async fn run_batch_in_session( // pinned path (`exec_query_pooled_txn`) for this; the batch path // bypasses `exec_query`, so it makes the same routing decision here. // The routing keys on whether the SELECT must run on this shared - // connection, with two cases (plus the cancel-guard note): + // connection, with three cases (plus the cancel-guard note): // // * `in_transaction` (the user's `BEGIN` has pinned the backend on // this shared connection — either a session-pinned batch reusing a @@ -295,8 +295,12 @@ async fn run_batch_in_session( // client above instead of the one-shot (#173). A `SET`/`CREATE TEMP // TABLE` here flips the flag; result-set statements and transaction // verbs don't, and `SET LOCAL` deliberately doesn't (transaction- - // scoped, not session-local). Sticky once set. - if sets_session_local_state(query) { + // scoped, not session-local). Sticky once set. Gated on success, like + // `in_transaction` above: a `SET`/`CREATE TEMP TABLE` the server + // rejected established no session state, so it must not steer a later + // SELECT onto the shared client (which, outside a transaction on a + // real pooler, runs an unpinned `prepare`). + if outcome.is_ok() && sets_session_local_state(query) { session_state_set = true; } @@ -446,8 +450,11 @@ async fn exec_query( // client — `execute_typed` is unnamed and needs no pinning. The batch // path (`run_batch_in_session`) mirrors this routing for its own queries // rather than calling here: a pooler result-set query with no open - // transaction routes to `exec_query_pooled_txn`, and one inside an open - // transaction pre-resolves on the shared client (see `run_batch_in_session`). + // transaction routes to `exec_query_pooled_txn`, *unless* a prior + // statement in the batch set session-local state (`session_state_set`, + // #173), in which case it pre-resolves on the shared client (see + // `run_batch_in_session`); one inside an open transaction likewise + // pre-resolves on the shared client. if conn_params.is_transaction_pooler() && returns_result_set(query) { return exec_query_pooled_txn(conn_params, query, limit, page, schema, request_id).await; } @@ -872,19 +879,24 @@ async fn exec_query_on_client( // unnamed `query_typed`, but the outcome — types resolved before the // data fetch, no mid-stream interleave — is the same), which worked. // - // Gated on `pre_resolve_types`, set by the caller for the two cases where - // a plain `prepare` is safe: + // Gated on `pre_resolve_types`, set by the caller for the cases where a + // plain `prepare` is safe: // * Direct connections (`!is_transaction_pooler`, the #166 path): no // backend reassignment to collide with, so a named statement is fine. // * Session-pinned pooler connections inside an open transaction (the // #169 batch path, `in_transaction`): the user's `BEGIN` pins the // backend, so a named statement can't be orphaned by reassignment. - // The remaining pooler case — a result-set query with NO open - // transaction — can't safely `prepare` on the shared client (the named - // statement would collide across the `s0`→`s1` reassignment that killed - // the removed pre-warm, commit ca33b98), so `run_batch_in_session` routes - // it to the #164 one-shot pinned path instead of setting - // `pre_resolve_types`. + // * A transaction-pooler bare batch whose earlier statement set + // session-local state (the #173 batch path, `session_state_set`, no + // open transaction): the SELECT must stay on the shared client to see + // the `SET`/temp table. This re-runs the `prepare` pattern the removed + // pre-warm (ca33b98) used, so on a *real* PgBouncer it can 42P05 if the + // pooler reassigns the backend during the prepare's typeinfo round-trip + // — but the prepare is non-fatal (swallowed below), and on the + // direct+`pooler_mode=transaction` config the PR targets there's no + // reassignment, so it's safe there. The remaining no-open-transaction + // pooler case with NO prior session state still routes to the #164 + // one-shot pinned path (see `run_batch_in_session`), not here. // // The pre-resolve is non-fatal on failure: a prepare error is swallowed // (`.ok()`) and the `query_typed` below runs anyway — for builtin types @@ -1205,6 +1217,11 @@ fn returns_result_set(query: &str) -> bool { /// branch already keeps the SELECT on the shared client. Either way a /// bare-batch SELECT on a fresh one-shot connection would never have seen /// it, so excluding it keeps the cheaper, pooler-safe one-shot path. +/// - `SET TRANSACTION ...` / `SET CONSTRAINTS ...` are transaction-control +/// verbs (not GUCs): outside a `BEGIN` they're a no-op or set only the +/// next transaction's characteristics, not session state a later SELECT +/// depends on. Excluding them avoids a false positive that would route a +/// later SELECT to the shared client unnecessarily. /// - `CREATE TEMP TABLE` *inside* a `BEGIN` is likewise covered by the /// `in_transaction` branch; this classifier only decides the /// `!in_transaction` case. @@ -1213,33 +1230,50 @@ fn returns_result_set(query: &str) -> bool { /// connection, so the one-shot sees it — not session-local, excluded. /// /// Soundness over completeness: uncommon session-state forms -/// (`SELECT set_config(..., false)`, `RESET`, `SET SESSION AUTHORIZATION`) -/// may be missed and degrade gracefully to the one-shot (the current, pre-fix -/// behavior) — a false negative is never a *new* regression. A false positive -/// routes to the shared client with a pre-resolve, the same safe path the +/// (`SELECT set_config(..., false)`, `RESET`) may be missed and degrade +/// gracefully to the one-shot (the current, pre-fix behavior) — a false +/// negative is never a *new* regression. A false positive routes to the +/// shared client with a pre-resolve, the same safe path the /// `in_transaction` branch uses, so it is not dangerous, only occasionally -/// suboptimal. +/// suboptimal. (`SET SESSION AUTHORIZATION` is *not* missed — it starts +/// with `SET` and its second keyword is `SESSION`, not `LOCAL`, so it is +/// caught.) /// /// Shares [`leading_keywords`] with [`transaction_effect`], so comment /// stripping and string-literal handling match (a comment-headed `SET` still /// counts; a `SET` inside a string literal does not). pub fn sets_session_local_state(query: &str) -> bool { - let mut words = leading_keywords(query, 3).into_iter(); + // Four keywords cover the longest form: `CREATE [GLOBAL|LOCAL] TEMP TABLE` + // (`GLOBAL`/`LOCAL` are optional deprecated noise words before TEMP, both + // still accepted by PostgreSQL and still create a session-scoped temp + // table — verified live). + let mut words = leading_keywords(query, 4).into_iter(); let Some(first) = words.next() else { return false; }; let second = words.next(); let third = words.next(); + let fourth = words.next(); match first.as_str() { - // Any `SET` that is not `SET LOCAL` is session-scoped (`SET `, - // `SET ROLE`, `SET SESSION ...`). - "SET" => second.as_deref() != Some("LOCAL"), - // `CREATE TEMP[ORARY] TABLE` lives for the session and is - // connection-scoped; permanent `CREATE TABLE` is visible everywhere. - "CREATE" => { - matches!(second.as_deref(), Some("TEMP") | Some("TEMPORARY")) - && third.as_deref() == Some("TABLE") - } + // Any `SET` that is neither `SET LOCAL` nor a transaction-control verb + // is session-scoped (`SET `, `SET ROLE`, `SET SESSION ...`). + "SET" => !matches!( + second.as_deref(), + Some("LOCAL") | Some("TRANSACTION") | Some("CONSTRAINTS") + ), + // `CREATE [GLOBAL|LOCAL] TEMP[ORARY] TABLE` lives for the session and + // is connection-scoped; the optional `GLOBAL`/`LOCAL` qualifier (both + // deprecated, both accepted) is skipped. Permanent `CREATE TABLE` / + // `CREATE UNLOGGED TABLE` / `CREATE TYPE` are visible everywhere, so + // they don't pin the SELECT to the shared connection. + "CREATE" => match second.as_deref() { + Some("TEMP") | Some("TEMPORARY") => third.as_deref() == Some("TABLE"), + Some("GLOBAL") | Some("LOCAL") => { + matches!(third.as_deref(), Some("TEMP") | Some("TEMPORARY")) + && fourth.as_deref() == Some("TABLE") + } + _ => false, + }, _ => false, } } diff --git a/src/handlers/query_tests.rs b/src/handlers/query_tests.rs index b9dcc54..7a57aed 100644 --- a/src/handlers/query_tests.rs +++ b/src/handlers/query_tests.rs @@ -524,10 +524,14 @@ fn sets_session_local_state_excludes_set_local() { // scoped and the `in_transaction` branch handles it when there is an open // transaction. Either way it is not state a *bare-batch* SELECT on a fresh // one-shot connection would have seen, so excluding it keeps the routing - // on the (cheaper, pooler-safe) one-shot path. + // on the (cheaper, pooler-safe) one-shot path. `SET TRANSACTION` / + // `SET CONSTRAINTS` are likewise transaction-control verbs, not GUCs — + // outside a `BEGIN` they set no session state a later SELECT depends on. for query in [ "SET LOCAL statement_timeout = 5000", "SET LOCAL search_path TO analytics", + "SET TRANSACTION ISOLATION LEVEL READ COMMITTED", + "SET CONSTRAINTS ALL IMMEDIATE", ] { assert!( !sets_session_local_state(query), @@ -536,6 +540,27 @@ fn sets_session_local_state_excludes_set_local() { } } +#[test] +fn sets_session_local_state_detects_global_local_temp_table() { + // PostgreSQL accepts an optional deprecated `GLOBAL`/`LOCAL` qualifier + // before `TEMP`/`TEMPORARY` (both still create a session-scoped temp table + // — verified against a live database). Without the 4-keyword lookahead + // the classifier would miss these and a later `SELECT * FROM t` would hit + // the exact #173 `relation does not exist` regression. + for query in [ + "CREATE LOCAL TEMP TABLE t (id int)", + "CREATE LOCAL TEMPORARY TABLE t AS SELECT 1 AS id", + "CREATE GLOBAL TEMP TABLE t (id int)", + "CREATE GLOBAL TEMPORARY TABLE t (id int)", + ] { + assert!( + sets_session_local_state(query), + "{query} should count as session-local state (the GLOBAL/LOCAL \ + qualifier is deprecated noise before TEMP)" + ); + } +} + #[test] fn sets_session_local_state_excludes_permanent_create_and_other_ddl() { // A permanent `CREATE TABLE` (or `CREATE TYPE`/`CREATE INDEX`) is visible to