Skip to main content

kftray_commons/utils/
db_mode.rs

1use std::sync::Arc;
2use std::sync::{
3    LazyLock,
4    Mutex,
5};
6
7use sqlx::SqlitePool;
8
9use crate::db::{
10    create_db_table,
11    get_db_pool,
12};
13
14#[derive(Debug, Clone, PartialEq, Eq, Hash, Default, Copy)]
15pub enum DatabaseMode {
16    #[default]
17    File,
18    Memory,
19}
20
21pub struct DatabaseContext {
22    pub pool: Arc<SqlitePool>,
23    pub mode: DatabaseMode,
24}
25
26pub struct DatabaseManager;
27
28static MEMORY_DB_POOL: LazyLock<Mutex<Option<Arc<SqlitePool>>>> =
29    LazyLock::new(|| Mutex::new(None));
30
31impl DatabaseManager {
32    pub async fn get_context(mode: DatabaseMode) -> Result<DatabaseContext, String> {
33        match mode {
34            DatabaseMode::File => {
35                let pool = get_db_pool().await.map_err(|e| e.to_string())?;
36                Ok(DatabaseContext { pool, mode })
37            }
38            DatabaseMode::Memory => {
39                {
40                    let pool_guard = MEMORY_DB_POOL.lock().unwrap();
41                    if let Some(pool) = pool_guard.as_ref() {
42                        return Ok(DatabaseContext {
43                            pool: pool.clone(),
44                            mode,
45                        });
46                    }
47                }
48
49                let connection_string = "sqlite::memory:";
50
51                let pool = Arc::new(
52                    SqlitePool::connect(connection_string)
53                        .await
54                        .map_err(|e| e.to_string())?,
55                );
56                create_db_table(&pool).await.map_err(|e| e.to_string())?;
57                if let Err(error) =
58                    crate::utils::settings::establish_expose_history_baseline_at_init(&pool, mode)
59                        .await
60                {
61                    // Bookkeeping only: expose::kubernetes::ensure_expose_history_baseline
62                    // re-establishes it lazily, using the snapshot taken above
63                    // of which config ids already existed, so a
64                    // configuration inserted after this point is never
65                    // mistaken for one that predates ingress history.
66                    log::warn!("Failed to establish the expose history baseline: {error}");
67                }
68                crate::utils::migration::migrate_configs(Some(&pool))
69                    .await
70                    .map_err(|e| e.to_string())?;
71
72                {
73                    let mut pool_guard = MEMORY_DB_POOL.lock().unwrap();
74                    *pool_guard = Some(pool.clone());
75                }
76
77                Ok(DatabaseContext { pool, mode })
78            }
79        }
80    }
81
82    pub fn cleanup_memory_pools() {
83        let mut pool_guard = MEMORY_DB_POOL.lock().unwrap();
84        *pool_guard = None;
85    }
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91
92    #[tokio::test]
93    async fn test_database_mode_default() {
94        let mode = DatabaseMode::default();
95        assert_eq!(mode, DatabaseMode::File);
96    }
97
98    #[tokio::test]
99    async fn test_database_context_memory() {
100        let context = DatabaseManager::get_context(DatabaseMode::Memory)
101            .await
102            .unwrap();
103        assert_eq!(context.mode, DatabaseMode::Memory);
104        assert!(!context.pool.is_closed());
105    }
106
107    #[tokio::test]
108    async fn test_database_context_file() {
109        let context = DatabaseManager::get_context(DatabaseMode::File).await;
110        if let Ok(ctx) = context {
111            assert_eq!(ctx.mode, DatabaseMode::File);
112        }
113    }
114}