devplacepy/devplacepy/database/core.py

# retoor <retoor@molodetz.nl>

import dataset
import logging
from pathlib import Path
from sqlalchemy import or_
from sqlalchemy.pool import NullPool
from collections import defaultdict
from datetime import datetime, timedelta, timezone
from devplacepy.cache import TTLCache
from devplacepy.config import (
    AQUALITY_NEWS_GRADING_MODEL,
    AQUALITY_NEWS_GRADING_URL,
    DATABASE_URL,
    DEFAULT_CORRECTION_PROMPT,
    DEFAULT_MODIFIER_PROMPT,
    INTERNAL_GATEWAY_URL,
    ensure_data_dirs,
)

logger = logging.getLogger(__name__)

ensure_data_dirs()
if DATABASE_URL.startswith("sqlite:///"):
    _db_file = DATABASE_URL[len("sqlite:///") :]
    if _db_file and _db_file != ":memory:":
        Path(_db_file).parent.mkdir(parents=True, exist_ok=True)

db = dataset.connect(
    DATABASE_URL,
    engine_kwargs={
        "connect_args": {
            "timeout": 30,
            "check_same_thread": False,
        },
        "poolclass": NullPool,
    },
    on_connect_statements=[
        "PRAGMA journal_mode=WAL",
        "PRAGMA synchronous=NORMAL",
        "PRAGMA busy_timeout=30000",
        "PRAGMA cache_size=-8000",
        "PRAGMA temp_store=MEMORY",
        "PRAGMA mmap_size=268435456",
    ],
)


def refresh_snapshot() -> None:
    connection = db.executable
    if connection.in_transaction() and not db.in_transaction:
        connection.commit()


_local_cache_versions: dict = {}


_cache_version_cache = TTLCache(ttl=1)


_cache_state_ready = False


def _ensure_cache_state() -> None:
    global _cache_state_ready
    if _cache_state_ready:
        return
    with db:
        db.query(
            "CREATE TABLE IF NOT EXISTS cache_state "
            "(name TEXT PRIMARY KEY, version INTEGER NOT NULL DEFAULT 0)"
        )
    _cache_state_ready = True


def get_cache_version(name: str) -> int:
    cached = _cache_version_cache.get(name)
    if cached is not None:
        return cached
    try:
        _ensure_cache_state()
        with db:
            rows = list(db.query("SELECT name, version FROM cache_state"))
        versions = {row["name"]: int(row["version"]) for row in rows}
    except Exception as e:
        logger.warning(f"Could not read cache version {name}: {e}")
        return 0
    for key, version in versions.items():
        _cache_version_cache.set(key, version)
    version = versions.get(name, 0)
    _cache_version_cache.set(name, version)
    return version


def bump_cache_version(name: str) -> None:
    try:
        _ensure_cache_state()
        with db:
            db.query(
                "INSERT OR IGNORE INTO cache_state (name, version) VALUES (:name, 0)",
                name=name,
            )
            db.query(
                "UPDATE cache_state SET version = version + 1 WHERE name = :name",
                name=name,
            )
        _cache_version_cache.pop(name)
    except Exception as e:
        logger.warning(f"Could not bump cache version {name}: {e}")


def sync_local_cache(name: str, cache) -> None:
    current = get_cache_version(name)
    if name not in _local_cache_versions:
        _local_cache_versions[name] = current
        return
    if _local_cache_versions[name] != current:
        cache.clear()
        _local_cache_versions[name] = current


def _index(db, table, name, columns, *, where=None, unique=False):
    try:
        if table in db.tables:
            cols = ", ".join(columns)
            kind = "UNIQUE INDEX" if unique else "INDEX"
            clause = f" WHERE {where}" if where else ""
            with db:
                db.query(
                    f"CREATE {kind} IF NOT EXISTS {name} ON {table} ({cols}){clause}"
                )
    except Exception as e:
        logger.warning(f"Could not create index {name} on {table}: {e}")


def _drop_index(db, name):
    try:
        with db:
            db.query(f"DROP INDEX IF EXISTS {name}")
    except Exception as e:
        logger.warning(f"Could not drop index {name}: {e}")


def _uid_index(db, table):
    if table not in db.tables or "uid" not in get_table(table).columns:
        return
    name = f"idx_{table}_uid"
    try:
        with db:
            db.query(f"CREATE UNIQUE INDEX IF NOT EXISTS {name} ON {table} (uid)")
    except Exception as e:
        logger.warning(f"Unique uid index on {table} failed ({e}); using non-unique")
        _index(db, table, name, ["uid"])


def get_table(name):
    return db[name]


def _in_clause(uids, prefix="p"):
    placeholders = ", ".join(f":{prefix}{i}" for i in range(len(uids)))
    params = {f"{prefix}{i}": uid for i, uid in enumerate(uids)}
    return placeholders, params


def _now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()