qml_rs/storage/
database_init.rs1#[cfg(feature = "postgres")]
7use crate::storage::{PostgresConfig, PostgresStorage};
8use crate::storage::error::StorageError;
9#[cfg(feature = "postgres")]
10use std::time::Duration;
11
12#[cfg(feature = "postgres")]
14#[derive(Debug, Clone)]
15pub struct DatabaseInitializer {
16 config: PostgresConfig,
17 retry_attempts: u32,
18 retry_delay: Duration,
19 health_check_enabled: bool,
20}
21
22#[cfg(feature = "postgres")]
23impl DatabaseInitializer {
24 pub fn new() -> Self {
26 Self {
27 config: PostgresConfig::new(),
28 retry_attempts: 3,
29 retry_delay: Duration::from_secs(2),
30 health_check_enabled: true,
31 }
32 }
33
34 pub fn with_database_url<S: Into<String>>(mut self, url: S) -> Self {
36 self.config = self.config.with_database_url(url);
37 self
38 }
39
40 pub fn with_auto_migrate(mut self, enabled: bool) -> Self {
42 self.config = self.config.with_auto_migrate(enabled);
43 self
44 }
45
46 pub fn with_pool_config(mut self, max_connections: u32, min_connections: u32) -> Self {
48 self.config = self.config
49 .with_max_connections(max_connections)
50 .with_min_connections(min_connections);
51 self
52 }
53
54 pub fn with_retry_config(mut self, attempts: u32, delay: Duration) -> Self {
56 self.retry_attempts = attempts;
57 self.retry_delay = delay;
58 self
59 }
60
61 pub fn with_health_checks(mut self, enabled: bool) -> Self {
63 self.health_check_enabled = enabled;
64 self
65 }
66
67 pub async fn initialize(self) -> Result<PostgresStorage, DatabaseInitError> {
69 let mut last_error = None;
70
71 for attempt in 1..=self.retry_attempts {
72 tracing::info!("Database initialization attempt {} of {}", attempt, self.retry_attempts);
73
74 match self.try_initialize().await {
75 Ok(storage) => {
76 tracing::info!("Database initialized successfully on attempt {}", attempt);
77
78 if self.health_check_enabled {
79 self.perform_health_check(&storage).await?;
80 }
81
82 return Ok(storage);
83 }
84 Err(e) => {
85 tracing::warn!("Initialization attempt {} failed: {}", attempt, e);
86 last_error = Some(e);
87
88 if attempt < self.retry_attempts {
89 tracing::info!("Retrying in {:?}...", self.retry_delay);
90 tokio::time::sleep(self.retry_delay).await;
91 }
92 }
93 }
94 }
95
96 Err(last_error.unwrap_or(DatabaseInitError::Unknown))
97 }
98
99 async fn try_initialize(&self) -> Result<PostgresStorage, DatabaseInitError> {
101 PostgresStorage::new(self.config.clone())
102 .await
103 .map_err(DatabaseInitError::StorageError)
104 }
105
106 async fn perform_health_check(&self, storage: &PostgresStorage) -> Result<(), DatabaseInitError> {
108 tracing::debug!("Performing database health check...");
109
110 match storage.schema_exists().await {
112 Ok(true) => {
113 tracing::debug!("Schema health check passed");
114 }
115 Ok(false) => {
116 return Err(DatabaseInitError::HealthCheck(
117 "Schema does not exist after initialization".to_string()
118 ));
119 }
120 Err(e) => {
121 return Err(DatabaseInitError::HealthCheck(
122 format!("Schema check failed: {}", e)
123 ));
124 }
125 }
126
127 match storage.pool().acquire().await {
129 Ok(_) => {
130 tracing::debug!("Connectivity health check passed");
131 Ok(())
132 }
133 Err(e) => {
134 Err(DatabaseInitError::HealthCheck(
135 format!("Connectivity check failed: {}", e)
136 ))
137 }
138 }
139 }
140}
141
142#[cfg(feature = "postgres")]
143impl Default for DatabaseInitializer {
144 fn default() -> Self {
145 Self::new()
146 }
147}
148
149#[derive(Debug, thiserror::Error)]
151pub enum DatabaseInitError {
152 #[error("Storage initialization failed: {0}")]
153 StorageError(#[from] StorageError),
154
155 #[error("Health check failed: {0}")]
156 HealthCheck(String),
157
158 #[error("Configuration error: {0}")]
159 Configuration(String),
160
161 #[error("Unknown initialization error")]
162 Unknown,
163}
164
165#[cfg(feature = "postgres")]
167pub mod quick_init {
168 use super::*;
169
170 pub async fn development(database_url: String) -> Result<PostgresStorage, DatabaseInitError> {
172 DatabaseInitializer::new()
173 .with_database_url(database_url)
174 .with_auto_migrate(true)
175 .with_pool_config(10, 1)
176 .initialize()
177 .await
178 }
179
180 pub async fn production(database_url: String) -> Result<PostgresStorage, DatabaseInitError> {
182 DatabaseInitializer::new()
183 .with_database_url(database_url)
184 .with_auto_migrate(false)
185 .with_pool_config(50, 5)
186 .with_retry_config(5, Duration::from_secs(3))
187 .initialize()
188 .await
189 }
190
191 pub async fn testing(database_url: String) -> Result<PostgresStorage, DatabaseInitError> {
193 DatabaseInitializer::new()
194 .with_database_url(database_url)
195 .with_auto_migrate(true)
196 .with_pool_config(2, 1)
197 .with_health_checks(false)
198 .initialize()
199 .await
200 }
201}
202
203#[cfg(test)]
204mod tests {
205 #[cfg(feature = "postgres")]
206 use super::{DatabaseInitializer, Duration};
207
208 #[tokio::test]
209 #[cfg(feature = "postgres")]
210 async fn test_database_initializer_builder() {
211 let initializer = DatabaseInitializer::new()
212 .with_database_url("postgresql://test")
213 .with_auto_migrate(false)
214 .with_pool_config(5, 1)
215 .with_retry_config(2, Duration::from_millis(100));
216
217 assert_eq!(initializer.config.database_url, "postgresql://test");
218 assert!(!initializer.config.auto_migrate);
219 assert_eq!(initializer.config.max_connections, 5);
220 assert_eq!(initializer.retry_attempts, 2);
221 }
222}