diff --git a/src/lib.rs b/src/lib.rs index 9621884..6bdf17d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -75,6 +75,20 @@ //! It's not currently possible to use `Tx` for a dynamic number of databases. Feel free to open an //! issue if you have a requirement for this. //! +//! ## Accessing the pool +//! +//! Note that [`State`] implements [`FromRef`](axum_core::extract::FromRef) into the inner SQLx pool. Therefore, +//! if you still need to access the database pool at some handler, you can use axum's `State` +//! extractor normally. +//! +//! ``` +//! use axum::extract::State; +//! +//! async fn this_still_works(State(pool): State) { +//! /* ... */ +//! } +//! ``` +//! //! # Examples //! //! See [`examples/`][examples] in the repo for more examples. diff --git a/src/state.rs b/src/state.rs index 1e4d019..fecf879 100644 --- a/src/state.rs +++ b/src/state.rs @@ -1,3 +1,5 @@ +use axum_core::extract::FromRef; + use crate::Marker; /// Application state that enables the [`Tx`] extractor. @@ -34,3 +36,9 @@ impl Clone for State { } } } + +impl FromRef> for sqlx::Pool { + fn from_ref(input: &State) -> Self { + input.pool.clone() + } +} diff --git a/tests/lib.rs b/tests/lib.rs index accec5a..16ea510 100644 --- a/tests/lib.rs +++ b/tests/lib.rs @@ -208,6 +208,35 @@ async fn substates() { assert!(response.status().is_success()); } +#[tokio::test] +async fn extract_pool_from_state() { + let pool = sqlx::SqlitePool::connect("sqlite::memory:").await.unwrap(); + + let (state, layer) = Tx::setup(pool); + + let app = axum::Router::new() + .route( + "/", + axum::routing::get( + |axum::extract::State(_pool): axum::extract::State| async move {}, + ), + ) + .layer(layer) + .with_state(state); + + let response = app + .oneshot( + http::Request::builder() + .uri("/") + .body(axum::body::Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + + assert!(response.status().is_success()); +} + #[tokio::test] async fn missing_layer() { let pool = sqlx::SqlitePool::connect("sqlite::memory:").await.unwrap();