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::time::Duration;
20
21pub use sqlx;
22use sqlx::{
23    AssertSqlSafe, PgPool, Row,
24    migrate::{MigrateError, Migrator},
25    postgres::PgPoolOptions,
26};
27
28/// Pool settings.
29#[derive(Debug, Clone, PartialEq, Eq)]
30pub struct PoolConfig {
31    /// Connection string.
32    pub uri: String,
33    /// Connections kept open.
34    pub min_connections: u32,
35    /// Largest number of connections.
36    pub max_connections: u32,
37    /// How long to wait for a connection, in milliseconds.
38    pub connect_timeout_ms: u64,
39    /// How long an idle connection stays open, in milliseconds.
40    pub idle_timeout_ms: u64,
41}
42
43impl PoolConfig {
44    /// Settings for `uri` with a small pool.
45    #[must_use]
46    pub fn new(uri: impl Into<String>) -> Self {
47        Self {
48            uri: uri.into(),
49            min_connections: 1,
50            max_connections: 10,
51            connect_timeout_ms: 5_000,
52            idle_timeout_ms: 600_000,
53        }
54    }
55
56    /// The sqlx options these settings describe.
57    #[must_use]
58    pub fn options(&self) -> PgPoolOptions {
59        PgPoolOptions::new()
60            .min_connections(self.min_connections)
61            .max_connections(self.max_connections)
62            .acquire_timeout(Duration::from_millis(self.connect_timeout_ms))
63            .idle_timeout(Duration::from_millis(self.idle_timeout_ms))
64    }
65}
66
67/// Open a pool and prove it works with one connection.
68///
69/// # Errors
70///
71/// Fails when the database cannot be reached.
72pub async fn connect(config: &PoolConfig) -> Result<PgPool, sqlx::Error> {
73    config.options().connect(&config.uri).await
74}
75
76/// Check that the database answers.
77///
78/// # Errors
79///
80/// Fails when it does not.
81pub async fn ping(pool: &PgPool) -> Result<(), sqlx::Error> {
82    sqlx::query("SELECT 1").execute(pool).await?;
83    Ok(())
84}
85
86/// Whether an error is a violated unique constraint, which a service usually
87/// reports as `409` rather than `500`.
88#[must_use]
89pub fn is_unique_violation(error: &sqlx::Error) -> bool {
90    error
91        .as_database_error()
92        .is_some_and(sqlx::error::DatabaseError::is_unique_violation)
93}
94
95/// Quote a table or type name for use in SQL.
96///
97/// Identifiers cannot be bound as parameters, so this is the one place names
98/// reach SQL as text. Only lowercase ASCII letters, digits, and underscores
99/// are accepted, starting with a letter or underscore. That is stricter than
100/// Postgres, and it leaves nothing for an attacker to work with.
101///
102/// ```
103/// assert_eq!(axa_sqlx::quote_identifier("users").as_deref(), Some("\"users\""));
104/// assert_eq!(axa_sqlx::quote_identifier("users; DROP TABLE x"), None);
105/// ```
106#[must_use]
107pub fn quote_identifier(name: &str) -> Option<String> {
108    let mut chars = name.chars();
109    let first = chars.next()?;
110    let valid = (first.is_ascii_lowercase() || first == '_')
111        && chars.all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '_');
112    valid.then(|| format!("\"{name}\""))
113}
114
115/// Quote every name, failing on the first one [`quote_identifier`] refuses.
116fn quote_all<S: AsRef<str>>(names: &[S]) -> Result<Vec<String>, sqlx::Error> {
117    names
118        .iter()
119        .map(|name| {
120            quote_identifier(name.as_ref()).ok_or_else(|| {
121                sqlx::Error::Protocol(format!("invalid identifier: {}", name.as_ref()))
122            })
123        })
124        .collect()
125}
126
127/// Apply pending migrations. Sqlx holds a lock while it does, so services
128/// booting together take turns.
129///
130/// # Errors
131///
132/// Fails when a migration fails, or an applied one no longer matches.
133pub async fn migrate(pool: &PgPool, migrator: &Migrator) -> Result<(), MigrateError> {
134    migrator.run(pool).await
135}
136
137/// The newest applied versions, newest first. Empty on a database that has
138/// never been migrated.
139async fn applied_versions(pool: &PgPool) -> Result<Vec<i64>, sqlx::Error> {
140    let exists: bool = sqlx::query("SELECT to_regclass('_sqlx_migrations') IS NOT NULL AS found")
141        .fetch_one(pool)
142        .await?
143        .try_get("found")?;
144    if !exists {
145        return Ok(Vec::new());
146    }
147    let rows =
148        sqlx::query("SELECT version FROM _sqlx_migrations WHERE success ORDER BY version DESC")
149            .fetch_all(pool)
150            .await?;
151    rows.iter().map(|row| row.try_get("version")).collect()
152}
153
154/// Revert the newest `steps` migrations. Each needs a `.down.sql` file.
155///
156/// # Errors
157///
158/// Fails on a database error or a failed down migration.
159pub async fn revert(pool: &PgPool, migrator: &Migrator, steps: u32) -> Result<(), MigrateError> {
160    let applied = applied_versions(pool).await?;
161    let steps = usize::try_from(steps).unwrap_or(usize::MAX);
162    // `undo` reverts everything newer than its target, so the target is the
163    // version just past the last one to revert, or 0 to revert them all.
164    let target = applied.get(steps).copied().unwrap_or(0);
165    migrator.undo(pool, target).await
166}
167
168/// One migration and whether it has run.
169#[derive(Debug, Clone, PartialEq, Eq)]
170pub struct MigrationStatus {
171    /// Its version.
172    pub version: i64,
173    /// Its description.
174    pub description: String,
175    /// Whether it has been applied.
176    pub applied: bool,
177}
178
179/// Every migration in `migrator`, oldest first, with whether it has run.
180///
181/// # Errors
182///
183/// Fails on a database error.
184pub async fn status(
185    pool: &PgPool,
186    migrator: &Migrator,
187) -> Result<Vec<MigrationStatus>, sqlx::Error> {
188    let applied = applied_versions(pool).await?;
189    Ok(migrator
190        .iter()
191        .filter(|migration| !migration.migration_type.is_down_migration())
192        .map(|migration| MigrationStatus {
193            version: migration.version,
194            description: migration.description.to_string(),
195            applied: applied.contains(&migration.version),
196        })
197        .collect())
198}
199
200/// Empty `tables`, restarting their sequences. Each name goes through
201/// [`quote_identifier`].
202///
203/// # Errors
204///
205/// Fails on an invalid name or a database error.
206pub async fn truncate<S: AsRef<str>>(pool: &PgPool, tables: &[S]) -> Result<(), sqlx::Error> {
207    if tables.is_empty() {
208        return Ok(());
209    }
210    let tables = quote_all(tables)?.join(", ");
211    let sql = format!("TRUNCATE {tables} RESTART IDENTITY CASCADE");
212    sqlx::query(AssertSqlSafe(sql)).execute(pool).await?;
213    Ok(())
214}
215
216/// Drop every table and enum type in `public`, then migrate from nothing.
217/// This gives each test an empty schema, whatever tables it touched.
218///
219/// **This destroys all data.** It exists for tests, and a service must never
220/// reach it outside the `test` environment.
221///
222/// It drops tables and types by name rather than dropping the schema, because
223/// `DROP SCHEMA public CASCADE` would also remove the extensions a database
224/// was prepared with, such as `citext`. Functions are left in place, so a
225/// migration that creates one should use `CREATE OR REPLACE`.
226///
227/// # Errors
228///
229/// Fails on a database error, or a name [`quote_identifier`] refuses.
230pub async fn dangerously_recreate(pool: &PgPool, migrator: &Migrator) -> Result<(), MigrateError> {
231    let tables: Vec<String> =
232        sqlx::query("SELECT tablename::text AS name FROM pg_tables WHERE schemaname = 'public'")
233            .fetch_all(pool)
234            .await?
235            .iter()
236            .map(|row| row.try_get("name"))
237            .collect::<Result<_, _>>()?;
238    // Types owned by an extension (`deptype = 'e'`) belong to that extension.
239    let types: Vec<String> = sqlx::query(
240        "SELECT t.typname::text AS name FROM pg_type t
241         JOIN pg_namespace n ON n.oid = t.typnamespace
242         WHERE t.typtype = 'e' AND n.nspname = 'public'
243           AND NOT EXISTS (
244               SELECT 1 FROM pg_depend d WHERE d.objid = t.oid AND d.deptype = 'e'
245           )",
246    )
247    .fetch_all(pool)
248    .await?
249    .iter()
250    .map(|row| row.try_get("name"))
251    .collect::<Result<_, _>>()?;
252
253    if !tables.is_empty() {
254        let names = quote_all(&tables)?.join(", ");
255        let sql = format!("DROP TABLE {names} CASCADE");
256        sqlx::query(AssertSqlSafe(sql)).execute(pool).await?;
257    }
258    for name in quote_all(&types)? {
259        let sql = format!("DROP TYPE {name} CASCADE");
260        sqlx::query(AssertSqlSafe(sql)).execute(pool).await?;
261    }
262    migrate(pool, migrator).await
263}
264
265#[cfg(test)]
266mod tests {
267    use super::*;
268
269    #[test]
270    fn identifiers_are_quoted() {
271        assert_eq!(quote_identifier("users").as_deref(), Some("\"users\""));
272        assert_eq!(quote_identifier("_a1").as_deref(), Some("\"_a1\""));
273    }
274
275    #[test]
276    fn unsafe_identifiers_are_refused() {
277        for name in [
278            "", "1users", "Users", "a b", "a\"b", "a;b", "a.b", "users--", "é",
279        ] {
280            assert_eq!(quote_identifier(name), None, "{name:?}");
281        }
282    }
283
284    #[test]
285    fn quote_all_fails_on_any_bad_name() {
286        assert!(quote_all(&["good", "also_good"]).is_ok());
287        assert!(quote_all(&["good", "bad name"]).is_err());
288    }
289
290    #[test]
291    fn pool_settings_reach_the_options() {
292        let config = PoolConfig {
293            min_connections: 2,
294            max_connections: 7,
295            connect_timeout_ms: 1_500,
296            idle_timeout_ms: 9_000,
297            ..PoolConfig::new("postgres://x")
298        };
299        let options = config.options();
300        assert_eq!(options.get_min_connections(), 2);
301        assert_eq!(options.get_max_connections(), 7);
302        assert_eq!(options.get_acquire_timeout(), Duration::from_millis(1_500));
303        assert_eq!(options.get_idle_timeout(), Some(Duration::from_secs(9)));
304    }
305
306    #[test]
307    fn errors_that_are_not_from_the_database_are_not_unique_violations() {
308        assert!(!is_unique_violation(&sqlx::Error::RowNotFound));
309    }
310}
axa_sqlx/lib.rs

axa_sqlx/
lib.rs
