axum-postgres-tx 0.3.0

use std::{
    marker::PhantomData,
    task::{Context, Poll},
};

use axum_core::{
    extract::Request,
    response::{IntoResponse, Response},
};
#[cfg(feature = "bb8")]
use bb8_postgres::tokio_postgres;
#[cfg(feature = "deadpool")]
use deadpool_postgres::tokio_postgres;
use futures_core::future::BoxFuture;

use super::{Pool, extension::Extension};

/// A [`tower_layer::Layer`] that adds request-scoped PostgreSQL transaction
/// management to an axum router.
///
/// Insert this layer into your router to enable the [`Tx`](super::tx::Tx)
/// extractor. Each request will have its own transaction, which is committed
/// automatically when the response is sent.
///
/// # Examples
///
/// With `bb8-postgres`:
///
/// ```rust,no_run
/// use axum::Router;
/// use bb8_postgres::{PostgresConnectionManager, tokio_postgres::NoTls};
///
/// async fn example() {
///     let manager = PostgresConnectionManager::new("host=localhost user=postgres", NoTls);
///     let pool = bb8_postgres::bb8::Pool::builder().build(manager).await.unwrap();
///     let router = Router::new().layer(axum_postgres_tx::layer::Layer::from(pool));
/// }
/// ```
///
/// With `deadpool-postgres`:
///
/// ```rust,no_run
/// use axum::Router;
/// use deadpool_postgres::{Config, Runtime};
/// use tokio_postgres::NoTls;
///
/// async fn example() {
///     let mut cfg = Config::new();
///     cfg.dbname = Some("postgres".to_string());
///     let pool = cfg.create_pool(Some(Runtime::Tokio1), NoTls).unwrap();
///     let router = Router::new().layer(axum_postgres_tx::layer::Layer::from(pool));
/// }
/// ```
pub struct Layer<E> {
    pool: Pool,
    _error: PhantomData<E>,
}

impl<E> Clone for Layer<E> {
    fn clone(&self) -> Self {
        Self {
            pool: self.pool.clone(),
            _error: self._error,
        }
    }
}

impl From<Pool> for Layer<super::Error> {
    fn from(value: Pool) -> Self {
        Self {
            pool: value,
            _error: PhantomData,
        }
    }
}

impl<S, E> tower_layer::Layer<S> for Layer<E> {
    type Service = Service<S, E>;
    fn layer(&self, inner: S) -> Self::Service {
        Service {
            pool: self.pool.clone(),
            inner,
            _error: self._error,
        }
    }
}

/// A [`tower_service::Service`] that manages request-scoped PostgreSQL
/// transactions.
///
/// Created by [`Layer`]. You do not typically need to construct this directly.
pub struct Service<S, E> {
    pool: Pool,
    inner: S,
    _error: PhantomData<E>,
}

impl<S: Clone, E> Clone for Service<S, E> {
    fn clone(&self) -> Self {
        Self {
            pool: self.pool.clone(),
            inner: self.inner.clone(),
            _error: self._error,
        }
    }
}

impl<S, E> tower_service::Service<Request> for Service<S, E>
where
    S: tower_service::Service<Request, Response = Response> + Send + 'static,
    S::Future: Send + 'static,
    E: From<tokio_postgres::Error> + IntoResponse,
{
    type Response = S::Response;
    type Error = S::Error;
    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, mut req: Request) -> Self::Future {
        let ext = Extension::from(self.pool.clone());
        req.extensions_mut().insert(ext.clone());

        let res = self.inner.call(req);

        Box::pin(async move {
            let res = res.await?;

            if !res.status().is_server_error()
                && !res.status().is_client_error()
                && let Err(err) = ext.commit().await
            {
                return Ok(E::from(err).into_response());
            }

            Ok(res)
        })
    }
}