runlight_sqlx/
postgres.rs

1//! Postgres, on the app's pool. Each statement runs on the pool on its own, so the store's writes never
2//! join a transaction the app has open; a transaction or a lock holds one connection until it ends.
3
4use std::sync::Arc;
5
6use runlight::BoxFuture;
7use runlight::store::{Db, DbError, Dialect, Hold, Param, Row, fill_placeholders, postgres_literal};
8use sqlx::pool::PoolConnection;
9use sqlx::{AssertSqlSafe, PgPool, Postgres};
10use tokio::sync::Mutex;
11
12use crate::rows::{db_error, pg_row};
13
14/// Arbitrary but fixed, so every Runlight process takes the same lock to create tables.
15const MIGRATION_LOCK: i64 = 7_331_906;
16
17/// The store's driver over a Postgres pool. Values go in the statement's text as literals of unknown
18/// type, so Postgres infers each one's type as it does for node-postgres's untyped parameters.
19#[derive(Clone)]
20pub struct PostgresDb {
21    pool: PgPool,
22}
23
24impl PostgresDb {
25    /// The store's driver over a pool.
26    pub fn new(pool: PgPool) -> PostgresDb {
27        PostgresDb { pool }
28    }
29}
30
31fn text(sql: &str, params: &[Param]) -> Result<AssertSqlSafe<String>, DbError> {
32    Ok(AssertSqlSafe(fill_placeholders(sql, params, false, postgres_literal)?))
33}
34
35impl Db for PostgresDb {
36    fn dialect(&self) -> Dialect {
37        Dialect::Postgres
38    }
39
40    fn all<'a>(&'a self, sql: &'a str, params: Vec<Param>) -> BoxFuture<'a, Result<Vec<Row>, DbError>> {
41        Box::pin(async move {
42            let rows = sqlx::raw_sql(text(sql, &params)?).fetch_all(&self.pool).await.map_err(db_error)?;
43            rows.iter().map(pg_row).collect()
44        })
45    }
46
47    fn run<'a>(&'a self, sql: &'a str, params: Vec<Param>) -> BoxFuture<'a, Result<u64, DbError>> {
48        Box::pin(async move {
49            Ok(sqlx::raw_sql(text(sql, &params)?).execute(&self.pool).await.map_err(db_error)?.rows_affected())
50        })
51    }
52
53    fn hold(&self, kind: Hold) -> BoxFuture<'_, Result<Arc<dyn Db>, DbError>> {
54        Box::pin(async move {
55            // In its slot from the start, so a request cancelled part way closes the connection rather than
56            // give it back to the pool in a transaction or with the lock taken.
57            let mut slot = Slot { conn: Some(self.pool.acquire().await.map_err(db_error)?), nested: false };
58            let conn = slot.conn.as_mut().ok_or_else(ended)?;
59            match kind {
60                Hold::Transaction => {
61                    sqlx::raw_sql(AssertSqlSafe("BEGIN".to_string())).execute(&mut **conn).await.map_err(db_error)?;
62                }
63                Hold::Exclusive => loop {
64                    // Asked for again and again rather than waited on: a waiting statement would hold up an
65                    // index being built CONCURRENTLY by whoever has the lock, and the two would wait on each
66                    // other for good.
67                    let rows =
68                        sqlx::raw_sql(AssertSqlSafe(format!("SELECT pg_try_advisory_lock({MIGRATION_LOCK}) AS ok")))
69                            .fetch_all(&mut **conn)
70                            .await
71                            .map_err(db_error)?;
72                    let ok = rows.first().map(pg_row).transpose()?.is_some_and(|r| r.num("ok") == 1.0);
73                    if ok {
74                        break;
75                    }
76                    tokio::time::sleep(std::time::Duration::from_millis(100)).await;
77                },
78            }
79            Ok(Arc::new(Held { conn: Arc::new(Mutex::new(slot)), kind, role: Role::Owner }) as Arc<dyn Db>)
80        })
81    }
82}
83
84/// The held connection. Let go of with it still in (a request cancelled before `finish`), it closes the
85/// connection, so no transaction or lock goes back to the pool.
86struct Slot {
87    conn: Option<PoolConnection<Postgres>>,
88    /// A transaction begun inside the lock that has not ended.
89    nested: bool,
90}
91
92impl Slot {
93    fn close(&mut self) {
94        if let Some(conn) = self.conn.take() {
95            drop(conn.detach());
96        }
97    }
98}
99
100impl Drop for Slot {
101    fn drop(&mut self) {
102        self.close();
103    }
104}
105
106/// One connection held for a transaction or a lock.
107struct Held {
108    conn: Arc<Mutex<Slot>>,
109    kind: Hold,
110    role: Role,
111}
112
113impl Drop for Held {
114    fn drop(&mut self) {
115        // Its owner gone without `finish`: closed now rather than when the last joined handle goes.
116        if self.role == Role::Owner
117            && let Ok(mut slot) = self.conn.try_lock()
118        {
119            slot.close();
120        }
121    }
122}
123
124/// What finishing a held connection does.
125#[derive(Clone, Copy, PartialEq, Eq)]
126enum Role {
127    /// Ends the transaction or lets the lock go, and gives the connection back.
128    Owner,
129    /// A transaction begun inside a lock: ends the transaction, and keeps the connection.
130    Nested,
131    /// Joined what was already held: ends nothing.
132    Joined,
133}
134
135fn ended() -> DbError {
136    DbError("Runlight: the held connection has been let go".into())
137}
138
139impl Db for Held {
140    fn dialect(&self) -> Dialect {
141        Dialect::Postgres
142    }
143
144    fn all<'a>(&'a self, sql: &'a str, params: Vec<Param>) -> BoxFuture<'a, Result<Vec<Row>, DbError>> {
145        Box::pin(async move {
146            let mut slot = self.conn.lock().await;
147            let conn = slot.conn.as_mut().ok_or_else(ended)?;
148            let rows = sqlx::raw_sql(text(sql, &params)?).fetch_all(&mut **conn).await.map_err(db_error)?;
149            rows.iter().map(pg_row).collect()
150        })
151    }
152
153    fn run<'a>(&'a self, sql: &'a str, params: Vec<Param>) -> BoxFuture<'a, Result<u64, DbError>> {
154        Box::pin(async move {
155            let mut slot = self.conn.lock().await;
156            let conn = slot.conn.as_mut().ok_or_else(ended)?;
157            Ok(sqlx::raw_sql(text(sql, &params)?).execute(&mut **conn).await.map_err(db_error)?.rows_affected())
158        })
159    }
160
161    fn hold(&self, kind: Hold) -> BoxFuture<'_, Result<Arc<dyn Db>, DbError>> {
162        Box::pin(async move {
163            // A transaction inside the lock runs on the locked connection; anything else joins what is held.
164            if kind == Hold::Transaction && self.kind == Hold::Exclusive {
165                let mut slot = self.conn.lock().await;
166                let slot = &mut *slot;
167                let conn = slot.conn.as_mut().ok_or_else(ended)?;
168                // Marked before it is sent: a BEGIN cancelled part way may still have run.
169                slot.nested = true;
170                sqlx::raw_sql(AssertSqlSafe("BEGIN".to_string())).execute(&mut **conn).await.map_err(db_error)?;
171                return Ok(Arc::new(Held { conn: self.conn.clone(), kind, role: Role::Nested }) as Arc<dyn Db>);
172            }
173            Ok(Arc::new(Held { conn: self.conn.clone(), kind: self.kind, role: Role::Joined }) as Arc<dyn Db>)
174        })
175    }
176
177    fn finish(&self, ok: bool) -> BoxFuture<'_, Result<(), DbError>> {
178        Box::pin(async move {
179            let mut slot = self.conn.lock().await;
180            if self.role == Role::Joined {
181                return Ok(());
182            }
183            if self.role == Role::Nested {
184                let slot = &mut *slot;
185                let Some(conn) = slot.conn.as_mut() else { return Ok(()) };
186                let end = if ok { "COMMIT" } else { "ROLLBACK" };
187                sqlx::raw_sql(AssertSqlSafe(end.to_string())).execute(&mut **conn).await.map_err(db_error)?;
188                slot.nested = false;
189                return Ok(());
190            }
191            // A transaction inside the lock that never ended: closing the connection rolls it back and lets
192            // the lock go.
193            if slot.nested {
194                slot.close();
195                return Ok(());
196            }
197            let Some(mut conn) = slot.conn.take() else { return Ok(()) };
198            let result = match self.kind {
199                Hold::Transaction => {
200                    let end = if ok { "COMMIT" } else { "ROLLBACK" };
201                    sqlx::raw_sql(AssertSqlSafe(end.to_string()))
202                        .execute(&mut *conn)
203                        .await
204                        .map(|_| ())
205                        .map_err(db_error)
206                }
207                // A connection whose unlock failed may still have the lock, so it is closed below, and a
208                // lost connection ends its session, and the lock with it.
209                Hold::Exclusive => sqlx::raw_sql(AssertSqlSafe(format!("SELECT pg_advisory_unlock({MIGRATION_LOCK})")))
210                    .execute(&mut *conn)
211                    .await
212                    .map(|_| ())
213                    .map_err(db_error),
214            };
215            if result.is_err() {
216                // A connection still in its transaction or holding the lock must never go back to the pool.
217                drop(conn.detach());
218            }
219            if self.kind == Hold::Exclusive { Ok(()) } else { result }
220        })
221    }
222}
runlight_sqlx/postgres.rs

runlight_sqlx/
postgres.rs
