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}