Menu
Public repository / master
AkurAI-Drive / backend/src/db/mod.rs
Latest change beff27fc3a062a220cde556e7bb28da49f6d04ea - Drive login: recover oidc.rs (IdP PKCE + introspection) from deployed binary 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)
}
/// Drop a pre-SSO table that lacks `column`; it holds only transient login
/// state, so it is recreated empty by [`create_tables`].
fn drop_if_missing(conn: &Connection, table: &str, column: &str) {
let has_column = conn
.prepare(&format!("PRAGMA table_info({table})"))
.and_then(|mut stmt| {
let names = stmt.query_map([], |row| row.get::<_, String>(1))?;
names.collect::<Result<Vec<_>, _>>()
})
.map(|names| names.is_empty() || names.iter().any(|n| n == column))
.unwrap_or(true);
if !has_column {
tracing::warn!("dropping pre-SSO `{table}` table (missing `{column}`)");
conn.execute_batch(&format!("DROP TABLE {table};")).ok();
}
}
fn create_tables(conn: &Connection) {
drop_if_missing(conn, "sessions", "id_token");
drop_if_missing(conn, "oauth_states", "verifier");
conn.execute_batch(
"
CREATE TABLE IF NOT EXISTS sessions (
id TEXT PRIMARY KEY,
sub TEXT NOT NULL,
email TEXT NOT NULL,
role TEXT NOT NULL, -- owner | admin | member
organization_id TEXT,
tenant_id TEXT,
groups TEXT NOT NULL, -- JSON array of group names
id_token TEXT NOT NULL, -- kept for RP-initiated logout
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);
-- Server-side half of an in-flight login: the nonce and PKCE verifier
-- plus the same-origin `return_to`, keyed on the OAuth `state`.
-- Single use, TTL 600 s.
CREATE TABLE IF NOT EXISTS oauth_states (
state TEXT PRIMARY KEY,
nonce TEXT NOT NULL,
verifier TEXT NOT NULL,
return_to TEXT NOT NULL,
created INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS oauth_states_created_idx ON oauth_states(created);
",
)
.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();
}
#[cfg(test)]
mod tests {
use super::*;
fn columns(conn: &Connection, table: &str) -> Vec<String> {
let mut stmt = conn
.prepare(&format!("PRAGMA table_info({table})"))
.unwrap();
stmt.query_map([], |r| r.get(1))
.unwrap()
.map(Result::unwrap)
.collect()
}
#[test]
fn upgrades_pre_sso_login_tables() {
let conn = Connection::open_in_memory().unwrap();
conn.execute_batch(
"CREATE TABLE sessions (id TEXT PRIMARY KEY, sub TEXT, email TEXT, groups TEXT,
tenant_id TEXT, created INTEGER, expires INTEGER);
INSERT INTO sessions VALUES ('x','u','e','[]',NULL,0,0);
CREATE TABLE oauth_states (state TEXT PRIMARY KEY, return_to TEXT, created INTEGER);
CREATE TABLE 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);
INSERT INTO storage VALUES ('s','u','o','e','p','b','r',0);",
)
.unwrap();
create_tables(&conn);
assert!(columns(&conn, "sessions").contains(&"id_token".to_string()));
assert!(columns(&conn, "oauth_states").contains(&"verifier".to_string()));
let kept: i64 = conn
.query_row("SELECT COUNT(*) FROM storage", [], |r| r.get(0))
.unwrap();
assert_eq!(kept, 1, "storage rows survive the login-table upgrade");
// Idempotent on an already-current schema.
create_tables(&conn);
}
}