AkurAI-Drive / backend/src/db/mod.rs

Latest change 19126b373cd637ce1c055d17f2198a09424fe140 - fix(drive): cargo fmt + collapsible-if clippy fix (AKURAI-DRIVE-6) by AkurAI Agent

use crate::config;
use rusqlite::Connection;
use std::sync::LazyLock;
use std::sync::Mutex;

static DB: LazyLock<Mutex<Connection>> = LazyLock::new(|| {
    let cfg = config::get();
    if let Some(parent) = std::path::Path::new(&cfg.db_path).parent() {
        std::fs::create_dir_all(parent).ok();
    }
    let conn = Connection::open(&cfg.db_path).expect("failed to open SQLite database");
    conn.execute_batch("PRAGMA journal_mode = WAL; PRAGMA foreign_keys = ON;")
        .expect("failed to set pragmas");
    create_tables(&conn);
    Mutex::new(conn)
});

/// Run a closure with exclusive access to the single shared connection.
pub fn with_db<F, R>(f: F) -> R
where
    F: FnOnce(&Connection) -> R,
{
    let conn = DB.lock().unwrap();
    f(&conn)
}

fn create_tables(conn: &Connection) {
    conn.execute_batch(
        "
        CREATE TABLE IF NOT EXISTS sessions (
            id         TEXT PRIMARY KEY,
            sub        TEXT NOT NULL,
            email      TEXT NOT NULL,
            groups     TEXT NOT NULL,        -- JSON array of group names
            tenant_id  TEXT,
            created    INTEGER NOT NULL,
            expires    INTEGER NOT NULL
        );
        CREATE INDEX IF NOT EXISTS sessions_expires_idx ON sessions(expires);

        CREATE TABLE IF NOT EXISTS storage (
            id         TEXT PRIMARY KEY,
            owner_sub  TEXT NOT NULL,
            org        TEXT NOT NULL,
            email      TEXT NOT NULL,
            plan       TEXT NOT NULL,
            bucket     TEXT NOT NULL,
            region     TEXT NOT NULL,
            created         INTEGER NOT NULL,
            organization_id TEXT,
            workspace_id    TEXT
        );
        CREATE INDEX IF NOT EXISTS storage_owner_idx ON storage(owner_sub);

        CREATE TABLE IF NOT EXISTS telemetry_events (
            event_id        TEXT PRIMARY KEY,
            schema_version  INTEGER NOT NULL,
            organization_id TEXT NOT NULL,
            workspace_id    TEXT NOT NULL,
            installation_id TEXT NOT NULL,
            app_id          TEXT NOT NULL,
            event_name      TEXT NOT NULL,
            occurred_at     INTEGER NOT NULL,
            attributes      TEXT NOT NULL
        );
        CREATE INDEX IF NOT EXISTS telemetry_events_org_time_idx
            ON telemetry_events(organization_id, occurred_at);

        CREATE TABLE IF NOT EXISTS organization_daily_facts (
            organization_id    TEXT NOT NULL,
            workspace_id       TEXT NOT NULL,
            day                INTEGER NOT NULL,
            active_day         INTEGER NOT NULL DEFAULT 0,
            first_upload_at    INTEGER,
            first_retrieval_at INTEGER,
            storage_used_bytes INTEGER,
            quota_bytes        INTEGER,
            quota_pressure     TEXT,
            PRIMARY KEY(organization_id, workspace_id, day)
        );
        CREATE INDEX IF NOT EXISTS organization_daily_facts_org_day_idx
            ON organization_daily_facts(organization_id, day);


        CREATE TABLE IF NOT EXISTS device_tokens (
            id         TEXT PRIMARY KEY,
            token      TEXT NOT NULL UNIQUE,
            sub        TEXT NOT NULL,
            email      TEXT NOT NULL,
            created    INTEGER NOT NULL,
            expires    INTEGER NOT NULL,
            last_used  INTEGER NOT NULL
        );
        CREATE INDEX IF NOT EXISTS device_tokens_expires_idx ON device_tokens(expires);

        -- Carries the post-login `return_to` URL across the OAuth round-trip, keyed
        -- on the OAuth `state` (so /auth/device can resume after a fresh login).
        CREATE TABLE IF NOT EXISTS oauth_states (
            state      TEXT PRIMARY KEY,
            return_to  TEXT NOT NULL,
            created    INTEGER NOT NULL
        );
        ",
    )
    .expect("failed to create tables");

    // Backfill only opaque workspace IDs; user-entered organization labels are
    // never promoted into canonical telemetry IDs. The LIMIT bounds startup work.
    conn.execute("ALTER TABLE storage ADD COLUMN organization_id TEXT", [])
        .ok();
    conn.execute("ALTER TABLE storage ADD COLUMN workspace_id TEXT", [])
        .ok();
    conn.execute(
        "UPDATE storage SET workspace_id = id
         WHERE workspace_id IS NULL
           AND rowid IN (SELECT rowid FROM storage WHERE workspace_id IS NULL LIMIT 1000)",
        [],
    )
    .ok();
}