Skip to main content

aro_fletch/
database.rs

1//! Extension trait for registering a fletch database connection with [`App`].
2//!
3//! # Example
4//!
5//! ```ignore
6//! use aro_web::App;
7//! use aro_fletch::AppDatabaseExt;
8//!
9//! let app = App::new()
10//!     .database::<sqlx::Sqlite>("sqlite::memory:")
11//!     .await?
12//!     .migrate::<sqlx::Sqlite>("migrations/")
13//!     .await?
14//!     .health_check_db::<sqlx::Sqlite>()
15//!     .routes(router)
16//!     .build();
17//! ```
18
19use aro_web::App;
20use axum::Router;
21use axum::routing::get;
22use sqlx::Database;
23use std::sync::Arc;
24
25use crate::health::health_db_handler;
26
27/// Extension trait that adds database connection and migration methods to [`App`].
28pub trait AppDatabaseExt: Sized {
29    /// Connect to a database using the given URL and register the pool.
30    ///
31    /// The [`fletch_orm::Pool<DB>`] is registered as a dependency and can be
32    /// extracted in actions via `Dep<fletch_orm::Pool<DB>>`.
33    fn database<DB: Database>(
34        self,
35        url: &str,
36    ) -> impl std::future::Future<Output = Result<Self, fletch_orm::FletchError>> + Send;
37
38    /// Connect to a database using the given URL and a pre-configured
39    /// [`fletch_orm::PoolBuilder`], then register the pool.
40    ///
41    /// This allows customizing pool settings such as max connections,
42    /// idle timeout, and acquire timeout.
43    fn database_with<DB: Database>(
44        self,
45        url: &str,
46        builder: fletch_orm::PoolBuilder<DB>,
47    ) -> impl std::future::Future<Output = Result<Self, fletch_orm::FletchError>> + Send;
48
49    /// Register a `GET /health/db` route that checks database connectivity.
50    ///
51    /// Returns 200 with pool stats when healthy, or 503 when unreachable.
52    fn health_check_db<DB: Database>(self) -> Self;
53
54    /// Run SQLx migrations from the given directory path.
55    ///
56    /// This registers a startup hook that runs all pending migrations
57    /// when the app starts (via [`serve()`](App::serve) or
58    /// [`build_with_hooks()`](App::build_with_hooks)).
59    ///
60    /// The `path` should point to a directory containing `.sql` migration files
61    /// following SQLx naming conventions (e.g. `20230101000000_create_users.sql`).
62    fn migrate<DB>(
63        self,
64        path: &str,
65    ) -> impl std::future::Future<Output = Result<Self, sqlx::migrate::MigrateError>> + Send
66    where
67        DB: Database,
68        <DB as Database>::Connection: sqlx::migrate::Migrate;
69}
70
71impl AppDatabaseExt for App {
72    async fn database<DB: Database>(self, url: &str) -> Result<Self, fletch_orm::FletchError> {
73        let pool = fletch_orm::Pool::<DB>::connect(url).await?;
74        Ok(self.register::<fletch_orm::Pool<DB>>(Arc::new(pool)))
75    }
76
77    async fn database_with<DB: Database>(
78        self,
79        url: &str,
80        builder: fletch_orm::PoolBuilder<DB>,
81    ) -> Result<Self, fletch_orm::FletchError> {
82        let pool = builder.connect(url).await?;
83        Ok(self.register::<fletch_orm::Pool<DB>>(Arc::new(pool)))
84    }
85
86    fn health_check_db<DB: Database>(self) -> Self {
87        let router = Router::new().route("/health/db", get(health_db_handler::<DB>));
88        self.routes_with_state(router)
89    }
90
91    async fn migrate<DB>(self, path: &str) -> Result<Self, sqlx::migrate::MigrateError>
92    where
93        DB: Database,
94        <DB as Database>::Connection: sqlx::migrate::Migrate,
95    {
96        let migrator = sqlx::migrate::Migrator::new(std::path::Path::new(path)).await?;
97        Ok(self.on_startup(move |state| {
98            Box::pin(async move {
99                let Some(pool) = state.get::<fletch_orm::Pool<DB>>() else {
100                    return Err(
101                        "fletch_orm::Pool not registered — call .database() before .migrate()"
102                            .into(),
103                    );
104                };
105                migrator.run(pool.inner()).await?;
106                Ok(())
107            })
108        }))
109    }
110}