vacuum-wall/lib/db.py

"""Abstract database layer for Vacuum Wall.

Provides query ID constants and an abstract Database baseclass so subsystems
interact with the database through opaque query identifiers, never raw SQL.
Backend implementations (SQLite, PostgreSQL) provide the actual SQL.

Usage:
    from lib.db import Q_INSERT_USER, Database, get_db

    class SQLiteBackend(Database):
        QUERY_MAP = {
            Q_INSERT_USER: "INSERT INTO users (username, password_hash) VALUES (?, ?)",
            ...
        }
"""

from __future__ import annotations

import logging
import os
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Any, ClassVar

logger = logging.getLogger(__name__)

PROJECT_DIR = Path(__file__).resolve().parent.parent

# ---------------------------------------------------------------------------
# Query ID constants — single source of truth for all database operations
# ---------------------------------------------------------------------------

Q_INIT_TABLES = "init_tables"
Q_INSERT_USER = "insert_user"
Q_SELECT_USER_BY_NAME = "select_user_by_name"
Q_SELECT_USER_BY_ID = "select_user_by_id"
Q_UPDATE_PASSWORD = "update_password"
Q_DELETE_USER = "delete_user"
Q_UPSERT_PERMISSION = "upsert_permission"
Q_SELECT_PERMISSIONS = "select_permissions"
Q_DELETE_PERMISSIONS = "delete_permissions"
Q_INSERT_BLACKLIST = "insert_blacklist"
Q_SELECT_BLACKLIST = "select_blacklist_jti"
Q_DELETE_EXPIRED_BLACKLIST = "delete_expired_blacklist"
Q_UPSERT_REFRESH_TOKEN = "upsert_refresh_token"
Q_SELECT_REFRESH_TOKEN = "select_refresh_token"
Q_DELETE_REFRESH_TOKEN = "delete_refresh_token"
Q_INSERT_WEBAUTHN = "insert_webauthn"
Q_SELECT_WEBAUTHN_USER = "select_webauthn_user"
Q_SELECT_WEBAUTHN_ID = "select_webauthn_id"
Q_SELECT_WEBAUTHN_COUNTS = "select_webauthn_counts"
Q_DELETE_WEBAUTHN = "delete_webauthn"
Q_UPDATE_WEBAUTHN_SIGN_COUNT = "update_webauthn_sign_count"
Q_SELECT_ALL_USERS = "select_all_users"
Q_SELECT_USER_JWT_SECRET = "select_user_jwt_secret"
Q_UPDATE_JWT_SECRET = "update_jwt_secret"

# ---------------------------------------------------------------------------
# Schema DDL
# ---------------------------------------------------------------------------

INIT_SQL = """
CREATE TABLE IF NOT EXISTS users (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    username TEXT UNIQUE NOT NULL,
    password_hash TEXT NOT NULL,
    jwt_secret TEXT NOT NULL,
    created_at INTEGER NOT NULL DEFAULT (unixepoch())
    );

CREATE TABLE IF NOT EXISTS permissions (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    username TEXT NOT NULL REFERENCES users(username) ON DELETE CASCADE,
    subsystem TEXT NOT NULL,
    level TEXT NOT NULL CHECK (level IN ('read', 'rw')),
    UNIQUE(username, subsystem)
);

CREATE TABLE IF NOT EXISTS token_blacklist (
    jti TEXT PRIMARY KEY,
    token_type TEXT NOT NULL,
    expires INTEGER NOT NULL
);

CREATE TABLE IF NOT EXISTS refresh_tokens (
    username TEXT UNIQUE NOT NULL,
    jti TEXT NOT NULL,
    issued_at INTEGER NOT NULL
);

CREATE TABLE IF NOT EXISTS webauthn_creds (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    username TEXT NOT NULL REFERENCES users(username) ON DELETE CASCADE,
    credential_id TEXT NOT NULL,
    public_key TEXT NOT NULL,
    sign_count INTEGER NOT NULL DEFAULT 0,
    name TEXT NOT NULL DEFAULT '',
    transports TEXT NOT NULL DEFAULT '[]',
    UNIQUE(username, credential_id)
);

CREATE TABLE IF NOT EXISTS init_sequence (
    seq INTEGER PRIMARY KEY
);
"""


class Transaction:
    """Context manager for database transactions.

    Provides BEGIN/COMMIT/ROLLBACK semantics. Auto-commit is suppressed
    inside the transaction block.

    Usage:
        with db.in_transaction() as tx:
            tx.run(Q_INSERT_USER, ("user1", "hash"))
            tx.run(Q_UPSERT_PERMISSION, ("user1", "firewall", "rw"))
    """

    def __init__(self, parent: Database) -> None:
        self._parent = parent

    def __enter__(self) -> Transaction:
        self._parent._begin()
        self._parent._suppress_auto_commit()
        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> bool:
        try:
            if exc_type is None:
                self._parent._commit()
            else:
                self._parent._rollback()
        finally:
            self._parent._restore_auto_commit()
        return False

    def query(self, query_id: str, params: tuple = ()) -> list[dict]:
        return self._parent.query(query_id, params)

    def run(self, query_id: str, params: tuple = ()) -> int:
        return self._parent.run(query_id, params)

    def run_one(self, query_id: str, params: tuple = ()) -> int | dict:
        return self._parent.run_one(query_id, params)


