Lines
71.79 %
Functions
47.06 %
Branches
100 %
use crate::config::ConfigError;
use cfg_if::cfg_if;
use sqlx::migrate::MigrateError;
use sqlx::pool::PoolConnection;
use sqlx::{PgPool, Postgres};
use thiserror::Error;
cfg_if! {
if #[cfg(test)] {
use std::cell::Cell;
} else if #[cfg(feature = "test-utils")] {
use sqlx::postgres::PgPoolOptions;
use std::env::var;
use std::sync::LazyLock;
use std::time::Duration;
use tokio::sync::OnceCell;
static DB_URL: LazyLock<String> = LazyLock::new(|| {
var("DATABASE_URL")
.unwrap_or_else(|_| panic!("{}", String::from(t!("DATABASE_URL is not provided"))))
});
} else {
}
#[derive(Debug, Error)]
pub enum DBError {
#[error("Database error: {0}")]
Sqlx(#[from] sqlx::Error),
#[error("DB migration error: {0}")]
Migration(#[from] MigrateError),
#[error("Configuration access error")]
Config(#[from] ConfigError),
#[error("DATABASE_URL is not provided")]
MissingUrl,
#[error("The database role lacks CREATEDB privilege")]
NoCreateDb,
#[error("Failed to generate an authentication keypair")]
KeyGen,
pub async fn migrate_db() -> Result<(), DBError> {
Ok(sqlx::migrate!("../migrations")
.run(get_pool().await?)
.await?)
pub async fn get_connection() -> Result<PoolConnection<Postgres>, DBError> {
Ok(get_pool().await?.acquire().await?)
/// Runs a multi-statement SQL script against the admin pool via the simple
/// query protocol. Executed against the shared pool reference (not a
/// `&mut PgConnection`) so the future stays `Send`/spawnable — the `&mut`
/// executor form trips a higher-ranked-lifetime bound once `boot()` is spawned.
pub async fn execute_raw(sql: &str) -> Result<(), DBError> {
sqlx::raw_sql(sql).execute(get_pool().await?).await?;
Ok(())
// Server's own tests rely on `local_db_sqlx_test` (supp_macro) injecting
// a sqlx::test pool into DB_POOL before commands run. Forgetting that
// setup is a programming error — panic loudly rather than silently
// sharing a real DATABASE_URL pool across tests, which would break the
// per-test isolation sqlx::test guarantees.
thread_local!(pub static DB_POOL: Cell<*const PgPool> = const {
Cell::new(std::ptr::null())
/// True when the current thread has installed a test pool. Server-side
/// callers consult this to decide between the test pool and the
/// per-user production pool.
#[must_use]
pub fn test_pool_is_set() -> bool {
DB_POOL.with(|c| !c.get().is_null())
async fn get_pool() -> Result<&'static PgPool, DBError> {
let p = DB_POOL.with(|c| c.get());
assert!(!p.is_null(), "DB_POOL must be set; see local_db_sqlx_test macro");
unsafe { Ok(&*p) }
// Downstream consumers (the workspace `tests-integration` crate, and
// any web/sshd test that wants to drive a `server::command::*` flow
// against an isolated DB) install a sqlx::test pool via DB_POOL. When
// it is not installed we fall back to the production `DATABASE_URL`
// pool — that's the path web's existing tests take when compiled with
// `--all-features`, since they never set DB_POOL and instead let the
// production pool resolve via JWT'd handlers.
/// True when the current thread has installed a test pool. Used by
/// `User::get_connection` to choose between the test override and the
static FALLBACK_POOL: OnceCell<PgPool> = OnceCell::const_new();
if !p.is_null() {
return unsafe { Ok(&*p) };
Ok(FALLBACK_POOL
.get_or_init(|| async {
log::debug!("Fallback pool initialization");
let options = PgPoolOptions::new()
.max_connections(10)
.acquire_timeout(Duration::from_secs(10));
options.connect(&DB_URL).await.unwrap()
}).await)
// Production: lazy-init from DATABASE_URL.
static DB_POOL: OnceCell<PgPool> = OnceCell::const_new();
Ok(DB_POOL
log::debug!("Pool initialization");
// Provisioning helpers. Available outside server's own `cfg(test)` runs (which
// use isolated sqlx::test pools and never provision); the per-user DSN is
// derived from the admin `DATABASE_URL`.
// Provisioning helpers. The per-user DSN is derived from the admin
// `DATABASE_URL`. Fully-qualified paths keep these independent of the per-arm
// imports in the `cfg_if!` above so they compile under every cfg (server's own
// `cfg(test)` never provisions, but the pure helpers stay unit-testable).
/// The admin `DATABASE_URL` the server was configured with.
///
/// # Errors
/// [`DBError::MissingUrl`] if the environment variable is unset.
pub fn admin_database_url() -> Result<String, DBError> {
std::env::var("DATABASE_URL").map_err(|_| DBError::MissingUrl)
/// Runs the full migration set against the database at `url`. Used by
/// provisioning to bring a freshly-created per-user database up to the
/// (DDL-only) schema.
/// [`DBError::Sqlx`] on connect failure, [`DBError::Migration`] if a migration
/// fails.
pub async fn migrate_url(url: &str) -> Result<(), DBError> {
let pool = sqlx::postgres::PgPoolOptions::new()
.max_connections(1)
.acquire_timeout(std::time::Duration::from_secs(10))
.connect(url)
.await?;
let result = sqlx::migrate!("../migrations").run(&pool).await;
// Close the pool on BOTH paths before returning. A still-open connection to
// the per-user DB would otherwise block a compensating `DROP DATABASE` on
// the migration-failure path.
pool.close().await;
result.map_err(DBError::from)
#[cfg(test)]
mod db_tests {
use sqlx::PgPool;
/// Context for keeping environment intact
static CONTEXT: OnceCell<()> = OnceCell::const_new();
async fn setup() {
CONTEXT
#[cfg(feature = "testlog")]
let _ = env_logger::builder()
.is_test(true)
.filter_level(log::LevelFilter::Trace)
.try_init();
})
.await;
#[sqlx::test(migrations = "../migrations")]
async fn migrations_create_schema_without_seed(pool: PgPool) -> sqlx::Result<()> {
setup().await;
let mut conn = pool.acquire().await?;
// The migration set is DDL-only: it creates the `config` table but no
// longer seeds it (seeding moved to `bootstrap::seed`). So a freshly
// migrated DB has the table present and empty.
let config_rows: i64 = sqlx::query_scalar("SELECT count(*) FROM config")
.fetch_one(&mut *conn)
assert_eq!(config_rows, 0, "migration set must not seed config rows");