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}