server/
db.rs

1use crate::config::ConfigError;
2use cfg_if::cfg_if;
3use sqlx::migrate::MigrateError;
4use sqlx::pool::PoolConnection;
5use sqlx::{PgPool, Postgres};
6use thiserror::Error;
7cfg_if! {
8    if #[cfg(test)] {
9    use std::cell::RefCell;
10    } else if #[cfg(feature = "test-utils")] {
11    use std::cell::RefCell;
12    use sqlx::postgres::PgPoolOptions;
13    use std::env::var;
14    use std::sync::LazyLock;
15    use std::time::Duration;
16    use tokio::sync::OnceCell;
17
18    static DB_URL: LazyLock<Option<String>> = LazyLock::new(|| var("DATABASE_URL").ok());
19    } else {
20    use sqlx::postgres::PgPoolOptions;
21    use std::env::var;
22    use std::sync::LazyLock;
23    use std::time::Duration;
24    use tokio::sync::OnceCell;
25
26    static DB_URL: LazyLock<Option<String>> = LazyLock::new(|| var("DATABASE_URL").ok());
27    }
28}
29
30#[derive(Debug, Error)]
31pub enum DBError {
32    #[error("Database error: {0}")]
33    Sqlx(#[from] sqlx::Error),
34    #[error("DB migration error: {0}")]
35    Migration(#[from] MigrateError),
36    #[error("Configuration access error")]
37    Config(#[from] ConfigError),
38    #[error("DATABASE_URL is not provided")]
39    MissingUrl,
40    #[error("The database role lacks CREATEDB privilege")]
41    NoCreateDb,
42    #[error("Failed to generate an authentication keypair")]
43    KeyGen,
44    #[error("No connection-pool slot became free in time")]
45    PoolBudgetExhausted,
46    #[error("This account is undergoing maintenance and is temporarily unavailable")]
47    TenantQuarantined,
48    #[error("This account is disabled")]
49    TenantDisabled,
50    #[error("A fleet migration worker did not finish: {0}")]
51    FleetWorkerLost(String),
52    #[error("{0}")]
53    GlobalSchemaStale(String),
54    #[error("Refusing to build a connection string from database name {0:?}")]
55    UnsafeDatabaseName(String),
56    #[cfg(test)]
57    #[error("DB_POOL must be set; see local_db_sqlx_test macro")]
58    TestPoolUnset,
59}
60
61/// Connections the global/admin pool may open, per process.
62///
63/// Read on every authenticated request (`jwt_auth` re-reads the user row), so
64/// this is the web tier's effective concurrency ceiling, not a background
65/// allowance. Budget per process is this plus the tenant registry's own
66/// ceiling; see ADR 0032.
67///
68/// Server's own `cfg(test)` arm asserts an injected `sqlx::test` pool and never
69/// builds one, so the constant does not exist there.
70#[cfg(not(test))]
71pub(crate) const ADMIN_MAX_CONNECTIONS: u32 = 5;
72
73/// Shared migration set. Exposed to tests (server's own and `tests-integration`)
74/// so every `#[sqlx::test(migrator = "...")]` site names this once instead of
75/// embedding all migrations per test.
76#[cfg(any(test, feature = "test-utils"))]
77pub static MIGRATOR: sqlx::migrate::Migrator = sqlx::migrate!("../migrations");
78
79pub async fn migrate_db() -> Result<(), DBError> {
80    let pool = get_pool().await?;
81    Ok(sqlx::migrate!("../migrations").run(&pool).await?)
82}
83
84pub async fn get_connection() -> Result<PoolConnection<Postgres>, DBError> {
85    Ok(get_pool().await?.acquire().await?)
86}
87
88/// Runs a multi-statement SQL script against the admin pool via the simple
89/// query protocol. Executed against the shared pool (not a
90/// `&mut PgConnection`) so the future stays `Send`/spawnable — the `&mut`
91/// executor form trips a higher-ranked-lifetime bound once `boot()` is spawned.
92pub async fn execute_raw(sql: &str) -> Result<(), DBError> {
93    let pool = get_pool().await?;
94    sqlx::raw_sql(sqlx::AssertSqlSafe(sql))
95        .execute(&pool)
96        .await?;
97    Ok(())
98}
99
100cfg_if! {
101    if #[cfg(test)] {
102    // Server's own tests rely on `local_db_sqlx_test` (supp_macro) injecting
103    // a sqlx::test pool into DB_POOL before commands run. Forgetting that
104    // setup is a programming error — panic loudly rather than silently
105    // sharing a real DATABASE_URL pool across tests, which would break the
106    // per-test isolation sqlx::test guarantees.
107    //
108    // The slot owns a pool handle rather than pointing at the test's own, so a
109    // task still running after the test body sees a closed pool instead of
110    // freed memory.
111    thread_local!(pub static DB_POOL: RefCell<Option<PgPool>> = const {
112        RefCell::new(None)
113    });
114
115    /// True when the current thread has installed a test pool. Server-side
116    /// callers consult this to decide between the test pool and the
117    /// per-user production pool.
118    #[must_use]
119    pub fn test_pool_is_set() -> bool {
120        DB_POOL.with_borrow(Option::is_some)
121    }
122
123    async fn get_pool() -> Result<PgPool, DBError> {
124        DB_POOL
125            .with_borrow(Clone::clone)
126            .ok_or(DBError::TestPoolUnset)
127    }
128
129    } else if #[cfg(feature = "test-utils")] {
130    // Downstream consumers (the workspace `tests-integration` crate, and
131    // any web/sshd test that wants to drive a `server::command::*` flow
132    // against an isolated DB) install a sqlx::test pool via DB_POOL. When
133    // it is not installed we fall back to the production `DATABASE_URL`
134    // pool — that's the path web's existing tests take when compiled with
135    // `--all-features`, since they never set DB_POOL and instead let the
136    // production pool resolve via JWT'd handlers.
137    //
138    // The slot owns a pool handle for the same reason as the `cfg(test)` arm.
139    thread_local!(pub static DB_POOL: RefCell<Option<PgPool>> = const {
140        RefCell::new(None)
141    });
142
143    /// True when the current thread has installed a test pool. Used by
144    /// `User::get_connection` to choose between the test override and the
145    /// per-user production pool.
146    #[must_use]
147    pub fn test_pool_is_set() -> bool {
148        DB_POOL.with_borrow(Option::is_some)
149    }
150
151    static FALLBACK_POOL: OnceCell<PgPool> = OnceCell::const_new();
152
153    async fn get_pool() -> Result<PgPool, DBError> {
154        if let Some(pool) = DB_POOL.with_borrow(Clone::clone) {
155            return Ok(pool);
156        }
157        let url = DB_URL.as_deref().ok_or(DBError::MissingUrl)?;
158        FALLBACK_POOL
159            .get_or_try_init(|| async {
160                log::debug!("Fallback pool initialization");
161                let options = PgPoolOptions::new()
162                    .max_connections(ADMIN_MAX_CONNECTIONS)
163                    .acquire_timeout(Duration::from_secs(10));
164                Ok(options.connect(url).await?)
165            })
166            .await
167            .cloned()
168    }
169
170    } else {
171    // Production: lazy-init from DATABASE_URL.
172    static DB_POOL: OnceCell<PgPool> = OnceCell::const_new();
173
174    async fn get_pool() -> Result<PgPool, DBError> {
175        let url = DB_URL.as_deref().ok_or(DBError::MissingUrl)?;
176        DB_POOL
177            .get_or_try_init(|| async {
178                log::debug!("Pool initialization");
179                let options = PgPoolOptions::new()
180                    .max_connections(ADMIN_MAX_CONNECTIONS)
181                    .acquire_timeout(Duration::from_secs(10));
182                Ok(options.connect(url).await?)
183            })
184            .await
185            .cloned()
186    }
187    }
188}
189
190// Provisioning helpers. The per-user DSN is derived from the admin
191// `DATABASE_URL`. Fully-qualified paths keep these independent of the per-arm
192// imports in the `cfg_if!` above so they compile under every cfg (server's own
193// `cfg(test)` never provisions, but the pure helpers stay unit-testable).
194
195/// The admin `DATABASE_URL` the server was configured with.
196///
197/// # Errors
198/// [`DBError::MissingUrl`] if the environment variable is unset.
199pub fn admin_database_url() -> Result<String, DBError> {
200    std::env::var("DATABASE_URL").map_err(|_| DBError::MissingUrl)
201}
202
203/// Runs the migration set on an existing connection.
204///
205/// The fleet migrator already holds one per tenant for the claim lock, so
206/// opening a second just to migrate would double its connection cost — four
207/// workers would spend eight tenant connections against a budget of five
208/// (ADR 0032).
209///
210/// # Errors
211/// [`DBError::Migration`] if a migration fails or the applied set has drifted.
212pub async fn migrate_on(conn: &mut sqlx::PgConnection) -> Result<(), DBError> {
213    sqlx::migrate!("../migrations")
214        .run(conn)
215        .await
216        .map_err(DBError::from)
217}
218
219/// Runs the full migration set against the database at `url`. Used by
220/// provisioning to bring a freshly-created per-user database up to the
221/// (DDL-only) schema.
222///
223/// # Errors
224/// [`DBError::Sqlx`] on connect failure, [`DBError::Migration`] if a migration
225/// fails.
226pub async fn migrate_url(url: &str) -> Result<(), DBError> {
227    let pool = sqlx::postgres::PgPoolOptions::new()
228        .max_connections(1)
229        .acquire_timeout(std::time::Duration::from_secs(10))
230        .connect(url)
231        .await?;
232    let result = sqlx::migrate!("../migrations").run(&pool).await;
233    // Close the pool on BOTH paths before returning. A still-open connection to
234    // the per-user DB would otherwise block a compensating `DROP DATABASE` on
235    // the migration-failure path.
236    pool.close().await;
237    result.map_err(DBError::from)
238}
239
240#[cfg(test)]
241mod db_tests;
server/db.rs

server/
db.rs
