from __future__ import annotations import sqlite3 from pathlib import Path from ..security import hash_lookup_token, new_bearer_secret BUSY_TIMEOUT_MS = 10_000 CACHE_SIZE_KIB = 20_000 MMAP_SIZE_BYTES = 256 * 1024 * 1024 WAL_AUTOCHECKPOINT_PAGES = 1_000 SCHEMA_VERSION = 3 def connect_database(database_path: Path) -> sqlite3.Connection: if str(database_path) != ":memory:": database_path.parent.mkdir(parents=True, exist_ok=True) conn = sqlite3.connect(database_path, check_same_thread=False) conn.row_factory = sqlite3.Row configure_connection(conn) return conn def configure_connection(conn: sqlite3.Connection) -> None: conn.execute("PRAGMA foreign_keys = ON") conn.execute(f"PRAGMA busy_timeout = {BUSY_TIMEOUT_MS}") conn.execute("PRAGMA temp_store = MEMORY") conn.execute(f"PRAGMA cache_size = -{CACHE_SIZE_KIB}") conn.execute(f"PRAGMA mmap_size = {MMAP_SIZE_BYTES}") conn.execute("PRAGMA synchronous = NORMAL") conn.execute(f"PRAGMA wal_autocheckpoint = {WAL_AUTOCHECKPOINT_PAGES}") def initialize_schema(conn: sqlite3.Connection) -> None: configure_connection(conn) schema_version = conn.execute("PRAGMA user_version").fetchone()[0] conn.execute("PRAGMA journal_mode = WAL") conn.execute(""" CREATE TABLE IF NOT EXISTS records ( hostname TEXT PRIMARY KEY, ipv4 TEXT, ipv6 TEXT, updated_at TEXT NOT NULL ) """) conn.execute(""" CREATE TABLE IF NOT EXISTS update_events ( id INTEGER PRIMARY KEY AUTOINCREMENT, hostname TEXT NOT NULL, ipv4 TEXT, ipv6 TEXT, status TEXT NOT NULL, detail TEXT NOT NULL, source_ip TEXT NOT NULL, created_at TEXT NOT NULL ) """) conn.execute("CREATE INDEX IF NOT EXISTS idx_update_events_created_at ON update_events(created_at)") conn.execute( "CREATE INDEX IF NOT EXISTS idx_update_events_hostname_created_at ON update_events(hostname, created_at DESC)" ) conn.execute(""" CREATE TABLE IF NOT EXISTS update_event_rollups ( id INTEGER PRIMARY KEY AUTOINCREMENT, hostname TEXT NOT NULL, ipv4 TEXT, ipv6 TEXT, status TEXT NOT NULL, detail TEXT NOT NULL, source_ip TEXT NOT NULL, first_seen_at TEXT NOT NULL, last_seen_at TEXT NOT NULL, count INTEGER NOT NULL, UNIQUE(hostname, ipv4, ipv6, status, detail, source_ip) ) """) conn.execute( "CREATE INDEX IF NOT EXISTS idx_update_event_rollups_last_seen_at ON update_event_rollups(last_seen_at)" ) conn.execute(""" CREATE INDEX IF NOT EXISTS idx_update_event_rollups_hostname_last_seen_at ON update_event_rollups(hostname, last_seen_at DESC) """) conn.execute(""" CREATE TABLE IF NOT EXISTS cleanup_runs ( id INTEGER PRIMARY KEY AUTOINCREMENT, created_at TEXT NOT NULL, unused_accounts_deleted INTEGER NOT NULL ) """) _migrate_cleanup_runs_schema(conn) conn.execute("CREATE INDEX IF NOT EXISTS idx_cleanup_runs_created_at ON cleanup_runs(created_at)") conn.execute(""" CREATE TABLE IF NOT EXISTS scheduled_tasks ( name TEXT PRIMARY KEY, interval_seconds INTEGER NOT NULL, next_run_at TEXT NOT NULL, lease_owner TEXT, lease_expires_at TEXT, last_started_at TEXT, last_finished_at TEXT, last_status TEXT NOT NULL DEFAULT 'pending', last_error TEXT NOT NULL DEFAULT '', updated_at TEXT NOT NULL ) """) conn.execute("CREATE INDEX IF NOT EXISTS idx_scheduled_tasks_next_run_at ON scheduled_tasks(next_run_at)") conn.execute(""" CREATE INDEX IF NOT EXISTS idx_scheduled_tasks_due ON scheduled_tasks(next_run_at, lease_expires_at) """) conn.execute(""" CREATE TABLE IF NOT EXISTS rate_limits ( key TEXT NOT NULL, window INTEGER NOT NULL, count INTEGER NOT NULL, PRIMARY KEY(key, window) ) """) conn.execute("CREATE INDEX IF NOT EXISTS idx_rate_limits_window ON rate_limits(window)") conn.execute(""" CREATE TABLE IF NOT EXISTS service_metrics ( name TEXT PRIMARY KEY, value REAL NOT NULL, updated_at TEXT NOT NULL ) """) conn.execute(""" CREATE TABLE IF NOT EXISTS accounts ( hostname TEXT PRIMARY KEY, update_slug TEXT NOT NULL UNIQUE, management_hash TEXT UNIQUE, dashboard_password_hash TEXT, created_at TEXT NOT NULL, accepted_at TEXT, disabled INTEGER NOT NULL DEFAULT 0 ) """) _ensure_column(conn, "accounts", "update_slug", "TEXT") _ensure_column(conn, "accounts", "management_hash", "TEXT") _ensure_column(conn, "accounts", "dashboard_password_hash", "TEXT") _ensure_column(conn, "accounts", "accepted_at", "TEXT") if schema_version < 2: conn.execute("UPDATE accounts SET accepted_at = created_at WHERE accepted_at IS NULL") for row in conn.execute("SELECT hostname FROM accounts WHERE update_slug IS NULL OR update_slug = ''"): conn.execute("UPDATE accounts SET update_slug = ? WHERE hostname = ?", (new_bearer_secret(), row[0])) for row in conn.execute("SELECT hostname FROM accounts WHERE management_hash IS NULL OR management_hash = ''"): conn.execute( "UPDATE accounts SET management_hash = ? WHERE hostname = ?", (hash_lookup_token(new_bearer_secret()), row[0]), ) _migrate_accounts_schema(conn) conn.execute("CREATE UNIQUE INDEX IF NOT EXISTS idx_accounts_update_slug ON accounts(update_slug)") conn.execute("CREATE UNIQUE INDEX IF NOT EXISTS idx_accounts_management_hash ON accounts(management_hash)") conn.execute("CREATE INDEX IF NOT EXISTS idx_accounts_accepted_at ON accounts(accepted_at)") conn.execute("CREATE INDEX IF NOT EXISTS idx_accounts_created_at ON accounts(created_at)") if schema_version < SCHEMA_VERSION: conn.execute(f"PRAGMA user_version = {SCHEMA_VERSION}") conn.execute("PRAGMA optimize") def _ensure_column(conn: sqlite3.Connection, table: str, column: str, definition: str) -> None: columns = {row[1] for row in conn.execute(f"PRAGMA table_info({table})")} if column not in columns: conn.execute(f"ALTER TABLE {table} ADD COLUMN {column} {definition}") def _migrate_accounts_schema(conn: sqlite3.Connection) -> None: columns = {row[1] for row in conn.execute("PRAGMA table_info(accounts)")} if "username" not in columns and "token_hash" not in columns: return conn.execute("DROP INDEX IF EXISTS idx_accounts_update_slug") conn.execute("DROP INDEX IF EXISTS idx_accounts_management_hash") conn.execute("DROP INDEX IF EXISTS idx_accounts_accepted_at") conn.execute("DROP INDEX IF EXISTS idx_accounts_created_at") conn.execute("DROP INDEX IF EXISTS idx_accounts_hostname_username") conn.execute(""" CREATE TABLE accounts_new ( hostname TEXT PRIMARY KEY, update_slug TEXT NOT NULL UNIQUE, management_hash TEXT UNIQUE, dashboard_password_hash TEXT, created_at TEXT NOT NULL, accepted_at TEXT, disabled INTEGER NOT NULL DEFAULT 0 ) """) conn.execute(""" INSERT INTO accounts_new( hostname, update_slug, management_hash, dashboard_password_hash, created_at, accepted_at, disabled ) SELECT hostname, update_slug, management_hash, dashboard_password_hash, created_at, accepted_at, disabled FROM accounts """) conn.execute("DROP TABLE accounts") conn.execute("ALTER TABLE accounts_new RENAME TO accounts") def _migrate_cleanup_runs_schema(conn: sqlite3.Connection) -> None: columns = {row[1] for row in conn.execute("PRAGMA table_info(cleanup_runs)")} if columns <= {"id", "created_at", "unused_accounts_deleted"}: return conn.execute("DROP INDEX IF EXISTS idx_cleanup_runs_created_at") conn.execute(""" CREATE TABLE cleanup_runs_new ( id INTEGER PRIMARY KEY AUTOINCREMENT, created_at TEXT NOT NULL, unused_accounts_deleted INTEGER NOT NULL ) """) conn.execute(""" INSERT INTO cleanup_runs_new(id, created_at, unused_accounts_deleted) SELECT id, created_at, unused_accounts_deleted FROM cleanup_runs """) conn.execute("DROP TABLE cleanup_runs") conn.execute("ALTER TABLE cleanup_runs_new RENAME TO cleanup_runs")