Skip to main content

pylon_pgcon/
lib.rs

1//
2// This source file is part of the Pylon open source project.
3//
4// Copyright (c) 2026 Jaldis B.V.
5//
6// Licensed under the MIT OR Apache-2.0 license (the "License");
7// you may not use this file except in compliance with the License.
8// You may obtain a copy of the License at
9//
10//     https://opensource.org/licenses/MIT
11//     https://www.apache.org/licenses/LICENSE-2.0
12//
13// Unless required by applicable law or agreed to in writing, software
14// distributed under the License is distributed on an "AS IS" BASIS,
15// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16// See the License for the specific language governing permissions and
17// limitations under the License.
18//
19
20//! Postgres connection pooling and query execution, via `tokio-postgres` +
21//! `deadpool-postgres`. No dependency on `pylon-core` or PyO3 — usable as a
22//! plain Rust Postgres client crate on its own; `pylon-core` depends on
23//! *this* crate for execution, not the other way around.
24//!
25//! `query_raw` returns raw `tokio_postgres::Row`s; `wire` decodes the
26//! composite `result` column those rows carry into `pylon_value::DecodedValue`
27//! (the shared decode target `pylon-cache` also stores).
28
29pub mod error;
30pub mod listener;
31pub mod numeric;
32pub mod wire;
33
34pub use error::{Error, Result};
35pub use listener::PgListener;
36pub use wire::{ExtensionOids, decode_value};
37
38use pylon_value::DecodedValue;
39use std::sync::{Arc, RwLock};
40
41/// Captures a column's raw wire bytes regardless of its declared Postgres
42/// type. `tokio_postgres`'s own `&[u8]` `FromSql` impl only accepts
43/// `BYTEA` — the `result` column pylon-core emits is always `record`
44/// (OID 2249), so a plain `row.get::<_, &[u8]>(0)` would panic on the type
45/// check. This wrapper's `accepts` is unconditionally `true`, matching the
46/// well-established pattern for pulling raw bytes out of any column.
47struct RawBytes<'a>(&'a [u8]);
48
49impl<'a> postgres_types::FromSql<'a> for RawBytes<'a> {
50    fn from_sql(
51        _ty: &postgres_types::Type,
52        raw: &'a [u8],
53    ) -> std::result::Result<Self, Box<dyn std::error::Error + Sync + Send>> {
54        Ok(RawBytes(raw))
55    }
56
57    fn accepts(_ty: &postgres_types::Type) -> bool {
58        true
59    }
60}
61
62/// Wraps a `DecodedValue` for binding as a query parameter. `accepts` is
63/// unconditionally `true` (mirroring `RawBytes` above) because the target
64/// `Type` isn't known until `Statement::params()` reports it — see
65/// `wire::encode_value`, which does the actual type-directed encoding.
66#[derive(Debug)]
67struct BoundParam<'a>(&'a DecodedValue);
68
69impl postgres_types::ToSql for BoundParam<'_> {
70    fn to_sql(
71        &self,
72        ty: &postgres_types::Type,
73        out: &mut bytes::BytesMut,
74    ) -> std::result::Result<postgres_types::IsNull, Box<dyn std::error::Error + Sync + Send>> {
75        Ok(wire::encode_value(self.0, ty, out)?)
76    }
77
78    fn accepts(_ty: &postgres_types::Type) -> bool {
79        true
80    }
81
82    postgres_types::to_sql_checked!();
83}
84
85// `deadpool_postgres::Pool` is `Arc`-backed internally, so cloning a
86// `PgPool` is cheap and shares the same underlying pool — needed at the
87// pyo3 boundary, where a lock guard over the process-global pool slot
88// can't be held across an `.await` (it isn't `Send`), so callers clone the
89// pool out from under the lock first.
90#[derive(Clone, Debug)]
91pub struct PgPool {
92    pool: deadpool_postgres::Pool,
93    /// Shared with every clone of this pool, so a refresh reaches all of
94    /// them — consumers hold the pool behind an `Arc`, with only `&self`
95    /// available at the moments the registry needs replacing.
96    types: Arc<RwLock<Arc<ExtensionOids>>>,
97    /// Label for this pool in observability output. Set via
98    /// `PgPool::set_name`; `"default"` until then.
99    name: String,
100}
101
102/// The one failure a stale type registry causes, and so the only one worth
103/// re-discovering types for and retrying. Any other error would survive a
104/// retry, which would also run the statement twice.
105pub(crate) fn is_unknown_oid(err: &Error) -> bool {
106    matches!(err, Error::UnknownTypeOid { .. })
107}
108
109/// Notified with `(pool name, time spent waiting)` every time a caller
110/// acquires a pooled connection.
111///
112/// A hook rather than a direct metrics call because the dependency only goes
113/// one way: `pylon-workers` (which owns the Prometheus registry) depends on
114/// this crate, so this crate can't call into it. Whoever owns the registry
115/// installs an observer at startup; nobody installing one costs a single
116/// `OnceLock` read per checkout.
117pub type PoolWaitObserver = fn(&str, std::time::Duration);
118
119static POOL_WAIT_OBSERVER: std::sync::OnceLock<PoolWaitObserver> = std::sync::OnceLock::new();
120
121/// Installs the pool-wait observer. The first call wins; later ones are
122/// ignored, so a process that initialises metrics twice can't double-count.
123pub fn set_pool_wait_observer(observer: PoolWaitObserver) {
124    let _ = POOL_WAIT_OBSERVER.set(observer);
125}
126
127/// Snapshot of a pool's connection accounting — deadpool's own `Status`,
128/// re-exported under this crate's name so callers (the Prometheus gauge
129/// sampler in `pylon-py`) don't need a direct `deadpool_postgres`
130/// dependency just to read four numbers.
131pub struct PoolStatus {
132    /// Connections actually established right now (idle + checked out).
133    pub size: usize,
134    /// Idle connections available for immediate checkout right now.
135    pub available: usize,
136    /// Callers currently blocked in `.get()` waiting for a connection.
137    pub waiting: usize,
138    pub max_size: usize,
139}
140
141impl PgPool {
142    /// Connects using a `postgresql://` DSN, matching the DSN Python's
143    /// `DatabaseConfig`/`_build_dsn` already produces today. No TLS support
144    /// yet — no SSL/TLS surface exists anywhere in the project currently
145    /// (confirmed by a full-repo grep during planning), so this isn't a
146    /// regression; it's simply not needed until it is.
147    ///
148    /// Unlike `deadpool_postgres::Pool::builder(..).build()` on its own —
149    /// which only validates the DSN and is otherwise lazy, deferring the
150    /// first real connection attempt to whenever a caller first acquires
151    /// one — this eagerly acquires and immediately releases one connection
152    /// before returning, so a bad host/port/database/credentials fails
153    /// right here. That eager-connect
154    /// behavior, which the caller (`Client.ensure_connected`) depends on to
155    /// raise `ConnectionFailedError`/`ConnectionTimeoutError` immediately
156    /// rather than silently deferring the failure to the first query.
157    pub async fn connect(dsn: &str, max_size: usize) -> Result<Self> {
158        let pg_config = session_config(dsn)?;
159        let manager = deadpool_postgres::Manager::new(pg_config, tokio_postgres::NoTls);
160        let pool = deadpool_postgres::Pool::builder(manager)
161            .max_size(max_size)
162            .create_timeout(Some(std::time::Duration::from_secs(10)))
163            .runtime(deadpool_postgres::Runtime::Tokio1)
164            .build()?;
165        let client = pool.get().await?;
166        // Same round trip that proves the pool can connect also classifies
167        // this database's enum/domain/extension type OIDs, which are
168        // assigned locally and so can't be known statically.
169        let types = discover_types(&client).await?;
170        drop(client);
171        Ok(Self {
172            pool,
173            types: Arc::new(RwLock::new(Arc::new(types))),
174            name: "default".to_string(),
175        })
176    }
177
178    /// The type OIDs currently discovered for this database. Pass this to
179    /// `query_typed`/`query_composite` rather than `ExtensionOids::default()`
180    /// — without it, a `vector`, enum, or domain column has no decoder.
181    /// Owned rather than borrowed because `refresh_types` can replace it;
182    /// `&pool.types()` coerces to the `&ExtensionOids` those methods take.
183    pub fn types(&self) -> Arc<ExtensionOids> {
184        self.types.read().unwrap().clone()
185    }
186
187    /// Re-runs type discovery. Needed after a migration creates an enum,
188    /// domain or extension type, and after a restore re-assigns every local
189    /// OID underneath a running pool — the registry is a connect-time
190    /// snapshot and a pool outlives both.
191    ///
192    /// Clears the prepared-statement caches unconditionally, since DDL also
193    /// invalidates a cached plan on its own (`0A000 cached plan must not
194    /// change result type`).
195    pub async fn refresh_types(&self) -> Result<()> {
196        let client = self.checkout().await?;
197        let fresh = discover_types(&client).await?;
198        *self.types.write().unwrap() = Arc::new(fresh);
199        drop(client);
200        self.clear_statement_caches();
201        Ok(())
202    }
203
204    /// Re-runs discovery on a held connection after a decode hit an unknown
205    /// OID. Clears the statement caches only when the registry really
206    /// changed: a type nothing can decode reaches this on every query, and
207    /// emptying every connection's cache that often is its own problem.
208    async fn heal_types(&self, client: &tokio_postgres::Client) -> Result<Arc<ExtensionOids>> {
209        let fresh = Arc::new(discover_types(client).await?);
210        let changed = {
211            let mut slot = self.types.write().unwrap();
212            let changed = *slot != fresh;
213            if changed {
214                *slot = fresh.clone();
215            }
216            changed
217        };
218        if changed {
219            self.clear_statement_caches();
220        }
221        Ok(fresh)
222    }
223
224    /// Drops every pooled connection's prepared-statement cache.
225    ///
226    /// Must be called after any DDL that could change a statement's result
227    /// type or parameter types — `refresh_types` already does, so migration
228    /// paths get this for free. Cheap: it only empties the cache maps, it
229    /// does not touch the connections.
230    pub fn clear_statement_caches(&self) {
231        self.pool.manager().statement_caches.clear();
232    }
233
234    /// Labels this pool in observability output — see `PoolWaitObserver`.
235    pub fn set_name(&mut self, name: impl Into<String>) {
236        self.name = name.into();
237    }
238
239    /// Acquires a pooled connection, reporting how long that took to the
240    /// pool-wait observer if one is installed.
241    ///
242    /// Every query path goes through this rather than `self.pool.get()`
243    /// directly, so "time spent waiting for a connection" is measured in one
244    /// place and can't drift between call sites.
245    async fn checkout(&self) -> Result<deadpool_postgres::Object> {
246        let started = std::time::Instant::now();
247        let client = self.pool.get().await?;
248        if let Some(observe) = POOL_WAIT_OBSERVER.get() {
249            observe(&self.name, started.elapsed());
250        }
251        Ok(client)
252    }
253
254    /// Current connection accounting — for the Prometheus gauge sampler
255    /// (`pylon-py`'s `record_pool_metrics`), not used on any query path.
256    pub fn status(&self) -> PoolStatus {
257        let s = self.pool.status();
258        PoolStatus {
259            size: s.size,
260            available: s.available,
261            waiting: s.waiting,
262            max_size: s.max_size,
263        }
264    }
265
266    /// Executes `sql` with no parameters and returns the raw rows.
267    /// Composite/record decoding is a later phase — this only proves the
268    /// pool can connect and round-trip a query end to end.
269    pub async fn query_raw(&self, sql: &str) -> Result<Vec<tokio_postgres::Row>> {
270        let client = self.checkout().await?;
271        let rows = client.query(sql, &[]).await?;
272        Ok(rows)
273    }
274
275    /// Runs `sql` (expected to produce exactly one column, matching
276    /// pylon-core's `SELECT (...) AS result` emission) and decodes that
277    /// column of every row via `wire::decode_value`, using its actual
278    /// declared Postgres type (not assumed to be `record` — a bare scalar
279    /// `result` column decodes just as well through the same path).
280    pub async fn query_composite(&self, sql: &str, ext: &ExtensionOids) -> Result<Vec<DecodedValue>> {
281        let client = self.checkout().await?;
282        let rows = client.query(sql, &[]).await?;
283        let decode = |ext: &ExtensionOids| -> Result<Vec<DecodedValue>> {
284            rows.iter().map(|r| decode_result_column(r, ext)).collect()
285        };
286        match decode(ext) {
287            Err(e) if is_unknown_oid(&e) => {
288                let healed = self.heal_types(&client).await?;
289                decode(&healed)
290            }
291            other => other,
292        }
293    }
294
295    /// Runs `sql` with bound `params`, matched positionally to `$1, $2, ...`
296    /// — the same convention `pylon-core`'s `param_names` already assumes.
297    /// No caller-supplied parameter types: `prepare` asks Postgres itself
298    /// to analyze the SQL and report each placeholder's expected `Type`
299    /// (`Statement::params()`), which drives `wire::encode_value`'s
300    /// encoding directly. Decodes the single result column exactly like
301    /// `query_composite`.
302    pub async fn query_typed(
303        &self,
304        sql: &str,
305        params: &[DecodedValue],
306        ext: &ExtensionOids,
307    ) -> Result<Vec<DecodedValue>> {
308        let client = self.checkout().await?;
309        match query_typed_on(&client, sql, params, ext).await {
310            Err(e) if is_unknown_oid(&e) => {
311                let healed = self.heal_types(&client).await?;
312                query_typed_on(&client, sql, params, &healed).await
313            }
314            other => other,
315        }
316    }
317
318    /// Like `query_typed`, but decodes every column of every row by name
319    /// instead of assuming a single `(...) AS result` column — for
320    /// hand-written admin SQL (CLI commands, not `pylon-core`-emitted
321    /// query bodies) that reads named columns directly.
322    pub async fn query_typed_named(
323        &self,
324        sql: &str,
325        params: &[DecodedValue],
326        ext: &ExtensionOids,
327    ) -> Result<Vec<DecodedValue>> {
328        let client = self.checkout().await?;
329        match query_typed_named_on(&client, sql, params, ext).await {
330            Err(e) if is_unknown_oid(&e) => {
331                let healed = self.heal_types(&client).await?;
332                query_typed_named_on(&client, sql, params, &healed).await
333            }
334            other => other,
335        }
336    }
337
338    /// `query_typed`, with the session globals database triggers read
339    /// (`pylon.globals`) set first on the same connection. Each mutating
340    /// statement sets them anew, so a pooled connection never hands one
341    /// statement's globals to the next.
342    pub async fn query_typed_with_globals(
343        &self,
344        sql: &str,
345        params: &[DecodedValue],
346        ext: &ExtensionOids,
347        globals: &str,
348    ) -> Result<Vec<DecodedValue>> {
349        let client = self.checkout().await?;
350        set_globals_on(&client, globals, false).await?;
351        match query_typed_on(&client, sql, params, ext).await {
352            Err(e) if is_unknown_oid(&e) => {
353                let healed = self.heal_types(&client).await?;
354                query_typed_on(&client, sql, params, &healed).await
355            }
356            other => other,
357        }
358    }
359
360    /// `execute_typed`, with `pylon.globals` set — see `query_typed_with_globals`.
361    pub async fn execute_typed_with_globals(&self, sql: &str, params: &[DecodedValue], globals: &str) -> Result<u64> {
362        let client = self.checkout().await?;
363        set_globals_on(&client, globals, false).await?;
364        execute_typed_on(&client, sql, params).await
365    }
366
367    /// Runs `sql` with bound `params` (same convention as `query_typed`)
368    /// and discards the result, returning the number of rows affected —
369    /// for `INSERT`/`UPDATE`/`DELETE` where the caller has no `RETURNING`
370    /// clause to decode.
371    pub async fn execute_typed(&self, sql: &str, params: &[DecodedValue]) -> Result<u64> {
372        let client = self.checkout().await?;
373        execute_typed_on(&client, sql, params).await
374    }
375
376    /// Runs `EXPLAIN (ANALYZE, FORMAT JSON, VERBOSE) sql` with bound
377    /// `params` and returns the raw JSON output text verbatim — for
378    /// `analyze <query>` (see `pylon_core::analyze`), which parses this text
379    /// itself to correlate plan nodes back to the query's own shape.
380    /// Actually *runs* the query (`ANALYZE`), same as Postgres's own
381    /// `EXPLAIN ANALYZE`, not just its planner estimate.
382    pub async fn query_explain(&self, sql: &str, params: &[DecodedValue]) -> Result<String> {
383        let client = self.checkout().await?;
384        query_explain_on(&client, sql, params).await
385    }
386
387    /// Acquires one pooled connection and starts an explicit transaction at
388    /// the given isolation level (`"read_uncommitted"`, `"read_committed"`,
389    /// `"repeatable_read"`, or `"serializable"` — matching
390    /// `AsyncTransaction`'s existing accepted values in `client.py`, itself
391    /// the four standard SQL isolation levels). The returned
392    /// `PgTransaction` owns the connection until `commit`/`rollback`
393    /// consumes it.
394    pub async fn begin(&self, isolation: &str) -> Result<PgTransaction> {
395        let client = self.checkout().await?;
396        let level = match isolation {
397            "read_uncommitted" => "READ UNCOMMITTED",
398            "read_committed" => "READ COMMITTED",
399            "repeatable_read" => "REPEATABLE READ",
400            "serializable" => "SERIALIZABLE",
401            other => return Err(Error::message(format!("unknown isolation level: {other:?}"))),
402        };
403        client.batch_execute(&format!("BEGIN ISOLATION LEVEL {level}")).await?;
404        Ok(PgTransaction {
405            client,
406            types: self.types(),
407        })
408    }
409
410    /// Starts a transaction with no explicit isolation level — whatever
411    /// Postgres's own session/database default is applies. Used by
412    /// migration execution, which (unlike `begin`) never needs a specific
413    /// isolation level — this matches a plain `conn.transaction()`
414    /// (no `isolation=` kwarg) that the old Python migration executor used.
415    pub async fn begin_default(&self) -> Result<PgTransaction> {
416        let client = self.checkout().await?;
417        client.batch_execute("BEGIN").await?;
418        Ok(PgTransaction {
419            client,
420            types: self.types(),
421        })
422    }
423
424    /// Runs `sql` via the simple query protocol — no bind parameters, but
425    /// (unlike `query_typed`/`execute_typed`, which prepare via the extended
426    /// protocol and so accept exactly one statement) able to run several
427    /// `;`-separated statements in one call. Matches
428    /// `Connection.execute(sql)` called with no arguments, which migration
429    /// DDL steps rely on (a step's body is whatever raw SQL text sits
430    /// between `-- pylon:step` markers, often more than one statement).
431    pub async fn batch_execute(&self, sql: &str) -> Result<()> {
432        let client = self.checkout().await?;
433        client.batch_execute(sql).await?;
434        Ok(())
435    }
436
437    /// Checks out one pooled connection and hands back a handle the caller
438    /// holds across several calls, with no transaction started — for
439    /// state that's scoped to a single session rather than a single
440    /// statement or transaction, the way Postgres advisory locks
441    /// (`pg_advisory_lock`/`pg_advisory_unlock`) are: they're released by
442    /// an explicit unlock (or the session ending), not by a transaction
443    /// boundary, so acquiring and releasing one has to happen on the same
444    /// held connection — calling `PgPool::batch_execute` twice wouldn't
445    /// work, since each call may checkout a different pooled connection.
446    pub async fn connection(&self) -> Result<PgConnection> {
447        let client = self.checkout().await?;
448        Ok(PgConnection {
449            client,
450            types: self.types(),
451        })
452    }
453}
454
455/// One connection checked out of the pool and held by the caller, with no
456/// transaction open — see `PgPool::connection`.
457#[derive(Debug)]
458pub struct PgConnection {
459    client: deadpool_postgres::Object,
460    types: Arc<ExtensionOids>,
461}
462
463impl PgConnection {
464    /// The registry the owning pool held at checkout — a snapshot, since a
465    /// held connection lives for one unit of work.
466    pub fn types(&self) -> &ExtensionOids {
467        &self.types
468    }
469
470    pub async fn query_typed(
471        &self,
472        sql: &str,
473        params: &[DecodedValue],
474        ext: &ExtensionOids,
475    ) -> Result<Vec<DecodedValue>> {
476        query_typed_on(&self.client, sql, params, ext).await
477    }
478
479    pub async fn execute_typed(&self, sql: &str, params: &[DecodedValue]) -> Result<u64> {
480        execute_typed_on(&self.client, sql, params).await
481    }
482
483    pub async fn batch_execute(&self, sql: &str) -> Result<()> {
484        self.client.batch_execute(sql).await?;
485        Ok(())
486    }
487}
488
489/// Binds `params` to an already-prepared `stmt` and decodes the single
490/// `result` column of every row.
491///
492/// Split out from the `*_on` entry points below so the two kinds of
493/// connection this crate holds can each prepare the way that suits them —
494/// a pooled `deadpool_postgres::Object` through its per-connection
495/// statement cache, a bare `tokio_postgres::Client` (the LISTEN/NOTIFY
496/// connection, which is not pooled) through a plain `prepare` — without
497/// either duplicating the binding and decoding below.
498async fn query_with_stmt(
499    client: &tokio_postgres::Client,
500    stmt: &tokio_postgres::Statement,
501    params: &[DecodedValue],
502    ext: &ExtensionOids,
503) -> Result<Vec<DecodedValue>> {
504    let bound: Vec<BoundParam<'_>> = params.iter().map(BoundParam).collect();
505    let param_refs: Vec<&(dyn postgres_types::ToSql + Sync)> =
506        bound.iter().map(|p| p as &(dyn postgres_types::ToSql + Sync)).collect();
507    let rows = client.query(stmt, &param_refs).await?;
508    rows.iter().map(|row| decode_result_column(row, ext)).collect()
509}
510
511/// Like `query_with_stmt`, but decodes every column of every row by name
512/// (`decode_row_named`) instead of assuming column 0 is the whole result —
513/// see `decode_row_named`'s doc comment.
514async fn query_named_with_stmt(
515    client: &tokio_postgres::Client,
516    stmt: &tokio_postgres::Statement,
517    params: &[DecodedValue],
518    ext: &ExtensionOids,
519) -> Result<Vec<DecodedValue>> {
520    let bound: Vec<BoundParam<'_>> = params.iter().map(BoundParam).collect();
521    let param_refs: Vec<&(dyn postgres_types::ToSql + Sync)> =
522        bound.iter().map(|p| p as &(dyn postgres_types::ToSql + Sync)).collect();
523    let rows = client.query(stmt, &param_refs).await?;
524    rows.iter().map(|row| decode_row_named(row, ext)).collect()
525}
526
527async fn execute_with_stmt(
528    client: &tokio_postgres::Client,
529    stmt: &tokio_postgres::Statement,
530    params: &[DecodedValue],
531) -> Result<u64> {
532    let bound: Vec<BoundParam<'_>> = params.iter().map(BoundParam).collect();
533    let param_refs: Vec<&(dyn postgres_types::ToSql + Sync)> =
534        bound.iter().map(|p| p as &(dyn postgres_types::ToSql + Sync)).collect();
535    Ok(client.execute(stmt, &param_refs).await?)
536}
537
538/// Prepares through the pooled connection's own statement cache, so a query
539/// run twice on the same connection costs one round trip rather than two.
540///
541/// `tokio_postgres::Client::prepare` always issues a PARSE/DESCRIBE round
542/// trip; on a query whose execution is itself a single round trip that
543/// doubles the wire cost. `deadpool`'s cache is keyed on the SQL text and
544/// lives as long as the connection, which is exactly the lifetime a
545/// prepared statement has server-side. Any statement it hands back can go
546/// stale if the schema changes underneath it — see
547/// `PgPool::clear_statement_caches`, which every migration path calls.
548pub(crate) async fn query_typed_on(
549    client: &deadpool_postgres::Object,
550    sql: &str,
551    params: &[DecodedValue],
552    ext: &ExtensionOids,
553) -> Result<Vec<DecodedValue>> {
554    let stmt = client.prepare_cached(sql).await?;
555    query_with_stmt(client, &stmt, params, ext).await
556}
557
558/// See `query_typed_on` — same caching, `decode_row_named` decoding.
559pub(crate) async fn query_typed_named_on(
560    client: &deadpool_postgres::Object,
561    sql: &str,
562    params: &[DecodedValue],
563    ext: &ExtensionOids,
564) -> Result<Vec<DecodedValue>> {
565    let stmt = client.prepare_cached(sql).await?;
566    query_named_with_stmt(client, &stmt, params, ext).await
567}
568
569/// Sets `pylon.globals`, which a database trigger reads its session globals
570/// from; `local` scopes it to the current transaction.
571async fn set_globals_on(client: &deadpool_postgres::Object, globals: &str, local: bool) -> Result<()> {
572    let sql = if local {
573        "SELECT set_config('pylon.globals', $1, true)"
574    } else {
575        "SELECT set_config('pylon.globals', $1, false)"
576    };
577    execute_typed_on(client, sql, &[DecodedValue::Str(globals.to_string())]).await?;
578    Ok(())
579}
580
581/// See `query_typed_on` — same caching, result discarded.
582pub(crate) async fn execute_typed_on(
583    client: &deadpool_postgres::Object,
584    sql: &str,
585    params: &[DecodedValue],
586) -> Result<u64> {
587    let stmt = client.prepare_cached(sql).await?;
588    execute_with_stmt(client, &stmt, params).await
589}
590
591/// The unpooled counterparts, for the LISTEN/NOTIFY connection — a bare
592/// `tokio_postgres::Client` with no `deadpool` wrapper around it and so no
593/// statement cache to reach. Its query volume is a handful of worker
594/// statements, not the query hot path, so a plain `prepare` is fine here.
595pub(crate) async fn query_typed_on_raw(
596    client: &tokio_postgres::Client,
597    sql: &str,
598    params: &[DecodedValue],
599    ext: &ExtensionOids,
600) -> Result<Vec<DecodedValue>> {
601    let stmt = client.prepare(sql).await?;
602    query_with_stmt(client, &stmt, params, ext).await
603}
604
605pub(crate) async fn query_typed_named_on_raw(
606    client: &tokio_postgres::Client,
607    sql: &str,
608    params: &[DecodedValue],
609    ext: &ExtensionOids,
610) -> Result<Vec<DecodedValue>> {
611    let stmt = client.prepare(sql).await?;
612    query_named_with_stmt(client, &stmt, params, ext).await
613}
614
615pub(crate) async fn execute_typed_on_raw(
616    client: &tokio_postgres::Client,
617    sql: &str,
618    params: &[DecodedValue],
619) -> Result<u64> {
620    let stmt = client.prepare(sql).await?;
621    execute_with_stmt(client, &stmt, params).await
622}
623
624/// `EXPLAIN (FORMAT JSON)` returns exactly one row with one column (named
625/// `QUERY PLAN`, typed `json`) holding the entire plan as JSON text — reused
626/// via `RawBytes` (see its own doc comment) rather than depending on the
627/// `with-serde_json-1` feature, since `json`'s wire format is just its plain
628/// UTF-8 text (unlike `jsonb`, which prefixes a version byte).
629pub(crate) async fn query_explain_on(
630    client: &deadpool_postgres::Object,
631    sql: &str,
632    params: &[DecodedValue],
633) -> Result<String> {
634    let wrapped = format!("EXPLAIN (ANALYZE, FORMAT JSON, VERBOSE) {sql}");
635    let stmt = client.prepare_cached(&wrapped).await?;
636    let bound: Vec<BoundParam<'_>> = params.iter().map(BoundParam).collect();
637    let param_refs: Vec<&(dyn postgres_types::ToSql + Sync)> =
638        bound.iter().map(|p| p as &(dyn postgres_types::ToSql + Sync)).collect();
639    let rows = client.query(&stmt, &param_refs).await?;
640    let row = rows
641        .into_iter()
642        .next()
643        .ok_or_else(|| Error::message("EXPLAIN produced no output row".to_string()))?;
644    let RawBytes(bytes) = row.try_get::<_, RawBytes<'_>>(0)?;
645    Ok(String::from_utf8_lossy(bytes).into_owned())
646}
647
648/// An explicit transaction on a single connection checked out of the pool.
649/// `deadpool-postgres`'s default recycling method (`Fast`) does *not* run
650/// any reset query when a connection is returned to the pool — unlike
651/// pool release, which always issues `ROLLBACK` itself if the
652/// released connection still has an open transaction. That safety net has
653/// to be reproduced here explicitly, or a connection released mid- or
654/// aborted-transaction would silently corrupt the next borrower's session.
655/// Hence `commit` rolls back on its own failure before returning the error,
656/// and both `commit`/`rollback` consume `self` so the connection is only
657/// ever returned to the pool (via `Drop`) once it is guaranteed to be back
658/// in a clean, non-transactional state.
659#[derive(Debug)]
660pub struct PgTransaction {
661    client: deadpool_postgres::Object,
662    types: Arc<ExtensionOids>,
663}
664
665impl PgTransaction {
666    /// The registry the owning pool held when this transaction began — a
667    /// snapshot, as in `PgConnection::types`.
668    pub fn types(&self) -> &ExtensionOids {
669        &self.types
670    }
671
672    /// `query_typed`, with `pylon.globals` set for the rest of this
673    /// transaction — see `PgPool::query_typed_with_globals`.
674    pub async fn query_typed_with_globals(
675        &self,
676        sql: &str,
677        params: &[DecodedValue],
678        ext: &ExtensionOids,
679        globals: &str,
680    ) -> Result<Vec<DecodedValue>> {
681        set_globals_on(&self.client, globals, true).await?;
682        query_typed_on(&self.client, sql, params, ext).await
683    }
684
685    /// `execute_typed`, with `pylon.globals` set — see `query_typed_with_globals`.
686    pub async fn execute_typed_with_globals(&self, sql: &str, params: &[DecodedValue], globals: &str) -> Result<u64> {
687        set_globals_on(&self.client, globals, true).await?;
688        execute_typed_on(&self.client, sql, params).await
689    }
690
691    pub async fn query_typed(
692        &self,
693        sql: &str,
694        params: &[DecodedValue],
695        ext: &ExtensionOids,
696    ) -> Result<Vec<DecodedValue>> {
697        query_typed_on(&self.client, sql, params, ext).await
698    }
699
700    pub async fn execute_typed(&self, sql: &str, params: &[DecodedValue]) -> Result<u64> {
701        execute_typed_on(&self.client, sql, params).await
702    }
703
704    /// Runs `sql` via the simple query protocol — see `PgPool::batch_execute`
705    /// for why migration DDL steps need this instead of `execute_typed`.
706    pub async fn batch_execute(&self, sql: &str) -> Result<()> {
707        self.client.batch_execute(sql).await?;
708        Ok(())
709    }
710
711    /// Establishes a named savepoint inside this transaction. Used by
712    /// migration execution's dev-mode rebase: a step runs inside a
713    /// savepoint so an "already exists" error (from `watch` having applied
714    /// the same DDL earlier) can be rolled back to just that step instead
715    /// of aborting the whole outer transaction.
716    pub async fn savepoint(&self, name: &str) -> Result<()> {
717        self.client
718            .batch_execute(&format!("SAVEPOINT {}", listener::quote_ident(name)))
719            .await?;
720        Ok(())
721    }
722
723    pub async fn release_savepoint(&self, name: &str) -> Result<()> {
724        self.client
725            .batch_execute(&format!("RELEASE SAVEPOINT {}", listener::quote_ident(name)))
726            .await?;
727        Ok(())
728    }
729
730    pub async fn rollback_to_savepoint(&self, name: &str) -> Result<()> {
731        self.client
732            .batch_execute(&format!("ROLLBACK TO SAVEPOINT {}", listener::quote_ident(name)))
733            .await?;
734        Ok(())
735    }
736
737    /// Commits the transaction. On failure (e.g. a serialization failure or
738    /// deadlock detected at COMMIT time), best-effort rolls back first so
739    /// the connection isn't returned to the pool still aborted — the
740    /// original commit error is what's returned either way.
741    pub async fn commit(self) -> Result<()> {
742        match self.client.batch_execute("COMMIT").await {
743            Ok(()) => Ok(()),
744            Err(e) => {
745                let _ = self.client.batch_execute("ROLLBACK").await;
746                Err(e.into())
747            }
748        }
749    }
750
751    pub async fn rollback(self) -> Result<()> {
752        self.client.batch_execute("ROLLBACK").await?;
753        Ok(())
754    }
755}
756
757/// Decodes a row's column 0 (the `result` column pylon-core's SQL always
758/// projects) into a `DecodedValue`, using its actual declared Postgres type.
759/// Column 0 itself can be SQL NULL at the top level (not just a NULL
760/// *field within* a composite, which `decode_record`/`decode_array`
761/// already handle) — `RawBytes` has no `from_sql_null` override, so a
762/// direct `row.try_get::<_, RawBytes>(0)` errors on a null column; going
763/// through `Option<RawBytes>` (which `postgres_types` implements generically
764/// for any `T: FromSql`, yielding `None` for SQL NULL) avoids that.
765/// Parses `dsn` and pins the session settings every Pylon connection needs.
766///
767/// `timezone` decides what a `timestamptz` renders as and which day
768/// `date_part('day', …)` reports; `intervalstyle` decides whether an interval
769/// renders as `01:30:00` or `PT1H30M`. Left to the server's configuration,
770/// the same query would give different answers against two databases, and
771/// `<str>` of a datetime or a duration would stop matching the ISO 8601 the
772/// rest of the stack assumes. Sent as startup-packet options rather than a
773/// `SET` after connecting, so a pooled connection cannot be handed out
774/// before they apply.
775fn session_config(dsn: &str) -> Result<tokio_postgres::Config> {
776    let mut config: tokio_postgres::Config = dsn.parse()?;
777    let pinned = "-c timezone=UTC -c intervalstyle=iso_8601";
778    let options = match config.get_options() {
779        Some(existing) => format!("{existing} {pinned}"),
780        None => pinned.to_string(),
781    };
782    config.options(options);
783    Ok(config)
784}
785
786/// Runs `wire::TYPE_DISCOVERY_SQL` on `client` and builds the registry from
787/// it. Shared by `PgPool::connect` and `PgListener::connect` so a pooled
788/// connection and a dedicated listener classify types identically.
789pub(crate) async fn discover_types(client: &tokio_postgres::Client) -> Result<ExtensionOids> {
790    let rows = client.query(wire::TYPE_DISCOVERY_SQL, &[]).await?;
791    Ok(ExtensionOids::from_discovery_rows(rows.iter().map(|r| {
792        let oid: i64 = r.get(0);
793        let typtype: String = r.get(1);
794        let base_oid: i64 = r.get(2);
795        let typname: String = r.get(3);
796        (oid as u32, typtype, base_oid as u32, typname)
797    })))
798}
799
800fn decode_result_column(row: &tokio_postgres::Row, ext: &ExtensionOids) -> Result<DecodedValue> {
801    let oid = row.columns()[0].type_().oid();
802    match row.try_get::<_, Option<RawBytes>>(0)? {
803        None => Ok(DecodedValue::Null),
804        Some(RawBytes(bytes)) => wire::decode_value(oid, bytes, ext),
805    }
806}
807
808/// Decodes *every* column of `row`, keyed by name, as a `DecodedValue::Object`
809/// — the general named-column decode path, unlike
810/// `decode_result_column` (which only ever decodes column 0, matching
811/// pylon-core's own single-composite-column SQL emission convention). For
812/// hand-written queries with several named columns a caller accesses by
813/// name (`row["col"]`) — `pylon.worker`/`pylon.vector`/`pylon.search`'s
814/// index-outbox queries, not PyQL-compiled SQL.
815fn decode_row_named(row: &tokio_postgres::Row, ext: &ExtensionOids) -> Result<DecodedValue> {
816    let mut fields = Vec::with_capacity(row.columns().len());
817    for (i, col) in row.columns().iter().enumerate() {
818        let oid = col.type_().oid();
819        let value = match row.try_get::<_, Option<RawBytes>>(i)? {
820            None => DecodedValue::Null,
821            Some(RawBytes(bytes)) => wire::decode_value(oid, bytes, ext)?,
822        };
823        fields.push((col.name().to_string(), value));
824    }
825    Ok(DecodedValue::Object(fields))
826}
827
828#[cfg(test)]
829mod tests {
830    use super::*;
831
832    /// Real Postgres required — `PYLON_PGCON_TEST_DSN` must be set (no
833    /// hardcoded fallback; point it at a disposable database, never a real
834    /// one). Not run by default (`cargo test -- --ignored` to opt in) so
835    /// the default test run stays hermetic.
836    fn test_dsn() -> String {
837        std::env::var("PYLON_PGCON_TEST_DSN").expect("PYLON_PGCON_TEST_DSN must be set to run live-Postgres tests")
838    }
839
840    #[tokio::test]
841    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
842    async fn connects_and_round_trips_a_scalar_query() {
843        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
844        let rows = pool.query_raw("SELECT 1 + 1").await.unwrap();
845        assert_eq!(rows.len(), 1);
846        let value: i32 = rows[0].get(0);
847        assert_eq!(value, 2);
848    }
849
850    #[tokio::test]
851    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
852    async fn pool_is_reused_across_multiple_queries() {
853        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
854        for i in 0..5 {
855            let rows = pool.query_raw(&format!("SELECT {i}")).await.unwrap();
856            let value: i32 = rows[0].get(0);
857            assert_eq!(value, i);
858        }
859    }
860
861    #[tokio::test]
862    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
863    async fn invalid_dsn_fails_to_connect() {
864        let result = PgPool::connect("not-a-valid-dsn", 5).await;
865        assert!(result.is_err());
866    }
867
868    #[tokio::test]
869    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
870    async fn connect_fails_eagerly_against_a_nonexistent_database() {
871        // A well-formed DSN pointing at a database that doesn't exist must
872        // fail right here, not lazily on the first query — matching
873        // the eager-connect behavior, which
874        // `Client.ensure_connected()` depends on to map this straight to
875        // `ConnectionFailedError` instead of silently deferring the
876        // failure past `ensure_connected()` returning successfully.
877        // Swap out whatever the real database name is (not a hardcoded
878        // substring — that silently no-ops and leaves this pointed at the
879        // real, existing test database if the DSN's db name ever changes)
880        // for one guaranteed not to exist.
881        let dsn = test_dsn();
882        let (prefix, _db) = dsn.rsplit_once('/').expect("DSN must have a database path segment");
883        let bad_dsn = format!("{prefix}/pgcon_definitely_does_not_exist");
884        let result = PgPool::connect(&bad_dsn, 5).await;
885        assert!(result.is_err());
886    }
887
888    #[tokio::test]
889    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
890    async fn bad_sql_returns_an_error_not_a_panic() {
891        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
892        let result = pool.query_raw("SELECT this is not valid sql").await;
893        assert!(result.is_err());
894    }
895
896    // ── query_composite against real Postgres wire data ─────────────────
897    //
898    // The wire::tests module already checks decode_value against
899    // hand-crafted byte buffers — these tests instead run real SQL through
900    // a real connection, so any mismatch between my assumptions about
901    // Postgres's binary format and what Postgres actually sends shows up
902    // here, not just in self-consistent hand-rolled fixtures.
903
904    #[tokio::test]
905    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
906    async fn decodes_a_bare_scalar_result_column() {
907        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
908        let rows = pool
909            .query_composite("SELECT 42::int8 AS result", &ExtensionOids::default())
910            .await
911            .unwrap();
912        assert_eq!(rows, vec![DecodedValue::I64(42)]);
913    }
914
915    #[tokio::test]
916    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
917    async fn decodes_a_composite_matching_pylon_cores_own_emission_shape() {
918        // Mirrors exactly what `sql/mod.rs::emit_bound_select` emits:
919        // `SELECT (type_disc, col1, col2, ...) AS result`.
920        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
921        let sql = "SELECT ('Person'::text, 'Alice'::text, 30::int8, NULL::text) AS result";
922        let rows = pool.query_composite(sql, &ExtensionOids::default()).await.unwrap();
923        assert_eq!(
924            rows,
925            vec![DecodedValue::Composite(vec![
926                DecodedValue::Str("Person".into()),
927                DecodedValue::Str("Alice".into()),
928                DecodedValue::I64(30),
929                DecodedValue::Null,
930            ])]
931        );
932    }
933
934    #[tokio::test]
935    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
936    async fn decodes_nested_composite_and_array_of_composite_for_real() {
937        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
938        let sql = "SELECT (\
939            'Product'::text, \
940            ROW('Tag'::text, 'sale'::text), \
941            ARRAY[ROW(1::int8), ROW(2::int8)]::record[]\
942        ) AS result";
943        let rows = pool.query_composite(sql, &ExtensionOids::default()).await.unwrap();
944        assert_eq!(
945            rows,
946            vec![DecodedValue::Composite(vec![
947                DecodedValue::Str("Product".into()),
948                DecodedValue::Composite(vec![DecodedValue::Str("Tag".into()), DecodedValue::Str("sale".into())]),
949                DecodedValue::Array(vec![
950                    DecodedValue::Composite(vec![DecodedValue::I64(1)]),
951                    DecodedValue::Composite(vec![DecodedValue::I64(2)]),
952                ]),
953            ])]
954        );
955    }
956
957    #[tokio::test]
958    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
959    async fn decodes_array_of_text_for_real() {
960        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
961        let sql = "SELECT (ARRAY['a', 'b', NULL]::text[]) AS result";
962        let rows = pool.query_composite(sql, &ExtensionOids::default()).await.unwrap();
963        assert_eq!(
964            rows,
965            vec![DecodedValue::Array(vec![
966                DecodedValue::Str("a".into()),
967                DecodedValue::Str("b".into()),
968                DecodedValue::Null,
969            ])]
970        );
971    }
972
973    #[tokio::test]
974    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
975    async fn decodes_numeric_and_jsonb_and_uuid_for_real() {
976        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
977        let sql = "SELECT (\
978            12.50::numeric, \
979            '{\"a\": 1, \"b\": [1,2]}'::jsonb, \
980            '11111111-1111-1111-1111-111111111111'::uuid\
981        ) AS result";
982        let rows = pool.query_composite(sql, &ExtensionOids::default()).await.unwrap();
983        let DecodedValue::Composite(fields) = &rows[0] else {
984            panic!("expected Composite")
985        };
986        assert_eq!(fields[0], DecodedValue::Decimal("12.50".to_string()));
987        assert_eq!(
988            fields[1],
989            DecodedValue::Object(vec![
990                ("a".into(), DecodedValue::I64(1)),
991                (
992                    "b".into(),
993                    DecodedValue::Array(vec![DecodedValue::I64(1), DecodedValue::I64(2)])
994                ),
995            ])
996        );
997        assert_eq!(fields[2], DecodedValue::Uuid([0x11; 16]));
998    }
999
1000    #[tokio::test]
1001    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1002    async fn decodes_bytea_for_real() {
1003        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1004        let rows = pool
1005            .query_composite("SELECT '\\xdeadbeef'::bytea AS result", &ExtensionOids::default())
1006            .await
1007            .unwrap();
1008        assert_eq!(rows, vec![DecodedValue::Bytes(vec![0xde, 0xad, 0xbe, 0xef])]);
1009    }
1010
1011    #[tokio::test]
1012    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1013    async fn decodes_enum_cast_to_text_for_real() {
1014        // pylon-core always ::text-casts enum-typed columns (sql/mod.rs's
1015        // emit_scalar) specifically so the runtime-assigned enum OID never
1016        // needs to be known statically — this is the actually-exercised path.
1017        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1018        let sql = "DO $$ BEGIN CREATE TYPE pgcon_test_enum AS ENUM ('a', 'b'); \
1019                   EXCEPTION WHEN duplicate_object THEN NULL; END $$;";
1020        pool.query_raw(sql).await.ok();
1021        let rows = pool
1022            .query_composite(
1023                "SELECT ('a'::pgcon_test_enum::text) AS result",
1024                &ExtensionOids::default(),
1025            )
1026            .await
1027            .unwrap();
1028        assert_eq!(rows, vec![DecodedValue::Str("a".to_string())]);
1029    }
1030
1031    #[tokio::test]
1032    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1033    async fn decodes_a_real_pgvector_value_through_connect_time_discovery() {
1034        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1035        if pool.query_raw("CREATE EXTENSION IF NOT EXISTS vector").await.is_err() {
1036            eprintln!("skipping: pgvector not installable on this server");
1037            return;
1038        }
1039        // The extension may have been created after this pool connected.
1040        pool.refresh_types().await.unwrap();
1041        assert!(
1042            pool.types().vector.is_some(),
1043            "discovery should have found the vector OID"
1044        );
1045
1046        // Nested in a record, exactly as pylon-core emits every result.
1047        let rows = pool
1048            .query_composite("SELECT ('doc', '[1.5,2.5]'::vector) AS result", &pool.types())
1049            .await
1050            .unwrap();
1051
1052        let DecodedValue::Composite(fields) = &rows[0] else {
1053            panic!("expected Composite, got {:?}", rows[0])
1054        };
1055        assert_eq!(fields[0], DecodedValue::Str("doc".to_string()));
1056        assert_eq!(
1057            fields[1],
1058            DecodedValue::Array(vec![DecodedValue::F64(1.5), DecodedValue::F64(2.5)]),
1059            "a vector must decode to its floats, not to mojibake or an error"
1060        );
1061    }
1062
1063    #[tokio::test]
1064    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1065    async fn decodes_a_real_enum_and_domain_through_connect_time_discovery() {
1066        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1067        pool.query_raw(
1068            "DO $$ BEGIN CREATE TYPE pgcon_disc_enum AS ENUM ('x', 'y'); \
1069             EXCEPTION WHEN duplicate_object THEN NULL; END $$;",
1070        )
1071        .await
1072        .ok();
1073        pool.query_raw(
1074            "DO $$ BEGIN CREATE DOMAIN pgcon_disc_domain AS int8 CHECK (VALUE > 0); \
1075             EXCEPTION WHEN duplicate_object THEN NULL; END $$;",
1076        )
1077        .await
1078        .ok();
1079
1080        pool.refresh_types().await.unwrap();
1081
1082        // Neither is ::text-cast here, so both arrive with their real
1083        // database-assigned OID and can only decode via the registry.
1084        let rows = pool
1085            .query_composite(
1086                "SELECT ('x'::pgcon_disc_enum, 42::pgcon_disc_domain) AS result",
1087                &pool.types(),
1088            )
1089            .await
1090            .unwrap();
1091
1092        let DecodedValue::Composite(fields) = &rows[0] else {
1093            panic!("expected Composite, got {:?}", rows[0])
1094        };
1095        assert_eq!(fields[0], DecodedValue::Str("x".to_string()));
1096        assert_eq!(
1097            fields[1],
1098            DecodedValue::I64(42),
1099            "a domain must decode as its base type, not as text"
1100        );
1101    }
1102
1103    // ── a registry that goes stale underneath a running pool ────────────
1104
1105    /// Read as `int8`, so this needs no registry of its own.
1106    async fn live_type_oid(pool: &PgPool, typname: &str) -> u32 {
1107        let rows = pool
1108            .query_typed(
1109                "SELECT (oid::int8) AS result FROM pg_type WHERE typname = $1",
1110                &[DecodedValue::Str(typname.to_string())],
1111                &ExtensionOids::default(),
1112            )
1113            .await
1114            .unwrap();
1115        match rows.first() {
1116            Some(DecodedValue::I64(oid)) => *oid as u32,
1117            other => panic!("expected one int8 oid for {typname}, got {other:?}"),
1118        }
1119    }
1120
1121    /// Recreating a type reproduces what a restore does to every OID at
1122    /// once: Postgres never reissues the one it just freed.
1123    #[tokio::test]
1124    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1125    async fn a_type_whose_oid_moved_under_a_live_pool_still_decodes() {
1126        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1127        pool.batch_execute(
1128            "DROP TYPE IF EXISTS pgcon_moved_enum; \
1129             CREATE TYPE pgcon_moved_enum AS ENUM ('a', 'b')",
1130        )
1131        .await
1132        .unwrap();
1133        // Stand in for a pool that connected while this type existed.
1134        pool.refresh_types().await.unwrap();
1135        let stale = pool.types();
1136        let before = live_type_oid(&pool, "pgcon_moved_enum").await;
1137        assert!(
1138            stale.enums.contains(&before),
1139            "the pool must start out knowing the pre-move OID, or this test proves nothing"
1140        );
1141
1142        pool.batch_execute(
1143            "DROP TYPE pgcon_moved_enum; \
1144             CREATE TYPE pgcon_moved_enum AS ENUM ('a', 'b')",
1145        )
1146        .await
1147        .unwrap();
1148        let after = live_type_oid(&pool, "pgcon_moved_enum").await;
1149        assert_ne!(before, after, "recreating a type must assign it a fresh OID");
1150
1151        // Before: the registry the pool is holding can't decode the new OID.
1152        let err = decode_value(after, b"a", &stale).unwrap_err();
1153        assert!(
1154            matches!(err, Error::UnknownTypeOid { oid } if oid == after),
1155            "expected UnknownTypeOid for the moved OID, got {err}"
1156        );
1157
1158        // After: `pool.types()` is still that stale registry when it's
1159        // passed, and the query succeeds anyway.
1160        let rows = pool
1161            .query_composite("SELECT ('a'::pgcon_moved_enum) AS result", &pool.types())
1162            .await
1163            .unwrap();
1164        assert_eq!(rows, vec![DecodedValue::Str("a".to_string())]);
1165        assert!(
1166            pool.types().enums.contains(&after),
1167            "the heal must leave the new OID in the pool's registry, not just decode one query"
1168        );
1169    }
1170
1171    /// Consumers hold clones, so a refresh that reached only the handle it
1172    /// was called on would leave the clone serving queries stale.
1173    #[tokio::test]
1174    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1175    async fn a_refresh_through_one_handle_is_visible_through_every_clone() {
1176        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1177        pool.batch_execute("DROP TYPE IF EXISTS pgcon_clone_enum")
1178            .await
1179            .unwrap();
1180        pool.refresh_types().await.unwrap();
1181
1182        let clone = pool.clone();
1183        pool.batch_execute("CREATE TYPE pgcon_clone_enum AS ENUM ('a')")
1184            .await
1185            .unwrap();
1186        let oid = live_type_oid(&pool, "pgcon_clone_enum").await;
1187        assert!(!clone.types().enums.contains(&oid), "the clone must start out stale");
1188
1189        clone.refresh_types().await.unwrap();
1190
1191        assert!(pool.types().enums.contains(&oid), "the original must see the refresh");
1192        let tx = pool.begin_default().await.unwrap();
1193        assert!(
1194            tx.types().enums.contains(&oid),
1195            "a transaction opened after the refresh must inherit the refreshed registry"
1196        );
1197        tx.rollback().await.unwrap();
1198        pool.batch_execute("DROP TYPE pgcon_clone_enum").await.unwrap();
1199    }
1200
1201    // ── query_typed: bound-parameter round trips against real Postgres ──
1202    //
1203    // Each test binds a DecodedValue as $1, has Postgres echo it straight
1204    // back out (so both encode_value AND decode_value are exercised in one
1205    // pass — a mismatch in either direction fails the assertion), matching
1206    // exactly how a real PyQL query binds a param and gets a result back.
1207
1208    async fn round_trip(pool: &PgPool, pg_type: &str, param: DecodedValue) -> DecodedValue {
1209        let sql = format!("SELECT ($1::{pg_type}) AS result");
1210        let rows = pool
1211            .query_typed(&sql, &[param], &ExtensionOids::default())
1212            .await
1213            .unwrap();
1214        rows.into_iter().next().unwrap()
1215    }
1216
1217    // ── query_explain against real Postgres ─────────────────────────────
1218
1219    #[tokio::test]
1220    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1221    async fn query_explain_returns_parseable_json_with_a_plan_node() {
1222        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1223        let raw = pool.query_explain("SELECT 1 + 1", &[]).await.unwrap();
1224        let parsed: serde_json::Value = serde_json::from_str(&raw).unwrap();
1225        assert!(parsed[0]["Plan"]["Node Type"].is_string());
1226        // ANALYZE was requested, so the plan must carry real execution stats.
1227        assert!(parsed[0]["Plan"]["Actual Total Time"].is_number());
1228    }
1229
1230    #[tokio::test]
1231    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1232    async fn query_explain_binds_params_the_same_way_query_typed_does() {
1233        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1234        let raw = pool
1235            .query_explain("SELECT $1::int8 + 1", &[DecodedValue::I64(41)])
1236            .await
1237            .unwrap();
1238        let parsed: serde_json::Value = serde_json::from_str(&raw).unwrap();
1239        assert!(parsed[0]["Plan"]["Node Type"].is_string());
1240    }
1241
1242    #[tokio::test]
1243    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1244    async fn round_trips_bool_param() {
1245        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1246        assert_eq!(
1247            round_trip(&pool, "bool", DecodedValue::Bool(true)).await,
1248            DecodedValue::Bool(true)
1249        );
1250        assert_eq!(
1251            round_trip(&pool, "bool", DecodedValue::Bool(false)).await,
1252            DecodedValue::Bool(false)
1253        );
1254    }
1255
1256    #[tokio::test]
1257    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1258    async fn round_trips_integer_params_at_every_width() {
1259        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1260        assert_eq!(
1261            round_trip(&pool, "int2", DecodedValue::I64(30)).await,
1262            DecodedValue::I64(30)
1263        );
1264        assert_eq!(
1265            round_trip(&pool, "int4", DecodedValue::I64(70_000)).await,
1266            DecodedValue::I64(70_000)
1267        );
1268        assert_eq!(
1269            round_trip(&pool, "int8", DecodedValue::I64(9_223_372_036_854_775_807)).await,
1270            DecodedValue::I64(9_223_372_036_854_775_807)
1271        );
1272        assert_eq!(
1273            round_trip(&pool, "int8", DecodedValue::I64(-1)).await,
1274            DecodedValue::I64(-1)
1275        );
1276    }
1277
1278    #[tokio::test]
1279    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1280    async fn round_trips_float_params() {
1281        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1282        assert_eq!(
1283            round_trip(&pool, "float4", DecodedValue::F64(1.5)).await,
1284            DecodedValue::F64(1.5)
1285        );
1286        assert_eq!(
1287            round_trip(&pool, "float8", DecodedValue::F64(2.25)).await,
1288            DecodedValue::F64(2.25)
1289        );
1290    }
1291
1292    #[tokio::test]
1293    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1294    async fn round_trips_text_param() {
1295        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1296        assert_eq!(
1297            round_trip(&pool, "text", DecodedValue::Str("héllo 🎉".to_string())).await,
1298            DecodedValue::Str("héllo 🎉".to_string())
1299        );
1300    }
1301
1302    #[tokio::test]
1303    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1304    async fn round_trips_bytea_param() {
1305        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1306        assert_eq!(
1307            round_trip(&pool, "bytea", DecodedValue::Bytes(vec![1, 2, 3, 255])).await,
1308            DecodedValue::Bytes(vec![1, 2, 3, 255])
1309        );
1310    }
1311
1312    #[tokio::test]
1313    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1314    async fn round_trips_uuid_param() {
1315        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1316        let bytes = [0x11u8; 16];
1317        assert_eq!(
1318            round_trip(&pool, "uuid", DecodedValue::Uuid(bytes)).await,
1319            DecodedValue::Uuid(bytes)
1320        );
1321    }
1322
1323    #[tokio::test]
1324    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1325    async fn binds_a_plain_string_as_a_uuid_param() {
1326        // A JSON API request body (see `pylon-server`'s /api/<connection>/query
1327        // handler) necessarily carries a UUID query parameter as plain text —
1328        // there's no JSON "uuid" type — so it arrives as `DecodedValue::Str`,
1329        // not `::Uuid`. Regression test for a real bug: binding that Str
1330        // directly against a `uuid`-typed parameter used to send raw UTF-8
1331        // text bytes for a binary-format column, which Postgres rejected
1332        // with "incorrect binary data format in bind parameter 1".
1333        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1334        let as_string = DecodedValue::Str("11111111-1111-1111-1111-111111111111".to_string());
1335        assert_eq!(
1336            round_trip(&pool, "uuid", as_string).await,
1337            DecodedValue::Uuid([0x11; 16])
1338        );
1339    }
1340
1341    #[tokio::test]
1342    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1343    async fn binds_a_plain_string_as_a_jsonb_param() {
1344        // Regression test for a real bug: `migration apply`'s db_state
1345        // snapshot update binds already-serialized JSON text (a plain Rust
1346        // `String`, not a `DecodedValue::Object`) as `$1::jsonb` — the same
1347        // class of bug as the UUID case above (a Str value bound against a
1348        // non-text binary-format target type).
1349        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1350        let as_string = DecodedValue::Str(r#"{"a":1,"b":[1,2]}"#.to_string());
1351        assert_eq!(
1352            round_trip(&pool, "jsonb", as_string).await,
1353            DecodedValue::Object(vec![
1354                ("a".into(), DecodedValue::I64(1)),
1355                (
1356                    "b".into(),
1357                    DecodedValue::Array(vec![DecodedValue::I64(1), DecodedValue::I64(2)])
1358                ),
1359            ])
1360        );
1361    }
1362
1363    #[tokio::test]
1364    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1365    async fn round_trips_numeric_param() {
1366        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1367        assert_eq!(
1368            round_trip(&pool, "numeric", DecodedValue::Decimal("12.50".to_string())).await,
1369            DecodedValue::Decimal("12.50".to_string())
1370        );
1371        assert_eq!(
1372            round_trip(&pool, "numeric", DecodedValue::Decimal("-9999.001".to_string())).await,
1373            DecodedValue::Decimal("-9999.001".to_string())
1374        );
1375
1376        // Values a fixed 96-bit/28-place carrier cannot hold. Each of these
1377        // used to come back scaled by a power of ten, or refuse to serialize,
1378        // with `numeric` itself perfectly able to store them. The first is
1379        // what Python's `Decimal(0.0000004)` stringifies to.
1380        for text in [
1381            "0.00000039999999999999998189924473035450347424557548947632312774658203125",
1382            "0.00000012222222222222222222222222222222222222222",
1383            "12222222222222222222222222222222000000000",
1384            "100000000000000000000000000000",
1385            "123456789012345678901234567890.123456789012345678901234567890",
1386        ] {
1387            assert_eq!(
1388                round_trip(&pool, "numeric", DecodedValue::Decimal(text.to_string())).await,
1389                DecodedValue::Decimal(text.to_string()),
1390                "{text}"
1391            );
1392        }
1393
1394        // A JSON request body has no decimal type, so a numeric parameter can
1395        // arrive as a number or as text — both reach the same encoder.
1396        assert_eq!(
1397            round_trip(&pool, "numeric", DecodedValue::Str("0.0000004".to_string())).await,
1398            DecodedValue::Decimal("0.0000004".to_string())
1399        );
1400        assert_eq!(
1401            round_trip(&pool, "numeric", DecodedValue::I64(-12345)).await,
1402            DecodedValue::Decimal("-12345".to_string())
1403        );
1404        assert_eq!(
1405            round_trip(&pool, "numeric", DecodedValue::F64(0.1)).await,
1406            DecodedValue::Decimal("0.1".to_string())
1407        );
1408    }
1409
1410    #[tokio::test]
1411    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1412    async fn round_trips_null_param() {
1413        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1414        assert_eq!(round_trip(&pool, "int8", DecodedValue::Null).await, DecodedValue::Null);
1415        assert_eq!(round_trip(&pool, "text", DecodedValue::Null).await, DecodedValue::Null);
1416    }
1417
1418    #[tokio::test]
1419    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1420    async fn round_trips_array_param() {
1421        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1422        let param = DecodedValue::Array(vec![
1423            DecodedValue::Str("a".into()),
1424            DecodedValue::Str("b".into()),
1425            DecodedValue::Null,
1426        ]);
1427        assert_eq!(round_trip(&pool, "text[]", param.clone()).await, param);
1428    }
1429
1430    #[tokio::test]
1431    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1432    async fn round_trips_int_array_param() {
1433        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1434        let param = DecodedValue::Array(vec![DecodedValue::I64(1), DecodedValue::I64(2), DecodedValue::I64(3)]);
1435        assert_eq!(round_trip(&pool, "int8[]", param.clone()).await, param);
1436    }
1437
1438    #[tokio::test]
1439    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1440    async fn round_trips_jsonb_object_param() {
1441        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1442        let param = DecodedValue::Object(vec![
1443            ("a".into(), DecodedValue::I64(1)),
1444            ("b".into(), DecodedValue::Str("two".into())),
1445            (
1446                "c".into(),
1447                DecodedValue::Array(vec![DecodedValue::I64(1), DecodedValue::I64(2)]),
1448            ),
1449        ]);
1450        assert_eq!(round_trip(&pool, "jsonb", param.clone()).await, param);
1451    }
1452
1453    #[tokio::test]
1454    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1455    async fn query_typed_matches_pylon_cores_own_param_binding_convention() {
1456        // $1, $2, ... positional, matching multiple params in one query —
1457        // the same shape a real PyQL query with several kwargs produces.
1458        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1459        let sql = "SELECT ($1::text, $2::int8, $3::bool) AS result";
1460        let params = vec![
1461            DecodedValue::Str("Alice".into()),
1462            DecodedValue::I64(30),
1463            DecodedValue::Bool(true),
1464        ];
1465        let rows = pool.query_typed(sql, &params, &ExtensionOids::default()).await.unwrap();
1466        assert_eq!(
1467            rows,
1468            vec![DecodedValue::Composite(vec![
1469                DecodedValue::Str("Alice".into()),
1470                DecodedValue::I64(30),
1471                DecodedValue::Bool(true),
1472            ])]
1473        );
1474    }
1475
1476    #[tokio::test]
1477    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1478    async fn wrong_param_count_returns_an_error_not_a_panic() {
1479        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1480        let result = pool
1481            .query_typed(
1482                "SELECT $1::int8, $2::int8",
1483                &[DecodedValue::I64(1)],
1484                &ExtensionOids::default(),
1485            )
1486            .await;
1487        assert!(result.is_err());
1488    }
1489
1490    #[tokio::test]
1491    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1492    async fn execute_typed_runs_a_mutation_and_reports_affected_rows() {
1493        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1494        pool.query_raw("CREATE TEMP TABLE IF NOT EXISTS pgcon_execute_test (id int8, name text)")
1495            .await
1496            .unwrap();
1497
1498        let inserted = pool
1499            .execute_typed(
1500                "INSERT INTO pgcon_execute_test (id, name) VALUES ($1::int8, $2::text)",
1501                &[DecodedValue::I64(1), DecodedValue::Str("alice".into())],
1502            )
1503            .await
1504            .unwrap();
1505        assert_eq!(inserted, 1);
1506
1507        let updated = pool
1508            .execute_typed(
1509                "UPDATE pgcon_execute_test SET name = $1::text WHERE id = $2::int8",
1510                &[DecodedValue::Str("bob".into()), DecodedValue::I64(1)],
1511            )
1512            .await
1513            .unwrap();
1514        assert_eq!(updated, 1);
1515
1516        let rows = pool
1517            .query_composite(
1518                "SELECT (name) AS result FROM pgcon_execute_test",
1519                &ExtensionOids::default(),
1520            )
1521            .await
1522            .unwrap();
1523        assert_eq!(rows, vec![DecodedValue::Str("bob".to_string())]);
1524    }
1525
1526    // ── Error::sqlstate() against real Postgres constraint violations ──
1527    //
1528    // Error mapping to Pylon's Python exception hierarchy (a later phase)
1529    // classifies on these codes (`"23505"` unique violation, etc.)
1530    // already do today — verified against a real server response, not
1531    // assumed from the SQLSTATE spec alone.
1532
1533    #[tokio::test]
1534    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1535    async fn unique_violation_reports_23505() {
1536        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1537        pool.query_raw("CREATE TEMP TABLE pgcon_unique_test (id int8 PRIMARY KEY)")
1538            .await
1539            .unwrap();
1540        pool.execute_typed(
1541            "INSERT INTO pgcon_unique_test (id) VALUES ($1::int8)",
1542            &[DecodedValue::I64(1)],
1543        )
1544        .await
1545        .unwrap();
1546
1547        let err = pool
1548            .execute_typed(
1549                "INSERT INTO pgcon_unique_test (id) VALUES ($1::int8)",
1550                &[DecodedValue::I64(1)],
1551            )
1552            .await
1553            .unwrap_err();
1554        assert_eq!(err.sqlstate(), Some(&tokio_postgres::error::SqlState::UNIQUE_VIOLATION));
1555        assert_eq!(err.sqlstate().unwrap().code(), "23505");
1556    }
1557
1558    #[tokio::test]
1559    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1560    async fn foreign_key_violation_reports_23503() {
1561        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1562        pool.query_raw("CREATE TEMP TABLE pgcon_fk_parent (id int8 PRIMARY KEY)")
1563            .await
1564            .unwrap();
1565        pool.query_raw("CREATE TEMP TABLE pgcon_fk_child (parent_id int8 REFERENCES pgcon_fk_parent(id))")
1566            .await
1567            .unwrap();
1568
1569        let err = pool
1570            .execute_typed(
1571                "INSERT INTO pgcon_fk_child (parent_id) VALUES ($1::int8)",
1572                &[DecodedValue::I64(999)],
1573            )
1574            .await
1575            .unwrap_err();
1576        assert_eq!(
1577            err.sqlstate(),
1578            Some(&tokio_postgres::error::SqlState::FOREIGN_KEY_VIOLATION)
1579        );
1580    }
1581
1582    #[tokio::test]
1583    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1584    async fn check_violation_reports_23514() {
1585        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1586        pool.query_raw("CREATE TEMP TABLE pgcon_check_test (age int8 CHECK (age >= 0))")
1587            .await
1588            .unwrap();
1589
1590        let err = pool
1591            .execute_typed(
1592                "INSERT INTO pgcon_check_test (age) VALUES ($1::int8)",
1593                &[DecodedValue::I64(-1)],
1594            )
1595            .await
1596            .unwrap_err();
1597        assert_eq!(err.sqlstate(), Some(&tokio_postgres::error::SqlState::CHECK_VIOLATION));
1598        // Temp tables live in a session-specific `pg_temp_N` schema, so only
1599        // the table name (not the exact schema) is asserted here.
1600        assert_eq!(err.violated_table().map(|(_, table)| table), Some("pgcon_check_test"));
1601        assert_eq!(err.violated_scalar(), None);
1602    }
1603
1604    #[tokio::test]
1605    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1606    async fn domain_check_violation_reports_the_domain_name_not_the_constraint_name() {
1607        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1608        pool.query_raw(
1609            "DO $$ BEGIN CREATE DOMAIN pgcon_rating AS int8 CHECK (VALUE BETWEEN 1 AND 5); \
1610             EXCEPTION WHEN duplicate_object THEN NULL; END $$",
1611        )
1612        .await
1613        .unwrap();
1614        pool.query_raw("CREATE TEMP TABLE pgcon_domain_check_test (rating pgcon_rating)")
1615            .await
1616            .unwrap();
1617
1618        let err = pool
1619            .execute_typed(
1620                "INSERT INTO pgcon_domain_check_test (rating) VALUES ($1::pgcon_rating)",
1621                &[DecodedValue::I64(99)],
1622            )
1623            .await
1624            .unwrap_err();
1625        assert_eq!(err.sqlstate(), Some(&tokio_postgres::error::SqlState::CHECK_VIOLATION));
1626        assert_eq!(err.violated_scalar(), Some(("public", "pgcon_rating")));
1627    }
1628
1629    #[tokio::test]
1630    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1631    async fn syntax_error_has_no_sqlstate_matching_constraint_codes() {
1632        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1633        let err = pool.query_raw("SELECT this is not valid sql").await.unwrap_err();
1634        assert_ne!(err.sqlstate(), Some(&tokio_postgres::error::SqlState::UNIQUE_VIOLATION));
1635    }
1636
1637    #[tokio::test]
1638    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1639    async fn connection_pool_error_has_no_sqlstate() {
1640        // A bad DSN never reaches Postgres at all — no SQLSTATE to report,
1641        // unlike a real server-side rejection.
1642        let result = PgPool::connect("not-a-valid-dsn", 5).await;
1643        let err = result.unwrap_err();
1644        assert_eq!(err.sqlstate(), None);
1645    }
1646
1647    // ── PgTransaction: begin/commit/rollback on real Postgres ───────────
1648
1649    #[tokio::test]
1650    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1651    async fn committed_transaction_persists_its_writes() {
1652        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1653        pool.query_raw("CREATE TEMP TABLE pgcon_tx_commit_test (id int8 PRIMARY KEY)")
1654            .await
1655            .unwrap();
1656
1657        let tx = pool.begin("serializable").await.unwrap();
1658        tx.execute_typed(
1659            "INSERT INTO pgcon_tx_commit_test (id) VALUES ($1::int8)",
1660            &[DecodedValue::I64(1)],
1661        )
1662        .await
1663        .unwrap();
1664        tx.commit().await.unwrap();
1665
1666        let rows = pool
1667            .query_composite(
1668                "SELECT (id) AS result FROM pgcon_tx_commit_test",
1669                &ExtensionOids::default(),
1670            )
1671            .await
1672            .unwrap();
1673        assert_eq!(rows, vec![DecodedValue::I64(1)]);
1674    }
1675
1676    #[tokio::test]
1677    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1678    async fn rolled_back_transaction_discards_its_writes() {
1679        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1680        pool.query_raw("CREATE TEMP TABLE pgcon_tx_rollback_test (id int8 PRIMARY KEY)")
1681            .await
1682            .unwrap();
1683
1684        let tx = pool.begin("serializable").await.unwrap();
1685        tx.execute_typed(
1686            "INSERT INTO pgcon_tx_rollback_test (id) VALUES ($1::int8)",
1687            &[DecodedValue::I64(1)],
1688        )
1689        .await
1690        .unwrap();
1691        tx.rollback().await.unwrap();
1692
1693        let rows = pool
1694            .query_composite(
1695                "SELECT (id) AS result FROM pgcon_tx_rollback_test",
1696                &ExtensionOids::default(),
1697            )
1698            .await
1699            .unwrap();
1700        assert!(rows.is_empty());
1701    }
1702
1703    #[tokio::test]
1704    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1705    async fn begin_actually_sets_the_requested_isolation_level() {
1706        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1707        for (level, expected) in [
1708            ("read_committed", "read committed"),
1709            ("repeatable_read", "repeatable read"),
1710            ("serializable", "serializable"),
1711        ] {
1712            let tx = pool.begin(level).await.unwrap();
1713            let rows = tx
1714                .query_typed(
1715                    "SELECT (current_setting('transaction_isolation')) AS result",
1716                    &[],
1717                    &ExtensionOids::default(),
1718                )
1719                .await
1720                .unwrap();
1721            assert_eq!(rows, vec![DecodedValue::Str(expected.to_string())]);
1722            tx.rollback().await.unwrap();
1723        }
1724    }
1725
1726    #[tokio::test]
1727    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1728    async fn begin_rejects_an_unknown_isolation_level() {
1729        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1730        let result = pool.begin("not_a_real_level").await;
1731        assert!(result.is_err());
1732    }
1733
1734    #[tokio::test]
1735    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1736    async fn a_pooled_connection_is_reusable_after_commit_and_after_rollback() {
1737        // Guards the exact hazard begin()/PgTransaction's doc comment
1738        // describes: deadpool's default Fast recycling does nothing to a
1739        // connection returned mid-transaction, so if commit/rollback ever
1740        // failed to leave the session clean, this small pool (max_size 1)
1741        // would hang forever on the second `begin()` waiting for a
1742        // connection that never becomes usable again.
1743        let pool = PgPool::connect(&test_dsn(), 1).await.unwrap();
1744
1745        let tx = pool.begin("serializable").await.unwrap();
1746        tx.commit().await.unwrap();
1747
1748        let tx = pool.begin("serializable").await.unwrap();
1749        tx.rollback().await.unwrap();
1750
1751        let rows = pool.query_raw("SELECT 1").await.unwrap();
1752        assert_eq!(rows.len(), 1);
1753    }
1754
1755    #[tokio::test]
1756    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1757    async fn failed_commit_leaves_the_connection_reusable() {
1758        // Forces a real 40001 serialization failure at COMMIT time (the
1759        // same interleaving as `serializable_transactions_conflict_with_40001`
1760        // below), on a pool sized to exactly the two connections both
1761        // transactions occupy, then drains the pool with fresh queries to
1762        // prove every connection — including the one that failed COMMIT —
1763        // comes back healthy. Without `commit`'s best-effort
1764        // ROLLBACK-on-failure, the failed connection would still be
1765        // aborted server-side and the next query to land on it would
1766        // immediately fail with "current transaction is aborted".
1767        let pool = PgPool::connect(&test_dsn(), 2).await.unwrap();
1768        pool.query_raw("DROP TABLE IF EXISTS pgcon_tx_failed_commit_test")
1769            .await
1770            .unwrap();
1771        pool.query_raw("CREATE TABLE pgcon_tx_failed_commit_test (class int8, value int8)")
1772            .await
1773            .unwrap();
1774        pool.execute_typed(
1775            "INSERT INTO pgcon_tx_failed_commit_test (class, value) VALUES ($1::int8, $2::int8), ($3::int8, $4::int8)",
1776            &[
1777                DecodedValue::I64(1),
1778                DecodedValue::I64(10),
1779                DecodedValue::I64(2),
1780                DecodedValue::I64(20),
1781            ],
1782        )
1783        .await
1784        .unwrap();
1785
1786        let tx1 = pool.begin("serializable").await.unwrap();
1787        let tx2 = pool.begin("serializable").await.unwrap();
1788
1789        tx1.query_typed(
1790            "SELECT (sum(value)) AS result FROM pgcon_tx_failed_commit_test WHERE class = 1::int8",
1791            &[],
1792            &ExtensionOids::default(),
1793        )
1794        .await
1795        .unwrap();
1796        tx2.query_typed(
1797            "SELECT (sum(value)) AS result FROM pgcon_tx_failed_commit_test WHERE class = 2::int8",
1798            &[],
1799            &ExtensionOids::default(),
1800        )
1801        .await
1802        .unwrap();
1803        tx1.execute_typed(
1804            "INSERT INTO pgcon_tx_failed_commit_test (class, value) VALUES (2::int8, $1::int8)",
1805            &[DecodedValue::I64(10)],
1806        )
1807        .await
1808        .unwrap();
1809        tx2.execute_typed(
1810            "INSERT INTO pgcon_tx_failed_commit_test (class, value) VALUES (1::int8, $1::int8)",
1811            &[DecodedValue::I64(20)],
1812        )
1813        .await
1814        .unwrap();
1815
1816        tx1.commit().await.unwrap();
1817        let commit_result = tx2.commit().await;
1818        assert!(commit_result.is_err());
1819
1820        // Both pooled connections are back now (tx1 released on success,
1821        // tx2 released on Drop after the failed commit) — round-trip each.
1822        for _ in 0..2 {
1823            let rows = pool.query_raw("SELECT 1").await.unwrap();
1824            assert_eq!(rows.len(), 1);
1825        }
1826    }
1827
1828    #[tokio::test]
1829    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1830    async fn serializable_transactions_conflict_with_40001() {
1831        // The canonical serialization-anomaly example from the Postgres
1832        // docs (13.2.3): two SERIALIZABLE transactions each read one
1833        // class's total, then insert a row into the *other* class based on
1834        // what they read. Run concurrently with each SELECT completing
1835        // before either INSERT, this is guaranteed to leave one commit
1836        // rejected with 40001 — this is the exact SQLSTATE `pgcon_err`
1837        // (pylon-py/src/pgcon.rs) maps to `TransactionSerializationError`.
1838        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1839        pool.query_raw("DROP TABLE IF EXISTS pgcon_serialization_test")
1840            .await
1841            .unwrap();
1842        pool.query_raw("CREATE TABLE pgcon_serialization_test (class int8, value int8)")
1843            .await
1844            .unwrap();
1845        pool.execute_typed(
1846            "INSERT INTO pgcon_serialization_test (class, value) VALUES ($1::int8, $2::int8), ($3::int8, $4::int8)",
1847            &[
1848                DecodedValue::I64(1),
1849                DecodedValue::I64(10),
1850                DecodedValue::I64(2),
1851                DecodedValue::I64(20),
1852            ],
1853        )
1854        .await
1855        .unwrap();
1856
1857        let tx1 = pool.begin("serializable").await.unwrap();
1858        let tx2 = pool.begin("serializable").await.unwrap();
1859
1860        tx1.query_typed(
1861            "SELECT (sum(value)) AS result FROM pgcon_serialization_test WHERE class = 1::int8",
1862            &[],
1863            &ExtensionOids::default(),
1864        )
1865        .await
1866        .unwrap();
1867        tx2.query_typed(
1868            "SELECT (sum(value)) AS result FROM pgcon_serialization_test WHERE class = 2::int8",
1869            &[],
1870            &ExtensionOids::default(),
1871        )
1872        .await
1873        .unwrap();
1874
1875        tx1.execute_typed(
1876            "INSERT INTO pgcon_serialization_test (class, value) VALUES (2::int8, $1::int8)",
1877            &[DecodedValue::I64(10)],
1878        )
1879        .await
1880        .unwrap();
1881        tx2.execute_typed(
1882            "INSERT INTO pgcon_serialization_test (class, value) VALUES (1::int8, $1::int8)",
1883            &[DecodedValue::I64(20)],
1884        )
1885        .await
1886        .unwrap();
1887
1888        tx1.commit().await.unwrap();
1889        let err = tx2.commit().await.unwrap_err();
1890        assert_eq!(
1891            err.sqlstate(),
1892            Some(&tokio_postgres::error::SqlState::T_R_SERIALIZATION_FAILURE)
1893        );
1894    }
1895
1896    #[tokio::test]
1897    #[ignore = "requires a live Postgres via PYLON_PGCON_TEST_DSN"]
1898    async fn concurrent_transactions_deadlock_with_40p01() {
1899        // Classic reproducible deadlock: two transactions lock two rows in
1900        // opposite order. tx1 locks row 1 then blocks on row 2; tx2 locks
1901        // row 2 then blocks on row 1 — Postgres's deadlock detector aborts
1902        // one of them with 40P01, `pgcon_err`'s other mapped SQLSTATE
1903        // (-> `TransactionDeadlockError`).
1904        let pool = PgPool::connect(&test_dsn(), 5).await.unwrap();
1905        pool.query_raw("DROP TABLE IF EXISTS pgcon_deadlock_test")
1906            .await
1907            .unwrap();
1908        pool.query_raw("CREATE TABLE pgcon_deadlock_test (id int8 PRIMARY KEY, value int8)")
1909            .await
1910            .unwrap();
1911        pool.execute_typed(
1912            "INSERT INTO pgcon_deadlock_test (id, value) VALUES ($1::int8, $2::int8), ($3::int8, $4::int8)",
1913            &[
1914                DecodedValue::I64(1),
1915                DecodedValue::I64(0),
1916                DecodedValue::I64(2),
1917                DecodedValue::I64(0),
1918            ],
1919        )
1920        .await
1921        .unwrap();
1922
1923        let tx1 = pool.begin("read_committed").await.unwrap();
1924        let tx2 = pool.begin("read_committed").await.unwrap();
1925
1926        tx1.execute_typed("UPDATE pgcon_deadlock_test SET value = 1::int8 WHERE id = 1::int8", &[])
1927            .await
1928            .unwrap();
1929        tx2.execute_typed("UPDATE pgcon_deadlock_test SET value = 2::int8 WHERE id = 2::int8", &[])
1930            .await
1931            .unwrap();
1932
1933        // Now each blocks on the row the other is holding — issue both
1934        // concurrently and let Postgres's deadlock detector break the tie.
1935        let (r1, r2) = tokio::join!(
1936            tx1.execute_typed("UPDATE pgcon_deadlock_test SET value = 3::int8 WHERE id = 2::int8", &[]),
1937            tx2.execute_typed("UPDATE pgcon_deadlock_test SET value = 4::int8 WHERE id = 1::int8", &[]),
1938        );
1939
1940        let results = [r1, r2];
1941        let deadlock_errors: Vec<_> = results
1942            .iter()
1943            .filter(|r| matches!(r, Err(e) if e.sqlstate() == Some(&tokio_postgres::error::SqlState::T_R_DEADLOCK_DETECTED)))
1944            .collect();
1945        assert_eq!(
1946            deadlock_errors.len(),
1947            1,
1948            "expected exactly one side to be aborted with 40P01, got {results:?}"
1949        );
1950    }
1951}
1952
1953#[cfg(test)]
1954mod error_message_tests {
1955    use super::*;
1956
1957    /// A server-side failure must render what the server said, not
1958    /// `tokio_postgres::Error`'s generic `"db error"` — the regression that
1959    /// made `IndexWorker(Vector): claim_batch failed: db error` impossible
1960    /// to act on.
1961    #[tokio::test]
1962    #[ignore = "requires a live database"]
1963    async fn a_server_error_renders_the_server_message() {
1964        let pool = PgPool::connect(&std::env::var("PYLON_PGCON_TEST_DSN").unwrap(), 2)
1965            .await
1966            .unwrap();
1967        // Deliberately needs no fixture: a missing relation is a
1968        // server-side error on any database, so this cannot rot when some
1969        // other test's table comes or goes.
1970        let err = pool
1971            .query_typed("SELECT 1 FROM pylon_no_such_table", &[], &pool.types())
1972            .await
1973            .unwrap_err();
1974        let rendered = err.to_string();
1975        assert!(
1976            rendered.contains("pylon_no_such_table"),
1977            "expected the server's own message, got {rendered:?}"
1978        );
1979        assert_ne!(rendered, "db error");
1980    }
1981}