siderite_backends/
postgres.rs

1//! PostgreSQL backend built on SQLx.
2//!
3//! Every [`Value`] variant has a native PostgreSQL type (`Uuid` → `uuid`,
4//! `Decimal` → `numeric`, `Timestamp` → `timestamptz`, `Json` → `jsonb`, ...),
5//! so nothing is stored in a canonical text form. Rows are decoded by
6//! PostgreSQL type name into the matching native variant.
7//!
8//! Connections are pinned to the `UTC` time zone so that `timestamptz`
9//! date-part lookups (`EXTRACT`) agree with the values the ORM writes.
10
11use crate::connection_init::ConnectionInit;
12use crate::shared::{
13    ConnSlot, TxSlot, affected_result, map_error, returning_result, with_conn, with_tx,
14};
15use crate::sql::{CompiledQuery, Postgres, compile, compile_write};
16use async_trait::async_trait;
17use siderite_orm::{
18    Backend, BackendCapabilities, BackendError, ExecResult, Executor, IsolationLevel, OrmError,
19    QueryError, QueryPlan, QueryResult, Row, Transaction, Value, WritePlan,
20};
21use sqlx::encode::IsNull;
22use sqlx::error::BoxDynError;
23use sqlx::postgres::types::Oid;
24use sqlx::postgres::{
25    PgArgumentBuffer, PgArguments, PgPool, PgPoolOptions, PgRow, PgTypeInfo, PgValueRef,
26};
27use sqlx::types::chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime, Utc};
28use sqlx::types::uuid::Uuid;
29use sqlx::types::{Decimal, JsonValue};
30use sqlx::{Column as _, Executor as _, Row as _, TypeInfo as _, ValueRef as _};
31
32/// PostgreSQL adapter executing compiled plans on a connection pool.
33#[derive(Debug, Clone)]
34pub struct PgBackend {
35    pool: PgPool,
36}
37
38impl PgBackend {
39    /// Connect to `url` (`postgres://user:pass@host/db`) with a pool of up to
40    /// ten connections, each set to the UTC time zone.
41    ///
42    /// A connection is **not** checked when it is taken from the pool: SQLx's
43    /// `test_before_acquire` costs a round trip per query. The price is that
44    /// the first query on a connection the server has closed (a restart, an
45    /// idle timeout) fails with a connection error instead of being retried
46    /// on a fresh connection. To have connections checked, use
47    /// [`connect_with`](Self::connect_with) with
48    /// `PgPoolOptions::new().test_before_acquire(true)`.
49    ///
50    /// # Errors
51    /// [`BackendError::Connection`] if the pool cannot be created.
52    pub async fn connect(url: &str) -> Result<Self, BackendError> {
53        let options = PgPoolOptions::new()
54            .max_connections(10)
55            .test_before_acquire(false);
56        Self::connect_with(url, options).await
57    }
58
59    /// Connect with pool settings and mandatory adapter initialization.
60    ///
61    /// Replaces any after_connect callback already stored in options.
62    /// SQLx does not expose that callback for chaining; pass custom setup to
63    /// [`Self::connect_with_init`] instead. Other pool options are retained.
64    ///
65    /// Args:
66    ///     url: Database URL.
67    ///     options: Pool capacity, timeouts and checkout policy.
68    ///
69    /// Returns:
70    ///     Adapter whose fresh sessions receive mandatory initialization.
71    ///
72    /// # Errors
73    /// Connection or initialization failure within the acquire deadline.
74    pub async fn connect_with(url: &str, options: PgPoolOptions) -> Result<Self, BackendError> {
75        let init: ConnectionInit<sqlx::Postgres> =
76            std::sync::Arc::new(|_, _| Box::pin(async { Ok(()) }));
77        Self::connect_with_init(url, options, init).await
78    }
79
80    /// Compose explicit custom setup with mandatory session settings.
81    ///
82    /// Args:
83    ///     url: Database URL.
84    ///     options: Pool settings; its stored after_connect is replaced.
85    ///     init: Custom hook run first on every new/reconnected session.
86    ///
87    /// Returns:
88    ///     Adapter after custom setup and required settings succeed.
89    ///
90    /// # Errors
91    /// Connection or initialization failure within the acquire deadline.
92    pub async fn connect_with_init(
93        url: &str,
94        options: PgPoolOptions,
95        init: ConnectionInit<sqlx::Postgres>,
96    ) -> Result<Self, BackendError> {
97        let pool = options
98            .after_connect(move |conn, meta| {
99                let init = init.clone();
100                Box::pin(async move {
101                    init(conn, meta).await?;
102                    conn.execute("SET TIME ZONE 'UTC'").await?;
103                    Ok(())
104                })
105            })
106            .connect(url)
107            .await
108            .map_err(|e| BackendError::Connection(e.to_string()))?;
109        Ok(Self { pool })
110    }
111
112    /// Wrap an existing pool. Its connections should use the UTC time zone
113    /// (see the module docs).
114    pub fn from_pool(pool: PgPool) -> Self {
115        Self { pool }
116    }
117}
118
119#[async_trait]
120impl Executor for PgBackend {
121    fn capabilities(&self) -> BackendCapabilities {
122        BackendCapabilities::postgres()
123    }
124
125    async fn fetch(&self, plan: &QueryPlan) -> Result<QueryResult, OrmError> {
126        let compiled = compile(plan, &Postgres)?;
127        fetch_rows(&self.pool, &compiled.sql, compiled.params).await
128    }
129
130    async fn execute(&self, plan: &WritePlan) -> Result<ExecResult, OrmError> {
131        let compiled = compile_write(plan, &Postgres)?;
132        run_write(&self.pool, compiled, !plan.returning().is_empty()).await
133    }
134
135    async fn fetch_raw(&self, sql: &str, params: Vec<Value>) -> Result<QueryResult, OrmError> {
136        fetch_rows(&self.pool, sql, params).await
137    }
138
139    async fn execute_raw(&self, sql: &str, params: Vec<Value>) -> Result<u64, OrmError> {
140        execute_rows(&self.pool, sql, params).await
141    }
142
143    async fn execute_script(&self, sql: &str) -> Result<(), OrmError> {
144        script_result(sqlx::Executor::execute(&self.pool, sql).await)
145    }
146}
147
148#[async_trait]
149impl Backend for PgBackend {
150    fn read_parameter_count(&self, plan: &QueryPlan) -> Result<Option<usize>, OrmError> {
151        Ok(Some(compile(plan, &Postgres)?.params.len()))
152    }
153
154    async fn begin(
155        &self,
156        isolation: Option<IsolationLevel>,
157    ) -> Result<Box<dyn Transaction>, OrmError> {
158        if let Some(level) = isolation {
159            self.capabilities()
160                .require(siderite_orm::Feature::Isolation(level))?;
161        }
162        let mut tx = self.pool.begin().await.map_err(map_error)?;
163        if let Some(level) = isolation {
164            // Must be the first statement of the transaction.
165            let statement = format!("SET TRANSACTION ISOLATION LEVEL {}", isolation_sql(level));
166            sqlx::query(&statement)
167                .execute(&mut *tx)
168                .await
169                .map_err(map_error)?;
170        }
171        Ok(Box::new(PgTransaction(TxSlot::new(tx))))
172    }
173
174    async fn begin_schema(&self, transactional: bool) -> Result<Box<dyn Transaction>, OrmError> {
175        if transactional {
176            return self.begin(None).await;
177        }
178        let conn = self.pool.acquire().await.map_err(map_error)?;
179        Ok(Box::new(PgHeld(ConnSlot::new(conn))))
180    }
181}
182
183fn isolation_sql(level: IsolationLevel) -> &'static str {
184    match level {
185        IsolationLevel::ReadCommitted => "READ COMMITTED",
186        IsolationLevel::RepeatableRead => "REPEATABLE READ",
187        IsolationLevel::Serializable => "SERIALIZABLE",
188    }
189}
190
191/// An open PostgreSQL transaction. Dropped without commit, it rolls back.
192struct PgTransaction(TxSlot<sqlx::Postgres>);
193
194#[async_trait]
195impl Executor for PgTransaction {
196    fn capabilities(&self) -> BackendCapabilities {
197        BackendCapabilities::postgres()
198    }
199
200    async fn fetch(&self, plan: &QueryPlan) -> Result<QueryResult, OrmError> {
201        let compiled = compile(plan, &Postgres)?;
202        with_tx!(self.0, conn => fetch_rows(conn, &compiled.sql, compiled.params).await)
203    }
204
205    async fn execute(&self, plan: &WritePlan) -> Result<ExecResult, OrmError> {
206        let compiled = compile_write(plan, &Postgres)?;
207        let returning = !plan.returning().is_empty();
208        with_tx!(self.0, conn => run_write(conn, compiled, returning).await)
209    }
210
211    async fn fetch_raw(&self, sql: &str, params: Vec<Value>) -> Result<QueryResult, OrmError> {
212        with_tx!(self.0, conn => fetch_rows(conn, sql, params).await)
213    }
214
215    async fn execute_raw(&self, sql: &str, params: Vec<Value>) -> Result<u64, OrmError> {
216        with_tx!(self.0, conn => execute_rows(conn, sql, params).await)
217    }
218
219    async fn execute_script(&self, sql: &str) -> Result<(), OrmError> {
220        with_tx!(self.0, conn => script_result(sqlx::Executor::execute(conn, sql).await))
221    }
222}
223
224#[async_trait]
225impl Transaction for PgTransaction {
226    async fn commit(&self) -> Result<(), OrmError> {
227        self.0.commit().await
228    }
229
230    async fn rollback(&self) -> Result<(), OrmError> {
231        self.0.rollback().await
232    }
233}
234
235/// Pooled connection held for a migration lock (no SQL transaction).
236///
237/// Commit and rollback close the connection instead of returning it to the
238/// pool: a session lock (`pg_advisory_lock`) that was not released must never
239/// be handed to another caller. Releasing the lock and then closing costs one
240/// reconnect per migrate, which is negligible for a deploy-time operation.
241struct PgHeld(ConnSlot<sqlx::Postgres>);
242
243#[async_trait]
244impl Executor for PgHeld {
245    fn capabilities(&self) -> BackendCapabilities {
246        BackendCapabilities::postgres()
247    }
248
249    async fn fetch(&self, plan: &QueryPlan) -> Result<QueryResult, OrmError> {
250        let compiled = compile(plan, &Postgres)?;
251        with_conn!(self.0, conn => fetch_rows(conn, &compiled.sql, compiled.params).await)
252    }
253
254    async fn execute(&self, plan: &WritePlan) -> Result<ExecResult, OrmError> {
255        let compiled = compile_write(plan, &Postgres)?;
256        let returning = !plan.returning().is_empty();
257        with_conn!(self.0, conn => run_write(conn, compiled, returning).await)
258    }
259
260    async fn fetch_raw(&self, sql: &str, params: Vec<Value>) -> Result<QueryResult, OrmError> {
261        with_conn!(self.0, conn => fetch_rows(conn, sql, params).await)
262    }
263
264    async fn execute_raw(&self, sql: &str, params: Vec<Value>) -> Result<u64, OrmError> {
265        with_conn!(self.0, conn => execute_rows(conn, sql, params).await)
266    }
267
268    async fn execute_script(&self, sql: &str) -> Result<(), OrmError> {
269        with_conn!(self.0, conn => script_result(sqlx::Executor::execute(conn, sql).await))
270    }
271}
272
273#[async_trait]
274impl Transaction for PgHeld {
275    async fn commit(&self) -> Result<(), OrmError> {
276        let mut conn = self.0.take().await?;
277        conn.close_on_drop();
278        drop(conn);
279        Ok(())
280    }
281
282    async fn rollback(&self) -> Result<(), OrmError> {
283        let mut conn = self.0.take().await?;
284        conn.close_on_drop();
285        drop(conn);
286        Ok(())
287    }
288}
289
290/// A `NULL` parameter of unspecified type (OID 0): the server infers it from
291/// context, so it fits `text`, `uuid`, `timestamptz`, `jsonb`, ... columns.
292/// (A typed `None::<i64>` would be rejected for all but integer columns.)
293struct UntypedNull;
294
295impl sqlx::Type<sqlx::Postgres> for UntypedNull {
296    fn type_info() -> PgTypeInfo {
297        PgTypeInfo::with_oid(Oid(0))
298    }
299}
300
301impl sqlx::Encode<'_, sqlx::Postgres> for UntypedNull {
302    fn encode_by_ref(&self, _buf: &mut PgArgumentBuffer) -> Result<IsNull, BoxDynError> {
303        Ok(IsNull::Yes)
304    }
305}
306
307fn bind_all(sql: &str, params: Vec<Value>) -> sqlx::query::Query<'_, sqlx::Postgres, PgArguments> {
308    params
309        .into_iter()
310        .fold(sqlx::query(sql), |query, param| match param {
311            Value::Null => query.bind(UntypedNull),
312            Value::Bool(v) => query.bind(v),
313            Value::Int(v) => query.bind(v),
314            Value::Float(v) => query.bind(v),
315            Value::Decimal(v) => query.bind(v),
316            Value::Text(v) => query.bind(v),
317            Value::Bytes(v) => query.bind(v),
318            Value::Uuid(v) => query.bind(v),
319            Value::Date(v) => query.bind(v),
320            Value::Time(v) => query.bind(v),
321            Value::Timestamp(v) => query.bind(v),
322            Value::Json(v) => query.bind(v),
323        })
324}
325
326async fn fetch_rows<'c, E>(ex: E, sql: &str, params: Vec<Value>) -> Result<QueryResult, OrmError>
327where
328    E: sqlx::Executor<'c, Database = sqlx::Postgres>,
329{
330    let rows = bind_all(sql, params)
331        .fetch_all(ex)
332        .await
333        .map_err(map_error)?;
334    let rows = rows.iter().map(decode_row).collect::<Result<_, _>>()?;
335    Ok(QueryResult { rows })
336}
337
338async fn execute_rows<'c, E>(ex: E, sql: &str, params: Vec<Value>) -> Result<u64, OrmError>
339where
340    E: sqlx::Executor<'c, Database = sqlx::Postgres>,
341{
342    let done = bind_all(sql, params).execute(ex).await.map_err(map_error)?;
343    Ok(done.rows_affected())
344}
345
346async fn run_write<'c, E>(
347    ex: E,
348    compiled: CompiledQuery,
349    returning: bool,
350) -> Result<ExecResult, OrmError>
351where
352    E: sqlx::Executor<'c, Database = sqlx::Postgres>,
353{
354    if returning {
355        let rows = fetch_rows(ex, &compiled.sql, compiled.params).await?.rows;
356        Ok(returning_result(rows))
357    } else {
358        Ok(affected_result(
359            execute_rows(ex, &compiled.sql, compiled.params).await?,
360        ))
361    }
362}
363
364fn script_result(done: Result<sqlx::postgres::PgQueryResult, sqlx::Error>) -> Result<(), OrmError> {
365    done.map(|_| ()).map_err(|e| map_error(e).into())
366}
367
368fn decode_row(row: &PgRow) -> Result<Row, QueryError> {
369    row.columns()
370        .iter()
371        .enumerate()
372        .map(|(i, col)| {
373            let name = col.name().to_owned();
374            let raw = row.try_get_raw(i).map_err(|e| QueryError::Decode {
375                column: name.clone(),
376                reason: e.to_string(),
377            })?;
378            let value = if raw.is_null() {
379                Value::Null
380            } else {
381                decode_value(row, i, raw).map_err(|reason| QueryError::Decode {
382                    column: name.clone(),
383                    reason,
384                })?
385            };
386            Ok((name, value))
387        })
388        .collect::<Result<Vec<_>, _>>()
389        .map(Row::new)
390}
391
392/// Decode a non-null column by its PostgreSQL type name.
393fn decode_value(row: &PgRow, i: usize, raw: PgValueRef<'_>) -> Result<Value, String> {
394    /// `try_get` as `$t`, converted with `$into`.
395    macro_rules! get {
396        ($t:ty => $into:expr) => {
397            row.try_get::<$t, _>(i)
398                .map($into)
399                .map_err(|e| e.to_string())
400        };
401    }
402    /// A one-dimensional array as a JSON array.
403    macro_rules! array {
404        ($t:ty) => {
405            get!(Vec<$t> => |items| Value::Json(JsonValue::Array(items.into_iter().map(JsonValue::from).collect())))
406        };
407    }
408    let type_info = raw.type_info();
409    match type_info.name() {
410        "BOOL" => get!(bool => Value::Bool),
411        "INT2" => get!(i16 => |v| Value::Int(i64::from(v))),
412        "INT4" => get!(i32 => |v| Value::Int(i64::from(v))),
413        "INT8" => get!(i64 => Value::Int),
414        "FLOAT4" => get!(f32 => |v| Value::Float(f64::from(v))),
415        "FLOAT8" => get!(f64 => Value::Float),
416        "NUMERIC" => get!(Decimal => Value::Decimal),
417        "TEXT" | "VARCHAR" | "BPCHAR" | "NAME" => get!(String => Value::Text),
418        "BYTEA" => get!(Vec<u8> => Value::Bytes),
419        "UUID" => get!(Uuid => Value::Uuid),
420        "DATE" => get!(NaiveDate => Value::Date),
421        "TIME" => get!(NaiveTime => Value::Time),
422        "TIMESTAMPTZ" => get!(DateTime<Utc> => Value::Timestamp),
423        "TIMESTAMP" => get!(NaiveDateTime => |v| Value::Timestamp(v.and_utc())),
424        "JSON" | "JSONB" => get!(JsonValue => Value::Json),
425        "INT2[]" => array!(i16),
426        "INT4[]" => array!(i32),
427        "INT8[]" => array!(i64),
428        "FLOAT8[]" => array!(f64),
429        "BOOL[]" => array!(bool),
430        "TEXT[]" | "VARCHAR[]" => array!(String),
431        other => Err(format!("unsupported PostgreSQL type {other}")),
432    }
433}
siderite_backends/postgres.rs

siderite_backends/
postgres.rs
