reliar-store-postgres 0.8.0

//! [`Session`] — the crate-private connection context both stores run every statement through
//! (`docs/architecture/store-postgres-layout.md` Part II §7). Not generic, not public: `Session`
//! never leaves this crate, and neither store's
//! constructor does any I/O to build one (ADR 0047).

use std::time::Duration;

use sqlx::{Connection, PgConnection, PgPool};

use crate::error::FromDatabaseError;

/// The pooled connection plus the one policy knob every statement honours:
/// `statement_timeout`. Built once, by each store's constructor, from the pool and settings it
/// was handed — no query, no connection of its own. **There is no `pool()` accessor**: the pool
/// is reachable only through [`Self::run`], which is what makes "did this statement get the
/// timeout policy?" answerable by reading a signature.
#[derive(Clone, Debug)]
pub(crate) struct Session {
    pool: PgPool,

    statement_timeout: Duration,
}

impl Session {
    pub(crate) fn new(pool: PgPool, statement_timeout: Duration) -> Self {
        Self {
            pool,
            statement_timeout,
        }
    }

    /// Runs `op` on this session's own pool under its `statement_timeout` policy: with a zero
    /// timeout (the default) `op` runs on a pooled connection with no wrapping statements, so a
    /// single-statement `op` is still the one implicit-transaction statement ADR 0006 relies on;
    /// with a non-zero timeout it runs inside `BEGIN` / `SET LOCAL statement_timeout` / `COMMIT`.
    /// Acquires **one** connection for the whole closure — with a zero timeout, an `op` issuing
    /// several statements (e.g. `purge`'s three) now shares one pool checkout instead of one per
    /// statement; each is still its own implicit transaction, so no guarantee moves.
    pub(crate) async fn run<T, F>(&self, op: F) -> Result<T, sqlx::Error>
    where
        F: AsyncFnOnce(&mut PgConnection) -> Result<T, sqlx::Error>,
    {
        let mut conn = self.pool.acquire().await?;

        if self.statement_timeout.is_zero() {
            return op(&mut conn).await;
        }

        let timeout_ms = i64::try_from(self.statement_timeout.as_millis())
            .unwrap_or(i64::MAX)
            .to_string();
        let mut tx = conn.begin().await?;

        sqlx::query_scalar!(
            "SELECT set_config('statement_timeout', $1, true)",
            timeout_ms
        )
        .fetch_one(&mut *tx)
        .await?;

        let result = op(&mut tx).await;

        match result {
            Ok(value) => {
                tx.commit().await?;

                Ok(value)
            }
            Err(err) => {
                // Best effort: whichever error `op` returned is what the caller needs, and an
                // aborted transaction is rolled back by the connection's own teardown regardless.
                let _ = tx.rollback().await;

                Err(err)
            }
        }
    }

    /// This session's typed error for a `sqlx::Error` — the one call every operation fn ends its
    /// `Session::run` with (`.map_err(|e| session.map_err(e))?`). A method rather than an
    /// associated function so every call site reads `session.map_err(e)`, mirroring
    /// [`Self::run`]; today's mapping needs no session state, but the shape is the point.
    #[allow(
        clippy::unused_self,
        reason = "kept as a method for call-site symmetry with `run`"
    )]
    pub(crate) fn map_err<E: FromDatabaseError>(&self, err: sqlx::Error) -> E {
        E::from_database_error(err)
    }
}