kftray_commons/utils/
db.rs

1use std::fs::{
2    self,
3    File,
4};
5use std::io::Write;
6use std::path::Path;
7use std::sync::Arc;
8use std::sync::Mutex;
9
10use lazy_static::lazy_static;
11use log::{
12    error,
13    info,
14    warn,
15};
16use sqlx::SqlitePool;
17use tokio::sync::OnceCell;
18
19use crate::config_dir::{
20    get_db_file_path,
21    get_pod_manifest_path,
22};
23use crate::utils::db_mode::DatabaseMode;
24use crate::utils::manifests::{
25    create_expose_deployment_manifest,
26    create_expose_ingress_manifest,
27    create_expose_service_manifest,
28    create_proxy_deployment_manifest,
29    expose_deployment_manifest_exists,
30    expose_ingress_manifest_exists,
31    expose_service_manifest_exists,
32    migrate_expose_deployment_manifest_if_previous_default,
33    migrate_pod_manifest_if_previous_default,
34    migrate_proxy_deployment_manifest_if_previous_default,
35    proxy_deployment_manifest_exists,
36};
37
38lazy_static! {
39    static ref ENV_TEST_MUTEX: Mutex<()> = Mutex::new(());
40}
41
42pub async fn init() -> Result<(), Box<dyn std::error::Error>> {
43    if !db_file_exists() {
44        create_db_file()?;
45    }
46
47    if !pod_manifest_file_exists() {
48        create_server_config_manifest()?;
49    } else if let Err(error) = migrate_pod_manifest_if_previous_default() {
50        warn!("Failed to migrate pod manifest: {error}");
51    }
52
53    if !proxy_deployment_manifest_exists() {
54        info!("Creating proxy deployment manifest");
55        create_proxy_deployment_manifest()?;
56    } else if let Err(error) = migrate_proxy_deployment_manifest_if_previous_default() {
57        warn!("Failed to migrate proxy deployment manifest: {error}");
58    }
59
60    if !expose_deployment_manifest_exists() {
61        info!("Creating expose deployment manifest");
62        create_expose_deployment_manifest()?;
63    } else if let Err(error) = migrate_expose_deployment_manifest_if_previous_default() {
64        warn!("Failed to migrate expose deployment manifest: {error}");
65    }
66
67    if !expose_service_manifest_exists() {
68        info!("Creating expose service manifest");
69        create_expose_service_manifest()?;
70    }
71
72    if !expose_ingress_manifest_exists() {
73        info!("Creating expose ingress manifest");
74        create_expose_ingress_manifest()?;
75    }
76
77    let pool = get_db_pool().await.map_err(|e| e.to_string())?;
78    create_db_table(&pool).await?;
79    if let Err(error) =
80        crate::utils::settings::establish_expose_history_baseline_at_init(&pool, DatabaseMode::File)
81            .await
82    {
83        // Bookkeeping only: expose::kubernetes::ensure_expose_history_baseline
84        // re-establishes it lazily, using the snapshot taken above of
85        // which config ids already existed, so a configuration inserted
86        // after this point is never mistaken for one that predates ingress
87        // history. A transient SQLITE_BUSY from another kftray process
88        // sharing this file database must not stop this one from starting.
89        warn!("Failed to establish the expose history baseline: {error}");
90    }
91
92    Ok(())
93}
94
95pub static DB_POOL: OnceCell<Arc<SqlitePool>> = OnceCell::const_new();
96
97pub async fn get_db_pool() -> Result<Arc<SqlitePool>, String> {
98    DB_POOL
99        .get_or_try_init(|| async {
100            let db_dir = get_db_file_path().map_err(|e| {
101                error!("Failed to get DB file path: {e}");
102                e.to_string()
103            })?;
104            let db_dir_str = db_dir.to_str().ok_or("Invalid DB path")?;
105            info!("Database file path: {db_dir_str}");
106            let pool = SqlitePool::connect(db_dir_str).await.map_err(|e| {
107                error!("Failed to connect to DB: {e}");
108                e.to_string()
109            })?;
110            Ok(Arc::new(pool))
111        })
112        .await
113        .map(Arc::clone)
114}
115
116pub async fn create_db_table(pool: &SqlitePool) -> Result<(), sqlx::Error> {
117    info!("Creating database tables and triggers.");
118    let mut conn = pool.acquire().await.map_err(|e| {
119        error!("Failed to acquire connection: {e}");
120        e
121    })?;
122
123    sqlx::query("PRAGMA foreign_keys = ON;")
124        .execute(&mut *conn)
125        .await
126        .map_err(|e| {
127            error!("Failed to set PRAGMA foreign_keys: {e}");
128            e
129        })?;
130
131    sqlx::query(
132        "CREATE TABLE IF NOT EXISTS configs (
133            id INTEGER PRIMARY KEY,
134            data TEXT NOT NULL
135        )",
136    )
137    .execute(&mut *conn)
138    .await
139    .map_err(|e| {
140        error!("Failed to create configs table: {e}");
141        e
142    })?;
143
144    sqlx::query(
145        "CREATE TABLE IF NOT EXISTS config_state (
146            id INTEGER PRIMARY KEY,
147            config_id INTEGER NOT NULL,
148            is_running BOOLEAN NOT NULL DEFAULT false,
149            process_id INTEGER,
150            FOREIGN KEY(config_id) REFERENCES configs(id) ON DELETE CASCADE
151        )",
152    )
153    .execute(&mut *conn)
154    .await
155    .map_err(|e| {
156        error!("Failed to create config_state table: {e}");
157        e
158    })?;
159
160    sqlx::query(
161        "CREATE TRIGGER IF NOT EXISTS after_insert_config
162         AFTER INSERT ON configs
163         FOR EACH ROW
164         BEGIN
165             INSERT INTO config_state (config_id, is_running) VALUES (NEW.id, false);
166         END;",
167    )
168    .execute(&mut *conn)
169    .await
170    .map_err(|e| {
171        error!("Failed to create after_insert_config trigger: {e}");
172        e
173    })?;
174
175    sqlx::query(
176        "CREATE TRIGGER IF NOT EXISTS after_delete_config
177         AFTER DELETE ON configs
178         FOR EACH ROW
179         BEGIN
180             DELETE FROM config_state WHERE config_id = OLD.id;
181         END;",
182    )
183    .execute(&mut *conn)
184    .await
185    .map_err(|e| {
186        error!("Failed to create after_delete_config trigger: {e}");
187        e
188    })?;
189
190    sqlx::query(
191        "CREATE TABLE IF NOT EXISTS settings (
192            key TEXT PRIMARY KEY,
193            value TEXT NOT NULL,
194            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
195        )",
196    )
197    .execute(&mut *conn)
198    .await
199    .map_err(|e| {
200        error!("Failed to create settings table: {e}");
201        e
202    })?;
203
204    info!("Database tables and triggers created successfully.");
205    Ok(())
206}
207
208fn pod_manifest_file_exists() -> bool {
209    match get_pod_manifest_path() {
210        Ok(path) => {
211            let exists = path.exists();
212            if cfg!(test) {
213                println!(
214                    "pod_manifest_file_exists checking path: {}, exists: {}",
215                    path.display(),
216                    exists
217                );
218            }
219            exists
220        }
221        Err(e) => {
222            if cfg!(test) {
223                println!("pod_manifest_file_exists failed to get path: {e}");
224            }
225            false
226        }
227    }
228}
229
230fn create_server_config_manifest() -> Result<(), std::io::Error> {
231    let manifest_path = get_pod_manifest_path().map_err(std::io::Error::other)?;
232
233    let manifest_dir = manifest_path
234        .parent()
235        .ok_or_else(|| std::io::Error::other("Failed to get manifest directory"))?;
236
237    if !manifest_dir.exists() {
238        fs::create_dir_all(manifest_dir)?;
239    }
240
241    let manifest_json =
242        serde_json::to_string_pretty(&crate::utils::manifests::default_pod_manifest())?;
243
244    File::create(&manifest_path)?.write_all(manifest_json.as_bytes())
245}
246
247fn db_file_exists() -> bool {
248    match get_db_file_path() {
249        Ok(db_path) => {
250            let exists = db_path.exists();
251            if cfg!(test) {
252                println!(
253                    "db_file_exists checking path: {}, exists: {}",
254                    db_path.display(),
255                    exists
256                );
257            }
258            exists
259        }
260        Err(e) => {
261            if cfg!(test) {
262                println!("db_file_exists failed to get path: {e}");
263            }
264            false
265        }
266    }
267}
268
269fn create_db_file() -> Result<(), std::io::Error> {
270    let db_path = get_db_file_path().map_err(std::io::Error::other)?;
271
272    let db_dir = Path::new(&db_path)
273        .parent()
274        .expect("Failed to get db directory");
275
276    if !db_dir.exists() {
277        fs::create_dir_all(db_dir)?;
278    }
279
280    fs::File::create(db_path)?;
281
282    Ok(())
283}
284
285#[cfg(test)]
286mod tests {
287    use std::env;
288    use std::fs::{
289        self,
290        File,
291    };
292    use std::sync::Mutex;
293
294    use lazy_static::lazy_static;
295    use sqlx::SqlitePool;
296    use tempfile::tempdir;
297
298    use super::*;
299    use crate::config_dir::{
300        get_config_dir,
301        get_db_file_path,
302        get_pod_manifest_path,
303    };
304
305    lazy_static! {
306        static ref ENV_TEST_MUTEX: Mutex<()> = Mutex::new(());
307    }
308
309    struct StrictEnvGuard {
310        saved_vars: Vec<(String, Option<String>)>,
311    }
312
313    impl StrictEnvGuard {
314        fn new(keys: &[&str]) -> Self {
315            let saved_vars = keys
316                .iter()
317                .map(|&key| (key.to_string(), env::var(key).ok()))
318                .collect::<Vec<_>>();
319
320            for key in keys {
321                unsafe { env::remove_var(key) };
322            }
323
324            StrictEnvGuard { saved_vars }
325        }
326    }
327
328    impl Drop for StrictEnvGuard {
329        fn drop(&mut self) {
330            for (key, value) in self.saved_vars.drain(..) {
331                match value {
332                    Some(val) => unsafe { env::set_var(key, val) },
333
334                    None => unsafe { env::remove_var(key) },
335                }
336            }
337        }
338    }
339
340    #[test]
341    fn test_db_file_exists_and_create() {
342        let _lock = ENV_TEST_MUTEX.lock().unwrap();
343        let _guard = StrictEnvGuard::new(&["KFTRAY_CONFIG", "XDG_CONFIG_HOME", "HOME"]);
344
345        let temp_dir = tempdir().unwrap();
346        let temp_path = temp_dir.path();
347
348        let db_path = temp_path.join("configs.db");
349
350        unsafe { env::set_var("KFTRAY_CONFIG", temp_path.to_str().unwrap()) };
351
352        assert_eq!(
353            get_config_dir().unwrap().to_str().unwrap(),
354            temp_path.to_str().unwrap(),
355            "Config directory should match our temporary directory"
356        );
357
358        assert!(!db_path.exists(), "DB file should not exist initially");
359        assert!(
360            !db_file_exists(),
361            "db_file_exists() should return false initially"
362        );
363
364        File::create(&db_path).unwrap();
365        assert!(db_path.exists(), "DB file should exist after creation");
366
367        let db_exists = db_file_exists();
368
369        assert!(
370            db_exists,
371            "db_file_exists() should return true after file creation"
372        );
373    }
374
375    #[test]
376    fn test_pod_manifest_file_exists_and_create() {
377        let _lock = ENV_TEST_MUTEX.lock().unwrap();
378        let _guard = StrictEnvGuard::new(&["KFTRAY_CONFIG", "XDG_CONFIG_HOME", "HOME"]);
379
380        let temp_dir = tempdir().unwrap();
381        let test_dir = temp_dir.path();
382        std::fs::create_dir_all(test_dir).unwrap();
383        assert!(test_dir.exists(), "Test directory should exist");
384
385        let config_path = test_dir.to_str().unwrap();
386
387        unsafe { env::set_var("KFTRAY_CONFIG", config_path) };
388        println!("Set KFTRAY_CONFIG to: {config_path}");
389
390        assert!(env::var("HOME").is_err(), "HOME should not be set");
391        assert!(
392            env::var("XDG_CONFIG_HOME").is_err(),
393            "XDG_CONFIG_HOME should not be set"
394        );
395        assert_eq!(
396            env::var("KFTRAY_CONFIG").unwrap(),
397            config_path,
398            "KFTRAY_CONFIG should be set to test dir"
399        );
400
401        let expected_manifest_path = test_dir.join("proxy_manifest.json");
402        println!(
403            "Expected manifest path: {}",
404            expected_manifest_path.display()
405        );
406
407        assert!(
408            !expected_manifest_path.exists(),
409            "Manifest file should not exist initially"
410        );
411        assert!(
412            !pod_manifest_file_exists(),
413            "pod_manifest_file_exists() should return false initially"
414        );
415
416        println!("Creating manifest file...");
417        create_server_config_manifest().unwrap();
418
419        assert!(
420            expected_manifest_path.exists(),
421            "Manifest file should exist at: {}",
422            expected_manifest_path.display()
423        );
424
425        let func_result = pod_manifest_file_exists();
426        println!("pod_manifest_file_exists() returned: {func_result}");
427        assert!(func_result, "pod_manifest_file_exists() should return true");
428
429        let content = fs::read_to_string(&expected_manifest_path).unwrap();
430        assert!(
431            content.contains("apiVersion"),
432            "Manifest should contain apiVersion"
433        );
434        assert!(
435            content.contains("kftray-server"),
436            "Manifest should contain kftray-server"
437        );
438    }
439
440    #[test]
441    fn test_pod_manifest_create_directory() {
442        let _lock = ENV_TEST_MUTEX.lock().unwrap();
443        let _env_guard = StrictEnvGuard::new(&["KFTRAY_CONFIG", "XDG_CONFIG_HOME", "HOME"]);
444
445        let temp_dir = tempdir().unwrap();
446        let temp_dir_path = temp_dir.path().to_str().unwrap().to_string();
447
448        let manifest_dir = std::path::Path::new(&temp_dir_path).join("manifest_dir");
449        std::fs::create_dir_all(&manifest_dir).unwrap();
450
451        let manifest_dir_str = manifest_dir.to_str().unwrap();
452
453        unsafe { env::set_var("KFTRAY_CONFIG", manifest_dir_str) };
454
455        assert!(manifest_dir.exists(), "Directory should exist");
456
457        let manifest_path = manifest_dir.join("proxy_manifest.json");
458        let content = r#"{"apiVersion":"v1","kind":"Pod"}"#;
459        std::fs::write(&manifest_path, content).unwrap();
460
461        // Verify file exists
462        assert!(manifest_path.exists(), "Manifest file should exist at path");
463        println!("Created manifest file at: {}", manifest_path.display());
464
465        if let Ok(config_path) = get_config_dir() {
466            println!(
467                "Config dir from get_config_dir(): {}",
468                config_path.display()
469            );
470            let expected_manifest = config_path.join("proxy_manifest.json");
471            println!("Expected manifest path: {}", expected_manifest.display());
472            println!(
473                "File exists at expected path: {}",
474                expected_manifest.exists()
475            );
476        }
477
478        let start = std::time::Instant::now();
479        let timeout = std::time::Duration::from_secs(5);
480        let poll_interval = std::time::Duration::from_millis(50);
481
482        let mut result = false;
483        while start.elapsed() < timeout {
484            result = pod_manifest_file_exists();
485            if result {
486                break;
487            }
488            std::thread::sleep(poll_interval);
489        }
490
491        assert!(result, "pod_manifest_file_exists() should return true");
492
493        let file_content = std::fs::read_to_string(&manifest_path).unwrap();
494        assert!(file_content.contains("apiVersion"));
495    }
496
497    #[tokio::test]
498    async fn test_create_db_table() {
499        let pool = SqlitePool::connect("sqlite::memory:")
500            .await
501            .expect("Failed to connect to in-memory database");
502        let result = create_db_table(&pool).await;
503        assert!(result.is_ok());
504
505        let mut conn = pool.acquire().await.unwrap();
506        let result =
507            sqlx::query("SELECT name FROM sqlite_master WHERE type='table' AND name='configs'")
508                .fetch_optional(&mut *conn)
509                .await;
510
511        assert!(result.is_ok());
512        assert!(result.unwrap().is_some(), "configs table should exist");
513
514        let result = sqlx::query(
515            "SELECT name FROM sqlite_master WHERE type='table' AND name='config_state'",
516        )
517        .fetch_optional(&mut *conn)
518        .await;
519
520        assert!(result.is_ok());
521        assert!(result.unwrap().is_some(), "config_state table should exist");
522
523        let result = sqlx::query(
524            "SELECT name FROM sqlite_master WHERE type='trigger' AND name='after_insert_config'",
525        )
526        .fetch_optional(&mut *conn)
527        .await;
528
529        assert!(result.is_ok());
530        assert!(
531            result.unwrap().is_some(),
532            "after_insert_config trigger should exist"
533        );
534    }
535
536    #[test]
537    fn test_init_creates_files_and_db_direct() {
538        let _lock = ENV_TEST_MUTEX.lock().unwrap();
539        let _guard = StrictEnvGuard::new(&["KFTRAY_CONFIG", "XDG_CONFIG_HOME", "HOME"]);
540
541        let temp_dir = tempdir().unwrap();
542        let temp_path = temp_dir.path();
543
544        unsafe { env::set_var("KFTRAY_CONFIG", temp_path.to_str().unwrap()) };
545
546        let db_path = temp_path.join("configs.db");
547        let manifest_path = temp_path.join("proxy_manifest.json");
548
549        assert!(!db_path.exists(), "DB file should not exist initially");
550        assert!(
551            !manifest_path.exists(),
552            "Manifest file should not exist initially"
553        );
554
555        let db_result = create_db_file();
556        assert!(db_result.is_ok(), "create_db_file() should succeed");
557        assert!(db_path.exists(), "DB file should exist after creation");
558
559        let manifest_content = r#"{"apiVersion":"v1","kind":"Pod"}"#;
560        let manifest_result = fs::write(&manifest_path, manifest_content);
561        assert!(
562            manifest_result.is_ok(),
563            "Writing manifest file should succeed"
564        );
565        assert!(
566            manifest_path.exists(),
567            "Manifest file should exist after creation"
568        );
569
570        println!("Created DB file at: {}", db_path.to_str().unwrap());
571        println!("Does path exist? {}", db_path.exists());
572
573        if let Ok(cfg_path) = get_db_file_path() {
574            println!("Config DB path: {}", cfg_path.display());
575            assert!(
576                cfg_path.exists(),
577                "Path returned by get_db_file_path() must exist"
578            );
579            assert_eq!(
580                cfg_path, db_path,
581                "get_db_file_path() should return the expected path"
582            );
583            assert!(db_file_exists(), "db_file_exists() should return true");
584        } else {
585            panic!("get_db_file_path() failed unexpectedly");
586        }
587
588        if let Ok(manifest_cfg_path) = get_pod_manifest_path() {
589            println!("Config manifest path: {}", manifest_cfg_path.display());
590            assert!(
591                manifest_cfg_path.exists(),
592                "Path returned by get_pod_manifest_path() must exist"
593            );
594            assert_eq!(
595                manifest_cfg_path, manifest_path,
596                "get_pod_manifest_path() should return the expected path"
597            );
598            assert!(
599                pod_manifest_file_exists(),
600                "pod_manifest_file_exists() should return true"
601            );
602        } else {
603            panic!("get_pod_manifest_path() failed unexpectedly");
604        }
605    }
606}
kftray_commons/utils/db.rs

kftray_commons/utils/
db.rs
