grotti/grotti/db.py

"""PostgreSQL-Verbindung, Pool und Migrationen."""

from __future__ import annotations

import re
import threading
from contextlib import contextmanager
from pathlib import Path
from urllib.parse import urlparse, urlunparse

import psycopg
from psycopg import sql
from psycopg.rows import dict_row
from psycopg_pool import ConnectionPool

from grotti import logutil

log = logutil.get("grotti.db")

MIGRATIONS_DIR = Path(__file__).parent / "migrations"
SCHEMA_LEGACY = Path(__file__).parent / "schema.sql"

_pool: ConnectionPool | None = None
_pool_lock = threading.Lock()
_user_count_cache = {"n": None, "at": 0.0}
_USER_COUNT_TTL = 5.0


def database_name(url):
    name = urlparse(url).path.lstrip("/")
    if not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", name or ""):
        raise SystemExit(
            "Der Datenbankname in database_url darf nur Buchstaben, Zahlen und Unterstrich enthalten."
        )
    return name


def _migration_files():
    files = sorted(MIGRATIONS_DIR.glob("*.sql"))
    out = []
    for path in files:
        match = re.match(r"^(\d+)_", path.name)
        if not match:
            continue
        out.append((int(match.group(1)), path))
    return out


def _split_statements(text):
    """Einfache Trennung an Semikolon — Migrationen ohne Funktionen mit ; im Körper."""
    parts = []
    for statement in text.split(";"):
        cleaned = statement.strip()
        if cleaned:
            parts.append(cleaned)
    return parts


def _ensure_version_table(conn):
    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS schema_version (
            version integer PRIMARY KEY,
            name text NOT NULL,
            applied_at timestamptz NOT NULL DEFAULT now()
        )
        """
    )


def _applied_versions(conn):
    rows = conn.execute("SELECT version FROM schema_version ORDER BY version").fetchall()
    return {int(row["version"] if isinstance(row, dict) else row[0]) for row in rows}


def _legacy_schema_present(conn):
    row = conn.execute(
        """
        SELECT 1 FROM information_schema.tables
        WHERE table_schema = 'public' AND table_name = 'printers'
        """
    ).fetchone()
    return bool(row)


def apply_migrations(conn):
    """Wendet fehlende Migrationen an. Bestehende Installationen ohne Version bekommen 001 gutgeschrieben."""
    _ensure_version_table(conn)
    applied = _applied_versions(conn)
    migrations = _migration_files()
    if not migrations and SCHEMA_LEGACY.exists():
        # Fallback: altes schema.sql einmalig
        for statement in _split_statements(SCHEMA_LEGACY.read_text(encoding="utf-8")):
            conn.execute(statement)
        return

    if not applied and _legacy_schema_present(conn):
        # Schema war schon da (vor Migrationssystem) → Basis als angewandt markieren.
        base = next((m for m in migrations if m[0] == 1), None)
        if base:
            conn.execute(
                "INSERT INTO schema_version (version, name) VALUES (%s, %s) ON CONFLICT DO NOTHING",
                (1, base[1].name),
            )
            applied.add(1)

    for version, path in migrations:
        if version in applied:
            continue
        text = path.read_text(encoding="utf-8")
        for statement in _split_statements(text):
            conn.execute(statement)
        conn.execute(
            "INSERT INTO schema_version (version, name) VALUES (%s, %s)",
            (version, path.name),
        )


def prepare(url):
    name = database_name(url)
    admin = urlunparse(urlparse(url)._replace(path="/postgres"))
    try:
        with psycopg.connect(admin, autocommit=True, connect_timeout=5) as conn:
            exists = conn.execute(
                "SELECT 1 FROM pg_database WHERE datname = %s", (name,)
            ).fetchone()
            if not exists:
                conn.execute(sql.SQL("CREATE DATABASE {}").format(sql.Identifier(name)))
    except psycopg.Error as exc:
        log.warning("Datenbank anlegen: %s", exc)
    try:
        with psycopg.connect(url, row_factory=dict_row, connect_timeout=5) as conn:
            apply_migrations(conn)
            conn.commit()
    except psycopg.Error as exc:
        raise SystemExit(
            "PostgreSQL ist nicht erreichbar.\n"
            "In config.json bei database_url das Passwort eintragen, zum Beispiel\n"
            "postgresql://postgres:[email protected]:5432/grotti\n"
            "Die Datenbank grotti legt die App an, sobald die Anmeldung klappt.\n"
            + str(exc)
        ) from exc


def init_pool(url, min_size=2, max_size=16):
    global _pool
    with _pool_lock:
        if _pool is not None:
            return _pool
        _pool = ConnectionPool(
            conninfo=url,
            min_size=min_size,
            max_size=max_size,
            kwargs={"row_factory": dict_row, "connect_timeout": 5},
            open=True,
            name="grotti",
        )
        return _pool


def close_pool():
    global _pool
    with _pool_lock:
        if _pool is not None:
            _pool.close()
            _pool = None


def get_pool():
    if _pool is None:
        raise RuntimeError("Datenbank-Pool ist nicht gestartet.")
    return _pool


@contextmanager
def connect(url=None):
    """Verbindung aus dem Pool, sonst einmalig direkt (Tests / prepare)."""
    if _pool is not None:
        with _pool.connection() as conn:
            try:
                yield conn
            except Exception:
                conn.rollback()
                raise
        return
    if not url:
        raise RuntimeError("Keine database_url und kein Pool.")
    conn = psycopg.connect(url, row_factory=dict_row, connect_timeout=5)
    try:
        yield conn
    except Exception:
        conn.rollback()
        raise
    finally:
        conn.close()


def cached_user_count(conn):
    """Zählt Benutzer mit kurzem Cache — vermeidet count(*) bei jedem Request."""
    import time

    now = time.monotonic()
    cached = _user_count_cache["n"]
    if cached is not None and (now - _user_count_cache["at"]) < _USER_COUNT_TTL:
        return cached
    n = int(conn.execute("SELECT COUNT(*) AS n FROM users").fetchone()["n"])
    _user_count_cache["n"] = n
    _user_count_cache["at"] = now
    return n


def invalidate_user_count():
    _user_count_cache["n"] = None
    _user_count_cache["at"] = 0.0


def prune_run_samples(conn, keep_days=90):
    """Alte Temperaturproben löschen (Retention)."""
    days = max(7, int(keep_days))
    row = conn.execute(
        """
        DELETE FROM run_samples
        WHERE sampled_at < now() - (%s || ' days')::interval
        """,
        (str(days),),
    )
    return row.rowcount if row is not None else 0