lxc-streamutils/apps/server-rust/src/db.rs

//! PostgreSQL pool, migrations and low-level tenant-scoped queries.
//!
//! Multi-tenant tables use PostgreSQL row-level security in addition to owner
//! columns and composite foreign keys. Every tenant query must execute inside a
//! transaction after [`Db::set_tenant`], which uses `SET LOCAL` so pooled
//! connections cannot retain the previous request's identity.

use std::{fmt, str::FromStr};

use deadpool_postgres::{Manager, ManagerConfig, Object, Pool, RecyclingMethod, Runtime};
use serde_json::Value;
use tokio_postgres::{Config as PostgresConfig, NoTls, Transaction};
use uuid::Uuid;

/// Cloneable database handle intended to live directly in Axum's `AppState`.
///
/// HTTP handlers should acquire a pooled connection through [`Db::get`]. Any
/// query against an RLS-protected tenant table must run in a transaction after
/// calling [`Db::set_tenant`]. `SET LOCAL` is important: a session-level setting
/// could otherwise leak an identity when the connection returns to the pool.
#[derive(Clone)]
pub struct Db {
    pool: Pool,
}

impl Db {
    pub fn from_pool(pool: Pool) -> Self {
        Self { pool }
    }

    pub fn pool(&self) -> &Pool {
        &self.pool
    }

    pub fn connect(database_url: &str, max_size: usize) -> Result<Self, DbError> {
        let config = PostgresConfig::from_str(database_url)
            .map_err(|error| DbError::Configuration(error.to_string()))?;
        let manager = Manager::from_config(
            config,
            NoTls,
            ManagerConfig {
                recycling_method: RecyclingMethod::Fast,
            },
        );
        let pool = Pool::builder(manager)
            .runtime(Runtime::Tokio1)
            .max_size(max_size.max(1))
            .build()
            .map_err(|error| DbError::Configuration(error.to_string()))?;
        Ok(Self { pool })
    }

    pub async fn get(&self) -> Result<Object, DbError> {
        self.pool
            .get()
            .await
            .map_err(|error| DbError::Pool(error.to_string()))
    }

    /// Applies the additive multi-tenant migration. The legacy migrations must
    /// already have run because this migration deliberately links their tables.
    pub async fn migrate_multitenancy(&self) -> Result<(), DbError> {
        let client = self.get().await?;
        client
            .batch_execute(include_str!("../migrations/003_multitenancy.sql"))
            .await
            .map_err(DbError::Postgres)
    }

    pub async fn set_tenant(transaction: &Transaction<'_>, user_id: Uuid) -> Result<(), DbError> {
        transaction
            .query_one(
                "SELECT set_config('app.user_id', $1, true)",
                &[&user_id.to_string()],
            )
            .await
            .map_err(DbError::Postgres)?;
        Ok(())
    }

    /// Minimal, non-secret account inventory for the source supervisor at
    /// startup. Component rows are deliberately loaded in a second,
    /// tenant-scoped query so RLS remains effective.
    pub async fn list_active_tenants(&self) -> Result<Vec<ActiveTenant>, DbError> {
        let client = self.get().await?;
        let rows = client
            .query(
                "SELECT owner_user_id,room_id,source_id FROM list_active_live_sources()",
                &[],
            )
            .await?;
        Ok(rows
            .into_iter()
            .map(|row| ActiveTenant {
                user_id: row.get(0),
                room_id: row.get(1),
                source_id: row.get(2),
            })
            .collect())
    }

    pub async fn list_tenant_components(
        &self,
        user_id: Uuid,
    ) -> Result<Vec<ComponentRecord>, DbError> {
        let mut client = self.get().await?;
        let transaction = client.transaction().await?;
        Self::set_tenant(&transaction, user_id).await?;
        let rows = transaction
            .query(
                "SELECT component.id,source.id,component.kind,component.name,component.settings,\
                        component.settings_version,component.enabled \
                 FROM component_instances AS component \
                 JOIN live_sources AS source ON source.owner_user_id=component.owner_user_id \
                 WHERE component.owner_user_id=$1 ORDER BY component.created_at",
                &[&user_id],
            )
            .await?;
        transaction.commit().await?;
        Ok(rows
            .into_iter()
            .map(|row| ComponentRecord {
                id: row.get(0),
                owner_user_id: user_id,
                account_source_id: row.get(1),
                kind: row.get(2),
                name: row.get(3),
                settings: row.get(4),
                settings_version: row.get(5),
                enabled: row.get(6),
            })
            .collect())
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ActiveTenant {
    pub user_id: Uuid,
    pub room_id: String,
    pub source_id: Uuid,
}

#[derive(Clone, Debug, PartialEq)]
pub struct ComponentRecord {
    pub id: Uuid,
    pub owner_user_id: Uuid,
    /// Derived from the account's singleton live source; it is not component
    /// configuration and is never persisted on `component_instances`.
    pub account_source_id: Uuid,
    pub kind: String,
    pub name: String,
    pub settings: Value,
    pub settings_version: i32,
    pub enabled: bool,
}

#[derive(Debug)]
pub enum DbError {
    Configuration(String),
    Pool(String),
    Postgres(tokio_postgres::Error),
}

impl fmt::Display for DbError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Configuration(message) => write!(formatter, "database configuration: {message}"),
            Self::Pool(message) => write!(formatter, "database pool: {message}"),
            Self::Postgres(error) => write!(formatter, "database query: {error}"),
        }
    }
}

impl std::error::Error for DbError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            Self::Postgres(error) => Some(error),
            _ => None,
        }
    }
}

impl From<tokio_postgres::Error> for DbError {
    fn from(error: tokio_postgres::Error) -> Self {
        Self::Postgres(error)
    }
}