tenaxum 0.1.1

use crate::TenantId;
use sqlx::{PgPool, Postgres, Transaction};

/// Set `app.tenant_id` on an open transaction via `SET LOCAL`.
///
/// Useful when you already have a transaction in flight and need to scope
/// a sub-block to a specific tenant. Most call sites should prefer
/// [`PgPoolExt::begin_tenant`] which begins the transaction for you.
///
/// `SET LOCAL` is bounded to the current transaction, so when the tx ends
/// (commit or rollback), the connection returns to the pool with no
/// lingering tenant binding.
pub async fn set_tenant<'c>(
    tx: &mut Transaction<'c, Postgres>,
    tenant: TenantId,
) -> sqlx::Result<()> {
    sqlx::query("SELECT set_config('app.tenant_id', $1, true)")
        .bind(tenant.0.to_string())
        .execute(&mut **tx)
        .await?;
    Ok(())
}

/// Extension trait on [`sqlx::PgPool`] adding tenant-scoped transaction helpers.
#[async_trait::async_trait]
pub trait PgPoolExt {
    /// Begin a transaction and set `app.tenant_id` to the given tenant.
    ///
    /// The returned transaction is identical to one from [`PgPool::begin`];
    /// the only difference is that `SET LOCAL app.tenant_id = '<uuid>'` has
    /// already been run inside it.
    ///
    /// Pair with a Postgres RLS policy on your tenant-scoped tables:
    ///
    /// ```sql
    /// CREATE POLICY tenant_isolation ON your_table
    ///     USING       (tenant_id = current_setting('app.tenant_id', true)::uuid)
    ///     WITH CHECK  (tenant_id = current_setting('app.tenant_id', true)::uuid);
    /// ```
    ///
    /// Don't forget `ALTER TABLE your_table FORCE ROW LEVEL SECURITY` —
    /// without it, the table owner (which sqlx connects as) bypasses the
    /// policy and your isolation is silently broken.
    ///
    /// And — even with FORCE — your application must connect as a
    /// non-superuser role. Superusers bypass RLS unconditionally, FORCE
    /// or not. The default `POSTGRES_USER` in the official Postgres docker
    /// image is a superuser, which is the easiest way to ship a "tested,
    /// passes locally, leaks in prod" bug.
    async fn begin_tenant(&self, tenant: TenantId) -> sqlx::Result<Transaction<'_, Postgres>>;
}

#[async_trait::async_trait]
impl PgPoolExt for PgPool {
    async fn begin_tenant(&self, tenant: TenantId) -> sqlx::Result<Transaction<'_, Postgres>> {
        let mut tx = self.begin().await?;
        set_tenant(&mut tx, tenant).await?;
        Ok(tx)
    }
}