axum_postgres_tx/
layer.rs

1use std::{
2    marker::PhantomData,
3    task::{Context, Poll},
4};
5
6use axum_core::{
7    extract::Request,
8    response::{IntoResponse, Response},
9};
10#[cfg(feature = "bb8")]
11use bb8_postgres::tokio_postgres;
12#[cfg(feature = "deadpool")]
13use deadpool_postgres::tokio_postgres;
14use futures_core::future::BoxFuture;
15
16use super::{Pool, extension::Extension};
17
18/// A [`tower_layer::Layer`] that adds request-scoped PostgreSQL transaction
19/// management to an axum router.
20///
21/// Insert this layer into your router to enable the [`Tx`](super::tx::Tx)
22/// extractor. Each request will have its own transaction, which is committed
23/// automatically when the response is sent.
24///
25/// # Examples
26///
27/// With `bb8-postgres`:
28///
29/// ```rust,no_run
30/// use axum::Router;
31/// use bb8_postgres::{PostgresConnectionManager, tokio_postgres::NoTls};
32///
33/// async fn example() {
34///     let manager = PostgresConnectionManager::new("host=localhost user=postgres", NoTls);
35///     let pool = bb8_postgres::bb8::Pool::builder().build(manager).await.unwrap();
36///     let router = Router::new().layer(axum_postgres_tx::layer::Layer::from(pool));
37/// }
38/// ```
39///
40/// With `deadpool-postgres`:
41///
42/// ```rust,no_run
43/// use axum::Router;
44/// use deadpool_postgres::{Config, Runtime};
45/// use tokio_postgres::NoTls;
46///
47/// async fn example() {
48///     let mut cfg = Config::new();
49///     cfg.dbname = Some("postgres".to_string());
50///     let pool = cfg.create_pool(Some(Runtime::Tokio1), NoTls).unwrap();
51///     let router = Router::new().layer(axum_postgres_tx::layer::Layer::from(pool));
52/// }
53/// ```
54pub struct Layer<E> {
55    pool: Pool,
56    _error: PhantomData<E>,
57}
58
59impl<E> Clone for Layer<E> {
60    fn clone(&self) -> Self {
61        Self {
62            pool: self.pool.clone(),
63            _error: self._error,
64        }
65    }
66}
67
68impl From<Pool> for Layer<super::Error> {
69    fn from(value: Pool) -> Self {
70        Self {
71            pool: value,
72            _error: PhantomData,
73        }
74    }
75}
76
77impl<S, E> tower_layer::Layer<S> for Layer<E> {
78    type Service = Service<S, E>;
79    fn layer(&self, inner: S) -> Self::Service {
80        Service {
81            pool: self.pool.clone(),
82            inner,
83            _error: self._error,
84        }
85    }
86}
87
88/// A [`tower_service::Service`] that manages request-scoped PostgreSQL
89/// transactions.
90///
91/// Created by [`Layer`]. You do not typically need to construct this directly.
92pub struct Service<S, E> {
93    pool: Pool,
94    inner: S,
95    _error: PhantomData<E>,
96}
97
98impl<S: Clone, E> Clone for Service<S, E> {
99    fn clone(&self) -> Self {
100        Self {
101            pool: self.pool.clone(),
102            inner: self.inner.clone(),
103            _error: self._error,
104        }
105    }
106}
107
108impl<S, E> tower_service::Service<Request> for Service<S, E>
109where
110    S: tower_service::Service<Request, Response = Response> + Send + 'static,
111    S::Future: Send + 'static,
112    E: From<tokio_postgres::Error> + IntoResponse,
113{
114    type Response = S::Response;
115    type Error = S::Error;
116    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
117
118    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
119        self.inner.poll_ready(cx)
120    }
121
122    fn call(&mut self, mut req: Request) -> Self::Future {
123        let ext = Extension::from(self.pool.clone());
124        req.extensions_mut().insert(ext.clone());
125
126        let res = self.inner.call(req);
127
128        Box::pin(async move {
129            let res = res.await?;
130
131            if !res.status().is_server_error()
132                && !res.status().is_client_error()
133                && let Err(err) = ext.commit().await
134            {
135                return Ok(E::from(err).into_response());
136            }
137
138            Ok(res)
139        })
140    }
141}
axum_postgres_tx/layer.rs

axum_postgres_tx/
layer.rs
