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::Cell;
10    } else if #[cfg(feature = "test-utils")] {
11    use std::cell::Cell;
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}
45
46/// Shared migration set. Exposed to tests (server's own and `tests-integration`)
47/// so every `#[sqlx::test(migrator = "...")]` site names this once instead of
48/// embedding all migrations per test.
49#[cfg(any(test, feature = "test-utils"))]
50pub static MIGRATOR: sqlx::migrate::Migrator = sqlx::migrate!("../migrations");
51
52pub async fn migrate_db() -> Result<(), DBError> {
53    Ok(sqlx::migrate!("../migrations")
54        .run(get_pool().await?)
55        .await?)
56}
57
58pub async fn get_connection() -> Result<PoolConnection<Postgres>, DBError> {
59    Ok(get_pool().await?.acquire().await?)
60}
61
62/// Runs a multi-statement SQL script against the admin pool via the simple
63/// query protocol. Executed against the shared pool reference (not a
64/// `&mut PgConnection`) so the future stays `Send`/spawnable — the `&mut`
65/// executor form trips a higher-ranked-lifetime bound once `boot()` is spawned.
66pub async fn execute_raw(sql: &str) -> Result<(), DBError> {
67    sqlx::raw_sql(sqlx::AssertSqlSafe(sql))
68        .execute(get_pool().await?)
69        .await?;
70    Ok(())
71}
72
73cfg_if! {
74    if #[cfg(test)] {
75    // Server's own tests rely on `local_db_sqlx_test` (supp_macro) injecting
76    // a sqlx::test pool into DB_POOL before commands run. Forgetting that
77    // setup is a programming error — panic loudly rather than silently
78    // sharing a real DATABASE_URL pool across tests, which would break the
79    // per-test isolation sqlx::test guarantees.
80    thread_local!(pub static DB_POOL: Cell<*const PgPool> = const {
81        Cell::new(std::ptr::null())
82    });
83
84    /// True when the current thread has installed a test pool. Server-side
85    /// callers consult this to decide between the test pool and the
86    /// per-user production pool.
87    #[must_use]
88    pub fn test_pool_is_set() -> bool {
89        DB_POOL.with(|c| !c.get().is_null())
90    }
91
92    async fn get_pool() -> Result<&'static PgPool, DBError> {
93        let p = DB_POOL.with(|c| c.get());
94        assert!(!p.is_null(), "DB_POOL must be set; see local_db_sqlx_test macro");
95        unsafe { Ok(&*p) }
96    }
97
98    } else if #[cfg(feature = "test-utils")] {
99    // Downstream consumers (the workspace `tests-integration` crate, and
100    // any web/sshd test that wants to drive a `server::command::*` flow
101    // against an isolated DB) install a sqlx::test pool via DB_POOL. When
102    // it is not installed we fall back to the production `DATABASE_URL`
103    // pool — that's the path web's existing tests take when compiled with
104    // `--all-features`, since they never set DB_POOL and instead let the
105    // production pool resolve via JWT'd handlers.
106    thread_local!(pub static DB_POOL: Cell<*const PgPool> = const {
107        Cell::new(std::ptr::null())
108    });
109
110    /// True when the current thread has installed a test pool. Used by
111    /// `User::get_connection` to choose between the test override and the
112    /// per-user production pool.
113    #[must_use]
114    pub fn test_pool_is_set() -> bool {
115        DB_POOL.with(|c| !c.get().is_null())
116    }
117
118    static FALLBACK_POOL: OnceCell<PgPool> = OnceCell::const_new();
119
120    async fn get_pool() -> Result<&'static PgPool, DBError> {
121        let p = DB_POOL.with(|c| c.get());
122        if !p.is_null() {
123            return unsafe { Ok(&*p) };
124        }
125        let url = DB_URL.as_deref().ok_or(DBError::MissingUrl)?;
126        FALLBACK_POOL
127            .get_or_try_init(|| async {
128                log::debug!("Fallback pool initialization");
129                let options = PgPoolOptions::new()
130                    .max_connections(10)
131                    .acquire_timeout(Duration::from_secs(10));
132                Ok(options.connect(url).await?)
133            })
134            .await
135    }
136
137    } else {
138    // Production: lazy-init from DATABASE_URL.
139    static DB_POOL: OnceCell<PgPool> = OnceCell::const_new();
140
141    async fn get_pool() -> Result<&'static PgPool, DBError> {
142        let url = DB_URL.as_deref().ok_or(DBError::MissingUrl)?;
143        DB_POOL
144            .get_or_try_init(|| async {
145                log::debug!("Pool initialization");
146                let options = PgPoolOptions::new()
147                    .max_connections(10)
148                    .acquire_timeout(Duration::from_secs(10));
149                Ok(options.connect(url).await?)
150            })
151            .await
152    }
153    }
154}
155
156// Provisioning helpers. Available outside server's own `cfg(test)` runs (which
157// use isolated sqlx::test pools and never provision); the per-user DSN is
158// derived from the admin `DATABASE_URL`.
159// Provisioning helpers. The per-user DSN is derived from the admin
160// `DATABASE_URL`. Fully-qualified paths keep these independent of the per-arm
161// imports in the `cfg_if!` above so they compile under every cfg (server's own
162// `cfg(test)` never provisions, but the pure helpers stay unit-testable).
163
164/// The admin `DATABASE_URL` the server was configured with.
165///
166/// # Errors
167/// [`DBError::MissingUrl`] if the environment variable is unset.
168pub fn admin_database_url() -> Result<String, DBError> {
169    std::env::var("DATABASE_URL").map_err(|_| DBError::MissingUrl)
170}
171
172/// Runs the full migration set against the database at `url`. Used by
173/// provisioning to bring a freshly-created per-user database up to the
174/// (DDL-only) schema.
175///
176/// # Errors
177/// [`DBError::Sqlx`] on connect failure, [`DBError::Migration`] if a migration
178/// fails.
179pub async fn migrate_url(url: &str) -> Result<(), DBError> {
180    let pool = sqlx::postgres::PgPoolOptions::new()
181        .max_connections(1)
182        .acquire_timeout(std::time::Duration::from_secs(10))
183        .connect(url)
184        .await?;
185    let result = sqlx::migrate!("../migrations").run(&pool).await;
186    // Close the pool on BOTH paths before returning. A still-open connection to
187    // the per-user DB would otherwise block a compensating `DROP DATABASE` on
188    // the migration-failure path.
189    pool.close().await;
190    result.map_err(DBError::from)
191}
192
193#[cfg(test)]
194mod db_tests {
195    use sqlx::PgPool;
196    use tokio::sync::OnceCell;
197
198    /// Context for keeping environment intact
199    static CONTEXT: OnceCell<()> = OnceCell::const_new();
200
201    async fn setup() {
202        CONTEXT
203            .get_or_init(|| async {
204                #[cfg(feature = "testlog")]
205                let _ = env_logger::builder()
206                    .is_test(true)
207                    .filter_level(log::LevelFilter::Trace)
208                    .try_init();
209            })
210            .await;
211    }
212
213    #[sqlx::test(migrator = "server::db::MIGRATOR")]
214    async fn migrations_create_schema_without_seed(pool: PgPool) -> sqlx::Result<()> {
215        setup().await;
216
217        let mut conn = pool.acquire().await?;
218
219        // The migration set is DDL-only: it creates the `config` table but no
220        // longer seeds it (seeding moved to `bootstrap::seed`). So a freshly
221        // migrated DB has the table present and empty.
222        let config_rows: i64 = sqlx::query_scalar("SELECT count(*) FROM config")
223            .fetch_one(&mut *conn)
224            .await?;
225        assert_eq!(config_rows, 0, "migration set must not seed config rows");
226
227        Ok(())
228    }
229}
server/db.rs

server/
db.rs
