Skip to main content

aster_forge_db/
component.rs

1//! Runtime component integration for database handles.
2//!
3//! Product crates still own database configuration, migrations, repositories,
4//! and health-check semantics. Forge owns the repeated lifecycle mechanics for
5//! already prepared handles: registering the `database` runtime component,
6//! applying product-declared shutdown dependencies, and closing the handles
7//! exactly once during dependency-aware shutdown.
8
9use aster_forge_runtime::{
10    HealthCheckOptions, HealthCheckScopes, HealthComponentReport, RuntimeComponentBundle,
11    RuntimeComponentBundleRegistration, RuntimeComponentKind, RuntimeComponentRegistry,
12    runtime_component,
13};
14use sea_orm::DatabaseConnection;
15use std::time::Duration;
16
17use crate::DbHandles;
18
19/// Stable component name used for database handles.
20pub const DATABASE_COMPONENT: &str = "database";
21/// Stable shutdown phase name for database handle closing.
22pub const DATABASE_CONNECTIONS_SHUTDOWN_PHASE: &str = "database_connections";
23/// Stable health check name used for database ping checks.
24pub const DATABASE_HEALTH_CHECK: &str = "database";
25/// Default timeout for database readiness and diagnostics health checks.
26pub const DATABASE_HEALTH_CHECK_TIMEOUT: Duration = Duration::from_secs(5);
27
28/// Runtime component that closes database handles during graceful shutdown.
29pub struct DatabaseRuntimeComponent {
30    db_handles: DbHandles,
31    dependencies: &'static [&'static str],
32}
33
34impl DatabaseRuntimeComponent {
35    /// Creates a database runtime component from prepared handles.
36    pub const fn new(db_handles: DbHandles) -> Self {
37        Self {
38            db_handles,
39            dependencies: &[],
40        }
41    }
42
43    /// Declares components that must shut down before database handles close.
44    pub const fn depends_on_all(mut self, dependencies: &'static [&'static str]) -> Self {
45        self.dependencies = dependencies;
46        self
47    }
48}
49
50impl RuntimeComponentBundle for DatabaseRuntimeComponent {
51    fn register(self, registry: &mut RuntimeComponentRegistry) {
52        register_database_health_check(registry, self.db_handles.reader().clone());
53        register_database_shutdown(registry, self.db_handles, self.dependencies);
54    }
55}
56
57/// Runtime component that registers the standard database health check only.
58pub struct DatabaseHealthComponent {
59    db: DatabaseConnection,
60}
61
62impl DatabaseHealthComponent {
63    /// Creates a database health component from a prepared connection.
64    pub const fn new(db: DatabaseConnection) -> Self {
65        Self { db }
66    }
67}
68
69impl RuntimeComponentBundle for DatabaseHealthComponent {
70    fn register(self, registry: &mut RuntimeComponentRegistry) {
71        register_database_health_check(registry, self.db);
72    }
73}
74
75/// Creates the database runtime component used by product entrypoints.
76pub fn database_component(
77    db_handles: DbHandles,
78) -> RuntimeComponentBundleRegistration<DatabaseRuntimeComponent> {
79    runtime_component(DatabaseRuntimeComponent::new(db_handles))
80}
81
82/// Creates the database runtime component with shutdown dependencies.
83pub fn database_component_after(
84    db_handles: DbHandles,
85    dependencies: &'static [&'static str],
86) -> RuntimeComponentBundleRegistration<DatabaseRuntimeComponent> {
87    runtime_component(DatabaseRuntimeComponent::new(db_handles).depends_on_all(dependencies))
88}
89
90/// Creates the standard database health component.
91pub fn database_health_component(
92    db: DatabaseConnection,
93) -> RuntimeComponentBundleRegistration<DatabaseHealthComponent> {
94    runtime_component(DatabaseHealthComponent::new(db))
95}
96
97/// Registers database shutdown after product-declared dependency components.
98fn register_database_shutdown(
99    registry: &mut RuntimeComponentRegistry,
100    db_handles: DbHandles,
101    dependencies: &'static [&'static str],
102) {
103    registry
104        .component(DATABASE_COMPONENT)
105        .kind(RuntimeComponentKind::Database)
106        .depends_on_all(dependencies)
107        .shutdown_once(
108            DATABASE_CONNECTIONS_SHUTDOWN_PHASE,
109            None,
110            db_handles,
111            |db_handles| async move {
112                db_handles
113                    .close()
114                    .await
115                    .map_err(|error| error.to_string())?;
116                Ok(())
117            },
118        );
119}
120
121/// Registers a database readiness and diagnostics health check.
122fn register_database_health_check(registry: &mut RuntimeComponentRegistry, db: DatabaseConnection) {
123    registry.component_health_with_options(
124        DATABASE_COMPONENT,
125        RuntimeComponentKind::Database,
126        DATABASE_HEALTH_CHECK,
127        database_health_options(),
128        move || {
129            let db = db.clone();
130            async move { check_database_component(&db).await }
131        },
132    );
133}
134
135/// Returns the standard database health check options.
136pub fn database_health_options() -> HealthCheckOptions {
137    HealthCheckOptions::required(Some(DATABASE_HEALTH_CHECK_TIMEOUT))
138        .with_scopes(HealthCheckScopes::readiness_and_diagnostics())
139}
140
141/// Runs the standard database ping health check.
142pub async fn check_database_component(db: &DatabaseConnection) -> HealthComponentReport {
143    match ping_database(db).await {
144        Ok(()) => {
145            tracing::debug!("database health check succeeded");
146            HealthComponentReport::healthy(DATABASE_HEALTH_CHECK, "database ping succeeded")
147        }
148        Err(error) => {
149            tracing::debug!(error = %error, "database health check failed");
150            HealthComponentReport::unhealthy(
151                DATABASE_HEALTH_CHECK,
152                format!("database ping failed: {error}"),
153            )
154        }
155    }
156}
157
158/// Pings the database connection used by the standard health check.
159pub async fn ping_database(db: &DatabaseConnection) -> crate::Result<()> {
160    tracing::debug!("pinging database health check");
161    db.ping().await.map_err(crate::DbError::from)
162}
163
164#[cfg(test)]
165mod tests {
166    use aster_forge_runtime::{RuntimeComponentBundle, RuntimeComponentKind};
167
168    use super::{
169        DATABASE_COMPONENT, DATABASE_CONNECTIONS_SHUTDOWN_PHASE, DATABASE_HEALTH_CHECK,
170        check_database_component, database_component_after, database_health_component,
171    };
172    use aster_forge_runtime::{HealthCheckScope, HealthStatus};
173
174    #[tokio::test]
175    async fn database_component_registers_dependencies_and_shutdown() {
176        let db = sea_orm::Database::connect("sqlite::memory:")
177            .await
178            .expect("database runtime component test database should connect");
179        let db_handles = crate::DbHandles::single(db);
180
181        let registry = aster_forge_runtime::RuntimeComponentRegistry::configured(|registry| {
182            database_component_after(db_handles, &["background_tasks", "mail_outbox"])
183                .register(registry);
184        });
185
186        let descriptor = registry
187            .descriptor(DATABASE_COMPONENT)
188            .expect("database component should be registered");
189        assert_eq!(descriptor.kind, RuntimeComponentKind::Database);
190        assert_eq!(
191            descriptor.dependencies,
192            vec!["background_tasks", "mail_outbox"]
193        );
194        assert_eq!(
195            descriptor
196                .shutdown
197                .first()
198                .expect("database shutdown should be registered")
199                .phase_name,
200            DATABASE_CONNECTIONS_SHUTDOWN_PHASE
201        );
202        assert_eq!(descriptor.health_checks.len(), 1);
203    }
204
205    #[tokio::test]
206    async fn database_component_reports_ping_success_and_failure() {
207        let db = sea_orm::Database::connect("sqlite::memory:")
208            .await
209            .expect("database health test database should connect");
210
211        let healthy = check_database_component(&db).await;
212        assert_eq!(healthy.status, HealthStatus::Healthy);
213        assert_eq!(healthy.message, "database ping succeeded");
214
215        db.close_by_ref()
216            .await
217            .expect("database health test database should close");
218        let unhealthy = check_database_component(&db).await;
219        assert_eq!(unhealthy.status, HealthStatus::Unhealthy);
220        assert!(unhealthy.message.contains("database ping failed"));
221    }
222
223    #[tokio::test]
224    async fn database_health_component_registers_readiness_component() {
225        let db = sea_orm::Database::connect("sqlite::memory:")
226            .await
227            .expect("database readiness test database should connect");
228        let mut registry = aster_forge_runtime::RuntimeComponentRegistry::new();
229
230        database_health_component(db).register(&mut registry);
231
232        assert_eq!(registry.len(), 1);
233        let report = registry.run_health(HealthCheckScope::Readiness).await;
234        let component_names = report
235            .components
236            .iter()
237            .map(|component| component.name)
238            .collect::<Vec<_>>();
239        assert_eq!(component_names, vec![DATABASE_HEALTH_CHECK]);
240        assert_eq!(report.status(), HealthStatus::Healthy);
241    }
242
243    #[tokio::test]
244    async fn database_health_component_reports_healthy_status() {
245        let db = sea_orm::Database::connect("sqlite::memory:")
246            .await
247            .expect("database health component test database should connect");
248
249        let mut registry = aster_forge_runtime::RuntimeComponentRegistry::configured(|registry| {
250            database_health_component(db).register(registry);
251        });
252
253        let descriptor = registry
254            .descriptor(DATABASE_COMPONENT)
255            .expect("database component should be registered");
256        assert_eq!(descriptor.health_checks.len(), 1);
257        let report = registry.run_health(HealthCheckScope::Readiness).await;
258        assert_eq!(report.status(), HealthStatus::Healthy);
259    }
260}