app/storage/db.py

from __future__ import annotations

import sqlite3
from pathlib import Path


def connect(database_path: str) -> sqlite3.Connection:
    path = Path(database_path)
    path.parent.mkdir(parents=True, exist_ok=True)
    connection = sqlite3.connect(path, check_same_thread=False)
    connection.row_factory = sqlite3.Row
    connection.execute("PRAGMA journal_mode=WAL")
    connection.execute("PRAGMA foreign_keys=ON")
    return connection


def initialize_schema(connection: sqlite3.Connection) -> None:
    connection.executescript(
        """
        CREATE TABLE IF NOT EXISTS sessions (
            id TEXT PRIMARY KEY,
            status TEXT NOT NULL,
            question TEXT NOT NULL,
            model_name TEXT NOT NULL,
            owner_id TEXT NOT NULL,
            owner_name TEXT,
            error TEXT,
            created_at TEXT NOT NULL,
            updated_at TEXT NOT NULL
        );

        CREATE TABLE IF NOT EXISTS generation_results (
            session_id TEXT PRIMARY KEY REFERENCES sessions(id) ON DELETE CASCADE,
            answer TEXT NOT NULL,
            raw_generation_text TEXT NOT NULL,
            raw_trace_text TEXT NOT NULL,
            normalized_trace_text TEXT NOT NULL,
            sentences_json TEXT NOT NULL,
            generation_metadata_json TEXT NOT NULL
        );

        CREATE TABLE IF NOT EXISTS analysis_results (
            session_id TEXT PRIMARY KEY REFERENCES sessions(id) ON DELETE CASCADE,
            result_json TEXT NOT NULL
        );
        """
    )
    columns = {
        row["name"] for row in connection.execute("PRAGMA table_info(generation_results)").fetchall()
    }
    if "raw_generation_text" not in columns:
        connection.execute(
            "ALTER TABLE generation_results ADD COLUMN raw_generation_text TEXT NOT NULL DEFAULT ''"
        )
    session_columns = {row["name"] for row in connection.execute("PRAGMA table_info(sessions)").fetchall()}
    if "owner_id" not in session_columns:
        connection.execute("ALTER TABLE sessions ADD COLUMN owner_id TEXT NOT NULL DEFAULT 'legacy-user'")
    if "owner_name" not in session_columns:
        connection.execute("ALTER TABLE sessions ADD COLUMN owner_name TEXT")
    connection.commit()