Skip to main content

persistant_sqlx/
lib.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2//! PersistAnt's relational driver over sqlx (PostgreSQL first).
3//!
4//! [`PgBackend`] implements [`persistant_relational::RelationalBackend`]. The concrete
5//! [`sqlx::PgPool`] is available through [`PgBackend::pool`], so a program that wants
6//! compile-time-checked `sqlx::query!` keeps it (D1, C+): this crate does not hide sqlx.
7//!
8//! Failures are classified from the server's SQLSTATE into [`persistant_relational::RelError`].
9
10use persistant_relational::{Dialect, RelError, RelationalBackend, RelationalTx, Row, Value};
11use sqlx::postgres::{PgArguments, PgPoolOptions, PgRow};
12use sqlx::query::Query;
13use sqlx::{Column, Postgres, Row as SqlxRow, TypeInfo, ValueRef};
14
15/// Classify an sqlx error by SQLSTATE (https://www.postgresql.org/docs/current/errcodes-appendix.html).
16fn classify(e: sqlx::Error) -> RelError {
17    match &e {
18        sqlx::Error::Database(db) => {
19            let constraint = db.constraint().map(str::to_string);
20            match db.code().as_deref() {
21                Some("23505") => RelError::UniqueViolation { constraint },
22                Some("23503") => RelError::ForeignKeyViolation { constraint },
23                Some("23514") | Some("23502") => RelError::CheckViolation { constraint },
24                Some("40001") => RelError::Serialization,
25                Some("40P01") => RelError::Deadlock,
26                Some(c) if c.starts_with("08") || c.starts_with("57P") => {
27                    RelError::Unavailable(e.to_string())
28                }
29                _ => RelError::Backend(e.to_string()),
30            }
31        }
32        sqlx::Error::Io(_) | sqlx::Error::PoolTimedOut | sqlx::Error::PoolClosed => {
33            RelError::Unavailable(e.to_string())
34        }
35        _ => RelError::Backend(e.to_string()),
36    }
37}
38
39fn bind<'q>(
40    mut q: Query<'q, Postgres, PgArguments>,
41    params: &'q [Value],
42) -> Result<Query<'q, Postgres, PgArguments>, RelError> {
43    for p in params {
44        q = match p {
45            // A typed NULL: PostgreSQL needs a type to infer parameter types; text works for the
46            // statements this layer issues. A caller needing another NULL type casts in SQL.
47            Value::Null => q.bind(Option::<String>::None),
48            Value::Bool(b) => q.bind(*b),
49            Value::Int(n) => q.bind(*n),
50            Value::Float(f) => q.bind(*f),
51            Value::Text(s) => q.bind(s.as_str()),
52            Value::Bytes(b) => q.bind(b.as_slice()),
53            other => {
54                return Err(RelError::Unsupported(format!(
55                    "parameter {other:?} has no PostgreSQL binding in this driver"
56                )));
57            }
58        };
59    }
60    Ok(q)
61}
62
63fn decode(row: &PgRow) -> Result<Row, RelError> {
64    let mut values = Vec::with_capacity(row.columns().len());
65    for (i, col) in row.columns().iter().enumerate() {
66        let raw = row.try_get_raw(i).map_err(classify)?;
67        if raw.is_null() {
68            values.push(Value::Null);
69            continue;
70        }
71        let name = col.type_info().name().to_string();
72        let v = match name.as_str() {
73            "BOOL" => Value::Bool(row.try_get::<bool, _>(i).map_err(classify)?),
74            "INT2" => Value::Int(i64::from(row.try_get::<i16, _>(i).map_err(classify)?)),
75            "INT4" => Value::Int(i64::from(row.try_get::<i32, _>(i).map_err(classify)?)),
76            "INT8" => Value::Int(row.try_get::<i64, _>(i).map_err(classify)?),
77            "FLOAT4" => Value::Float(f64::from(row.try_get::<f32, _>(i).map_err(classify)?)),
78            "FLOAT8" => Value::Float(row.try_get::<f64, _>(i).map_err(classify)?),
79            "TEXT" | "VARCHAR" | "NAME" | "BPCHAR" => {
80                Value::Text(row.try_get::<String, _>(i).map_err(classify)?)
81            }
82            "BYTEA" => Value::Bytes(row.try_get::<Vec<u8>, _>(i).map_err(classify)?),
83            other => {
84                return Err(RelError::Unsupported(format!(
85                    "column {:?} has type {other}, which Value cannot carry; use PgBackend::pool",
86                    col.name()
87                )));
88            }
89        };
90        values.push(v);
91    }
92    Ok(Row::new(values))
93}
94
95/// A PostgreSQL database behind an sqlx pool.
96#[derive(Debug, Clone)]
97pub struct PgBackend {
98    pool: sqlx::PgPool,
99    version: String,
100}
101
102impl PgBackend {
103    /// Connect with default pool settings and read the server version.
104    pub async fn connect(url: &str) -> Result<Self, RelError> {
105        let pool = PgPoolOptions::new()
106            .max_connections(16)
107            .connect(url)
108            .await
109            .map_err(classify)?;
110        Self::from_pool(pool).await
111    }
112
113    /// Wrap an existing pool (a program that already owns one keeps it).
114    pub async fn from_pool(pool: sqlx::PgPool) -> Result<Self, RelError> {
115        let full: String = sqlx::query_scalar("SHOW server_version")
116            .fetch_one(&pool)
117            .await
118            .map_err(classify)?;
119        // "16.15 (Debian 16.15-1.pgdg13+2)" -> "16.15"
120        let version = full.split_whitespace().next().unwrap_or("").to_string();
121        Ok(PgBackend { pool, version })
122    }
123
124    /// The concrete pool: for `sqlx::query!` and anything this layer does not carry.
125    pub fn pool(&self) -> &sqlx::PgPool {
126        &self.pool
127    }
128}
129
130impl RelationalBackend for PgBackend {
131    type Tx<'a> = PgTx<'a>;
132
133    fn backend(&self) -> &str {
134        "postgres"
135    }
136
137    fn version(&self) -> &str {
138        &self.version
139    }
140
141    fn dialect(&self) -> Dialect {
142        Dialect::Postgres
143    }
144
145    async fn execute(&self, sql: &str, params: &[Value]) -> Result<u64, RelError> {
146        let q = bind(sqlx::query(sqlx::AssertSqlSafe(sql.to_string())), params)?;
147        q.execute(&self.pool)
148            .await
149            .map(|r| r.rows_affected())
150            .map_err(classify)
151    }
152
153    async fn query(&self, sql: &str, params: &[Value]) -> Result<Vec<Row>, RelError> {
154        let q = bind(sqlx::query(sqlx::AssertSqlSafe(sql.to_string())), params)?;
155        let rows = q.fetch_all(&self.pool).await.map_err(classify)?;
156        rows.iter().map(decode).collect()
157    }
158
159    async fn begin(&self) -> Result<PgTx<'_>, RelError> {
160        Ok(PgTx {
161            tx: self.pool.begin().await.map_err(classify)?,
162        })
163    }
164}
165
166/// A PostgreSQL transaction.
167#[derive(Debug)]
168pub struct PgTx<'a> {
169    tx: sqlx::Transaction<'a, Postgres>,
170}
171
172impl RelationalTx for PgTx<'_> {
173    async fn execute(&mut self, sql: &str, params: &[Value]) -> Result<u64, RelError> {
174        let q = bind(sqlx::query(sqlx::AssertSqlSafe(sql.to_string())), params)?;
175        q.execute(&mut *self.tx)
176            .await
177            .map(|r| r.rows_affected())
178            .map_err(classify)
179    }
180
181    async fn query(&mut self, sql: &str, params: &[Value]) -> Result<Vec<Row>, RelError> {
182        let q = bind(sqlx::query(sqlx::AssertSqlSafe(sql.to_string())), params)?;
183        let rows = q.fetch_all(&mut *self.tx).await.map_err(classify)?;
184        rows.iter().map(decode).collect()
185    }
186
187    async fn commit(self) -> Result<(), RelError> {
188        self.tx.commit().await.map_err(classify)
189    }
190
191    async fn rollback(self) -> Result<(), RelError> {
192        self.tx.rollback().await.map_err(classify)
193    }
194}