rustango 0.60.12

//! Database-mode tenant pool registry for non-Postgres backends.
//!
//! Postgres tenants — including schema-mode — keep using the existing
//! [`super::TenantPools`]. This module ships the parallel structure
//! for `BackendKind::MySql` and `BackendKind::Sqlite`, both of which
//! are database-mode only (schema-mode is rejected at validation time
//! by [`super::BackendKind::validate_storage_mode`]).
//!
//! ## Why a parallel type
//!
//! `TenantPools` has 111 call sites today; making it generic over
//! `sqlx::Database` would force trait bounds onto every caller and
//! break the v0.30 dep-flip stability we just settled into. The
//! database-mode-only shape is also much simpler — no `SET
//! search_path` dance, no schema cache, no scoped-pool builder —
//! so a smaller parallel struct is the right size for the smaller
//! responsibility.
//!
//! ## Lifecycle
//!
//! - At boot the host app picks one backend type for the entire
//!   process (`Cli::tenancy::<MySql>()` or
//!   `Cli::tenancy::<Sqlite>()` — Slice 6 wires these). A single
//!   `DatabasePools<DB>` lives in the request-router state.
//! - Per request, the `Tenant<DB>` extractor (Slice 3) reads the
//!   resolved `Org` from request extensions, dispatches to
//!   [`DatabasePools::acquire`], and hands the resulting connection
//!   to the handler.
//! - The cache is keyed on `org.slug`. Builds-on-miss; evictable via
//!   [`DatabasePools::invalidate`] when `database_url` changes.

use std::collections::HashMap;
use std::sync::Arc;

use rustango::sql::sqlx;
use sqlx::pool::{Pool, PoolConnection};
use sqlx::Database;
use tokio::sync::RwLock;

use super::error::TenancyError;
use super::org::{BackendKind, Org, StorageMode};
use super::pools::TenantPoolsConfig;
use super::secrets::{LiteralSecretsResolver, SecretsResolver};

/// One tenant's pool (database-mode only, parameterized by backend
/// type). Holds an `Arc<Pool<DB>>` so callers can clone cheaply for
/// `axum::extract::State`-style sharing.
#[derive(Debug)]
pub struct DatabasePool<DB: Database> {
    pool: Arc<Pool<DB>>,
}

impl<DB: Database> Clone for DatabasePool<DB> {
    fn clone(&self) -> Self {
        Self {
            pool: self.pool.clone(),
        }
    }
}

impl<DB: Database> DatabasePool<DB> {
    /// The underlying pool. Use `Pool::acquire` for connection
    /// checkout, or hand `&Pool<DB>` to sqlx query macros directly.
    #[must_use]
    pub fn pool(&self) -> &Pool<DB> {
        &self.pool
    }

    /// `Arc` clone — cheap enough that it's safe to pull out and
    /// stash in app state per request.
    #[must_use]
    pub fn pool_arc(&self) -> Arc<Pool<DB>> {
        self.pool.clone()
    }
}

/// Owned, tenant-scoped connection. Database-mode only — there is no
/// schema variant. Deref's to `PoolConnection<DB>` so the inner sqlx
/// connection is reachable for query macros / ORM helpers.
pub struct DatabaseConn<DB: Database> {
    inner: PoolConnection<DB>,
}

impl<DB: Database> std::ops::Deref for DatabaseConn<DB> {
    type Target = PoolConnection<DB>;
    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl<DB: Database> std::ops::DerefMut for DatabaseConn<DB> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner
    }
}

