aster_forge_db/
component.rs1use 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
19pub const DATABASE_COMPONENT: &str = "database";
21pub const DATABASE_CONNECTIONS_SHUTDOWN_PHASE: &str = "database_connections";
23pub const DATABASE_HEALTH_CHECK: &str = "database";
25pub const DATABASE_HEALTH_CHECK_TIMEOUT: Duration = Duration::from_secs(5);
27
28pub struct DatabaseRuntimeComponent {
30 db_handles: DbHandles,
31 dependencies: &'static [&'static str],
32}
33
34impl DatabaseRuntimeComponent {
35 pub const fn new(db_handles: DbHandles) -> Self {
37 Self {
38 db_handles,
39 dependencies: &[],
40 }
41 }
42
43 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
57pub struct DatabaseHealthComponent {
59 db: DatabaseConnection,
60}
61
62impl DatabaseHealthComponent {
63 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
75pub fn database_component(
77 db_handles: DbHandles,
78) -> RuntimeComponentBundleRegistration<DatabaseRuntimeComponent> {
79 runtime_component(DatabaseRuntimeComponent::new(db_handles))
80}
81
82pub 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
90pub fn database_health_component(
92 db: DatabaseConnection,
93) -> RuntimeComponentBundleRegistration<DatabaseHealthComponent> {
94 runtime_component(DatabaseHealthComponent::new(db))
95}
96
97fn 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
121fn 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
135pub fn database_health_options() -> HealthCheckOptions {
137 HealthCheckOptions::required(Some(DATABASE_HEALTH_CHECK_TIMEOUT))
138 .with_scopes(HealthCheckScopes::readiness_and_diagnostics())
139}
140
141pub 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
158pub 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}