axum-postgres-tx 0.3.0

#![deny(missing_docs, clippy::unwrap_used)]
#![warn(clippy::impl_trait_in_params)]

//! Request-scoped PostgreSQL transactions for axum.
//!
//! This crate provides a Tower [`Layer`](layer::Layer) and [`Service`](layer::Service)
//! that automatically manage request-scoped PostgreSQL transactions using
//! either `bb8-postgres` or `deadpool-postgres`. Handlers can extract a
//! [`Tx`](tx::Tx) to access the transaction; it is committed automatically
//! when the response is sent.
//!
//! # Examples
//!
//! ```rust,no_run
//! use axum::extract::State;
//! use axum_postgres_tx::Tx;
//!
//! async fn handler(State(pool): State<axum_postgres_tx::pool::Pool>, Tx(tx): Tx) -> String {
//!     let row = tx.query_one("SELECT 1", &[]).await.unwrap();
//!     format!("{}", row.get::<_, i32>(0))
//! }
//! ```

use axum_core::response::{IntoResponse, Response};
use http::StatusCode;
use thiserror::Error;

#[cfg(feature = "bb8")]
use bb8_postgres::{
    PostgresConnectionManager,
    bb8::{PooledConnection, RunError},
    tokio_postgres::{self, NoTls, Transaction},
};

#[cfg(feature = "deadpool")]
use deadpool_postgres::tokio_postgres::{self, Transaction};

mod extension;
/// Tower [`Layer`](tower_layer::Layer) and [`Service`](tower_service::Service)
/// for managing request-scoped PostgreSQL transactions.
pub mod layer;
/// Re-exported pool types for use in application state.
pub mod pool;
/// The [`Tx`](tx::Tx) extractor for accessing the current transaction.
pub mod tx;

#[cfg(feature = "bb8")]
type Pool = bb8_postgres::bb8::Pool<PostgresConnectionManager<NoTls>>;
#[cfg(feature = "deadpool")]
type Pool = deadpool_postgres::Pool;

/// Errors that can occur when using this crate.
#[derive(Debug, Error)]
pub enum Error {
    /// Failed to obtain a connection from the pool.
    #[cfg(feature = "bb8")]
    #[error(transparent)]
    Pool(#[from] RunError<tokio_postgres::Error>),
    /// Failed to obtain a connection from the pool.
    #[cfg(feature = "deadpool")]
    #[error(transparent)]
    Pool(#[from] deadpool::managed::PoolError<deadpool_postgres::tokio_postgres::Error>),
    /// A PostgreSQL error occurred while executing a query or transaction.
    #[error(transparent)]
    Db(#[from] tokio_postgres::Error),
    /// The [`Layer`](layer::Layer) was not applied to the router.
    ///
    /// This occurs when [`Tx`](tx::Tx) is extracted in a handler but the
    /// request was not processed by the transaction [`Layer`](layer::Layer).
    #[error("required extension not registered")]
    MissingExtension,
    /// [`Tx`](tx::Tx) was extracted multiple times in the same handler or middleware.
    ///
    /// Only one `Tx` extractor can be used per request. If you need to run
    /// multiple queries, use the same transaction for all of them.
    #[error("Tx extractor used multiple times in the same handler/middleware")]
    OverlappingExtractors,
}

impl IntoResponse for Error {
    fn into_response(self) -> Response {
        (StatusCode::INTERNAL_SERVER_ERROR, self.to_string()).into_response()
    }
}

/// A transaction paired with the connection it was started on.
///
/// This is a self-referential struct: `tx` borrows from `_conn`. The
/// `transmute` in [`OwnedTx::new`] extends the transaction's lifetime to
/// `'static` so both fields can be stored together. This is sound because:
///
/// 1. `_conn` is [`Box`]-ed, giving it a stable address that never moves.
/// 2. `tx` is declared **before** `_conn`, so Rust's drop order guarantees
///    the transaction is dropped while the connection is still alive.
/// 3. [`tokio_postgres::Transaction<'a>`] is covariant over `'a` (it holds
///    `&'a mut Client`), so a `Transaction<'short>` can be soundly
///    reinterpreted as `Transaction<'long>` when `'short: 'long`.
struct OwnedTx {
    // SAFETY INVARIANT: `tx` must be declared before `_conn` so that `tx` is
    // dropped while `_conn` is still alive.
    tx: Transaction<'static>,
    #[cfg(feature = "bb8")]
    _conn: Box<PooledConnection<'static, PostgresConnectionManager<NoTls>>>,
    #[cfg(feature = "deadpool")]
    _conn: Box<deadpool_postgres::Object>,
}

impl OwnedTx {
    /// Start a transaction on the given connection.
    ///
    /// The connection is boxed to ensure a stable address for the
    /// self-referential borrow. See the [`OwnedTx`] struct docs for the
    /// safety argument.
    #[cfg(feature = "bb8")]
    async fn new(
        conn: PooledConnection<'static, PostgresConnectionManager<NoTls>>,
        read_only: bool,
    ) -> Result<OwnedTx, tokio_postgres::Error> {
        let mut conn = Box::new(conn);
        let tx = conn
            .build_transaction()
            .read_only(read_only)
            .start()
            .await?;
        // SAFETY: `conn` is boxed (stable address). `tx` borrows `*conn` but
        // we need to store both in the same struct. This is sound because:
        //   1. Box guarantees the connection's address never changes.
        //   2. Field declaration order ensures `tx` drops before `_conn`.
        //   3. Transaction<'a> is covariant over 'a, so
        //      Transaction<'_> (where '_ is the borrow of *conn) can be
        //      reinterpreted as Transaction<'static> — the connection
        //      remains alive for the struct's entire lifetime.
        // The transaction MUST be dropped before the connection, which is
        // guaranteed by the field declaration order above.
        let tx = unsafe {
            std::mem::transmute::<
                tokio_postgres::Transaction<'_>,
                tokio_postgres::Transaction<'static>,
            >(tx)
        };
        Ok(OwnedTx { tx, _conn: conn })
    }

    /// Start a transaction on the given connection.
    ///
    /// The connection is boxed to ensure a stable address for the
    /// self-referential borrow. See the [`OwnedTx`] struct docs for the
    /// safety argument.
    #[cfg(feature = "deadpool")]
    async fn new(
        conn: deadpool_postgres::Object,
        read_only: bool,
    ) -> Result<OwnedTx, tokio_postgres::Error> {
        use std::ops::DerefMut;

        let mut conn = Box::new(conn);
        // Bypass deadpool's wrapper to get the raw tokio_postgres::Client.
        // Box<Object> -> Object (DerefMut) -> ClientWrapper (DerefMut) -> tokio_postgres::Client
        let client: &mut tokio_postgres::Client = conn.deref_mut().deref_mut().deref_mut();
        let tx = client
            .build_transaction()
            .read_only(read_only)
            .start()
            .await?;
        let tx = unsafe {
            std::mem::transmute::<
                tokio_postgres::Transaction<'_>,
                tokio_postgres::Transaction<'static>,
            >(tx)
        };
        Ok(OwnedTx { tx, _conn: conn })
    }
}