axa_sqlx/
lib.rs

1//! `PostgreSQL` pooling, migrations, and test resets over sqlx.
2//!
3//! This crate holds the operations every service repeats. It does not wrap
4//! queries: a service writes its own SQL with `sqlx::query!`, so a reviewer
5//! reads exactly what runs.
6//!
7//! ```no_run
8//! use axa_sqlx::{PoolConfig, connect, migrate, sqlx};
9//!
10//! static MIGRATOR: sqlx::migrate::Migrator = sqlx::migrate::Migrator::DEFAULT;
11//!
12//! # async fn run() -> Result<(), Box<dyn std::error::Error>> {
13//! let pool = connect(&PoolConfig::new("postgres://localhost/app")).await?;
14//! migrate(&pool, &MIGRATOR).await?;
15//! # Ok(())
16//! # }
17//! ```
18
19use std::{
20    ops::{Deref, DerefMut},
21    time::Duration,
22};
23
24pub use sqlx;
25use sqlx::{
26    AssertSqlSafe, PgConnection, PgPool, Postgres, Row,
27    migrate::{MigrateError, Migrator},
28    pool::PoolConnection,
29    postgres::PgPoolOptions,
30};
31
32/// Pool settings.
33#[derive(Clone, PartialEq, Eq)]
34pub struct PoolConfig {
35    /// Connection string.
36    pub uri: String,
37    /// Connections kept open.
38    pub min_connections: u32,
39    /// Largest number of connections.
40    pub max_connections: u32,
41    /// How long to wait for a connection, in milliseconds.
42    pub connect_timeout_ms: u64,
43    /// How long an idle connection stays open, in milliseconds.
44    pub idle_timeout_ms: u64,
45}
46
47impl std::fmt::Debug for PoolConfig {
48    /// Leaves the connection string out, because it can carry a password.
49    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
50        f.debug_struct("PoolConfig")
51            .field("min_connections", &self.min_connections)
52            .field("max_connections", &self.max_connections)
53            .field("connect_timeout_ms", &self.connect_timeout_ms)
54            .field("idle_timeout_ms", &self.idle_timeout_ms)
55            .finish_non_exhaustive()
56    }
57}
58
59impl PoolConfig {
60    /// Settings for `uri` with a small pool.
61    #[must_use]
62    pub fn new(uri: impl Into<String>) -> Self {
63        Self {
64            uri: uri.into(),
65            min_connections: 1,
66            max_connections: 10,
67            connect_timeout_ms: 5_000,
68            idle_timeout_ms: 600_000,
69        }
70    }
71
72    /// The sqlx options these settings describe.
73    #[must_use]
74    pub fn options(&self) -> PgPoolOptions {
75        PgPoolOptions::new()
76            .min_connections(self.min_connections)
77            .max_connections(self.max_connections)
78            .acquire_timeout(Duration::from_millis(self.connect_timeout_ms))
79            .idle_timeout(Duration::from_millis(self.idle_timeout_ms))
80    }
81}
82
83/// Open a pool and prove it works with one connection.
84///
85/// # Errors
86///
87/// Fails when the database cannot be reached.
88pub async fn connect(config: &PoolConfig) -> Result<PgPool, sqlx::Error> {
89    config.options().connect(&config.uri).await
90}
91
92/// Check that the database answers.
93///
94/// # Errors
95///
96/// Fails when it does not.
97pub async fn ping(pool: &PgPool) -> Result<(), sqlx::Error> {
98    sqlx::query("SELECT 1").execute(pool).await?;
99    Ok(())
100}
101
102/// Where a function gets its connection: the pool, or a connection the caller
103/// holds, usually an open transaction, so one function serves both.
104///
105/// `impl sqlx::Acquire` says the same, but its associated types stop the future
106/// being `Send`, which Axum handlers and `tokio::spawn` need.
107///
108/// ```no_run
109/// use axa_sqlx::{Db, sqlx};
110///
111/// async fn count(db: impl Into<Db<'_>>) -> Result<i64, sqlx::Error> {
112///     let mut conn = db.into().acquire().await?;
113///     sqlx::query_scalar("SELECT count(*) FROM items")
114///         .fetch_one(&mut *conn)
115///         .await
116/// }
117/// ```
118pub enum Db<'a> {
119    /// Take a connection from the pool for the length of the call.
120    Pool(&'a PgPool),
121    /// Use the caller's connection.
122    Conn(&'a mut PgConnection),
123}
124
125impl<'a> From<&'a PgPool> for Db<'a> {
126    fn from(pool: &'a PgPool) -> Self {
127        Self::Pool(pool)
128    }
129}
130
131impl<'a> From<&'a mut PgConnection> for Db<'a> {
132    fn from(conn: &'a mut PgConnection) -> Self {
133        Self::Conn(conn)
134    }
135}
136
137impl<'a> Db<'a> {
138    /// The connection to use.
139    ///
140    /// # Errors
141    ///
142    /// Fails when the pool has none to give.
143    pub async fn acquire(self) -> Result<Conn<'a>, sqlx::Error> {
144        match self {
145            Self::Pool(pool) => Ok(Conn::Pooled(pool.acquire().await?)),
146            Self::Conn(conn) => Ok(Conn::Borrowed(conn)),
147        }
148    }
149}
150
151/// A connection for the length of one call. It derefs to [`PgConnection`].
152pub enum Conn<'a> {
153    /// Taken from the pool, and given back when dropped.
154    Pooled(PoolConnection<Postgres>),
155    /// The caller's.
156    Borrowed(&'a mut PgConnection),
157}
158
159impl Deref for Conn<'_> {
160    type Target = PgConnection;
161
162    fn deref(&self) -> &PgConnection {
163        match self {
164            Self::Pooled(conn) => conn,
165            Self::Borrowed(conn) => conn,
166        }
167    }
168}
169
170impl DerefMut for Conn<'_> {
171    fn deref_mut(&mut self) -> &mut PgConnection {
172        match self {
173            Self::Pooled(conn) => conn,
174            Self::Borrowed(conn) => conn,
175        }
176    }
177}
178
179/// Whether an error is a violated unique constraint, which a service usually
180/// reports as `409` rather than `500`.
181#[must_use]
182pub fn is_unique_violation(error: &sqlx::Error) -> bool {
183    error
184        .as_database_error()
185        .is_some_and(sqlx::error::DatabaseError::is_unique_violation)
186}
187
188/// Quote a table or type name for use in SQL.
189///
190/// Identifiers cannot be bound as parameters, so this is the one place names
191/// reach SQL as text. Only lowercase ASCII letters, digits, and underscores
192/// are accepted, starting with a letter or underscore. That is stricter than
193/// Postgres, and it leaves nothing for an attacker to work with.
194///
195/// ```
196/// assert_eq!(axa_sqlx::quote_identifier("users").as_deref(), Some("\"users\""));
197/// assert_eq!(axa_sqlx::quote_identifier("users; DROP TABLE x"), None);
198/// ```
199#[must_use]
200pub fn quote_identifier(name: &str) -> Option<String> {
201    let mut chars = name.chars();
202    let first = chars.next()?;
203    let valid = (first.is_ascii_lowercase() || first == '_')
204        && chars.all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '_');
205    valid.then(|| format!("\"{name}\""))
206}
207
208/// Quote every name, failing on the first one [`quote_identifier`] refuses.
209fn quote_all<S: AsRef<str>>(names: &[S]) -> Result<Vec<String>, sqlx::Error> {
210    names
211        .iter()
212        .map(|name| {
213            quote_identifier(name.as_ref()).ok_or_else(|| {
214                sqlx::Error::Protocol(format!("invalid identifier: {}", name.as_ref()))
215            })
216        })
217        .collect()
218}
219
220/// Apply pending migrations. Sqlx holds a lock while it does, so services
221/// booting together take turns.
222///
223/// # Errors
224///
225/// Fails when a migration fails, or an applied one no longer matches.
226pub async fn migrate(pool: &PgPool, migrator: &Migrator) -> Result<(), MigrateError> {
227    migrator.run(pool).await
228}
229
230/// The newest applied versions, newest first. Empty on a database that has
231/// never been migrated.
232async fn applied_versions(pool: &PgPool) -> Result<Vec<i64>, sqlx::Error> {
233    let exists: bool = sqlx::query("SELECT to_regclass('_sqlx_migrations') IS NOT NULL AS found")
234        .fetch_one(pool)
235        .await?
236        .try_get("found")?;
237    if !exists {
238        return Ok(Vec::new());
239    }
240    let rows =
241        sqlx::query("SELECT version FROM _sqlx_migrations WHERE success ORDER BY version DESC")
242            .fetch_all(pool)
243            .await?;
244    rows.iter().map(|row| row.try_get("version")).collect()
245}
246
247/// Revert the newest `steps` migrations. Each needs a `.down.sql` file.
248///
249/// # Errors
250///
251/// Fails on a database error or a failed down migration.
252pub async fn revert(pool: &PgPool, migrator: &Migrator, steps: u32) -> Result<(), MigrateError> {
253    let applied = applied_versions(pool).await?;
254    let steps = usize::try_from(steps).unwrap_or(usize::MAX);
255    // `undo` reverts everything newer than its target, so the target is the
256    // version just past the last one to revert, or 0 to revert them all.
257    let target = applied.get(steps).copied().unwrap_or(0);
258    migrator.undo(pool, target).await
259}
260
261/// One migration and whether it has run.
262#[derive(Debug, Clone, PartialEq, Eq)]
263pub struct MigrationStatus {
264    /// Its version.
265    pub version: i64,
266    /// Its description.
267    pub description: String,
268    /// Whether it has been applied.
269    pub applied: bool,
270}
271
272/// Every migration in `migrator`, oldest first, with whether it has run.
273///
274/// # Errors
275///
276/// Fails on a database error.
277pub async fn status(
278    pool: &PgPool,
279    migrator: &Migrator,
280) -> Result<Vec<MigrationStatus>, sqlx::Error> {
281    let applied = applied_versions(pool).await?;
282    Ok(migrator
283        .iter()
284        .filter(|migration| !migration.migration_type.is_down_migration())
285        .map(|migration| MigrationStatus {
286            version: migration.version,
287            description: migration.description.to_string(),
288            applied: applied.contains(&migration.version),
289        })
290        .collect())
291}
292
293/// Empty `tables`, restarting their sequences. Each name goes through
294/// [`quote_identifier`].
295///
296/// # Errors
297///
298/// Fails on an invalid name or a database error.
299pub async fn truncate<S: AsRef<str>>(pool: &PgPool, tables: &[S]) -> Result<(), sqlx::Error> {
300    if tables.is_empty() {
301        return Ok(());
302    }
303    let tables = quote_all(tables)?.join(", ");
304    let sql = format!("TRUNCATE {tables} RESTART IDENTITY CASCADE");
305    sqlx::query(AssertSqlSafe(sql)).execute(pool).await?;
306    Ok(())
307}
308
309/// Drop every table and enum type in `public`, then migrate from nothing.
310/// This gives each test an empty schema, whatever tables it touched.
311///
312/// **This destroys all data.** It exists for tests, and a service must never
313/// reach it outside the `test` environment.
314///
315/// It drops tables and types by name rather than dropping the schema, because
316/// `DROP SCHEMA public CASCADE` would also remove the extensions a database
317/// was prepared with, such as `citext`. Functions are left in place, so a
318/// migration that creates one should use `CREATE OR REPLACE`.
319///
320/// # Errors
321///
322/// Fails on a database error, or a name [`quote_identifier`] refuses.
323pub async fn dangerously_recreate(pool: &PgPool, migrator: &Migrator) -> Result<(), MigrateError> {
324    let tables: Vec<String> =
325        sqlx::query("SELECT tablename::text AS name FROM pg_tables WHERE schemaname = 'public'")
326            .fetch_all(pool)
327            .await?
328            .iter()
329            .map(|row| row.try_get("name"))
330            .collect::<Result<_, _>>()?;
331    // Types owned by an extension (`deptype = 'e'`) belong to that extension.
332    let types: Vec<String> = sqlx::query(
333        "SELECT t.typname::text AS name FROM pg_type t
334         JOIN pg_namespace n ON n.oid = t.typnamespace
335         WHERE t.typtype = 'e' AND n.nspname = 'public'
336           AND NOT EXISTS (
337               SELECT 1 FROM pg_depend d WHERE d.objid = t.oid AND d.deptype = 'e'
338           )",
339    )
340    .fetch_all(pool)
341    .await?
342    .iter()
343    .map(|row| row.try_get("name"))
344    .collect::<Result<_, _>>()?;
345
346    if !tables.is_empty() {
347        let names = quote_all(&tables)?.join(", ");
348        let sql = format!("DROP TABLE {names} CASCADE");
349        sqlx::query(AssertSqlSafe(sql)).execute(pool).await?;
350    }
351    for name in quote_all(&types)? {
352        let sql = format!("DROP TYPE {name} CASCADE");
353        sqlx::query(AssertSqlSafe(sql)).execute(pool).await?;
354    }
355    migrate(pool, migrator).await
356}
357
358#[cfg(test)]
359mod tests {
360    use super::*;
361
362    #[test]
363    fn debug_text_leaves_out_the_connection_string() {
364        let text = format!("{:?}", PoolConfig::new("postgres://me:hunter2@db/app"));
365        assert!(!text.contains("hunter2"), "{text}");
366        assert!(text.contains("max_connections"), "{text}");
367    }
368
369    #[test]
370    fn identifiers_are_quoted() {
371        assert_eq!(quote_identifier("users").as_deref(), Some("\"users\""));
372        assert_eq!(quote_identifier("_a1").as_deref(), Some("\"_a1\""));
373    }
374
375    #[test]
376    fn unsafe_identifiers_are_refused() {
377        for name in [
378            "", "1users", "Users", "a b", "a\"b", "a;b", "a.b", "users--", "é",
379        ] {
380            assert_eq!(quote_identifier(name), None, "{name:?}");
381        }
382    }
383
384    #[test]
385    fn quote_all_fails_on_any_bad_name() {
386        assert!(quote_all(&["good", "also_good"]).is_ok());
387        assert!(quote_all(&["good", "bad name"]).is_err());
388    }
389
390    #[test]
391    fn pool_settings_reach_the_options() {
392        let config = PoolConfig {
393            min_connections: 2,
394            max_connections: 7,
395            connect_timeout_ms: 1_500,
396            idle_timeout_ms: 9_000,
397            ..PoolConfig::new("postgres://x")
398        };
399        let options = config.options();
400        assert_eq!(options.get_min_connections(), 2);
401        assert_eq!(options.get_max_connections(), 7);
402        assert_eq!(options.get_acquire_timeout(), Duration::from_millis(1_500));
403        assert_eq!(options.get_idle_timeout(), Some(Duration::from_secs(9)));
404    }
405
406    #[test]
407    fn errors_that_are_not_from_the_database_are_not_unique_violations() {
408        assert!(!is_unique_violation(&sqlx::Error::RowNotFound));
409    }
410}
axa_sqlx/lib.rs

axa_sqlx/
lib.rs
