database.py

import os
from datetime import datetime

import libsql_client

_state: dict = {}


async def init():
    _state["client"] = libsql_client.create_client(
        url=os.environ["TURSO_DATABASE_URL"],
        auth_token=os.environ["TURSO_AUTH_TOKEN"],
    )
    await _state["client"].execute("""
        CREATE TABLE IF NOT EXISTS inference_results (
            id          INTEGER PRIMARY KEY AUTOINCREMENT,
            filename    TEXT    NOT NULL,
            model_name  TEXT    NOT NULL,
            centiloid   REAL,
            raw_output  REAL,
            label       TEXT,
            fold        INTEGER,
            created_at  TEXT    NOT NULL,
            UNIQUE (filename, model_name)
        )
    """)
    await _ensure_fold_column()


async def _ensure_fold_column() -> None:
    rs = await _state["client"].execute("PRAGMA table_info(inference_results)")
    cols = [row[1] for row in rs.rows]
    if "fold" not in cols:
        await _state["client"].execute(
            "ALTER TABLE inference_results ADD COLUMN fold INTEGER"
        )


async def save_result(
    filename: str,
    model_name: str,
    centiloid: float,
    raw_output: float,
    label: str | None = None,
    fold: int | None = None,
) -> int:
    now = datetime.utcnow().isoformat()
    rs = await _state["client"].execute(
        "INSERT INTO inference_results"
        " (filename, model_name, centiloid, raw_output, label, fold, created_at)"
        " VALUES (?, ?, ?, ?, ?, ?, ?)"
        " ON CONFLICT(filename, model_name) DO UPDATE SET"
        "   centiloid  = excluded.centiloid,"
        "   raw_output = excluded.raw_output,"
        "   label      = excluded.label,"
        "   fold       = COALESCE(excluded.fold, inference_results.fold),"
        "   created_at = excluded.created_at",
        [filename, model_name, centiloid, raw_output, label, fold, now],
    )
    return rs.last_insert_rowid


async def get_results_page(
    limit: int,
    offset: int,
    fold: int | None = None,
) -> tuple[int, list[dict]]:
    where = ""
    params: list = []
    if fold is not None:
        where = " WHERE fold = ?"
        params.append(fold)

    count_rs = await _state["client"].execute(
        f"SELECT COUNT(*) AS count FROM inference_results{where}",
        params,
    )
    total = int(count_rs.rows[0][0])

    rs = await _state["client"].execute(
        "SELECT id, filename, model_name, centiloid, raw_output, label, fold, created_at"
        f" FROM inference_results{where} ORDER BY created_at DESC LIMIT ? OFFSET ?",
        [*params, limit, offset],
    )
    cols = [c.name if hasattr(c, "name") else c for c in rs.columns]
    return total, [dict(zip(cols, row)) for row in rs.rows]


async def get_done_pairs() -> list[dict]:
    rs = await _state["client"].execute(
        "SELECT filename, model_name FROM inference_results"
    )
    cols = [c.name if hasattr(c, "name") else c for c in rs.columns]
    return [dict(zip(cols, row)) for row in rs.rows]


async def get_folds() -> list[int]:
    rs = await _state["client"].execute(
        "SELECT DISTINCT fold FROM inference_results"
        " WHERE fold IS NOT NULL ORDER BY fold ASC"
    )
    return [int(row[0]) for row in rs.rows]


async def update_result_fold_by_id(row_id: int, fold: int | None) -> bool:
    rs = await _state["client"].execute(
        "UPDATE inference_results SET fold = ? WHERE id = ?",
        [fold, row_id],
    )
    return rs.rows_affected > 0


async def update_result_fold_by_pair(
    filename: str,
    model_name: str,
    fold: int | None,
) -> bool:
    rs = await _state["client"].execute(
        "UPDATE inference_results SET fold = ? WHERE filename = ? AND model_name = ?",
        [fold, filename, model_name],
    )
    return rs.rows_affected > 0


async def close():
    if "client" in _state:
        await _state["client"].close()