tenaxum 0.1.1

//! Pool-scoped tenant isolation: set `app.tenant_id` once when a connection
//! is checked out of the pool, reset it when it goes back. Every query made
//! during the request is auto-isolated, with no per-call-site boilerplate.
//!
//! This is the pattern most production codebases want. It pairs with:
//!
//! 1. The [`TENANT_ID`] task-local — set per-request by the
//!    [`tenant_scope`] middleware below.
//! 2. The [`with_tenant_hooks`] pool-builder helper — wires the
//!    `before_acquire` / `after_release` hooks that read the task-local
//!    and run `set_config` / `RESET` on the connection.
//!
//! ## Wiring
//!
//! ```no_run
//! # async fn run() -> sqlx::Result<()> {
//! use sqlx::postgres::PgPoolOptions;
//! use tenaxum::pool;
//!
//! let pool = pool::with_tenant_hooks(PgPoolOptions::new().max_connections(8))
//!     .connect("postgres://...").await?;
//! # Ok(()) }
//! ```
//!
//! Then in your Axum app:
//!
//! ```ignore
//! use axum::{Router, middleware};
//! use tenaxum::pool::tenant_scope;
//!
//! let app = Router::new()
//!     // ...your routes...
//!     .layer(middleware::from_fn(tenant_scope));
//! ```
//!
//! `tenant_scope` reads `Extension<TenantId>` (set by your auth layer) and
//! scopes the [`TENANT_ID`] task-local for the rest of the request. The
//! pool hooks then pick it up automatically on every connection checkout.
//!
//! ## When to use this vs [`crate::PgPoolExt::begin_tenant`]
//!
//! - **Pool hooks** (this module): for the common case where every
//!   handler in your app touches tenant-scoped data and you want
//!   isolation by default.
//! - **`begin_tenant`**: for explicit, scoped uses — background jobs
//!   that operate on a specific tenant outside of an HTTP request,
//!   one-off scripts, or admin paths where the pool hooks aren't wired.

use crate::TenantId;
use axum::{
    body::Body,
    http::Request,
    middleware::Next,
    response::Response,
};
use sqlx::postgres::PgPoolOptions;

tokio::task_local! {
    /// Per-request tenant binding. Set by [`tenant_scope`] middleware (or
    /// any code that opens a [`tokio::task_local::LocalKey::scope`] block
    /// over it). Read by the pool hooks installed via [`with_tenant_hooks`].
    pub static TENANT_ID: TenantId;
}

/// Apply  tenaxum pool hooks to a [`PgPoolOptions`] builder.
///
/// `before_acquire`: reads [`TENANT_ID`] task-local; if set, runs
/// `SELECT set_config('app.tenant_id', '<uuid>', false)` (session-level —
/// because the connection may be used outside an explicit transaction).
/// If unset, runs `RESET app.tenant_id` so a connection that previously
/// held a tenant binding can't leak it into the new request.
///
/// `after_release`: unconditionally runs `RESET app.tenant_id` before the
/// connection returns to the pool. Defense in depth — the next acquire
/// will set its own value, but if anything ever bypasses `before_acquire`,
/// the connection comes back clean.
pub fn with_tenant_hooks(opts: PgPoolOptions) -> PgPoolOptions {
    opts.before_acquire(|conn, _meta| {
        Box::pin(async move {
            match TENANT_ID.try_with(|t| *t) {
                Ok(t) => {
                    sqlx::query("SELECT set_config('app.tenant_id', $1, false)")
                        .bind(t.0.to_string())
                        .execute(&mut *conn)
                        .await?;
                }
                Err(_) => {
                    sqlx::query("RESET app.tenant_id")
                        .execute(&mut *conn)
                        .await?;
                }
            }
            Ok(true)
        })
    })
    .after_release(|conn, _meta| {
        Box::pin(async move {
            sqlx::query("RESET app.tenant_id")
                .execute(&mut *conn)
                .await
                .map(|_| true)
        })
    })
}

/// Axum middleware that scopes the [`TENANT_ID`] task-local for the
/// duration of the request.
///
/// Reads `Extension<TenantId>` off the request — your auth layer is
/// expected to have inserted it. If absent, the request runs with no
/// tenant binding (the pool hooks will run `RESET app.tenant_id` on
/// every checkout, which is the correct behavior for unauthenticated
/// or admin paths).
pub async fn tenant_scope(req: Request<Body>, next: Next) -> Response {
    match req.extensions().get::<TenantId>().copied() {
        Some(tenant) => TENANT_ID.scope(tenant, next.run(req)).await,
        None => next.run(req).await,
    }
}