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..e95e39b 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 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 @@ -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,19 @@ 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. 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; + } match outcome { Ok(result) => results.push(json!({ @@ -408,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; } @@ -834,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 @@ -1144,6 +1194,90 @@ 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. +/// - `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. +/// +/// 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`) 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. (`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 { + // 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 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, + } +} + /// 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..7a57aed 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,179 @@ 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. `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), + "{query} is transaction-scoped, not session-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 + // 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'" + )); +}