/// Lazy database-mode pool registry. Generic over the sqlx backend
/// (`Postgres` works too but the existing [`super::TenantPools`] is
/// strictly nicer for Postgres because it supports schema-mode;
/// keep this type for MySQL / SQLite).
pub struct DatabasePools<DB: Database> {
    config: TenantPoolsConfig,
    secrets: Arc<dyn SecretsResolver>,
    /// LRU stamp source, as in [`super::pools`].
    lru_tick: std::sync::atomic::AtomicU64,
    cache: RwLock<HashMap<String, super::pools::CachedPool<DB>>>,
    /// Single-flight for cache misses.
    building: super::pools::SingleFlight<Arc<Pool<DB>>>,
    /// The backend type this registry serves. Validated against each
    /// `org.backend_kind` on `pool_for_org` so a mis-routed tenant
    /// fails loudly instead of silently building the wrong pool.
    backend: BackendKind,
    /// Optional URL template for tenants that don't carry an explicit
    /// `database_url`. Useful for SQLite: set
    /// `"sqlite:./tenants/{slug}.db?mode=rwc"` and every SQLite tenant
    /// gets its own file under `./tenants/` keyed by slug. `None`
    /// means "require explicit `database_url` per org" — the v0.5
    /// default and the safe-by-default production setting.
    ///
    /// The `{slug}` placeholder is replaced with `org.slug`; no other
    /// placeholders are recognized. The template is only consulted
    /// when `org.database_url is None`.
    url_template: Option<String>,
}

impl<DB: Database> DatabasePools<DB> {
    /// Construct with the default literal secrets resolver — i.e.
    /// `Org.database_url` is treated as a verbatim connection string.
    /// Pair with [`super::EnvSecretsResolver`] /
    /// [`super::ChainSecretsResolver`] via [`Self::with_secrets`]
    /// when you want env-var or vault indirection.
    #[must_use]
    pub fn new(backend: BackendKind) -> Self {
        Self::with_secrets(backend, LiteralSecretsResolver)
    }

