Skip to main content

oxisql_postgres/
connection.rs

1//! PostgreSQL connection implementing [`oxisql_core::Connection`].
2//!
3//! # Interior mutability
4//!
5//! `tokio_postgres::Client::transaction` requires `&mut self`.  Since the
6//! `Connection` trait only offers `&self`, we wrap the client in an
7//! `Arc<Mutex<Client>>` and acquire an `OwnedMutexGuard` for each transaction
8//! scope.  This matches the same pattern used by `oxisql-embedded`.
9//!
10//! # Transaction semantics
11//!
12//! `PgTransaction` holds an `Option<OwnedMutexGuard<Client>>` for the duration
13//! of the transaction.  Callers should call [`Transaction::commit`] or
14//! [`Transaction::rollback`] to terminate the transaction explicitly.
15//!
16//! ## Automatic rollback on Drop
17//!
18//! If a `PgTransaction` is dropped without an explicit `commit` or `rollback`,
19//! the `Drop` implementation attempts to schedule a `ROLLBACK` command on the
20//! active Tokio runtime.  When no runtime is reachable (e.g., in a sync
21//! context), the guard is simply released and the server will roll back the
22//! open transaction on the next connection reset.
23//!
24//! Explicit termination is still recommended for clarity and to capture errors:
25//! ```rust,no_run
26//! # async fn example(conn: &oxisql_postgres::PgConnection) -> Result<(), oxisql_core::OxiSqlError> {
27//! use oxisql_core::Connection;
28//! let txn = conn.transaction().await?;
29//! // ... operations ...
30//! txn.rollback().await?;  // or .commit()
31//! # Ok(())
32//! # }
33//! ```
34
35use std::collections::HashMap;
36use std::sync::Arc;
37use std::time::Duration;
38
39use async_trait::async_trait;
40use tokio::sync::{Mutex, OwnedMutexGuard};
41use tokio_postgres::types::ToSql;
42use tokio_postgres::{NoTls, Statement};
43
44use oxisql_core::{
45    ColumnInfo, Connection, ForeignKeyInfo, IndexInfo, OxiSqlError, PreparedStatement, Row,
46    TableInfo, TableType, ToSqlValue, Transaction,
47};
48
49use crate::builder::TlsMode;
50use crate::copy;
51use crate::error::PgError;
52use crate::notify::{self, PgNotification};
53use crate::pipeline::PgPipeline;
54use crate::prepared::{PgPrepared, StmtCache};
55use crate::tls::PgTls;
56use crate::types::{pg_row_to_row, value_to_param, OwnedParam};
57
58// ── PgConnection ─────────────────────────────────────────────────────────────
59
60/// Metadata for a single output column from a prepared statement.
61///
62/// Obtained via [`PgConnection::describe`] — describes column names and
63/// PostgreSQL type names without executing the query.
64#[derive(Debug, Clone)]
65pub struct ColumnDescription {
66    /// The column name as reported by the server.
67    pub name: String,
68    /// The PostgreSQL type name (e.g. `"int4"`, `"text"`, `"uuid"`).
69    pub type_name: String,
70    /// Always `true`: PostgreSQL does not expose nullability in column
71    /// descriptions returned by the extended-query protocol.
72    pub nullable: bool,
73}
74
75/// An asynchronous PostgreSQL connection backed by `tokio-postgres` (Pure Rust,
76/// no `libpq`).
77///
78/// The inner `Client` is wrapped in `Arc<Mutex<_>>` to allow the `Connection`
79/// trait's `&self` transaction method to obtain exclusive mutable access.
80///
81/// `Clone` is cheap: both `Arc`s inside are reference-counted.
82///
83/// # Statement cache
84///
85/// A second `Arc<std::sync::Mutex<HashMap<…>>>` holds compiled
86/// `tokio_postgres::Statement` handles keyed on SQL text.  The
87/// `std::sync::Mutex` (not `tokio::sync::Mutex`) is intentional: we only hold
88/// the lock during purely synchronous HashMap lookups, never across an `await`.
89///
90/// # Reconnection
91///
92/// When a `PgConnection` is created via [`PgConnection::connect`], the original
93/// connection string and TLS mode are stored internally.  Call
94/// [`PgConnection::reconnect`] to obtain a brand-new `PgConnection` built from
95/// the same parameters — useful after a broken-pipe error.  The caller is
96/// responsible for replacing their handle with the returned value.
97#[derive(Clone)]
98pub struct PgConnection {
99    inner: Arc<Mutex<tokio_postgres::Client>>,
100    stmt_cache: StmtCache,
101    /// Original connection string passed to [`PgConnection::connect`].
102    ///
103    /// `None` when created via [`PgConnection::from_client`].
104    reconnect_uri: Option<String>,
105    /// TLS mode used for the initial connection.
106    ///
107    /// Stored alongside `reconnect_uri` so that [`PgConnection::reconnect`]
108    /// can re-establish an encrypted connection without silently downgrading
109    /// to plaintext.
110    reconnect_tls: TlsMode,
111    /// Sender side of the broadcast channel used for LISTEN/NOTIFY.
112    ///
113    /// `None` when the connection was created via [`PgConnection::from_client`]
114    /// (no `tokio_postgres::Connection` object was available to intercept
115    /// notifications).  `Some` when created via [`PgConnection::connect`].
116    notif_tx: Option<tokio::sync::broadcast::Sender<PgNotification>>,
117}
118
119impl PgConnection {
120    /// Construct a `PgConnection` from an existing `tokio_postgres::Client`.
121    ///
122    /// Useful when the `Client` was obtained outside of OxiSQL (e.g., from a
123    /// custom connection pool or test harness).  A fresh, empty statement cache
124    /// is created automatically.
125    ///
126    /// # Example
127    ///
128    /// ```rust,no_run
129    /// # #[tokio::main]
130    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
131    /// let (client, conn) = tokio_postgres::connect(
132    ///     "host=localhost user=postgres",
133    ///     tokio_postgres::NoTls,
134    /// ).await?;
135    /// tokio::spawn(async move { let _: Result<_, _> = conn.await; });
136    ///
137    /// let pg_conn = oxisql_postgres::PgConnection::from_client(client);
138    /// # Ok(())
139    /// # }
140    /// ```
141    pub fn from_client(client: tokio_postgres::Client) -> Self {
142        Self {
143            inner: Arc::new(Mutex::new(client)),
144            stmt_cache: Arc::new(std::sync::Mutex::new(std::collections::HashMap::new())),
145            reconnect_uri: None,
146            reconnect_tls: TlsMode::Disabled,
147            notif_tx: None,
148        }
149    }
150
151    /// Connect to a PostgreSQL server using the given connection string.
152    ///
153    /// `conn_str` follows the `tokio-postgres` format, e.g.:
154    /// `"host=localhost port=5432 user=postgres password=secret dbname=mydb"`.
155    ///
156    /// # TLS
157    ///
158    /// Pass [`TlsMode::Disabled`] for a plain-text connection (suitable for
159    /// `localhost` or trusted network).  Pass [`TlsMode::Rustls`] with a
160    /// `ClientConfig` built from OxiTLS for encrypted connections.
161    pub async fn connect(conn_str: &str, tls: TlsMode) -> Result<Self, PgError> {
162        // Preserve the TLS mode for reconnect before moving it into the match.
163        let tls_clone = tls.clone();
164        let (client, notif_tx) = match tls {
165            TlsMode::Disabled => {
166                let (client, connection) = tokio_postgres::connect(conn_str, NoTls).await?;
167                // Spawn the notification-forwarding connection driver.
168                let tx = notify::spawn_connection_driver(connection);
169                (client, tx)
170            }
171            TlsMode::Rustls(cfg) => {
172                let tls_connector = PgTls::new(cfg);
173                let (client, connection) = tokio_postgres::connect(conn_str, tls_connector).await?;
174                let tx = notify::spawn_connection_driver(connection);
175                (client, tx)
176            }
177        };
178        Ok(Self {
179            inner: Arc::new(Mutex::new(client)),
180            stmt_cache: Arc::new(std::sync::Mutex::new(HashMap::new())),
181            reconnect_uri: Some(conn_str.to_string()),
182            reconnect_tls: tls_clone,
183            notif_tx: Some(notif_tx),
184        })
185    }
186
187    /// Connect to a PostgreSQL server with a maximum wait duration.
188    ///
189    /// Wraps [`PgConnection::connect`] with [`tokio::time::timeout`].  If the
190    /// connection is not established within `timeout`, returns
191    /// [`PgError::Timeout`] rather than hanging indefinitely.
192    ///
193    /// # Errors
194    ///
195    /// Returns [`PgError::Timeout`] if `timeout` elapses before the server
196    /// acknowledges the connection.  All other error conditions are propagated
197    /// unchanged from [`PgConnection::connect`].
198    ///
199    /// # Example
200    ///
201    /// ```rust,no_run
202    /// # #[tokio::main]
203    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
204    /// use std::time::Duration;
205    /// use oxisql_postgres::{PgConnection, TlsMode};
206    ///
207    /// let conn = PgConnection::connect_with_timeout(
208    ///     "host=localhost user=postgres",
209    ///     TlsMode::Disabled,
210    ///     Duration::from_secs(5),
211    /// ).await?;
212    /// # Ok(())
213    /// # }
214    /// ```
215    pub async fn connect_with_timeout(
216        conn_str: &str,
217        tls: TlsMode,
218        timeout: Duration,
219    ) -> Result<Self, PgError> {
220        tokio::time::timeout(timeout, Self::connect(conn_str, tls))
221            .await
222            .map_err(|_| {
223                PgError::Timeout(format!(
224                    "connection timed out after {}ms",
225                    timeout.as_millis()
226                ))
227            })?
228    }
229
230    /// Connect with TLS that skips server certificate verification.
231    ///
232    /// # Security Warning
233    ///
234    /// This method accepts **any** server certificate without validation.  It
235    /// is vulnerable to man-in-the-middle attacks and should only be used in
236    /// development or testing environments.
237    ///
238    /// # Errors
239    ///
240    /// Returns [`PgError::Tls`] if the TLS config cannot be built, or a
241    /// [`PgError::Postgres`] / [`PgError::Timeout`] on connection failure.
242    ///
243    /// # Example
244    ///
245    /// ```rust,no_run
246    /// # #[tokio::main]
247    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
248    /// let conn = oxisql_postgres::PgConnection::connect_skip_verify(
249    ///     "host=localhost user=postgres",
250    /// ).await?;
251    /// # Ok(())
252    /// # }
253    /// ```
254    pub async fn connect_skip_verify(conn_str: &str) -> Result<Self, PgError> {
255        Self::connect(conn_str, TlsMode::skip_verify()?).await
256    }
257
258    /// Connect with TLS using a custom CA certificate in PEM format.
259    ///
260    /// The provided `ca_pem` bytes are parsed as a PEM bundle and added to the
261    /// trust store (in addition to the Mozilla WebPKI root certificates).
262    ///
263    /// # Errors
264    ///
265    /// Returns [`PgError::Tls`] if the PEM cannot be parsed or is empty, or a
266    /// [`PgError::Postgres`] / [`PgError::Timeout`] on connection failure.
267    ///
268    /// # Example
269    ///
270    /// ```rust,no_run
271    /// # #[tokio::main]
272    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
273    /// let ca_pem = std::fs::read("/path/to/ca.crt")?;
274    /// let conn = oxisql_postgres::PgConnection::connect_with_ca(
275    ///     "host=db.example.com user=postgres",
276    ///     ca_pem,
277    /// ).await?;
278    /// # Ok(())
279    /// # }
280    /// ```
281    pub async fn connect_with_ca(conn_str: &str, ca_pem: Vec<u8>) -> Result<Self, PgError> {
282        Self::connect(conn_str, TlsMode::with_ca_pem(ca_pem)?).await
283    }
284}
285
286// ── PgConnection — query analysis helpers ────────────────────────────────────
287
288impl PgConnection {
289    /// Return `true` if the SQL is a read-only query (SELECT, EXPLAIN, SHOW, …).
290    ///
291    /// Uses `oxisql_parse::parse_one` to parse the statement and
292    /// `oxisql_parse::is_read_only` for AST-level classification.  This can
293    /// be used to route read queries to a replica and write queries to a primary.
294    ///
295    /// Returns `false` when the SQL cannot be parsed or when the statement
296    /// modifies data or schema.
297    pub fn is_read_only_query(sql: &str) -> bool {
298        match oxisql_parse::parse_one(sql) {
299            Ok(stmt) => oxisql_parse::is_read_only(&stmt),
300            Err(_) => false,
301        }
302    }
303
304    /// Return the normalized canonical form of a SQL query for use as a cache key.
305    ///
306    /// Uses `oxisql_parse::normalize` to collapse whitespace, strip comments,
307    /// and upper-case keywords.  The result is stable across cosmetically
308    /// different but semantically equivalent SQL strings, making it suitable
309    /// as a prepared-statement cache key.
310    pub fn normalize_query(sql: &str) -> String {
311        oxisql_parse::normalize(sql)
312    }
313}
314
315// ── PgConnection — COPY / NOTIFY / pipeline methods ─────────────────────────
316
317impl PgConnection {
318    /// Bulk-insert rows into `table` using the PostgreSQL `COPY … FROM STDIN`
319    /// protocol (text/TSV format).
320    ///
321    /// Each element of `rows` must contain exactly `columns.len()` string
322    /// values.  Values are automatically escaped for the PostgreSQL text COPY
323    /// format.  Pass `"\\N"` (backslash-N) to represent a SQL `NULL`.
324    ///
325    /// Returns the number of rows inserted as reported by PostgreSQL.
326    ///
327    /// # Errors
328    ///
329    /// Returns [`PgError::Copy`] if `columns` is empty, if a row has the
330    /// wrong number of fields, or if the underlying COPY protocol fails.
331    pub async fn copy_in_text(
332        &self,
333        table: &str,
334        columns: &[&str],
335        rows: impl Iterator<Item = Vec<String>>,
336    ) -> Result<u64, PgError> {
337        copy::copy_in_text(&self.inner, table, columns, rows).await
338    }
339
340    /// Extract rows from `table` using the PostgreSQL `COPY … TO STDOUT`
341    /// protocol (text/TSV format).
342    ///
343    /// Returns a `Vec<Vec<String>>` where each inner vec contains the field
344    /// values for one row, with escape sequences decoded.  The literal
345    /// `"\\N"` in the returned data represents a SQL `NULL`.
346    ///
347    /// # Errors
348    ///
349    /// Returns [`PgError::Copy`] if `columns` is empty or if the underlying
350    /// COPY protocol fails.
351    pub async fn copy_out_text(
352        &self,
353        table: &str,
354        columns: &[&str],
355    ) -> Result<Vec<Vec<String>>, PgError> {
356        copy::copy_out_text(&self.inner, table, columns).await
357    }
358
359    /// Send a `NOTIFY` command to `channel` with an optional `payload`.
360    ///
361    /// The channel name must consist only of ASCII alphanumeric characters
362    /// and underscores.  The payload is automatically single-quote–escaped
363    /// before inclusion in the SQL string.
364    ///
365    /// # Errors
366    ///
367    /// Returns [`PgError::Notify`] if the channel name is invalid or if
368    /// the underlying query fails.
369    pub async fn notify(&self, channel: &str, payload: &str) -> Result<(), PgError> {
370        notify::notify(&self.inner, channel, payload).await
371    }
372
373    /// Register `LISTEN` on `channel` and return a [`notify::NotificationStream`].
374    ///
375    /// Requires that this connection was created via [`PgConnection::connect`].
376    /// Connections obtained via [`PgConnection::from_client`] return
377    /// [`PgError::Notify`] because no `tokio_postgres::Connection` driver is
378    /// available to intercept incoming notifications.
379    ///
380    /// # Errors
381    ///
382    /// Returns [`PgError::Notify`] if the channel name is invalid, if this
383    /// connection was created via `from_client`, or if the `LISTEN` command
384    /// fails.
385    pub async fn listen(&self, channel: &str) -> Result<notify::NotificationStream, PgError> {
386        notify::listen(&self.inner, &self.notif_tx, channel).await
387    }
388
389    /// Create a new [`PgPipeline`] for batched query dispatch.
390    ///
391    /// Queue queries with [`PgPipeline::add_execute`] and
392    /// [`PgPipeline::add_query`], then send all at once with
393    /// [`PgPipeline::finish`].
394    pub fn pipeline(&self) -> PgPipeline {
395        PgPipeline::new(Arc::clone(&self.inner))
396    }
397
398    /// Establish a fresh connection using the same URI and TLS mode as this
399    /// connection.
400    ///
401    /// Returns a brand-new [`PgConnection`] with an empty statement cache.
402    /// The caller should replace their handle with the returned value after a
403    /// broken-pipe or connection-lost error.
404    ///
405    /// # Errors
406    ///
407    /// Returns [`PgError::Notify`] (repurposed as a "no URI" sentinel) when
408    /// this connection was created via [`PgConnection::from_client`] and no
409    /// original connection string was recorded.
410    ///
411    /// Returns [`PgError::Postgres`] on connection failure.
412    pub async fn reconnect(&self) -> Result<PgConnection, PgError> {
413        match &self.reconnect_uri {
414            None => Err(PgError::Notify(
415                "cannot reconnect: connection was created via from_client (no URI stored)"
416                    .to_string(),
417            )),
418            Some(uri) => PgConnection::connect(uri, self.reconnect_tls.clone()).await,
419        }
420    }
421
422    /// Execute a query using the binary protocol for better throughput.
423    ///
424    /// Uses `tokio-postgres`'s `query_typed`, which sends explicit Postgres type
425    /// annotations alongside each parameter value.  The server then returns
426    /// result columns in **binary format** (format code 1), bypassing the
427    /// text-serialization step used by [`Connection::query`].  This reduces
428    /// per-row encoding/decoding overhead on both client and server, improving
429    /// throughput for large result sets.
430    ///
431    /// Parameter types are inferred from the [`ToSqlValue`] variants:
432    ///
433    /// | OxiSQL variant | Declared Postgres type |
434    /// |---|---|
435    /// | `Bool` | `BOOL` |
436    /// | `I64` | `INT8` |
437    /// | `F64` | `FLOAT8` |
438    /// | `Text` | `TEXT` |
439    /// | `Blob` | `BYTEA` |
440    /// | `Null` / extended types | `UNKNOWN` (server infers from context) |
441    ///
442    /// **Note**: parameters are declared as INT8/FLOAT8/TEXT rather than
443    /// narrower types (INT4, FLOAT4, etc.).  If a column has a narrower type
444    /// and no implicit cast is available, prefer [`Connection::query`] which
445    /// lets the server apply type inference after `prepare`.
446    ///
447    /// # Errors
448    ///
449    /// Returns [`PgError::Postgres`] on connection or server-side errors.
450    /// Returns [`PgError::TypeConversion`] if a result column cannot be decoded.
451    ///
452    /// # Example
453    ///
454    /// ```rust,no_run
455    /// # #[tokio::main]
456    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
457    /// use oxisql_postgres::{PgConnection, TlsMode};
458    ///
459    /// let conn = PgConnection::connect("host=localhost user=postgres", TlsMode::Disabled).await?;
460    /// let rows = conn.query_binary("SELECT id, name FROM users WHERE id = $1", &[&42_i64]).await?;
461    /// # Ok(())
462    /// # }
463    /// ```
464    pub async fn query_binary(
465        &self,
466        sql: &str,
467        params: &[&dyn oxisql_core::ToSqlValue],
468    ) -> Result<Vec<Row>, PgError> {
469        use tokio_postgres::types::Type;
470
471        let client = self.inner.lock().await;
472        let owned = build_params(params);
473
474        // Build (ref, Type) pairs.  Type::UNKNOWN lets the server use context
475        // to resolve the type when we cannot determine it exactly.
476        let typed_params: Vec<(&(dyn ToSql + Sync), Type)> = owned
477            .iter()
478            .map(|p| {
479                let ty = match p {
480                    OwnedParam::Bool(_) => Type::BOOL,
481                    OwnedParam::I64(_) => Type::INT8,
482                    OwnedParam::F64(_) => Type::FLOAT8,
483                    OwnedParam::Text(_) => Type::TEXT,
484                    OwnedParam::Blob(_) => Type::BYTEA,
485                    // NULL and extended types (sent as text): let the server infer.
486                    OwnedParam::Null => Type::UNKNOWN,
487                };
488                let dyn_ref: &(dyn ToSql + Sync) = p;
489                (dyn_ref, ty)
490            })
491            .collect();
492
493        let pg_rows = client
494            .query_typed(sql, typed_params.as_slice())
495            .await
496            .map_err(PgError::Postgres)?;
497
498        pg_rows
499            .into_iter()
500            .map(crate::types::pg_row_to_row)
501            .collect()
502    }
503
504    /// Return column metadata for the given SQL without executing it.
505    ///
506    /// Uses `client.prepare(sql)` to obtain the server-reported column
507    /// descriptions.  Each [`ColumnDescription`] carries the column name and
508    /// PostgreSQL type name (e.g. `"int4"`, `"text"`, `"uuid"`).
509    ///
510    /// `nullable` is always `true`; the PostgreSQL extended-query protocol
511    /// does not carry nullability information in statement descriptions.
512    ///
513    /// # Errors
514    ///
515    /// Returns [`PgError::Postgres`] if the server rejects the SQL (syntax
516    /// error, unknown table, etc.).
517    pub async fn describe(&self, sql: &str) -> Result<Vec<ColumnDescription>, PgError> {
518        let client = self.inner.lock().await;
519        let stmt = client.prepare(sql).await.map_err(PgError::Postgres)?;
520        Ok(stmt
521            .columns()
522            .iter()
523            .map(|col| ColumnDescription {
524                name: col.name().to_string(),
525                type_name: col.type_().name().to_string(),
526                nullable: true,
527            })
528            .collect())
529    }
530}
531
532// ── helper ────────────────────────────────────────────────────────────────────
533
534/// Convert `oxisql_core` params to `OwnedParam` values and return both vecs so
535/// the caller can build a `&[&(dyn ToSql + Sync)]` slice from the refs.
536fn build_params(params: &[&dyn ToSqlValue]) -> Vec<OwnedParam> {
537    params
538        .iter()
539        .map(|p| value_to_param(&p.to_value()))
540        .collect()
541}
542
543fn owned_refs(owned: &[OwnedParam]) -> Vec<&(dyn ToSql + Sync)> {
544    owned.iter().map(|p| p as &(dyn ToSql + Sync)).collect()
545}
546
547// ── Connection impl ───────────────────────────────────────────────────────────
548
549#[async_trait]
550impl Connection for PgConnection {
551    async fn execute(&self, sql: &str, params: &[&dyn ToSqlValue]) -> Result<u64, OxiSqlError> {
552        let client = self.inner.lock().await;
553        let owned = build_params(params);
554        let refs = owned_refs(&owned);
555        client
556            .execute(sql, refs.as_slice())
557            .await
558            .map_err(|e| OxiSqlError::Execution(e.to_string()))
559    }
560
561    async fn query(&self, sql: &str, params: &[&dyn ToSqlValue]) -> Result<Vec<Row>, OxiSqlError> {
562        let client = self.inner.lock().await;
563        let owned = build_params(params);
564        let refs = owned_refs(&owned);
565        let pg_rows = client
566            .query(sql, refs.as_slice())
567            .await
568            .map_err(|e| OxiSqlError::Execution(e.to_string()))?;
569        pg_rows
570            .into_iter()
571            .map(|r| pg_row_to_row(r).map_err(OxiSqlError::from))
572            .collect()
573    }
574
575    async fn transaction(&self) -> Result<Box<dyn Transaction + '_>, OxiSqlError> {
576        // Acquire an owned guard so `PgTransaction` can be `'static`.
577        let guard = Arc::clone(&self.inner).lock_owned().await;
578        // BEGIN is issued via batch_execute so we have explicit transaction control.
579        guard
580            .batch_execute("BEGIN")
581            .await
582            .map_err(|e| OxiSqlError::Execution(e.to_string()))?;
583        Ok(Box::new(PgTransaction {
584            guard: Some(guard),
585            done: false,
586        }))
587    }
588
589    async fn execute_batch(&self, sql: &str) -> Result<u64, OxiSqlError> {
590        let client = self.inner.lock().await;
591        client
592            .batch_execute(sql)
593            .await
594            .map_err(|e| OxiSqlError::Execution(e.to_string()))?;
595        // batch_execute does not return row counts
596        Ok(0)
597    }
598
599    async fn ping(&self) -> Result<(), OxiSqlError> {
600        let client = self.inner.lock().await;
601        client
602            .simple_query("")
603            .await
604            .map_err(|e| OxiSqlError::Other(e.to_string()))?;
605        Ok(())
606    }
607
608    async fn tables(&self) -> Result<Vec<TableInfo>, OxiSqlError> {
609        let rows = self
610            .query(
611                "SELECT table_name, table_schema, table_type \
612                 FROM information_schema.tables \
613                 WHERE table_schema NOT IN ('pg_catalog', 'information_schema') \
614                 ORDER BY table_schema, table_name",
615                &[],
616            )
617            .await?;
618        rows.into_iter()
619            .map(|row| {
620                let name = row
621                    .try_get::<String>("table_name")
622                    .map_err(|e| OxiSqlError::Other(e.to_string()))?;
623                let schema: Option<String> = row.try_get::<String>("table_schema").ok();
624                let type_str: String = row
625                    .try_get::<String>("table_type")
626                    .unwrap_or_else(|_| "TABLE".to_string());
627                Ok(TableInfo {
628                    name,
629                    schema,
630                    table_type: TableType::from(type_str.as_str()),
631                })
632            })
633            .collect()
634    }
635
636    async fn columns(&self, table: &str) -> Result<Vec<ColumnInfo>, OxiSqlError> {
637        let rows = self
638            .query(
639                "SELECT column_name, ordinal_position, data_type, is_nullable, column_default, \
640                        character_maximum_length, numeric_precision, numeric_scale \
641                 FROM information_schema.columns \
642                 WHERE table_name = $1 \
643                 ORDER BY ordinal_position",
644                &[&table],
645            )
646            .await?;
647        rows.into_iter()
648            .map(|row| {
649                let name = row
650                    .try_get::<String>("column_name")
651                    .map_err(|e| OxiSqlError::Other(e.to_string()))?;
652                let ordinal: i64 = row
653                    .try_get::<i64>("ordinal_position")
654                    .map_err(|e| OxiSqlError::Other(e.to_string()))?;
655                let data_type = row
656                    .try_get::<String>("data_type")
657                    .map_err(|e| OxiSqlError::Other(e.to_string()))?;
658                let nullable_str = row
659                    .try_get::<String>("is_nullable")
660                    .unwrap_or_else(|_| "YES".to_string());
661                let default = row.try_get::<String>("column_default").ok();
662                let max_len: Option<i64> = row.try_get::<i64>("character_maximum_length").ok();
663                let num_prec: Option<i64> = row.try_get::<i64>("numeric_precision").ok();
664                let num_scale: Option<i64> = row.try_get::<i64>("numeric_scale").ok();
665                Ok(ColumnInfo {
666                    name,
667                    ordinal_position: u32::try_from(ordinal).unwrap_or(0),
668                    data_type,
669                    nullable: !nullable_str.eq_ignore_ascii_case("NO"),
670                    default,
671                    max_length: max_len.and_then(|n| u64::try_from(n).ok()),
672                    numeric_precision: num_prec.and_then(|n| u32::try_from(n).ok()),
673                    numeric_scale: num_scale.and_then(|n| u32::try_from(n).ok()),
674                })
675            })
676            .collect()
677    }
678
679    async fn indexes(&self, table: &str) -> Result<Vec<IndexInfo>, OxiSqlError> {
680        let rows = self
681            .query(
682                "SELECT indexname, indexdef \
683                 FROM pg_indexes \
684                 WHERE tablename = $1 \
685                 ORDER BY indexname",
686                &[&table],
687            )
688            .await?;
689        rows.into_iter()
690            .map(|row| {
691                let name = row
692                    .try_get::<String>("indexname")
693                    .map_err(|e| OxiSqlError::Other(e.to_string()))?;
694                let def = row.try_get::<String>("indexdef").unwrap_or_default();
695                let unique = def.contains("UNIQUE");
696                let primary = name.ends_with("_pkey");
697                // Extract columns from the index definition: content between
698                // the last '(' and ')'.
699                let columns = if let (Some(start), Some(end)) = (def.rfind('('), def.rfind(')')) {
700                    def[start + 1..end]
701                        .split(',')
702                        .map(|s| s.trim().to_string())
703                        .collect()
704                } else {
705                    Vec::new()
706                };
707                Ok(IndexInfo {
708                    name,
709                    columns,
710                    unique,
711                    primary,
712                })
713            })
714            .collect()
715    }
716
717    async fn foreign_keys(&self, table: &str) -> Result<Vec<ForeignKeyInfo>, OxiSqlError> {
718        let rows = self
719            .query(
720                "SELECT tc.constraint_name, kcu.column_name, \
721                        ccu.table_name AS foreign_table, \
722                        ccu.column_name AS foreign_column \
723                 FROM information_schema.table_constraints tc \
724                 JOIN information_schema.key_column_usage kcu \
725                      ON tc.constraint_name = kcu.constraint_name \
726                      AND tc.table_schema = kcu.table_schema \
727                 JOIN information_schema.constraint_column_usage ccu \
728                      ON ccu.constraint_name = tc.constraint_name \
729                      AND ccu.table_schema = tc.table_schema \
730                 WHERE tc.constraint_type = 'FOREIGN KEY' AND tc.table_name = $1",
731                &[&table],
732            )
733            .await?;
734        rows.into_iter()
735            .map(|row| {
736                let constraint_name = row
737                    .try_get::<String>("constraint_name")
738                    .map_err(|e| OxiSqlError::Other(e.to_string()))?;
739                let column = row
740                    .try_get::<String>("column_name")
741                    .map_err(|e| OxiSqlError::Other(e.to_string()))?;
742                let foreign_table = row
743                    .try_get::<String>("foreign_table")
744                    .map_err(|e| OxiSqlError::Other(e.to_string()))?;
745                let foreign_column = row
746                    .try_get::<String>("foreign_column")
747                    .map_err(|e| OxiSqlError::Other(e.to_string()))?;
748                Ok(ForeignKeyInfo {
749                    constraint_name,
750                    column,
751                    foreign_table,
752                    foreign_column,
753                })
754            })
755            .collect()
756    }
757
758    /// Stream query results row-by-row from PostgreSQL.
759    ///
760    /// Acquires a lock on the client, runs `client.query()` (which fetches all
761    /// rows in a single round-trip), then yields each converted row individually.
762    /// This is equivalent to the default trait implementation but avoids going
763    /// through `async_trait`'s boxing.
764    fn query_stream<'a>(
765        &'a self,
766        sql: &'a str,
767        params: &'a [&'a dyn oxisql_core::ToSqlValue],
768    ) -> std::pin::Pin<
769        Box<
770            dyn futures::Stream<Item = Result<oxisql_core::Row, oxisql_core::OxiSqlError>>
771                + Send
772                + 'a,
773        >,
774    > {
775        use futures::StreamExt;
776        let fut = self.query(sql, params);
777        Box::pin(futures::stream::once(fut).flat_map(|result| match result {
778            Ok(rows) => futures::stream::iter(rows.into_iter().map(Ok)).left_stream(),
779            Err(e) => futures::stream::once(async move { Err(e) }).right_stream(),
780        }))
781    }
782
783    /// Compile a SQL statement for repeated execution.
784    ///
785    /// The compiled [`Statement`] is cached by SQL text: subsequent calls with
786    /// the same SQL return a clone of the cached handle without any server
787    /// round-trip.
788    async fn prepare(&self, sql: &str) -> Result<Box<dyn PreparedStatement + '_>, OxiSqlError> {
789        // --- cache look-up (synchronous lock, never held across await) ---
790        {
791            let cache = self
792                .stmt_cache
793                .lock()
794                .map_err(|_| OxiSqlError::Other("statement cache lock poisoned".into()))?;
795            if let Some(stmt) = cache.get(sql) {
796                return Ok(Box::new(PgPrepared {
797                    client: Arc::clone(&self.inner),
798                    stmt: stmt.clone(),
799                    sql_text: sql.to_string(),
800                }));
801            }
802        } // lock released before any await
803
804        // --- cache miss: prepare on the server ---
805        let stmt: Statement = {
806            let client = self.inner.lock().await;
807            client
808                .prepare(sql)
809                .await
810                .map_err(|e| OxiSqlError::Execution(e.to_string()))?
811        }; // client lock released
812
813        // --- insert into cache ---
814        {
815            let mut cache = self
816                .stmt_cache
817                .lock()
818                .map_err(|_| OxiSqlError::Other("statement cache lock poisoned".into()))?;
819            cache.insert(sql.to_string(), stmt.clone());
820        }
821
822        Ok(Box::new(PgPrepared {
823            client: Arc::clone(&self.inner),
824            stmt,
825            sql_text: sql.to_string(),
826        }))
827    }
828}
829
830// ── PgTransaction ─────────────────────────────────────────────────────────────
831
832/// A PostgreSQL transaction that holds exclusive access to the client for the
833/// duration of the transaction scope.
834///
835/// # Preferred cleanup
836///
837/// Call [`Transaction::commit`] or [`Transaction::rollback`] to end the
838/// transaction explicitly.  If neither is called before this value is dropped,
839/// the `Drop` implementation attempts to schedule a `ROLLBACK` on the current
840/// Tokio runtime (best-effort).  If no runtime is reachable the connection is
841/// simply released; the server will clean up via connection reset.
842pub struct PgTransaction {
843    /// `Some` until `commit`, `rollback`, or `Drop` consumes the guard.
844    guard: Option<OwnedMutexGuard<tokio_postgres::Client>>,
845    /// `true` once `commit` or `rollback` has been called.
846    done: bool,
847}
848
849impl Drop for PgTransaction {
850    fn drop(&mut self) {
851        if let Some(guard) = self.guard.take() {
852            if !self.done {
853                // Best-effort: schedule ROLLBACK on the active Tokio runtime.
854                // If no runtime is reachable, the guard drops here, releasing
855                // the mutex so the connection becomes usable again.  The server
856                // will roll back automatically on connection reset.
857                match tokio::runtime::Handle::try_current() {
858                    Ok(handle) => {
859                        handle.spawn(async move {
860                            let _ = guard.batch_execute("ROLLBACK").await;
861                        });
862                    }
863                    Err(_) => {
864                        // Not in a Tokio context — guard is dropped silently.
865                        drop(guard);
866                    }
867                }
868            }
869        }
870    }
871}
872
873/// Helper: borrow the active guard or return a "transaction consumed" error.
874fn guard_ref(
875    guard: &mut Option<OwnedMutexGuard<tokio_postgres::Client>>,
876) -> Result<&mut OwnedMutexGuard<tokio_postgres::Client>, OxiSqlError> {
877    guard.as_mut().ok_or_else(|| {
878        OxiSqlError::Execution("transaction already consumed (committed or rolled back)".into())
879    })
880}
881
882#[async_trait]
883impl Transaction for PgTransaction {
884    async fn execute(&mut self, sql: &str, params: &[&dyn ToSqlValue]) -> Result<u64, OxiSqlError> {
885        let owned = build_params(params);
886        let refs = owned_refs(&owned);
887        guard_ref(&mut self.guard)?
888            .execute(sql, refs.as_slice())
889            .await
890            .map_err(|e| OxiSqlError::Execution(e.to_string()))
891    }
892
893    async fn query(
894        &mut self,
895        sql: &str,
896        params: &[&dyn ToSqlValue],
897    ) -> Result<Vec<Row>, OxiSqlError> {
898        let owned = build_params(params);
899        let refs = owned_refs(&owned);
900        let pg_rows = guard_ref(&mut self.guard)?
901            .query(sql, refs.as_slice())
902            .await
903            .map_err(|e| OxiSqlError::Execution(e.to_string()))?;
904        pg_rows
905            .into_iter()
906            .map(|r| pg_row_to_row(r).map_err(OxiSqlError::from))
907            .collect()
908    }
909
910    async fn commit(mut self: Box<Self>) -> Result<(), OxiSqlError> {
911        self.done = true;
912        let guard = self
913            .guard
914            .take()
915            .ok_or_else(|| OxiSqlError::Execution("transaction already consumed".into()))?;
916        guard
917            .batch_execute("COMMIT")
918            .await
919            .map_err(|e| OxiSqlError::Execution(e.to_string()))
920    }
921
922    async fn rollback(mut self: Box<Self>) -> Result<(), OxiSqlError> {
923        self.done = true;
924        let guard = self
925            .guard
926            .take()
927            .ok_or_else(|| OxiSqlError::Execution("transaction already consumed".into()))?;
928        guard
929            .batch_execute("ROLLBACK")
930            .await
931            .map_err(|e| OxiSqlError::Execution(e.to_string()))
932    }
933
934    async fn savepoint(&mut self, name: &str) -> Result<(), OxiSqlError> {
935        validate_savepoint_name(name)?;
936        guard_ref(&mut self.guard)?
937            .batch_execute(&format!("SAVEPOINT {name}"))
938            .await
939            .map_err(|e| OxiSqlError::Execution(e.to_string()))
940    }
941
942    async fn release_savepoint(&mut self, name: &str) -> Result<(), OxiSqlError> {
943        validate_savepoint_name(name)?;
944        guard_ref(&mut self.guard)?
945            .batch_execute(&format!("RELEASE SAVEPOINT {name}"))
946            .await
947            .map_err(|e| OxiSqlError::Execution(e.to_string()))
948    }
949
950    async fn rollback_to_savepoint(&mut self, name: &str) -> Result<(), OxiSqlError> {
951        validate_savepoint_name(name)?;
952        guard_ref(&mut self.guard)?
953            .batch_execute(&format!("ROLLBACK TO SAVEPOINT {name}"))
954            .await
955            .map_err(|e| OxiSqlError::Execution(e.to_string()))
956    }
957}
958
959// ── PgTransaction — inherent savepoint/timeout API (returns PgError) ─────────
960
961impl PgTransaction {
962    /// Create a savepoint with the given name.
963    ///
964    /// The name must consist only of ASCII alphanumeric characters and
965    /// underscores (`[a-zA-Z0-9_]`).  This restriction prevents SQL injection
966    /// because savepoint names cannot be parameterized in the PostgreSQL
967    /// wire protocol.
968    ///
969    /// # Errors
970    ///
971    /// Returns [`PgError::Notify`] if the name is empty or contains unsafe
972    /// characters.  Returns [`PgError::Postgres`] on wire-level errors.
973    pub async fn savepoint_pg(&mut self, name: &str) -> Result<(), PgError> {
974        if !is_valid_savepoint_name(name) {
975            return Err(PgError::Notify(format!(
976                "invalid savepoint name {name:?}: only alphanumeric characters and underscores allowed"
977            )));
978        }
979        let guard = self
980            .guard
981            .as_ref()
982            .ok_or_else(|| PgError::Notify("transaction already consumed".into()))?;
983        guard
984            .batch_execute(&format!("SAVEPOINT {name}"))
985            .await
986            .map_err(PgError::Postgres)
987    }
988
989    /// Roll back to a savepoint without releasing it.
990    ///
991    /// The savepoint remains active after this call; you can roll back to it
992    /// again or release it with [`PgTransaction::release_savepoint_pg`].
993    ///
994    /// # Errors
995    ///
996    /// Returns [`PgError::Notify`] if the name is invalid or the transaction
997    /// has already been consumed.  Returns [`PgError::Postgres`] on server
998    /// errors (e.g. the named savepoint does not exist).
999    pub async fn rollback_to_savepoint_pg(&mut self, name: &str) -> Result<(), PgError> {
1000        if !is_valid_savepoint_name(name) {
1001            return Err(PgError::Notify(format!(
1002                "invalid savepoint name {name:?}: only alphanumeric characters and underscores allowed"
1003            )));
1004        }
1005        let guard = self
1006            .guard
1007            .as_ref()
1008            .ok_or_else(|| PgError::Notify("transaction already consumed".into()))?;
1009        guard
1010            .batch_execute(&format!("ROLLBACK TO SAVEPOINT {name}"))
1011            .await
1012            .map_err(PgError::Postgres)
1013    }
1014
1015    /// Release (remove) a savepoint.
1016    ///
1017    /// All savepoints created after the named savepoint are also removed.
1018    ///
1019    /// # Errors
1020    ///
1021    /// Returns [`PgError::Notify`] if the name is invalid or the transaction
1022    /// has already been consumed.  Returns [`PgError::Postgres`] on server
1023    /// errors (e.g. the named savepoint does not exist).
1024    pub async fn release_savepoint_pg(&mut self, name: &str) -> Result<(), PgError> {
1025        if !is_valid_savepoint_name(name) {
1026            return Err(PgError::Notify(format!(
1027                "invalid savepoint name {name:?}: only alphanumeric characters and underscores allowed"
1028            )));
1029        }
1030        let guard = self
1031            .guard
1032            .as_ref()
1033            .ok_or_else(|| PgError::Notify("transaction already consumed".into()))?;
1034        guard
1035            .batch_execute(&format!("RELEASE SAVEPOINT {name}"))
1036            .await
1037            .map_err(PgError::Postgres)
1038    }
1039}
1040
1041/// Fast path: check if a savepoint name contains only safe SQL identifier chars.
1042///
1043/// Accepts `[a-zA-Z0-9_]` and rejects empty strings.
1044fn is_valid_savepoint_name(name: &str) -> bool {
1045    !name.is_empty() && name.chars().all(|c| c.is_ascii_alphanumeric() || c == '_')
1046}
1047
1048/// Validate that a savepoint name contains only safe SQL identifier characters.
1049///
1050/// Accepts `[a-zA-Z0-9_]` — no spaces, no quotes, no SQL injection vectors.
1051fn validate_savepoint_name(name: &str) -> Result<(), OxiSqlError> {
1052    if name.is_empty() {
1053        return Err(OxiSqlError::Other(
1054            "savepoint name must not be empty".into(),
1055        ));
1056    }
1057    if !name.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') {
1058        return Err(OxiSqlError::Other(format!(
1059            "invalid savepoint name {name:?}: only alphanumeric characters and underscores allowed"
1060        )));
1061    }
1062    Ok(())
1063}
1064
1065// ── PgConnParts / parse_pg_conn_str ───────────────────────────────────────────
1066
1067/// Parsed fields extracted from a PostgreSQL connection string.
1068///
1069/// Obtained via [`parse_pg_conn_str`].
1070#[derive(Debug, Clone, PartialEq, Eq)]
1071pub struct PgConnParts {
1072    /// The host name or IP address (defaults to `"localhost"`).
1073    pub host: String,
1074    /// The server port (defaults to `5432`).
1075    pub port: u16,
1076    /// The target database name, if present.
1077    pub dbname: Option<String>,
1078    /// The login user name, if present.
1079    pub user: Option<String>,
1080}
1081
1082/// Parse a PostgreSQL connection string into its constituent parts.
1083///
1084/// Accepts two common formats:
1085///
1086/// **Key-value format** (the `libpq`/`tokio-postgres` native format):
1087/// ```text
1088/// host=myhost port=5433 dbname=testdb user=admin password=secret
1089/// ```
1090/// Unknown keys are silently ignored.  Missing `host` defaults to
1091/// `"localhost"`; missing `port` defaults to `5432`.
1092///
1093/// **URI format**:
1094/// ```text
1095/// postgres://user:password@host:port/dbname
1096/// postgresql://host/dbname
1097/// ```
1098/// The scheme must be `postgres://` or `postgresql://`.
1099///
1100/// # Errors
1101///
1102/// Returns [`OxiSqlError::Other`] when:
1103/// - A non-empty `port` value cannot be parsed as a `u16`.
1104/// - The URI scheme is unrecognised (neither `postgres://` nor `postgresql://`).
1105///
1106/// # Examples
1107///
1108/// ```rust
1109/// use oxisql_postgres::parse_pg_conn_str;
1110///
1111/// let p = parse_pg_conn_str("host=myhost port=5433 dbname=testdb user=admin").unwrap();
1112/// assert_eq!(p.host, "myhost");
1113/// assert_eq!(p.port, 5433);
1114/// assert_eq!(p.dbname, Some("testdb".to_string()));
1115/// assert_eq!(p.user, Some("admin".to_string()));
1116/// ```
1117pub fn parse_pg_conn_str(conn_str: &str) -> Result<PgConnParts, OxiSqlError> {
1118    let trimmed = conn_str.trim();
1119
1120    if trimmed.starts_with("postgres://") || trimmed.starts_with("postgresql://") {
1121        parse_pg_uri(trimmed)
1122    } else {
1123        parse_pg_kv(trimmed)
1124    }
1125}
1126
1127/// Parse the `key=value …` form of a Postgres connection string.
1128fn parse_pg_kv(conn_str: &str) -> Result<PgConnParts, OxiSqlError> {
1129    let mut parts = PgConnParts {
1130        host: "localhost".to_string(),
1131        port: 5432,
1132        dbname: None,
1133        user: None,
1134    };
1135    for token in conn_str.split_whitespace() {
1136        if let Some((k, v)) = token.split_once('=') {
1137            match k {
1138                "host" | "hostaddr" => parts.host = v.to_string(),
1139                "port" if !v.is_empty() => {
1140                    parts.port = v.parse::<u16>().map_err(|_| {
1141                        OxiSqlError::Other(format!(
1142                            "invalid port value in connection string: {v:?}"
1143                        ))
1144                    })?;
1145                }
1146                "dbname" => parts.dbname = Some(v.to_string()),
1147                "user" => parts.user = Some(v.to_string()),
1148                _ => {}
1149            }
1150        }
1151    }
1152    Ok(parts)
1153}
1154
1155/// Parse the `postgres://[user[:pass]@]host[:port][/dbname]` URI form.
1156fn parse_pg_uri(uri: &str) -> Result<PgConnParts, OxiSqlError> {
1157    // Strip scheme.
1158    let rest = uri
1159        .strip_prefix("postgresql://")
1160        .or_else(|| uri.strip_prefix("postgres://"))
1161        .ok_or_else(|| {
1162            OxiSqlError::Other(format!("unrecognised connection string scheme: {uri:?}"))
1163        })?;
1164
1165    // rest = [user[:pass]@]host[:port][/dbname][?options]
1166    // Split off query string first.
1167    let (rest, _query) = rest.split_once('?').unwrap_or((rest, ""));
1168
1169    // Split user-info from host-info at the last '@'.
1170    let (user_info, host_part) = if let Some(at_pos) = rest.rfind('@') {
1171        (Some(&rest[..at_pos]), &rest[at_pos + 1..])
1172    } else {
1173        (None, rest)
1174    };
1175
1176    // Extract user from user_info (ignore password).
1177    let user = user_info.and_then(|ui| {
1178        let cred = ui.split_once(':').map(|(u, _)| u).unwrap_or(ui);
1179        if cred.is_empty() {
1180            None
1181        } else {
1182            Some(cred.to_string())
1183        }
1184    });
1185
1186    // Split path (dbname) from host[:port].
1187    let (authority, dbname) = if let Some(slash_pos) = host_part.find('/') {
1188        let db = &host_part[slash_pos + 1..];
1189        (
1190            &host_part[..slash_pos],
1191            if db.is_empty() {
1192                None
1193            } else {
1194                Some(db.to_string())
1195            },
1196        )
1197    } else {
1198        (host_part, None)
1199    };
1200
1201    // Split host and port.
1202    let (host, port) = if authority.starts_with('[') {
1203        // IPv6 literal: [::1]:5432
1204        if let Some(bracket_end) = authority.find(']') {
1205            let h = &authority[1..bracket_end];
1206            let port_str = authority[bracket_end + 1..].trim_start_matches(':');
1207            let p = if port_str.is_empty() {
1208                5432
1209            } else {
1210                port_str
1211                    .parse::<u16>()
1212                    .map_err(|_| OxiSqlError::Other(format!("invalid port in URI: {port_str:?}")))?
1213            };
1214            (h.to_string(), p)
1215        } else {
1216            (authority.to_string(), 5432)
1217        }
1218    } else if let Some(colon_pos) = authority.rfind(':') {
1219        let h = &authority[..colon_pos];
1220        let port_str = &authority[colon_pos + 1..];
1221        let p = if port_str.is_empty() {
1222            5432
1223        } else {
1224            port_str
1225                .parse::<u16>()
1226                .map_err(|_| OxiSqlError::Other(format!("invalid port in URI: {port_str:?}")))?
1227        };
1228        (h.to_string(), p)
1229    } else {
1230        (authority.to_string(), 5432)
1231    };
1232
1233    let host = if host.is_empty() {
1234        "localhost".to_string()
1235    } else {
1236        host
1237    };
1238
1239    Ok(PgConnParts {
1240        host,
1241        port,
1242        dbname,
1243        user,
1244    })
1245}