AkurAI-Monitor / src/db.rs

Latest change d0b8d14afb524ab72435e5abd0bbd6049a205349 - Migrate storage layer from SQLite to MariaDB by Ólafur Búi Ólafsson

use crate::config;
use crate::schema;
use mysql::{Opts, OptsBuilder, Pool, PoolConstraints, PoolOpts, PooledConn};
use std::sync::LazyLock;

static POOL: LazyLock<Pool> = LazyLock::new(|| {
    let cfg = config::get();
    let opts = Opts::from_url(&cfg.db_url).expect("invalid MONITOR_DB_URL");
    let constraints = PoolConstraints::new(1, 8).expect("valid pool constraints");
    let opts =
        OptsBuilder::from_opts(opts).pool_opts(PoolOpts::default().with_constraints(constraints));
    let pool = Pool::new(opts).expect("failed to create MariaDB pool");
    let mut conn = pool.get_conn().expect("failed to connect to MariaDB");
    schema::create_tables(&mut conn);
    schema::seed_alert_rules(&mut conn);
    pool
});

/// Run `f` with a pooled MariaDB connection. Connections are returned to the
/// pool on drop; the pool is bounded (1..=8) so DB memory stays flat.
pub fn with_db<F, R>(f: F) -> R
where
    F: FnOnce(&mut PooledConn) -> R,
{
    let mut conn = POOL.get_conn().expect("failed to get MariaDB connection");
    f(&mut conn)
}

/// Test support: scratch databases on the dev MariaDB instance.
///
/// Tests are skipped (with a note on stderr) unless `MONITOR_TEST_DB_URL`
/// points at a user allowed to CREATE/DROP `akurai_test_%` databases —
/// on Titan/EC2 that is the dev container on port 3307.
#[cfg(test)]
pub mod testing {
    use mysql::prelude::Queryable;
    use mysql::{Opts, OptsBuilder, Pool, PoolConstraints, PoolOpts, PooledConn};
    use std::sync::atomic::{AtomicU32, Ordering};

    static COUNTER: AtomicU32 = AtomicU32::new(0);

    pub struct TestDb {
        pool: Pool,
        admin: Pool,
        pub name: String,
    }

    impl TestDb {
        /// Returns `None` (after logging) when no test database is configured.
        pub fn new() -> Option<TestDb> {
            let Ok(url) = std::env::var("MONITOR_TEST_DB_URL") else {
                eprintln!("skipping: MONITOR_TEST_DB_URL is not set");
                return None;
            };
            let opts = Opts::from_url(&url).expect("invalid MONITOR_TEST_DB_URL");
            let constraints = PoolConstraints::new(1, 2).expect("valid pool constraints");
            let small = PoolOpts::default().with_constraints(constraints);
            let admin = Pool::new(OptsBuilder::from_opts(opts.clone()).pool_opts(small.clone()))
                .expect("connect test MariaDB");
            let name = format!(
                "akurai_test_mon_{}_{}",
                std::process::id(),
                COUNTER.fetch_add(1, Ordering::Relaxed)
            );
            admin
                .get_conn()
                .expect("admin conn")
                .query_drop(format!("CREATE DATABASE `{name}`"))
                .expect("create scratch database");
            let scoped = OptsBuilder::from_opts(opts)
                .db_name(Some(name.clone()))
                .pool_opts(small);
            let pool = Pool::new(scoped).expect("connect scratch database");
            Some(TestDb { pool, admin, name })
        }

        pub fn conn(&self) -> PooledConn {
            self.pool.get_conn().expect("scratch connection")
        }
    }

    impl Drop for TestDb {
        fn drop(&mut self) {
            if let Ok(mut conn) = self.admin.get_conn() {
                conn.query_drop(format!("DROP DATABASE IF EXISTS `{}`", self.name))
                    .ok();
            }
        }
    }
}