    /// Next LRU stamp. See [`super::pools::CachedPool`].
    fn next_tick(&self) -> u64 {
        self.lru_tick
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed)
    }

    /// Construct with a user-supplied secrets resolver.
    #[must_use]
    pub fn with_secrets<R: SecretsResolver>(backend: BackendKind, secrets: R) -> Self {
        Self {
            config: TenantPoolsConfig::default(),
            secrets: Arc::new(secrets),
            lru_tick: std::sync::atomic::AtomicU64::new(0),
            cache: RwLock::new(HashMap::new()),
            building: super::pools::SingleFlight::default(),
            backend,
            url_template: None,
        }
    }

    /// Set a URL template used when an org has no explicit
    /// `database_url`. The `{slug}` placeholder is replaced with the
    /// org's slug at acquire time.
    ///
    /// Most useful for SQLite, where the natural per-tenant shape is
    /// one file per slug:
    ///
    /// ```ignore
    /// let pools: DatabasePools<sqlx::Sqlite> =
    ///     DatabasePools::new(BackendKind::Sqlite)
    ///         .with_url_template("sqlite:./tenants/{slug}.db?mode=rwc");
    /// // Orgs with database_url=None now auto-route to
    /// // ./tenants/<slug>.db; sqlite creates the file on first
    /// // write thanks to `mode=rwc`.
    /// ```
    ///
    /// Returns `self` for builder chaining.
    #[must_use]
    pub fn with_url_template(mut self, template: impl Into<String>) -> Self {
        self.url_template = Some(template.into());
        self
    }

    /// Replace the config. Returns `self` for builder ergonomics.
    #[must_use]
    pub fn config(mut self, config: TenantPoolsConfig) -> Self {
        self.config = config;
        self
    }

    /// Read access to the current config.
    #[must_use]
    pub fn pool_config(&self) -> &TenantPoolsConfig {
        &self.config
    }

    /// Which backend driver this registry serves. Used by the
    /// `Tenant<DB>` extractor (Slice 3) to verify the request's
    /// `Org.backend_kind` matches the process's configured backend.
    #[must_use]
    pub fn backend_kind(&self) -> BackendKind {
        self.backend
    }

    /// Resolve (or build-on-miss) the pool for `org`.
    ///
    /// Rejects orgs whose `backend_kind` doesn't match this
    /// registry's configured backend, or whose `storage_mode` is
    /// `Schema` (this struct is database-mode-only — schema-mode is a
    /// Postgres feature, [`super::TenantPools`] handles it). Built
    /// pools are cached keyed on `slug`.
    ///
    /// # Errors
    /// * [`TenancyError::Validation`] when `org.backend_kind` doesn't
    ///   match this registry, `storage_mode` isn't `Database`, or
    ///   `database_url` is missing.
    /// * [`TenancyError::Secrets`] when the database-URL reference
    ///   fails to resolve.
    /// * [`TenancyError::Driver`] for pool-build failure.
    pub async fn pool_for_org(&self, org: &Org) -> Result<DatabasePool<DB>, TenancyError> {
        // Backend-kind match — refuse to silently build the wrong
        // backend's pool for a mis-routed tenant.
        let want = BackendKind::parse(&org.backend_kind).map_err(|got| {
            TenancyError::Validation(format!(
                "org `{}` has unknown backend_kind `{got}` \
                 (expected `postgres`, `mysql`, or `sqlite`)",
                org.slug
            ))
        })?;
        if want != self.backend {
            return Err(TenancyError::Validation(format!(
                "org `{}` is configured for backend `{want}` \
                 but this DatabasePools is serving `{}`. \
                 The process boots one backend; route through the \
                 matching server instance.",
                org.slug, self.backend
            )));
        }

        let mode = StorageMode::parse(&org.storage_mode).map_err(|got| {
            TenancyError::Validation(format!(
                "org `{}` has unknown storage_mode `{got}` \
                 (expected `schema` or `database`)",
                org.slug
            ))
        })?;
        if mode != StorageMode::Database {
            return Err(TenancyError::Validation(format!(
                "org `{}` storage_mode is `{mode}` but DatabasePools \
                 only supports database-mode. \
                 (Schema-mode is Postgres-only — use TenantPools.)",
                org.slug
            )));
        }

        // Fast path, then one build per slug and source, as in `TenantPools`.
        let source = super::pools::PoolSource::database(org);
        if let Some(pool) = self.cached(&org.slug, &source).await {
            return Ok(DatabasePool { pool });
        }
        let build = self.build_for(org, source.clone());
        let pool = if self.config.max_cached_database_pools == 0 {
            build.await?
        } else {
            self.building.run(&org.slug, &source, build).await?
        };
        Ok(DatabasePool { pool })
    }

    async fn build_for(
        &self,
        org: &Org,
        source: super::pools::PoolSource,
    ) -> Result<Arc<sqlx::Pool<DB>>, TenancyError> {
        if let Some(pool) = self.cached(&org.slug, &source).await {
            return Ok(pool);
        }

        // Resolve the database_url — explicit `org.database_url`
        // wins; otherwise fall back to the configured URL template
        // with `{slug}` substituted. Templates make per-tenant SQLite
        // files trivial: one config line, N orgs.
        let url_ref = match org.database_url.as_deref() {
            Some(u) => u.to_owned(),
            None => match &self.url_template {
                Some(tpl) => tpl.replace("{slug}", &org.slug),
                None => {
                    return Err(TenancyError::Validation(format!(
                        "org `{}` is database-mode but has no database_url \
                         and no url_template is configured on DatabasePools \
                         (call `.with_url_template(...)` at boot)",
                        org.slug
                    )));
                }
            },
        };
        let resolved = self
            .secrets
            .resolve(&url_ref)
            .await
            .map_err(TenancyError::Secrets)?;

        // Build the pool. The sqlx::Pool::connect path is generic
        // over DB — every backend gets the same connection-limit
        // config from TenantPoolsConfig.
        let pool = self.build_pool(&resolved).await?;
        let pool_arc = Arc::new(pool);

        // Upgrade to write lock — guard against multiple
        // concurrent builds racing.
        let mut cache = self.cache.write().await;
        // Re-check under write lock (race-loser case).
        if let Some(existing) = super::pools::CachedPool::lookup(&cache, &org.slug, &source) {
            existing.touch(self.next_tick());
            return Ok(existing.pool.clone());
        }
        cache.remove(&org.slug);
        // Evict the most idle tenant rather than refuse. Same policy
        // as `TenantPools`, and the same type, so the two cannot drift
        // apart again (#1527).
        for slug in super::pools::evict_to_fit(&mut cache, self.config.max_cached_database_pools) {
            tracing::info!(
                target: "rustango::tenancy::pools",
                evicted = %slug,
                for_slug = %org.slug,
                cap = self.config.max_cached_database_pools,
                "evicted the least recently used tenant pool to make room",
            );
        }
        if self.config.max_cached_database_pools == 0 {
            return Ok(pool_arc);
        }
        cache.insert(
            org.slug.clone(),
            super::pools::CachedPool::new(pool_arc.clone(), source, self.next_tick()),
        );
        Ok(pool_arc)
    }

    async fn cached(
        &self,
        slug: &str,
        source: &super::pools::PoolSource,
    ) -> Option<Arc<sqlx::Pool<DB>>> {
        let cache = self.cache.read().await;
        let entry = super::pools::CachedPool::lookup(&cache, slug, source)?;
        entry.touch(self.next_tick());
        Some(entry.pool.clone())
    }

    /// Acquire a connection scoped to `org`. Thin wrapper around
    /// `pool_for_org` + `pool.acquire`.
    ///
    /// # Errors
    /// As [`Self::pool_for_org`] plus [`TenancyError::Driver`] for
    /// the underlying connection checkout.
    pub async fn acquire(&self, org: &Org) -> Result<DatabaseConn<DB>, TenancyError> {
        let pool = self.pool_for_org(org).await?;
        let inner = pool.pool().acquire().await?;
        Ok(DatabaseConn { inner })
    }

    /// Drop the cached pool for `org` (if any). Used by the operator
    /// console's `database_url` edit handler — the next request
    /// rebuilds the pool with the new URL.
    pub async fn invalidate(&self, slug: &str) {
        let mut cache = self.cache.write().await;
        cache.remove(slug);
    }

    /// Build a fresh pool for `url` with the same tuning as
    /// `TenantPools`. Unlike it, a missing SQLite file is not created.
    pub(crate) async fn build_pool(&self, url: &str) -> Result<Pool<DB>, TenancyError> {
        super::pools::tenant_pool_options::<DB>(
            &self.config,
            self.config.database_pool_max_connections,
            self.config.database_pool_min_connections,
        )
        .connect(url)
        .await
        .map_err(TenancyError::from)
    }
}

// Every test here needs sqlite or mysql; a PG-only build has none (#2313).
#[cfg(all(test, any(feature = "sqlite", feature = "mysql")))]
mod tests {
    use super::*;

    // Compile-time tests — these prove the API surface for each
    // backend exists. Functional tests require a live DB and ship in
    // Slice 5's `tenancy_sqlite_live.rs` / `tenancy_mysql_live.rs`.

    #[cfg(feature = "sqlite")]
    #[test]
    fn sqlite_pool_registry_constructible() {
        let _: DatabasePools<sqlx::Sqlite> = DatabasePools::new(BackendKind::Sqlite);
    }

    #[cfg(feature = "mysql")]
    #[test]
    fn mysql_pool_registry_constructible() {
        let _: DatabasePools<sqlx::MySql> = DatabasePools::new(BackendKind::MySql);
    }

    #[cfg(feature = "sqlite")]
    #[test]
    fn backend_kind_round_trips() {
        let p: DatabasePools<sqlx::Sqlite> = DatabasePools::new(BackendKind::Sqlite);
        assert_eq!(p.backend_kind(), BackendKind::Sqlite);
    }
}