class Database(ABC):
    """Abstract database interface.

    All subsystems interact with the database through this interface.
    Queries are identified by string IDs (e.g. Q_INSERT_USER) — never
    raw SQL strings.

    Connection is cached via the ``conn`` property. Prepared statements
    are auto-cached on first use.
    """

    QUERY_MAP: ClassVar[dict[str, str]] = {}

    def __init__(self, connection_string: str) -> None:
        self._connection_string = connection_string
        self._conn: Any = None
        self._prepared: dict[str, Any] = {}
        self._in_transaction = False

    @property
    def conn(self) -> Any:
        """Return the cached database connection, creating it lazily."""
        if self._conn is None:
            self._conn = self._connect(self._connection_string)
        return self._conn

    @abstractmethod
    def _connect(self, cs: str) -> Any: ...

    @abstractmethod
    def _prepare(self, sql: str) -> Any: ...

    @abstractmethod
    def _execute(self, stmt: Any, params: tuple) -> tuple[list[dict], int]:
        """Execute a prepared statement. Returns (rows, rowcount)."""
        ...

    @abstractmethod
    def _last_insert_id(self, stmt: Any) -> int | dict: ...

    @abstractmethod
    def _execute_direct(self, sql: str) -> None:
        """Execute raw SQL without prepared statements (for DDL, not implemented by base)."""

    def _begin(self) -> None:  # noqa: B027
        """Begin a transaction (not implemented by base)."""

    def _commit(self) -> None:  # noqa: B027
        """Commit a transaction (not implemented by base)."""

    def _rollback(self) -> None:  # noqa: B027
        """Rollback a transaction (not implemented by base)."""

    def _suppress_auto_commit(self) -> None:  # noqa: B027
        """Suppress auto-commit (not implemented by base)."""

    def _restore_auto_commit(self) -> None:  # noqa: B027
        """Restore auto-commit (not implemented by base)."""

    def query(self, query_id: str, params: tuple = ()) -> list[dict]:
        """Execute a SELECT query. Returns list of row dicts."""
        stmt = self._get_prepared(query_id)
        rows, _ = self._execute(stmt, params)
        return rows

    def run(self, query_id: str, params: tuple = ()) -> int:
        """Execute an INSERT/UPDATE/DELETE. Returns affected row count."""
        stmt = self._get_prepared(query_id)
        _, count = self._execute(stmt, params)
        return count

    def run_one(self, query_id: str, params: tuple = ()) -> int | dict:
        """Execute and return the last insert ID or row dict."""
        stmt = self._get_prepared(query_id)
        _, _ = self._execute(stmt, params)
        return self._last_insert_id(stmt)

    def in_transaction(self) -> Transaction:
        """Return a transaction context manager."""
        return Transaction(self)

    def init_tables(self) -> None:
        """Create schema tables if they don't exist."""
        self._execute_direct(INIT_SQL)

    def _get_prepared(self, query_id: str) -> Any:
        if query_id not in self._prepared:
            if query_id not in self.QUERY_MAP:
                raise KeyError(f"Unknown query ID: {query_id!r}")
            self._prepared[query_id] = self._prepare(self.QUERY_MAP[query_id])
        return self._prepared[query_id]


# ---------------------------------------------------------------------------
# Singleton accessor
# ---------------------------------------------------------------------------

_db_instance: Database | None = None


def _get_backend_name() -> str:
    return os.environ.get("VACUUM_WALL_DB_BACKEND", "sqlite")


def _get_db_path() -> str:
    return os.environ.get("VACUUM_WALL_DB_PATH", str(PROJECT_DIR / "data" / "auth.db"))


def get_db() -> Database:
    """Return the singleton Database instance.

    Creates the instance on first call using the backend specified by
    ``VACUUM_WALL_DB_BACKEND`` env var (default: sqlite).

    Call this at application startup to ensure the DB is initialized.
    """
    global _db_instance
    if _db_instance is None:
        backend = _get_backend_name()
        if backend == "sqlite":
            from lib.db_sqlite import SQLiteBackend

            path = _get_db_path()
            _db_instance = SQLiteBackend(path)
            Path(path).parent.mkdir(parents=True, exist_ok=True)
            _db_instance.init_tables()
        else:
            raise ValueError(f"Unknown database backend: {backend!r}")
    return _db_instance


def reset_db_for_test() -> None:
    """Reset the singleton — only for tests."""
    global _db_instance
    if _db_instance is not None:
        _db_instance = None