db.py

"""SQLite schema + helpers for kmap."""
from __future__ import annotations

import os
import sqlite3
from pathlib import Path

# Default to $KMAP_DB if set, else ./kmap.db in the user's current directory,
# else fall back to a db next to this package.
def _default_db_path() -> Path:
    env = os.environ.get("KMAP_DB")
    if env:
        return Path(env).expanduser().resolve()
    cwd = Path.cwd() / "kmap.db"
    if cwd.exists():
        return cwd
    return Path(__file__).parent.parent / "kmap.db"


DB_PATH = _default_db_path()

SCHEMA = """
CREATE TABLE IF NOT EXISTS programs (
    id         INTEGER PRIMARY KEY,
    path       TEXT UNIQUE NOT NULL,
    bug_id     TEXT,
    file_hash  TEXT,      -- so we only reparse on change
    n_calls    INTEGER,
    options    TEXT        -- JSON
);
CREATE INDEX IF NOT EXISTS idx_programs_bug ON programs(bug_id);

CREATE TABLE IF NOT EXISTS calls (
    program_id INTEGER NOT NULL,
    idx        INTEGER NOT NULL,
    name       TEXT NOT NULL,     -- "socket"
    variant    TEXT,              -- "nl_netfilter" or NULL
    raw        TEXT NOT NULL,     -- "socket$nl_netfilter"
    defines    TEXT,              -- "r0" or NULL
    PRIMARY KEY (program_id, idx),
    FOREIGN KEY (program_id) REFERENCES programs(id) ON DELETE CASCADE
);
CREATE INDEX IF NOT EXISTS idx_calls_name ON calls(name);
CREATE INDEX IF NOT EXISTS idx_calls_raw  ON calls(raw);

CREATE TABLE IF NOT EXISTS call_uses (
    program_id INTEGER NOT NULL,
    idx        INTEGER NOT NULL,
    resource   TEXT NOT NULL,
    FOREIGN KEY (program_id) REFERENCES programs(id) ON DELETE CASCADE
);

-- Per-program multiset of syscall names, for cheap similarity queries.
CREATE TABLE IF NOT EXISTS program_tokens (
    program_id INTEGER NOT NULL,
    token      TEXT NOT NULL,       -- raw syscall (with variant)
    count      INTEGER NOT NULL,
    PRIMARY KEY (program_id, token)
);
CREATE INDEX IF NOT EXISTS idx_tok_token ON program_tokens(token);

-- Syscall n-gram counts over whole corpus.
CREATE TABLE IF NOT EXISTS ngrams (
    n      INTEGER NOT NULL,
    gram   TEXT NOT NULL,           -- "open|ioctl|mmap"
    count  INTEGER NOT NULL,
    PRIMARY KEY (n, gram)
);

-- One row per bug_id, filled in by the metadata scraper.
CREATE TABLE IF NOT EXISTS bugs (
    bug_id     TEXT PRIMARY KEY,
    title      TEXT,
    crash_type TEXT,          -- KASAN/UBSAN/WARN/BUG/null-deref/lockdep/...
    subsystem  TEXT,          -- fs, net, bpf, ntfs, ... (first token of title prefix)
    status     TEXT,          -- open/fixed/invalid
    fetched_at INTEGER,
    raw_title  TEXT
);
"""


def connect(db_path: str | Path = DB_PATH) -> sqlite3.Connection:
    conn = sqlite3.connect(str(db_path))
    conn.execute("PRAGMA journal_mode = WAL;")
    conn.execute("PRAGMA synchronous = NORMAL;")
    conn.execute("PRAGMA foreign_keys = ON;")
    return conn


def init_db(db_path: str | Path = DB_PATH) -> None:
    conn = connect(db_path)
    try:
        conn.executescript(SCHEMA)
        conn.commit()
    finally:
        conn.close()