diff --git a/Cargo.lock b/Cargo.lock index 008872f..1b27f3c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -998,9 +998,9 @@ checksum = "eec689c0bc40ff2458a5977b6619cb718087084a18e02a131c599b62d05e1a5f" [[package]] name = "rustls-webpki" -version = "0.103.15" +version = "0.103.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f3c3cf1d8b1e7d4927e2d154c3fcb02979afb9939629c62cd9048d4f07b60ac2" +checksum = "4443b1b6d39d7c663139942cc4558262562d2c84123bb6282ef177009ca8b281" dependencies = [ "ring", "rustls-pki-types", @@ -1588,9 +1588,9 @@ dependencies = [ [[package]] name = "webpki-root-certs" -version = "1.0.9" +version = "1.0.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b96554aa2acc8ccdb7e1c9a58a7a68dd5d13bccc69cd124cb09406db612a1c9b" +checksum = "58ec80e73fd8baf57bf0d9615554c75348bda990d6d5bbeab964dd068319fecb" dependencies = [ "rustls-pki-types", ] diff --git a/src/client.rs b/src/client.rs index be0156d..47df0e1 100644 --- a/src/client.rs +++ b/src/client.rs @@ -73,7 +73,9 @@ pub(crate) fn format_pg_error(e: &tokio_postgres::Error) -> String { /// Format a `tokio_postgres::Error` from a query path, rewriting /// prepared-statement errors (SQLSTATE 26000 / 42P05) with actionable -/// guidance when the connection is through a transaction-mode pooler. +/// pooler guidance. The rewrite is unconditional on these SQLSTATEs; +/// `is_transaction_pooler` only selects the phrasing (definitive when the +/// host declared a transaction pooler, conditional when it didn't). /// /// Under PgBouncer `pool_mode = transaction`, tokio-postgres's internal /// `typeinfo_statement` (fired during `Row` deserialization for non-builtin @@ -85,25 +87,47 @@ pub(crate) fn format_pg_error(e: &tokio_postgres::Error) -> String { /// the cryptic error into that guidance so the user knows what to configure /// (#157, tabularis#1097). /// -/// For non-pooler connections or non-prepared-statement errors, falls -/// through to `format_pg_error` unchanged. -pub(crate) fn format_query_error(e: &tokio_postgres::Error, is_transaction_pooler: bool) -> String { +/// When `is_transaction_pooler` is already true, the guidance is definitive. +/// When it's false/unknown (the host did not send `pooler_mode`), the same +/// SQLSTATEs still suggest a transaction pooler the host didn't declare — so +/// guidance is phrased conditionally ("if this connection goes through a +/// transaction-mode pooler…") rather than suppressed. Suppressing it would +/// hide the only actionable hint from users on older hosts that can't send +/// `pooler_mode` yet (#160, #163). +/// The lead sentence of the pooler guidance [`format_query_error`] prepends +/// to a 26000/42P05 message. Extracted as a pure function so the conditional +/// phrasing — definitive when the host declared a transaction pooler, +/// conditional ("If this connection goes through…") when it didn't (#160, +/// #163) — is unit-testable without a live `tokio_postgres::Error` (whose +/// `DbError`/SQLSTATE can't be constructed without a real server response). +pub(crate) fn pooler_error_lead(is_transaction_pooler: bool) -> &'static str { if is_transaction_pooler { - if let Some(code) = e.code() { - if *code == tokio_postgres::error::SqlState::INVALID_SQL_STATEMENT_NAME - || *code == tokio_postgres::error::SqlState::DUPLICATE_PSTATEMENT - { - return format!( - "{}\n\n\ - This error occurs under transaction-mode poolers (e.g. PgBouncer) \ - when a named prepared statement is lost to connection reassignment. \ - For queries returning custom types (enums, composites, extension types), \ - set `max_prepared_statements > 0` in pgbouncer.ini — this enables \ - server-side prepared statement caching so named statements survive \ - reassignment. See: https://www.pgbouncer.org/config.html#max_prepared_statements", - format_pg_error(e) - ); - } + "This error occurs under transaction-mode poolers (e.g. PgBouncer) \ + when a named prepared statement is lost to connection reassignment." + } else { + "If this connection goes through a transaction-mode pooler \ + (e.g. PgBouncer or Supabase/Supavisor), this error occurs when a \ + named prepared statement is lost to connection reassignment." + } +} + +pub(crate) fn format_query_error(e: &tokio_postgres::Error, is_transaction_pooler: bool) -> String { + if let Some(code) = e.code() { + if *code == tokio_postgres::error::SqlState::INVALID_SQL_STATEMENT_NAME + || *code == tokio_postgres::error::SqlState::DUPLICATE_PSTATEMENT + { + let lead = pooler_error_lead(is_transaction_pooler); + return format!( + "{msg}\n\n\ + {lead} For queries returning custom types (enums, composites, extension \ + types), set `max_prepared_statements > 0` in pgbouncer.ini (or the \ + equivalent prepared-statement setting on your pooler) — this enables \ + server-side prepared statement caching so named statements survive \ + reassignment. In Tabularis, also mark the connection's Pooler Mode as \ + \"Transaction\" so the plugin uses its transaction-pooler-safe query \ + path. See: https://www.pgbouncer.org/config.html#max_prepared_statements", + msg = format_pg_error(e), + ); } } format_pg_error(e) @@ -530,6 +554,38 @@ pub(crate) async fn get_pool_client( } } +/// Acquire a one-shot client from a **transient** pool that is never inserted +/// into the global `POOLS` cache. The `pool` local is dropped at the end of +/// this function (only the `Object` is returned), so the underlying TCP +/// connection simply closes when the caller later drops the `Object`: +/// `Object`'s `Drop` upgrades a `WeakPool`; when the pool itself is already +/// gone the upgrade fails and the inner connection is dropped, never returned +/// to any pool. +/// +/// This exists for the transaction-pooler result-set path (#160 / #163): it +/// runs `DEALLOCATE ALL`, which clears server-side prepared statements but +/// cannot clear tokio-postgres's private cached `typeinfo_statement` handle +/// on the `InnerClient`. Reusing that same `Client` afterward would leave +/// its private cache pointing at a now-deallocated statement name. A +/// one-shot client guarantees no stale client-side type cache is ever reused +/// — the connection (and its caches) die when the operation ends. +/// +/// Builds via [`build_oneshot_pool`] (same as the cached path's +/// [`build_pool`] but without the startup-script preflight — the `post_create` +/// hook already validates the script on the real connection with the same +/// error shape, so the preflight would just open a second throwaway +/// connection per query). TLS, startup script, and `search_path` snapshot +/// hooks are otherwise identical to the cached path. +pub(crate) async fn oneshot_pool_client( + params: &ConnectionParams, +) -> Result { + let pool = build_oneshot_pool(params).await?; + match pool.get().await { + Ok(client) => Ok(client), + Err(e) => Err(format!("Connection failed: {}", format_pool_error(&e))), + } +} + /// Drop pools that currently have no checked-out connections. Called /// periodically so long-idle sessions don't linger for the plugin's /// lifetime — matches the sqlserver/dynamodb sibling plugins' pattern. @@ -574,7 +630,26 @@ pub(crate) fn make_tls_connect(params: &ConnectionParams) -> Result Result { + build_pool_impl(params, false).await +} + +/// Like [`build_pool`] but skips the startup-script preflight. See +/// [`build_pool`]'s `skip_preflight` note for why the one-shot path uses this. +async fn build_oneshot_pool(params: &ConnectionParams) -> Result { + build_pool_impl(params, true).await +} + +async fn build_pool_impl(params: &ConnectionParams, skip_preflight: bool) -> Result { let mut cfg = Config::new(); match params @@ -621,8 +696,10 @@ async fn build_pool(params: &ConnectionParams) -> Result { if needs_tls(params) { let tls_config = build_tls_connector(params)?; let tls = MakeRustlsConnect::new(tls_config); - if let Some(script) = script { - preflight_startup_script(&cfg, tls.clone(), script).await?; + if !skip_preflight { + if let Some(script) = script { + preflight_startup_script(&cfg, tls.clone(), script).await?; + } } let mut builder = cfg .builder(tls) @@ -641,8 +718,10 @@ async fn build_pool(params: &ConnectionParams) -> Result { .build() .map_err(|e| format!("Pool creation failed (TLS): {e}")) } else { - if let Some(script) = script { - preflight_startup_script(&cfg, NoTls, script).await?; + if !skip_preflight { + if let Some(script) = script { + preflight_startup_script(&cfg, NoTls, script).await?; + } } let mut builder = cfg .builder(NoTls) diff --git a/src/client_live_tests.rs b/src/client_live_tests.rs index 404a9e4..815f80c 100644 --- a/src/client_live_tests.rs +++ b/src/client_live_tests.rs @@ -144,3 +144,290 @@ async fn wrong_password_error_message_has_no_debug_dump() { let key = connection_key(&wrong); POOLS.lock().unwrap().remove(&key); } + +/// Regression guard for #160 / #163: a query against a table with a +/// non-builtin column type (a custom `enum` here) must return rows, not +/// hang. The failure mode under a transaction-mode pooler (#160/#163) is +/// tokio-postgres's internal cached named `typeinfo_statement` — fired for +/// non-builtin result column OIDs — colliding with stale server-side +/// prepared statements across pooler backend reassignment, or hanging on +/// Supavisor. +/// +/// This test runs against a *direct* Postgres (no pooler), so it validates +/// the happy path and that the enum type still decodes correctly through the +/// unnamed `query_typed` binary path. It does NOT reproduce the pooler +/// failure — that needs a transaction-mode pooler fixture (see +/// `pooled_txn_enum_column_returns_rows` below for the pinned-path coverage, +/// and the manual PgBouncer matrix documented in the PR body). +#[tokio::test] +#[ignore = "needs a live PostgreSQL (PG_PLUGIN_TEST_HOST=127.0.0.1 ... podman 54320)"] +async fn query_with_nonbuiltin_enum_column_returns_rows() { + use crate::client; + use crate::handlers::query::execute_query; + 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"); + + // Unique names so repeated runs (and concurrent test DB users) don't + // collide. `DROP ... IF EXISTS` is idempotent on both setup and teardown. + let type_name = "plugin_test_160_color"; + let table_name = "plugin_test_160_enum_table"; + + // Setup: create a custom enum type and a table with an enum column, then + // insert a row. `execute_typed` uses the unnamed protocol with no result + // columns, so setup itself can't hang on the type-resolution interleave. + for stmt in [ + format!("DROP TABLE IF EXISTS {table_name}"), + format!("DROP TYPE IF EXISTS {type_name}"), + format!("CREATE TYPE {type_name} AS ENUM ('red', 'green', 'blue')"), + format!("CREATE TABLE {table_name} (id int, color {type_name})"), + format!("INSERT INTO {table_name} (id, color) VALUES (1, 'red')"), + ] { + client::execute_typed(¶ms, &stmt, &[]) + .await + .unwrap_or_else(|e| panic!("setup statement failed ({stmt}): {e}")); + } + + // The actual regression check: SELECT * from the enum-typed table through + // the full `execute_query` handler path. Build the JSON params object the + // host would send (connection fields at the top level; `inner_params` + // falls back to the value itself when there's no nested `params` key). + let request_params = json!({ + "driver": "postgresql", + "host": params.host, + "port": params.port, + "database": params.database, + "username": params.username, + "password": params.password, + "query": format!("SELECT * FROM {table_name}"), + "limit": 100, + "page": 1, + }); + + // Fail fast instead of hanging: the bug manifests as a stall past 120s, + // so a 30s ceiling catches a regression well before the host-style + // timeout would, while leaving headroom for a slow podman container. + let response = tokio::time::timeout( + std::time::Duration::from_secs(30), + execute_query(json!(1), &request_params), + ) + .await + .expect( + "execute_query on an enum-typed table hung past 30s — the #160/#163 regression is present", + ); + + // Success: `result.rows` is a non-empty array. + assert!( + response.get("error").is_none(), + "execute_query should succeed for an enum-typed table; got: {response}" + ); + let result = response + .get("result") + .expect("ok_response carries a `result` object"); + let rows = result + .get("rows") + .and_then(|r| r.as_array()) + .expect("`result.rows` is an array"); + assert!( + !rows.is_empty(), + "the enum-typed table has one row; got empty rows: {result}" + ); + + // Teardown: drop the test objects, then evict the pool entry this test + // populated so it doesn't linger for the idle sweep or leak into other + // tests' cache assertions. + for stmt in [ + format!("DROP TABLE IF EXISTS {table_name}"), + format!("DROP TYPE IF EXISTS {type_name}"), + ] { + let _ = client::execute_typed(¶ms, &stmt, &[]).await; + } + let key = connection_key(¶ms); + POOLS.lock().unwrap().remove(&key); +} + +/// Exercises the transaction-pooler pinned path (`exec_query_pooled_txn`) +/// against a *direct* Postgres. Setting `pooler_mode: "transaction"` routes +/// result-set queries through `BEGIN; DEALLOCATE ALL; prepare; query_typed; +/// DEALLOCATE ALL; COMMIT` on a one-shot client. On a direct connection this +/// must still succeed — the sequence is valid SQL and the enum column must +/// decode correctly — proving the pinned path doesn't break the happy path. +/// +/// The actual pooler hang (#160/#163) needs a transaction-mode pooler +/// fixture; this test covers the direct-Postgres correctness of the new +/// path. The manual PgBouncer matrix in the PR body covers the pooler case. +#[tokio::test] +#[ignore = "needs a live PostgreSQL (PG_PLUGIN_TEST_HOST=127.0.0.1 ... podman 54320)"] +async fn pooled_txn_enum_column_returns_rows() { + use crate::client; + use crate::handlers::query::execute_query; + 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"); + + let type_name = "plugin_test_160b_color"; + let table_name = "plugin_test_160b_enum_table"; + + for stmt in [ + format!("DROP TABLE IF EXISTS {table_name}"), + format!("DROP TYPE IF EXISTS {type_name}"), + format!("CREATE TYPE {type_name} AS ENUM ('red', 'green', 'blue')"), + format!("CREATE TABLE {table_name} (id int, color {type_name})"), + format!("INSERT INTO {table_name} (id, color) VALUES (1, 'red')"), + ] { + client::execute_typed(¶ms, &stmt, &[]) + .await + .unwrap_or_else(|e| panic!("setup statement failed ({stmt}): {e}")); + } + + // `pooler_mode: "transaction"` routes the result-set query through the + // pinned one-shot transaction path even against a direct Postgres. + let request_params = json!({ + "driver": "postgresql", + "host": params.host, + "port": params.port, + "database": params.database, + "username": params.username, + "password": params.password, + "pooler_mode": "transaction", + "query": format!("SELECT * FROM {table_name}"), + "limit": 100, + "page": 1, + }); + + // Run twice. This proves the pinned path is repeatable on a direct + // connection — no client-side state (a leaked BEGIN, an uncleared + // DEALLOCATE, a cached typeinfo handle pointing at a dropped statement) + // from run 1 breaks run 2. It does NOT prove the one-shot client's + // *discard* prevents the stale-`typeinfo_statement` collision: that + // collision is pooler-only (it needs a backend reassignment between the + // prepare and a later reuse), and on direct PG reusing the client would + // also succeed. The discard property itself is covered by the manual + // PgBouncer matrix in the PR body, not by this test. A regression in the + // direct-Postgres repeatable path would surface here as a 26000/42P05 + // or a hang. + for run in 1..=2 { + let response = tokio::time::timeout( + std::time::Duration::from_secs(30), + execute_query(json!(run), &request_params), + ) + .await + .expect("pooled-txn enum query hung past 30s on direct Postgres"); + + assert!( + response.get("error").is_none(), + "pooled-txn run {run} should succeed; got: {response}" + ); + let result = response + .get("result") + .expect("ok_response carries a `result` object"); + let rows = result + .get("rows") + .and_then(|r| r.as_array()) + .expect("`result.rows` is an array"); + assert!( + !rows.is_empty(), + "pooled-txn run {run}: expected one row; got empty: {result}" + ); + // The enum value must decode correctly through the binary path. + assert_eq!( + rows[0].as_array().unwrap()[1], + json!("red"), + "pooled-txn run {run}: enum column should decode to \"red\"; got: {result}" + ); + } + + for stmt in [ + format!("DROP TABLE IF EXISTS {table_name}"), + format!("DROP TYPE IF EXISTS {type_name}"), + ] { + let _ = client::execute_typed(¶ms, &stmt, &[]).await; + } + let key = connection_key(¶ms); + POOLS.lock().unwrap().remove(&key); +} + +/// Regression guard for the routing decision at `exec_query` (#160/#163): +/// a *non-result-set* statement (DDL) on a `pooler_mode: "transaction"` +/// connection must stay on the shared cached-client path, NOT take the +/// pinned one-shot path. The guard is `is_transaction_pooler() && +/// returns_result_set(query)`; if the `returns_result_set` half were dropped, +/// this DDL would route to `exec_query_pooled_txn` → `run_pinned_body` → +/// `prepare("CREATE TABLE ...")`, which fails (DDL has no prepared-statement +/// form) — so a successful CREATE here proves the guard's FALSE side holds. +/// +/// `execute_query` returns a `result` with `affected_rows` for non-result-set +/// statements (the `execute_typed` branch of `exec_query_on_client`), never +/// rows. Run against direct Postgres so the assertion is about routing, not +/// pooler behavior. +#[tokio::test] +#[ignore = "needs a live PostgreSQL (PG_PLUGIN_TEST_HOST=127.0.0.1 ... podman 54320)"] +async fn pooled_txn_ddl_stays_off_the_pinned_path() { + use crate::client; + use crate::handlers::query::execute_query; + 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"); + + let table_name = "plugin_test_160c_ddl_table"; + // Idempotent setup: drop first so a leftover from a prior run doesn't + // make the CREATE a no-op. The DROP uses execute_typed directly, not the + // path under test, so it can't mask a routing regression. + let _ = + client::execute_typed(¶ms, &format!("DROP TABLE IF EXISTS {table_name}"), &[]).await; + + // `pooler_mode: "transaction"` + a DDL query. If the routing guard's + // `returns_result_set` half is intact, this stays on the shared cached + // client and the CREATE succeeds. If it were removed, this would take + // the pinned path and fail at `prepare("CREATE TABLE ...")`. + let request_params = json!({ + "driver": "postgresql", + "host": params.host, + "port": params.port, + "database": params.database, + "username": params.username, + "password": params.password, + "pooler_mode": "transaction", + "query": format!("CREATE TABLE {table_name} (id int)"), + "limit": 100, + "page": 1, + }); + + let response = tokio::time::timeout( + std::time::Duration::from_secs(30), + execute_query(json!(1), &request_params), + ) + .await + .expect("pooled-txn DDL query hung past 30s on direct Postgres"); + + assert!( + response.get("error").is_none(), + "DDL on a transaction-pooler connection should succeed (routing guard keeps it off the \ + pinned path); got: {response}" + ); + let result = response + .get("result") + .expect("ok_response carries a `result` object"); + // Non-result-set statements return rows: [] and affected_rows (0 for + // CREATE TABLE, which reports no affected rows in PostgreSQL). The key + // assertion is no error; assert the shape, not the count. + assert!( + result.get("rows").and_then(|r| r.as_array()).is_some(), + "non-result-set statements return a rows array: {result}" + ); + + // Teardown. + let _ = + client::execute_typed(¶ms, &format!("DROP TABLE IF EXISTS {table_name}"), &[]).await; + let key = connection_key(¶ms); + POOLS.lock().unwrap().remove(&key); +} diff --git a/src/client_tests.rs b/src/client_tests.rs index a0c466d..3255e64 100644 --- a/src/client_tests.rs +++ b/src/client_tests.rs @@ -5,8 +5,8 @@ use tokio::sync::Mutex; use super::{ build_pool_pub, build_tls_connector, cleanup_idle_pools, connection_key, get_or_create_pool, - get_pool_client, load_client_cert_from_pem, load_roots_from_pem, resolve_ssl_mode, - with_text_types, NoCertVerifier, VerifyCaCertVerifier, POOLS, + get_pool_client, load_client_cert_from_pem, load_roots_from_pem, pooler_error_lead, + resolve_ssl_mode, with_text_types, NoCertVerifier, VerifyCaCertVerifier, POOLS, }; use crate::models::ConnectionParams; use crate::settings; @@ -1301,3 +1301,41 @@ fn num_cpus_from_deadpool_default() -> usize { .map(|n| n.get() * 2) .unwrap_or(20) } + +/// `pooler_error_lead` selects the lead sentence of the pooler guidance +/// `format_query_error` prepends to a 26000/42P05 error (#160/#163). The +/// `false` branch is the new behavior this PR added (older hosts that don't +/// send `pooler_mode` still see conditional guidance instead of a bare +/// message); the `true` branch is definitive. `format_query_error` itself +/// needs a live `tokio_postgres::Error` to exercise (its `DbError`/SQLSTATE +/// can't be built without a server response), so the pure phrasing selection +/// is tested here and the SQLSTATE-gated assembly is covered by the live +/// `client_live_tests` error-message tests. +#[test] +fn pooler_error_lead_definitive_when_pooler_declared() { + let lead = pooler_error_lead(true); + assert!( + lead.starts_with("This error occurs under transaction-mode poolers"), + "definitive lead for a declared transaction pooler: {lead:?}" + ); + assert!( + !lead.starts_with("If this connection"), + "must not hedge when the host declared pooler_mode=transaction" + ); +} + +#[test] +fn pooler_error_lead_conditional_when_pooler_not_declared() { + // The PR's new behavior: even when is_transaction_pooler is false, the + // guidance is emitted (conditionally phrased) rather than suppressed — + // so users on older hosts that can't send pooler_mode still see the hint. + let lead = pooler_error_lead(false); + assert!( + lead.starts_with("If this connection goes through a transaction-mode pooler"), + "conditional lead for an undeclared pooler: {lead:?}" + ); + assert!( + lead.contains("Supabase/Supavisor"), + "conditional lead should name the pooler flavors users hit (#163): {lead:?}" + ); +} diff --git a/src/handlers/query.rs b/src/handlers/query.rs index bff79d0..51ca974 100644 --- a/src/handlers/query.rs +++ b/src/handlers/query.rs @@ -290,6 +290,23 @@ async fn exec_query( schema: Option<&str>, request_id: &Value, ) -> Result { + // Transaction-pooler result-set queries take a pinned one-shot path + // (#160 / #163): tokio-postgres's internal cached named `typeinfo_statement` + // — fired for non-builtin result column types (custom enums, composites, + // extension types like `citext`/Supabase `auth.uid`) — collides with stale + // server-side prepared statements across pooler backend reassignment, or + // hangs on Supavisor. Pinning the backend with an explicit transaction, + // clearing stale statements, pre-resolving column types, then running the + // #149-unnamed data fetch on a connection that is *not* returned to the + // pool closes that. See `exec_query_pooled_txn`. + // + // Non-result-set statements (DDL, writes) and the batch path + // (`run_batch_in_session`, which owns its own transaction) stay on the + // shared cached client — `execute_typed` is unnamed and needs no pinning. + if conn_params.is_transaction_pooler() && returns_result_set(query) { + return exec_query_pooled_txn(conn_params, query, limit, page, schema, request_id).await; + } + let pg_client = client::get_pool_client(conn_params).await?; // Register a cancel handle for this in-flight query (#126). The guard @@ -313,23 +330,252 @@ async fn exec_query( .await } -/// Build a `SET search_path TO` statement that prepends `schema` onto the -/// connection's existing `search_path` rather than replacing it (#51), so +/// Bounded lifetime of the transaction-pooler pinned path. Kept below the +/// host's plugin-call timeout (120s) so a wedged pooler connection surfaces +/// as an actionable error here instead of the opaque "timed out after 120s" +/// the host would otherwise emit. Long enough that a legitimately slow query +/// on a healthy pooler isn't cut short, short enough that a Supavisor-style +/// stall never reaches the host timeout (#160 / #163). +const POOLED_TXN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(90); + +/// Run a result-set query on a transaction-mode pooler through an explicit +/// pinned transaction on a one-shot client (#160 / #163). +/// +/// On a transaction-mode pooler, `BEGIN` pins the client to one server backend +/// for the whole sequence. While pinned: +/// +/// 1. `DEALLOCATE ALL` clears stale server-side prepared statements on that +/// backend (left behind by a previous assignment's `typeinfo_statement`). +/// 2. `search_path` is set locally (inside the transaction) if a schema was +/// selected. +/// 3. `prepare(&query)` (the original user query, before pagination) resolves +/// every result column's type via tokio's `typeinfo_statement` — *while the +/// backend is pinned*, so the named statement can't be orphaned +/// mid-resolution, and *before* the data fetch drains the response stream, +/// so tokio never interleaves a fresh named prepare+query into a mid-drain +/// unnamed batch (the Supavisor hang mechanism). Pagination's appended +/// `LIMIT`/`OFFSET` does not change result column types, so the warmed +/// cache applies to the paginated data fetch too. +/// 4. `query_typed(&final_query, &[])` runs the data fetch using the +/// unnamed protocol — #149's pooler-safe path is preserved (#1097 stays +/// fixed), and binary `Row` decoding / `extract_value` is unchanged. +/// 5. A final `DEALLOCATE ALL` clears the statements `prepare` created. +/// 6. `COMMIT` ends the pin. +/// +/// The client comes from [`client::oneshot_pool_client`] — a transient pool +/// not inserted into the global cache. Both the client and its pool drop at +/// the end, so the connection (and tokio-postgres's private cached +/// `typeinfo_statement` handle, which `DEALLOCATE ALL` can't clear) dies with +/// it. No stale client-side type cache is ever returned to the pool. +/// +/// The whole sequence is bounded by [`POOLED_TXN_TIMEOUT`]; a stall is an +/// actionable error, not a 120s host timeout. On any error after `BEGIN` we +/// best-effort `ROLLBACK` (the one-shot client drop closes the connection +/// regardless). +async fn exec_query_pooled_txn( + conn_params: &ConnectionParams, + query: &str, + limit: Option, + page: u32, + schema: Option<&str>, + request_id: &Value, +) -> Result { + // One-shot client: not returned to any pool, so its caches die with it. + let pg_client = client::oneshot_pool_client(conn_params).await?; + + // Register a cancel handle for this in-flight query (#126). The guard + // deregisters on drop — at every return path below, Ok or Err. + let _cancel_guard = + cancel::CancelGuard::register(request_id.as_u64(), &pg_client, conn_params.clone()); + + // Wrap the pinned work in a plugin-owned timeout so a Supavisor-style + // stall surfaces here instead of at the host's 120s call timeout. A + // timeout means the pooler wedged; the one-shot client drops and its + // connection (and stale caches) dies regardless of the COMMIT/ROLLBACK + // below. + match tokio::time::timeout( + POOLED_TXN_TIMEOUT, + pin_pooled_txn(&pg_client, query, limit, page, schema), + ) + .await + { + Ok(inner) => inner, + Err(_elapsed) => Err(format!( + "Query timed out after {}s on a transaction-mode pooler. \ + This is the known #160/#163 failure mode when the pooler can't \ + track prepared statements for non-builtin result column types \ + (custom enums, composites, extension types). For PgBouncer, set \ + `max_prepared_statements > 0` in pgbouncer.ini; for poolers \ + without that setting, use a direct/session connection or the \ + pooler's session mode. See: \ + https://www.pgbouncer.org/config.html#max_prepared_statements", + POOLED_TXN_TIMEOUT.as_secs() + )), + } +} + +/// The pinned-transaction body, split out so [`exec_query_pooled_txn`] can +/// wrap it in a timeout without holding the one-shot client's lifetime in a +/// `Future`. `pg_client` must already be acquired. +/// +/// Every error path after `BEGIN` except a `COMMIT` failure best-effort +/// rolls back before returning. A failed `COMMIT` doesn't need it: PostgreSQL +/// itself rolls the transaction back when `COMMIT` fails (e.g. on a deferred +/// constraint). In all cases the one-shot client drops on return and the +/// server aborts any leftover transaction when the connection closes, so the +/// explicit rollback is belt-and-suspenders — but explicit is cheaper than +/// relying on that. +async fn pin_pooled_txn( + pg_client: &PgClient, + query: &str, + limit: Option, + page: u32, + schema: Option<&str>, +) -> Result { + // Pin the backend for the whole sequence. + if let Err(e) = pg_client.batch_execute("BEGIN").await { + return Err(client::format_pg_error(&e)); + } + + let result = run_pinned_body(pg_client, query, limit, page, schema).await; + + match result { + Ok(value) => { + if let Err(e) = pg_client.batch_execute("COMMIT").await { + return Err(format!("COMMIT failed: {}", client::format_pg_error(&e))); + } + Ok(value) + } + Err(e) => { + // best-effort rollback; the one-shot client drops and closes + // regardless, so the server aborts the txn even if this fails. + let _ = pg_client.batch_execute("ROLLBACK").await; + Err(e) + } + } +} + +/// The pinned-transaction steps between `BEGIN` and `COMMIT`/`ROLLBACK`: +/// clear stale statements, set search_path, pre-resolve column types, run +/// the unnamed data fetch, clear statements again. Any error propagates up +/// to [`pin_pooled_txn`]'s rollback branch. +async fn run_pinned_body( + pg_client: &PgClient, + query: &str, + limit: Option, + page: u32, + schema: Option<&str>, +) -> Result { + // Clear stale server-side prepared statements left by a previous + // assignment's typeinfo query (or the connection's prior use). + pg_client + .batch_execute("DEALLOCATE ALL") + .await + .map_err(|e| format!("DEALLOCATE ALL failed: {}", client::format_pg_error(&e)))?; + + // Apply search_path inside the pinned transaction. `SET LOCAL` keeps the + // change scoped to this transaction so it can't leak across pooler + // reassignment — and since the backend is pinned for the transaction, + // the `SET LOCAL` is guaranteed to land on the same backend that runs + // the query. + if let Some(s) = schema { + apply_search_path_local(pg_client, s).await?; + } + + // Pre-resolve result column types while the backend is pinned. This is + // the core of the fix: `prepare`'s internal `get_type`/`typeinfo_statement` + // round-trip runs inside the pin, before the data fetch drains anything. + // A 26000 here (the pooler couldn't track the statement even while + // pinned — only possible if it rejects named statements outright) is + // surfaced via `format_query_error` so the user sees the actionable + // `max_prepared_statements` guidance. + let is_transaction_pooler = true; + if let Err(e) = pg_client.prepare(query).await { + return Err(client::format_query_error(&e, is_transaction_pooler)); + } + + // Data fetch: unnamed protocol (#149's pooler-safe path). Binary Row + // decoding + extract_value unchanged. + let result = exec_query_on_client(pg_client, query, limit, page, is_transaction_pooler).await; + + // Clear the statements `prepare` created before COMMIT. Errors here are + // non-fatal — the one-shot client drops and closes regardless. Log and + // continue to the result. + if let Err(e) = pg_client.batch_execute("DEALLOCATE ALL").await { + log::debug!( + "pooled-txn trailing DEALLOCATE failed: {}", + client::format_pg_error(&e) + ); + } + + result +} + +/// Read the connection's base `search_path` and set it locally (`SET LOCAL`) +/// for the current pinned transaction. Same base-snapshot logic as +/// [`apply_search_path`], but the `SET LOCAL` (rather than `SET`) keeps the +/// change scoped to this transaction so it can't leak across pooler +/// reassignment — and since the backend is pinned for the transaction, the +/// `SET LOCAL` is guaranteed to land on the same backend that runs the query. +async fn apply_search_path_local(pg_client: &PgClient, schema: &str) -> Result<(), String> { + let read = pg_client + .query_typed_one( + "SELECT current_setting($1, true), current_setting('search_path')", + &[( + &client::BASE_SEARCH_PATH_GUC as &(dyn ToSql + Sync), + Type::TEXT, + )], + ) + .await + .map_err(|e| { + format!( + "Failed to read search_path: {}", + client::format_pg_error(&e) + ) + })?; + let snapshot: Option = read.try_get(0).map_err(|e| { + format!( + "Failed to read search_path: {}", + client::format_pg_error(&e) + ) + })?; + let live: String = read.try_get(1).map_err(|e| { + format!( + "Failed to read search_path: {}", + client::format_pg_error(&e) + ) + })?; + let current = snapshot.unwrap_or(live); + let set_path = build_search_path_sql(schema, ¤t, true); + pg_client + .batch_execute(&set_path) + .await + .map_err(|e| format!("Failed to set search_path: {}", client::format_pg_error(&e))) +} + +/// Build a `SET [LOCAL] search_path TO` statement that prepends `schema` onto +/// the connection's existing `search_path` rather than replacing it (#51), so /// queries against objects in a different schema still resolve. `current` /// is the verbatim value of `current_setting('search_path')`, emitted /// unchanged — PostgreSQL accepts its own display form back as the right-hand /// side of `SET search_path TO`, including the `"$user"` placeholder (probed /// against a live database: `"$user"` set back verbatim still resolves to the /// user's schema; unquoting it is a syntax error). -fn build_search_path_sql(schema: &str, current: &str) -> String { +/// +/// `local` selects `SET LOCAL` (scoped to the current transaction) over `SET` +/// (session-scoped). The pinned one-shot transaction path uses `SET LOCAL` so +/// the change can't leak across pooler reassignment; the shared cached path +/// uses `SET`. +fn build_search_path_sql(schema: &str, current: &str, local: bool) -> String { + let kw = if local { + "SET LOCAL search_path TO" + } else { + "SET search_path TO" + }; if current.trim().is_empty() { - format!("SET search_path TO {}", quote_identifier(schema)) + format!("{kw} {}", quote_identifier(schema)) } else { - format!( - "SET search_path TO {}, {}", - quote_identifier(schema), - current - ) + format!("{kw} {}, {}", quote_identifier(schema), current) } } @@ -387,7 +633,7 @@ async fn apply_search_path(pg_client: &PgClient, schema: &str) -> Result<(), Str // Prefer the stable base snapshot; fall back to the live value if it's // unset so an older pool (no snapshot hook) keeps working. let current = snapshot.unwrap_or(live); - let set_path = build_search_path_sql(schema, ¤t); + let set_path = build_search_path_sql(schema, ¤t, false); pg_client .batch_execute(&set_path) .await diff --git a/src/handlers/query_tests.rs b/src/handlers/query_tests.rs index f994035..1f86ff8 100644 --- a/src/handlers/query_tests.rs +++ b/src/handlers/query_tests.rs @@ -365,7 +365,7 @@ fn build_search_path_sql_prepends_onto_existing_path() { // The selected schema resolves first; the connection's real path is kept // as fallback so queries against other schemas still resolve (#51). assert_eq!( - build_search_path_sql("analyst", "\"$user\", public"), + build_search_path_sql("analyst", "\"$user\", public", false), "SET search_path TO \"analyst\", \"$user\", public" ); } @@ -373,7 +373,7 @@ fn build_search_path_sql_prepends_onto_existing_path() { #[test] fn build_search_path_sql_emits_only_schema_when_current_is_empty() { assert_eq!( - build_search_path_sql("analyst", ""), + build_search_path_sql("analyst", "", false), "SET search_path TO \"analyst\"" ); } @@ -381,7 +381,7 @@ fn build_search_path_sql_emits_only_schema_when_current_is_empty() { #[test] fn build_search_path_sql_emits_only_schema_when_current_is_whitespace() { assert_eq!( - build_search_path_sql("analyst", " "), + build_search_path_sql("analyst", " ", false), "SET search_path TO \"analyst\"" ); } @@ -393,7 +393,7 @@ fn build_search_path_sql_preserves_user_placeholder_verbatim() { // live database). The helper must not unquote it — unquoted `$user` is a // syntax error in `SET search_path TO`. assert_eq!( - build_search_path_sql("analyst", "\"$user\", public"), + build_search_path_sql("analyst", "\"$user\", public", false), "SET search_path TO \"analyst\", \"$user\", public" ); } @@ -403,7 +403,54 @@ fn build_search_path_sql_quotes_schema_with_special_characters() { // A schema name needing quoting is escaped via `quote_identifier`, not // emitted raw. assert_eq!( - build_search_path_sql("a\"b", "public"), + build_search_path_sql("a\"b", "public", false), "SET search_path TO \"a\"\"b\", public" ); } + +// `local: true` selects `SET LOCAL` (transaction-scoped) over `SET` +// (session-scoped). The pinned one-shot transaction path (#160/#163) uses +// `SET LOCAL` so the change can't leak across pooler reassignment. These +// mirror the `SET` cases above so the local branch is exercised for every +// shape the shared cached path covers — a unit-test analog of +// `apply_search_path_local`, which is otherwise only reachable through the +// live pinned-path test (and only when a `schema` param is sent). +#[test] +fn build_search_path_sql_local_prepends_onto_existing_path() { + assert_eq!( + build_search_path_sql("analyst", "\"$user\", public", true), + "SET LOCAL search_path TO \"analyst\", \"$user\", public" + ); +} + +#[test] +fn build_search_path_sql_local_emits_only_schema_when_current_is_empty() { + assert_eq!( + build_search_path_sql("analyst", "", true), + "SET LOCAL search_path TO \"analyst\"" + ); +} + +#[test] +fn build_search_path_sql_local_emits_only_schema_when_current_is_whitespace() { + assert_eq!( + build_search_path_sql("analyst", " ", true), + "SET LOCAL search_path TO \"analyst\"" + ); +} + +#[test] +fn build_search_path_sql_local_preserves_user_placeholder_verbatim() { + assert_eq!( + build_search_path_sql("analyst", "\"$user\", public", true), + "SET LOCAL search_path TO \"analyst\", \"$user\", public" + ); +} + +#[test] +fn build_search_path_sql_local_quotes_schema_with_special_characters() { + assert_eq!( + build_search_path_sql("a\"b", "public", true), + "SET LOCAL search_path TO \"a\"\"b\", public" + ); +} diff --git a/src/models.rs b/src/models.rs index a2a9937..c9aae50 100644 --- a/src/models.rs +++ b/src/models.rs @@ -70,7 +70,7 @@ impl ConnectionParams { pub fn is_transaction_pooler(&self) -> bool { self.pooler_mode .as_deref() - .is_some_and(|m| m.eq_ignore_ascii_case("transaction")) + .is_some_and(|m| m.trim().eq_ignore_ascii_case("transaction")) } } diff --git a/src/models_tests.rs b/src/models_tests.rs index b76d0a9..c028b68 100644 --- a/src/models_tests.rs +++ b/src/models_tests.rs @@ -58,6 +58,19 @@ fn is_transaction_pooler_case_insensitive() { assert!(params_with_pooler_mode(Some("Transaction")).is_transaction_pooler()); } +/// Trimmed: a `pooler_mode` with surrounding whitespace (e.g. a hand-edited +/// `connections.json` with a trailing space) still routes to the pinned path +/// — matching how `startup_script` is trimmed in `build_pool`. Without the +/// trim, "transaction " would silently bypass the fix and the user would hit +/// the #160/#163 hang on the shared cached client (the very thing this PR +/// fixes), with no diagnostic pointing at the whitespace. +#[test] +fn is_transaction_pooler_trims_surrounding_whitespace() { + assert!(params_with_pooler_mode(Some("transaction ")).is_transaction_pooler()); + assert!(params_with_pooler_mode(Some(" transaction")).is_transaction_pooler()); + assert!(params_with_pooler_mode(Some(" transaction ")).is_transaction_pooler()); +} + /// Non-string JSON for `pooler_mode` (number, bool, null) is silently /// treated as None — not a transaction pooler. #[test]