diff --git a/src/client_live_tests.rs b/src/client_live_tests.rs index 89ffdaa..bf583ea 100644 --- a/src/client_live_tests.rs +++ b/src/client_live_tests.rs @@ -659,3 +659,430 @@ async fn wide_enum_table_returns_rows_via_batch_on_direct_connection() { let key = connection_key(¶ms); POOLS.lock().unwrap().remove(&key); } + +/// Regression guard for #169: an interactive transaction **session** on a +/// transaction-mode pooler (`pooler_mode: "transaction"` + `session_id` set) +/// must not hang on a table with a non-builtin column type and a large +/// result. The #166 fix (#168) only pre-resolves for non-session single +/// queries; the #164 pinned one-shot path only covers non-session single +/// queries. The batch/session path (`run_batch_in_session`) bypassed both, +/// so a `SELECT *` on a wide enum table inside an open session transaction +/// still deadlocked. +/// +/// This test reproduces the reporter's actual scenario: a two-batch session +/// (`BEGIN` in batch 1, `SELECT *` in batch 2) on a `pooler_mode: +/// "transaction"` connection, so the SELECT runs on the reused/pinned +/// connection. The wide-enum table (600 rows × 500-byte pad + the enum +/// column) triggers the mid-stream `typeinfo_statement` deadlock without +/// the pre-resolve. rc.11/#168 hangs past 120s on this; the fix pre-resolves +/// via `prepare` inside the user's open transaction (the backend is pinned, +/// so the named statement can't be orphaned) → returns cleanly. +/// +/// Run against direct Postgres with `pooler_mode: "transaction"` (the +/// `pooler_mode` flag routes through the pooler code paths even without a +/// real pooler; the deadlock is a tokio-postgres interleave, not +/// pooler-specific, so direct PG reproduces it). +#[tokio::test] +#[ignore = "needs a live PostgreSQL (PG_PLUGIN_TEST_HOST=127.0.0.1 ... podman 54320)"] +async fn wide_enum_table_returns_rows_via_session_pinned_batch_on_pooler() { + use crate::client; + use crate::handlers::query::{execute_query_batch, release_session}; + 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_169_session_mood"; + let table_name = "plugin_test_169_session_enum_table"; + let session_id = "test-169-session-pinned"; + + // Setup: same wide-table fixture as the #166 tests (600 rows × 500-byte + // pad + the enum column — enough to fill the channel buffer and trigger + // the mid-stream get_type deadlock without the pre-resolve). + for stmt in [ + format!("DROP TABLE IF EXISTS {table_name}"), + format!("DROP TYPE IF EXISTS {type_name}"), + format!("CREATE TYPE {type_name} AS ENUM ('happy', 'sad')"), + format!( + "CREATE TABLE {table_name} AS SELECT g AS id, 'happy'::{type_name} AS m, \ + repeat('x', 500) AS pad FROM generate_series(1, 600) g" + ), + ] { + client::execute_typed(¶ms, &stmt, &[]) + .await + .unwrap_or_else(|e| panic!("setup statement failed ({stmt}): {e}")); + } + + // Batch 1: BEGIN (opens a transaction, pins the connection to the session). + let begin_params = json!({ + "driver": "postgresql", + "host": params.host, + "port": params.port, + "database": params.database, + "username": params.username, + "password": params.password, + "pooler_mode": "transaction", + "queries": ["BEGIN"], + "session_id": session_id, + "limit": 100, + "page": 1, + }); + let begin_resp = execute_query_batch(json!(1), &begin_params).await; + assert!( + begin_resp.get("error").is_none(), + "BEGIN batch should succeed; got: {begin_resp}" + ); + + // Batch 2: SELECT * on the wide enum table, same session (reused/pinned + // connection). This is the hanging case on rc.11/#168: the batch path + // passes is_transaction_pooler=true to exec_query_on_client, skipping the + // pre-resolve → query_typed fires typeinfo_statement mid-stream on the + // wide result → deadlock. With the fix, in_transaction=true → the + // pre-resolve prepare fires inside the user's open transaction (backend + // pinned) → no mid-stream typeinfo_statement → returns cleanly. + let select_params = json!({ + "driver": "postgresql", + "host": params.host, + "port": params.port, + "database": params.database, + "username": params.username, + "password": params.password, + "pooler_mode": "transaction", + "queries": [format!("SELECT * FROM {table_name}")], + "session_id": session_id, + "limit": 50, + "page": 1, + }); + let select_resp = tokio::time::timeout( + std::time::Duration::from_secs(30), + execute_query_batch(json!(2), &select_params), + ) + .await + .expect( + "session-pinned SELECT on a wide enum table hung past 30s on a pooler-mode connection — \ + the #169 regression is present (the session-pinned pre-resolve is missing or broken)", + ); + + assert!( + select_resp.get("error").is_none(), + "session-pinned SELECT batch should succeed; got: {select_resp}" + ); + // Session-aware batch returns {results, in_transaction}. + let results = select_resp + .get("result") + .and_then(|r| r.get("results")) + .and_then(|r| r.as_array()) + .expect("session-aware batch returns result.results"); + assert_eq!( + results.len(), + 1, + "one query in the batch; got {select_resp}" + ); + let first = &results[0]; + assert!( + first.get("error").and_then(|e| e.as_str()).is_none(), + "the SELECT should not error; got: {first}" + ); + let rows = first + .get("result") + .and_then(|r| r.get("rows")) + .and_then(|r| r.as_array()) + .expect("the query's result.rows is an array"); + assert!( + !rows.is_empty(), + "the wide enum-typed table has 600 rows; got empty rows: {first}" + ); + assert_eq!( + rows[0].as_array().unwrap()[1], + json!("happy"), + "enum column should decode to \"happy\"; got: {first}" + ); + + // Batch 3: ROLLBACK (ends the transaction, releases the pin). + let rollback_params = json!({ + "driver": "postgresql", + "host": params.host, + "port": params.port, + "database": params.database, + "username": params.username, + "password": params.password, + "pooler_mode": "transaction", + "queries": ["ROLLBACK"], + "session_id": session_id, + "limit": 100, + "page": 1, + }); + let _ = execute_query_batch(json!(3), &rollback_params).await; + + // Release the session (frees the pinned connection). + let _ = release_session(json!(4), &json!({"session_id": session_id})).await; + + // Teardown. + 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 #169 (non-session sub-case): a non-session batch +/// (`session_id` absent) on a transaction-mode pooler must route result-set +/// queries to the one-shot pinned path (`exec_query_pooled_txn`), not the +/// shared cached client (where a plain `prepare` would collide across +/// backend reassignment). This is the batch-path equivalent of the existing +/// `pooled_txn_enum_column_returns_rows` test (which covers the single-query +/// `exec_query` routing). +/// +/// Without the fix, a non-session pooler batch's result-set query would run +/// `query_typed` on the shared cached client with no pre-resolve → the same +/// mid-stream `typeinfo_statement` deadlock on wide non-builtin results. With +/// the fix, `in_transaction=false` + `is_transaction_pooler()` + +/// `returns_result_set` routes to `exec_query_pooled_txn` (one-shot client, +/// BEGIN/prepare/query_typed/COMMIT) → returns cleanly. +#[tokio::test] +#[ignore = "needs a live PostgreSQL (PG_PLUGIN_TEST_HOST=127.0.0.1 ... podman 54320)"] +async fn wide_enum_table_returns_rows_via_non_session_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"); + + let type_name = "plugin_test_169_nosession_mood"; + let table_name = "plugin_test_169_nosession_enum_table"; + + // Setup: same wide-table fixture (600 rows × 500-byte pad + enum column). + for stmt in [ + format!("DROP TABLE IF EXISTS {table_name}"), + format!("DROP TYPE IF EXISTS {type_name}"), + format!("CREATE TYPE {type_name} AS ENUM ('happy', 'sad')"), + format!( + "CREATE TABLE {table_name} AS SELECT g AS id, 'happy'::{type_name} AS m, \ + repeat('x', 500) AS pad FROM generate_series(1, 600) g" + ), + ] { + client::execute_typed(¶ms, &stmt, &[]) + .await + .unwrap_or_else(|e| panic!("setup statement failed ({stmt}): {e}")); + } + + // Non-session batch (no `session_id`): `in_transaction=false` + + // `is_transaction_pooler()` + `returns_result_set` → routes to + // `exec_query_pooled_txn` (one-shot client + BEGIN/prepare/query_typed/ + // COMMIT). The wide enum table at limit 50 is the hanging case without + // the routing. + 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!("SELECT * FROM {table_name}")], + "limit": 50, + "page": 1, + }); + + let response = tokio::time::timeout( + std::time::Duration::from_secs(30), + execute_query_batch(json!(1), &request_params), + ) + .await + .expect( + "non-session batch SELECT on a wide enum table hung past 30s on a pooler-mode connection — \ + the #169 regression is present (the non-session batch routing is missing or broken)", + ); + + assert!( + response.get("error").is_none(), + "non-session batch should succeed; got: {response}" + ); + // Non-session batch: `ok_response(id, json!(results))` puts the bare + // results array under the JSON-RPC `result` field. + let results = response + .get("result") + .and_then(|r| r.as_array()) + .expect("non-session batch returns a bare `result` array"); + assert_eq!(results.len(), 1, "one query in the batch; got {response}"); + let first = &results[0]; + assert!( + first.get("error").and_then(|e| e.as_str()).is_none(), + "the batch query should not error; got: {first}" + ); + let rows = first + .get("result") + .and_then(|r| r.get("rows")) + .and_then(|r| r.as_array()) + .expect("the query's result.rows is an array"); + assert!( + !rows.is_empty(), + "the wide enum-typed table has 600 rows; got empty rows: {first}" + ); + assert_eq!( + rows[0].as_array().unwrap()[1], + json!("happy"), + "enum column should decode to \"happy\"; got: {first}" + ); + + // Teardown. + 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 empty-rows column-discovery branch on the #169 +/// session-pinned pooler batch path. `exec_query_on_client`'s empty-rows +/// block falls back to `prepare` for column discovery when +/// `pre_resolve_types` is true (query.rs:849) — a path that became reachable +/// for session-pinned pooler batches only after #169 passed +/// `pre_resolve_types=true` on the `in_transaction` branch. Before #169 a +/// pooler (`is_transaction_pooler=true`) always skipped `prepare`, so an +/// empty result returned `columns: []`; after #169 it returns the real +/// column names. Neither #169 test exercises this (both use 600-row tables), +/// so this test guards the branch: a `SELECT ... WHERE 1=0` in a +/// session-pinned transaction on a `pooler_mode: "transaction"` connection +/// must return the table's column names, not `[]`. +#[tokio::test] +#[ignore = "needs a live PostgreSQL (PG_PLUGIN_TEST_HOST=127.0.0.1 ... podman 54320)"] +async fn empty_result_returns_columns_via_session_pinned_batch_on_pooler() { + use crate::client; + use crate::handlers::query::{execute_query_batch, release_session}; + 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_169_empty_mood"; + let table_name = "plugin_test_169_empty_enum_table"; + let session_id = "test-169-empty-columns"; + + // Setup: the same enum-typed table, but the query below selects zero rows. + for stmt in [ + format!("DROP TABLE IF EXISTS {table_name}"), + format!("DROP TYPE IF EXISTS {type_name}"), + format!("CREATE TYPE {type_name} AS ENUM ('happy', 'sad')"), + format!( + "CREATE TABLE {table_name} AS SELECT g AS id, 'happy'::{type_name} AS m, \ + repeat('x', 500) AS pad FROM generate_series(1, 600) g" + ), + ] { + client::execute_typed(¶ms, &stmt, &[]) + .await + .unwrap_or_else(|e| panic!("setup statement failed ({stmt}): {e}")); + } + + // Batch 1: BEGIN (pins the connection; sets in_transaction=true for the + // next batch, so the SELECT routes to the pre-resolve-on-shared-client + // path with pre_resolve_types=true). + let begin_params = json!({ + "driver": "postgresql", + "host": params.host, + "port": params.port, + "database": params.database, + "username": params.username, + "password": params.password, + "pooler_mode": "transaction", + "queries": ["BEGIN"], + "session_id": session_id, + "limit": 100, + "page": 1, + }); + let begin_resp = execute_query_batch(json!(1), &begin_params).await; + assert!( + begin_resp.get("error").is_none(), + "BEGIN batch should succeed; got: {begin_resp}" + ); + + // Batch 2: a SELECT that returns ZERO rows. With the #169 fix this runs + // pre_resolve_types=true on the shared (session-pinned) client, so + // `prepare` succeeds and the empty-rows branch reuses `.columns()` → + // real column names. If the empty-rows fallback condition is reverted to + // `!is_transaction_pooler` (or the pre-resolve is skipped), this returns + // `columns: []` and the assertion fails. + let select_params = json!({ + "driver": "postgresql", + "host": params.host, + "port": params.port, + "database": params.database, + "username": params.username, + "password": params.password, + "pooler_mode": "transaction", + "queries": [format!("SELECT id, m, pad FROM {table_name} WHERE 1 = 0")], + "session_id": session_id, + "limit": 50, + "page": 1, + }); + let select_resp = execute_query_batch(json!(2), &select_params).await; + assert!( + select_resp.get("error").is_none(), + "empty SELECT batch should succeed; got: {select_resp}" + ); + let results = select_resp + .get("result") + .and_then(|r| r.get("results")) + .and_then(|r| r.as_array()) + .expect("session-aware batch returns result.results"); + let first = &results[0]; + assert!( + first.get("error").and_then(|e| e.as_str()).is_none(), + "the SELECT should not error; got: {first}" + ); + let columns = first + .get("result") + .and_then(|r| r.get("columns")) + .and_then(|r| r.as_array()) + .expect("the query's result.columns is an array"); + assert_eq!( + columns, + &["id", "m", "pad"] + .iter() + .map(|s| json!(s)) + .collect::>(), + "an empty result on a session-pinned pooler batch must return the \ + real column names (the pre-resolve ran); got `columns: {columns:?}` — \ + if this is `[]`, the empty-rows column-discovery fallback regressed", + ); + + // Batch 3: ROLLBACK (ends the transaction, releases the pin). + let rollback_params = json!({ + "driver": "postgresql", + "host": params.host, + "port": params.port, + "database": params.database, + "username": params.username, + "password": params.password, + "pooler_mode": "transaction", + "queries": ["ROLLBACK"], + "session_id": session_id, + "limit": 100, + "page": 1, + }); + let _ = execute_query_batch(json!(3), &rollback_params).await; + let _ = release_session(json!(4), &json!({"session_id": session_id})).await; + + // Teardown. + 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); +} diff --git a/src/handlers/query.rs b/src/handlers/query.rs index 9fe90ac..13fac97 100644 --- a/src/handlers/query.rs +++ b/src/handlers/query.rs @@ -146,8 +146,15 @@ async fn run_batch_in_session( // for a reused id can't target a later one. The cancel only needs the // token (captured into the action) and the connection params, never the // borrowed client. - let _cancel_guard = - cancel::CancelGuard::register(request_id.as_u64(), &pg_client, conn_params.clone()); + // Held across the whole batch so a host `cancel(request_id)` can stop the + // in-flight shared-client query. Re-armed after each one-shot routing + // below (the `!in_transaction` branch) — see there for why it can't be a + // single guard. + let mut cancel_guard = Some(cancel::CancelGuard::register( + request_id.as_u64(), + &pg_client, + conn_params.clone(), + )); // Applying search_path to a reused connection would run inside the open // transaction and change what the rest of it sees. @@ -158,14 +165,98 @@ async fn run_batch_in_session( let mut results: Vec = Vec::new(); for query in queries { let start = Instant::now(); - let outcome = exec_query_on_client( - &pg_client, - query, - limit, - page, - conn_params.is_transaction_pooler(), - ) - .await; + + // On a transaction-mode pooler, a result-set query with non-builtin + // column types needs its column types pre-resolved before the data + // fetch, or `query_typed`'s internal `typeinfo_statement` fires + // mid-stream and deadlocks on large results (#169, same root as + // #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`): + // + // * `in_transaction` (the user's `BEGIN` has pinned the backend on + // this shared connection — either a session-pinned batch reusing a + // prior `BEGIN`, or a bare batch whose earlier statement was + // `BEGIN`): a plain `prepare` inside that transaction is safe (no + // reassignment can orphan the named statement) and keeps the + // 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 + // 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). + // + // 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 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. + exec_query_on_client( + &pg_client, + query, + limit, + page, + conn_params.is_transaction_pooler(), + /* pre_resolve_types = */ true, + ) + .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). + // + // Cancel-guard coexistence: `exec_query_pooled_txn` registers + // its own `CancelGuard` under the SAME `request_id` we hold + // below. `cancel::insert_action` is a `HashMap::insert` + // (overwrite), and the one-shot's guard `Drop` removes the + // key entirely — so if we kept a single guard across the whole + // batch, the one-shot would evict it and every subsequent + // shared-client query would lose cancel coverage (#126's + // "a late cancel is never a no-op for an in-flight query" + // guarantee). Take+drop our guard so the one-shot owns the + // map entry while it runs, then re-register for the shared + // client before the next iteration. + drop(cancel_guard.take()); + let result = + exec_query_pooled_txn(conn_params, query, limit, page, schema, request_id) + .await; + cancel_guard = Some(cancel::CancelGuard::register( + request_id.as_u64(), + &pg_client, + conn_params.clone(), + )); + result + } + } else { + // Non-result-set (DDL, writes) or direct connection — the + // unnamed protocol is pooler-safe and needs no pre-resolve. + // `cancel_guard` stays registered for this shared client. + exec_query_on_client( + &pg_client, + query, + limit, + page, + conn_params.is_transaction_pooler(), + !conn_params.is_transaction_pooler(), + ) + .await + }; let elapsed_ms = start.elapsed().as_secs_f64() * 1000.0; in_transaction = @@ -184,9 +275,12 @@ async fn run_batch_in_session( })), } } - // `_cancel_guard` drops here, deregistering the cancel handle before the + // `cancel_guard` drops here, deregistering the cancel handle before the // connection is re-pinned or released below, so a cancel racing the end - // of the batch can't target a query that's already done. + // of the batch can't target a query that's already done. It is `Some` + // here except after a final one-shot query (which left it `None`); a `None` + // drop is a no-op, so the map is clean either way. + drop(cancel_guard); match slot.as_mut() { Some(slot) if in_transaction => slot.pin(pg_client), @@ -310,9 +404,12 @@ async fn exec_query( // (#166) — cheaper than the one-shot path (preserves pool reuse) and safe // (no backend reassignment to collide with). // - // 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. + // Non-result-set statements (DDL, writes) stay on the shared cached + // 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`). if conn_params.is_transaction_pooler() && returns_result_set(query) { return exec_query_pooled_txn(conn_params, query, limit, page, schema, request_id).await; } @@ -336,6 +433,7 @@ async fn exec_query( limit, page, conn_params.is_transaction_pooler(), + !conn_params.is_transaction_pooler(), ) .await } @@ -505,8 +603,17 @@ async fn run_pinned_body( } // 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; + // Decoding + extract_value unchanged. `pre_resolve_types = false` because + // `run_pinned_body` already pre-resolved via its own `prepare` above. + let result = exec_query_on_client( + pg_client, + query, + limit, + page, + is_transaction_pooler, + /* pre_resolve_types = */ false, + ) + .await; // Clear the statements `prepare` created before COMMIT. Errors here are // non-fatal — the one-shot client drops and closes regardless. Log and @@ -651,12 +758,22 @@ async fn apply_search_path(pg_client: &PgClient, schema: &str) -> Result<(), Str } /// Execute a query on an existing client (used by both single and batch execution). +/// +/// `is_transaction_pooler` selects error-message phrasing (definitive when +/// the host declared a transaction pooler, conditional when it didn't). +/// `pre_resolve_types` controls whether column types are pre-resolved via +/// `prepare` before the `query_typed` data fetch (#166/#168). The two are +/// independent: a session-pinned pooler connection passes +/// `is_transaction_pooler = true` (for correct error phrasing) AND +/// `pre_resolve_types = true` (because the backend is pinned by the user's +/// open transaction, so the `prepare` is safe) — see `run_batch_in_session`. async fn exec_query_on_client( pg_client: &PgClient, query: &str, limit: Option, page: u32, is_transaction_pooler: bool, + pre_resolve_types: bool, ) -> Result { // Check if the statement returns a result set if !returns_result_set(query) { @@ -694,16 +811,17 @@ async fn exec_query_on_client( }; let manual_limit = if is_paginated { None } else { limit }; - // Pre-resolve result column types on a *direct* connection (no pooler) - // before the data fetch (#166). tokio-postgres's `query_typed` resolves - // non-builtin column types (`get_type` → the internal named - // `typeinfo_statement`) *inside* its `RowDescription` loop, mid-stream, - // over the same channel delivering the query's `DataRow`s. The typeinfo - // reply is queued behind the pending rows, but `query_typed` won't read - // the rows until `get_type` returns → a deadlock once the row data fills - // the bounded `responses` channel + socket buffer (hence the size - // threshold — small results pass, large ones hang). This is not - // pooler-specific; it deadlocks on any connection. + // Pre-resolve result column types before the data fetch (#166 for direct + // connections, #169 for session-pinned pooler connections inside an open + // transaction). tokio-postgres's `query_typed` resolves non-builtin column + // types (`get_type` → the internal named `typeinfo_statement`) *inside* + // its `RowDescription` loop, mid-stream, over the same channel delivering + // the query's `DataRow`s. The typeinfo reply is queued behind the pending + // rows, but `query_typed` won't read the rows until `get_type` returns → + // a deadlock once the row data fills the bounded `responses` channel + + // socket buffer (hence the size threshold — small results pass, large + // ones hang). This is not pooler-specific; it deadlocks on any + // connection. // // `prepare` runs Parse+Describe+Sync with no Execute/no DataRows, so its // `get_type` interleave has an empty channel and completes cleanly, @@ -716,21 +834,32 @@ 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 to direct connections: on a transaction-mode pooler a plain - // `prepare`'s named statement collides across backend reassignment (the - // `s0`→`s1` shift that killed the removed pre-warm, commit ca33b98), so - // poolers keep the #164 pinned one-shot path instead. The pre-resolve is - // non-fatal on failure: a prepare error is swallowed (`.ok()`) and the - // `query_typed` below runs anyway — for builtin types `get_type` is always - // a `Type::from_oid` cache hit (no `typeinfo_statement`, no hang), and a - // direct prepare failure on a valid result-set query is rare. The - // `Statement` is kept alive (not dropped immediately) so its `.columns()` - // can be reused by the empty-rows branch below instead of preparing a - // second time. (The warmed type cache persists on the shared `InnerClient` - // regardless of the `Statement`'s lifetime — `Statement::Drop` sends a - // server `Close` but does not clear `CachedTypeInfo.types` — so holding - // it alive is for `.columns()` reuse, not cache persistence.) - let prepared = if !is_transaction_pooler { + // Gated on `pre_resolve_types`, set by the caller for the two 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`. + // + // The pre-resolve is non-fatal on failure: a prepare error is swallowed + // (`.ok()`) and the `query_typed` below runs anyway — for builtin types + // `get_type` is always a `Type::from_oid` cache hit (no + // `typeinfo_statement`, no hang), and a prepare failure on a valid + // result-set query is rare. The `Statement` is kept alive (not dropped + // immediately) so its `.columns()` can be reused by the empty-rows branch + // below instead of preparing a second time. (The warmed type cache + // persists on the shared `InnerClient` regardless of the `Statement`'s + // lifetime — `Statement::Drop` sends a server `Close` but does not clear + // `CachedTypeInfo.types` — so holding it alive is for `.columns()` reuse, + // not cache persistence.) + let prepared = if pre_resolve_types { pg_client.prepare(&final_query).await.ok() } else { None @@ -751,19 +880,22 @@ async fn exec_query_on_client( .map_err(|e| client::format_query_error(&e, is_transaction_pooler))?; if rows.is_empty() { - // Get columns from the statement if possible. On a direct connection - // the pre-resolve above already prepared `final_query`, so reuse its - // `.columns()` instead of preparing a second time. On a transaction- - // mode pooler, `prepare()` is skipped (its named statement would - // collide across backend reassignment) and `columns: []` is returned - // — the host handles empty columns gracefully (#157, tabularis#1097). + // Get columns from the statement if possible. When the pre-resolve + // ran and succeeded, reuse its `.columns()` instead of preparing a + // second time. When the pre-resolve didn't run (or failed), try once + // more for column discovery only if `prepare` is safe here (direct + // connections, or session-pinned pooler connections where the backend + // is pinned). On a transaction-mode pooler where `prepare` is unsafe + // (the one-shot path handles its own pre-resolve), return `columns: + // []` — the host handles empty columns gracefully (#157, + // tabularis#1097). let columns: Vec = if let Some(stmt) = prepared.as_ref() { stmt.columns() .iter() .map(|c| c.name().to_string()) .collect() - } else if !is_transaction_pooler { - // Pre-resolve didn't run or failed; try once more for column + } else if pre_resolve_types { + // Pre-resolve was attempted but failed; try once more for column // discovery (the pre-#166 behavior for direct empty result sets). if let Ok(stmt) = pg_client.prepare(&final_query).await { stmt.columns()