feat: implement FromRef for State into sqlx pool

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<sqlx::SqlitePool>) {
+//!     /* ... */
+//! }
+//! ```
+//!
 //! # Examples
 //!
 //! See [`examples/`][examples] in the repo for more examples.

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<DB: Marker> Clone for State<DB> {
         }
     }
 }
+
+impl<DB: Marker> FromRef<State<DB>> for sqlx::Pool<DB::Driver> {
+    fn from_ref(input: &State<DB>) -> Self {
+        input.pool.clone()
+    }
+}

tests/lib.rs

@@ -208,6 +208,35 @@
     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<sqlx::SqlitePool>| 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();
@@ -457,7 +486,7 @@
     (pool, Response { status, body })
 }
 
 struct MyExtractorError(axum_sqlx_tx::Error);
 
 impl From<axum_sqlx_tx::Error> for MyExtractorError {
     fn from(error: axum_sqlx_tx::Error) -> Self {
@@ -471,7 +500,7 @@
     }
 }
 
 struct MyLayerError(sqlx::Error);
 
 impl From<sqlx::Error> for MyLayerError {
     fn from(error: sqlx::Error) -> Self {