Skip to main content

bindizr_db/
lib.rs

1use std::sync::OnceLock;
2
3use sqlx::{MySql, Pool, Postgres, Sqlite, sqlite::SqlitePoolOptions};
4
5pub mod error;
6pub mod repository;
7mod schema;
8mod utils;
9
10pub use bindizr_core::model;
11pub(crate) use bindizr_core::{config, log_error, log_info};
12
13static DATABASE_POOL: OnceLock<DatabasePool> = OnceLock::new();
14static INITIALIZE_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
15
16#[derive(Debug)]
17pub enum DatabasePool {
18    MySQL(Pool<MySql>),
19    PostgreSQL(Pool<Postgres>),
20    SQLite(Pool<Sqlite>),
21}
22
23#[derive(Debug, Clone)]
24pub enum DatabaseType {
25    MySQL,
26    PostgreSQL,
27    SQLite,
28}
29
30pub async fn initialize() {
31    if is_initialized() {
32        return;
33    }
34
35    let _guard = INITIALIZE_LOCK.lock().await;
36
37    if is_initialized() {
38        return;
39    }
40
41    let bindizr_config = config::get_bindizr_config();
42
43    let database_type = match bindizr_config.database.database_type {
44        config::DatabaseType::Mysql => DatabaseType::MySQL,
45        config::DatabaseType::Postgresql => DatabaseType::PostgreSQL,
46        config::DatabaseType::Sqlite => DatabaseType::SQLite,
47    };
48
49    let database_url = match database_type {
50        DatabaseType::MySQL => bindizr_config.database.mysql.server_url.clone(),
51        DatabaseType::PostgreSQL => bindizr_config.database.postgresql.server_url.clone(),
52        DatabaseType::SQLite => utils::to_sqlite_url(&bindizr_config.database.sqlite.file_path)
53            .unwrap_or_else(|e| {
54                log_error!("{}", e);
55                std::process::exit(1);
56            }),
57    };
58
59    let pool = match database_type {
60        DatabaseType::MySQL => DatabasePool::new_mysql(&database_url).await,
61        DatabaseType::PostgreSQL => DatabasePool::new_postgres(&database_url).await,
62        DatabaseType::SQLite => DatabasePool::new_sqlite(&database_url).await,
63    };
64
65    if DATABASE_POOL.set(pool).is_err() {
66        return;
67    }
68
69    log_info!("Database pool initialized");
70}
71
72fn is_initialized() -> bool {
73    DATABASE_POOL.get().is_some()
74}
75
76pub fn get_pool() -> &'static DatabasePool {
77    DATABASE_POOL.get().expect("Database pool not initialized")
78}
79
80impl DatabasePool {
81    pub async fn new_mysql(url: &str) -> Self {
82        let pool = Pool::<MySql>::connect(url).await.unwrap_or_else(|e| {
83            log_error!("Failed to create MySQL database pool: {}", e);
84            std::process::exit(1);
85        });
86
87        let database_pool = DatabasePool::MySQL(pool);
88
89        // Create tables
90        if let Err(e) = database_pool.create_tables().await {
91            log_error!("Failed to create tables: {}", e);
92            std::process::exit(1);
93        }
94
95        database_pool
96    }
97
98    pub async fn new_postgres(url: &str) -> Self {
99        let pool = Pool::<Postgres>::connect(url).await.unwrap_or_else(|e| {
100            log_error!("Failed to create PostgreSQL database pool: {}", e);
101            std::process::exit(1);
102        });
103
104        let database_pool = DatabasePool::PostgreSQL(pool);
105
106        // Create tables
107        if let Err(e) = database_pool.create_tables().await {
108            log_error!("Failed to create tables: {}", e);
109            std::process::exit(1);
110        }
111
112        database_pool
113    }
114    pub async fn new_sqlite(url: &str) -> Self {
115        let pool = SqlitePoolOptions::new()
116            .after_connect(|conn, _meta| {
117                Box::pin(async move {
118                    // Enable foreign key constraints for SQLite
119                    sqlx::query("PRAGMA foreign_keys = ON")
120                        .execute(conn)
121                        .await
122                        .map(|_| ())
123                })
124            })
125            .connect(url)
126            .await
127            .unwrap_or_else(|e| {
128                log_error!("Failed to create SQLite database pool: {}", e);
129                std::process::exit(1);
130            });
131
132        let database_pool = DatabasePool::SQLite(pool);
133
134        // Create tables
135        if let Err(e) = database_pool.create_tables().await {
136            log_error!("Failed to create tables: {}", e);
137            std::process::exit(1);
138        }
139
140        database_pool
141    }
142
143    async fn create_tables(&self) -> Result<(), String> {
144        // Get table creation queries from schema module based on database type
145        let queries = match self {
146            DatabasePool::MySQL(_) => schema::get_mysql_table_creation_queries(),
147            DatabasePool::PostgreSQL(_) => schema::get_postgres_table_creation_queries(),
148            DatabasePool::SQLite(_) => schema::get_sqlite_table_creation_queries(),
149        };
150
151        match self {
152            DatabasePool::MySQL(pool) => {
153                for query in queries {
154                    let mut conn = pool.acquire().await.map_err(|e| {
155                        log_error!("Failed to acquire MySQL connection: {}", e);
156                        e.to_string()
157                    })?;
158                    sqlx::query(query).execute(&mut *conn).await.map_err(|e| {
159                        log_error!("Failed to execute query '{}': {}", query, e);
160                        e.to_string()
161                    })?;
162                }
163            }
164            DatabasePool::PostgreSQL(pool) => {
165                for query in queries {
166                    let mut conn = pool.acquire().await.map_err(|e| {
167                        log_error!("Failed to acquire PostgreSQL connection: {}", e);
168                        e.to_string()
169                    })?;
170                    sqlx::query(query).execute(&mut *conn).await.map_err(|e| {
171                        log_error!("Failed to execute query '{}': {}", query, e);
172                        e.to_string()
173                    })?;
174                }
175            }
176            DatabasePool::SQLite(pool) => {
177                for query in queries {
178                    let mut conn = pool.acquire().await.map_err(|e| {
179                        log_error!("Failed to acquire SQLite connection: {}", e);
180                        e.to_string()
181                    })?;
182                    sqlx::query(query).execute(&mut *conn).await.map_err(|e| {
183                        log_error!("Failed to execute query '{}': {}", query, e);
184                        e.to_string()
185                    })?;
186                }
187            }
188        }
189        Ok(())
190    }
191}
192
193// Repository convenience functions - returns trait objects for runtime dispatch
194pub fn get_zone_repository() -> Box<dyn repository::ZoneRepository> {
195    let pool = get_pool();
196    repository::RepositoryFactory::create_zone_repository(pool)
197}
198
199pub fn get_record_repository() -> Box<dyn repository::RecordRepository> {
200    let pool = get_pool();
201    repository::RepositoryFactory::create_record_repository(pool)
202}
203
204pub fn get_api_token_repository() -> Box<dyn repository::ApiTokenRepository> {
205    let pool = get_pool();
206    repository::RepositoryFactory::create_api_token_repository(pool)
207}
208
209pub fn get_zone_change_repository() -> Box<dyn repository::ZoneChangeRepository> {
210    let pool = get_pool();
211    repository::RepositoryFactory::create_zone_change_repository(pool)
212}
213
214pub fn get_zone_snapshot_repository() -> Box<dyn repository::ZoneSnapshotRepository> {
215    let pool = get_pool();
216    repository::RepositoryFactory::create_zone_snapshot_repository(pool)
217}
218
219pub fn get_catalog_zone_state_repository() -> Box<dyn repository::CatalogZoneStateRepository> {
220    let pool = get_pool();
221    repository::RepositoryFactory::create_catalog_zone_state_repository(pool)
222}