FabledScribe/src/scribe/services/db_maintenance.py

"""Basic Postgres maintenance — targeted VACUUM (ANALYZE).

Postgres autovacuum already reclaims dead tuples and refreshes planner stats
in the background. This adds a daily, off-hours top-up over the handful of
tables the retention/purge sweeps churn hardest (log retention, notification
purge, auth-token purge, trash purge, version pruning) — so the bloat those
bulk DELETEs leave behind is collected promptly and the planner keeps good
stats. It is deliberately narrow; the rest of the schema is left to autovacuum.

Security: table names cannot be parameterized in SQL, so VACUUM is only ever
issued against names from the hardcoded MAINTENANCE_TABLES allowlist — there is
no path from user input to the table name.
"""
from __future__ import annotations

import json
import logging
import time
from datetime import datetime, timezone

from sqlalchemy import text

from scribe.models import async_session, engine
from scribe.models.base import iso
from scribe.services.settings import get_admin_setting, set_admin_setting

logger = logging.getLogger(__name__)

# High-churn tables — exactly the ones the scheduled delete sweeps hit:
#   app_logs              ← log retention (hourly)
#   notifications         ← read-notification purge (hourly)
#   password_reset_tokens ← auth-token purge (daily)
#   invitation_tokens     ← auth-token purge (daily)
#   notes                 ← trash purge (daily) + general churn
#   note_versions         ← version-pinning prune (daily)
MAINTENANCE_TABLES: tuple[str, ...] = (
    "app_logs",
    "notifications",
    "password_reset_tokens",
    "invitation_tokens",
    "notes",
    "note_versions",
)

_LAST_RUN_KEY = "db_maintenance_last_run"


async def run_maintenance(tables: tuple[str, ...] | list[str] | None = None) -> dict:
    """Run VACUUM (ANALYZE) over the allowlisted high-churn tables.

    VACUUM cannot run inside a transaction block, so this uses a dedicated
    AUTOCOMMIT connection rather than the usual session. Each table is
    vacuumed independently; one failure is logged and does not abort the rest.

    Returns a summary dict (also persisted as the db_maintenance_last_run
    admin setting) of shape:
        {"started_at": iso, "elapsed_ms": int, "tables": [
            {"table": str, "ok": bool, "elapsed_ms": int, "error": str|None}
        ]}
    """
    # Only ever operate on the closed allowlist — never an arbitrary name.
    requested = tuple(tables) if tables is not None else MAINTENANCE_TABLES
    targets = [t for t in requested if t in MAINTENANCE_TABLES]

    started = time.monotonic()
    started_at = datetime.now(timezone.utc).isoformat()
    results: list[dict] = []

    async with engine.connect() as conn:
        # On AsyncConnection, execution_options() is a coroutine and MUST be
        # awaited (it returns the connection with AUTOCOMMIT set — VACUUM can't
        # run inside a transaction block).
        autocommit = await conn.execution_options(isolation_level="AUTOCOMMIT")
        for table in targets:
            t0 = time.monotonic()
            try:
                # table is from the allowlist above — safe to interpolate.
                await autocommit.exec_driver_sql(f"VACUUM (ANALYZE) {table}")
                results.append({
                    "table": table,
                    "ok": True,
                    "elapsed_ms": int((time.monotonic() - t0) * 1000),
                    "error": None,
                })
            except Exception as exc:  # noqa: BLE001 — record per-table, keep going
                logger.exception("VACUUM (ANALYZE) failed for %s", table)
                results.append({
                    "table": table,
                    "ok": False,
                    "elapsed_ms": int((time.monotonic() - t0) * 1000),
                    "error": str(exc),
                })

    summary = {
        "started_at": started_at,
        "elapsed_ms": int((time.monotonic() - started) * 1000),
        "tables": results,
    }
    ok = sum(1 for r in results if r["ok"])
    logger.info(
        "DB maintenance: vacuumed %d/%d table(s) in %dms",
        ok, len(results), summary["elapsed_ms"],
    )
    try:
        await set_admin_setting(_LAST_RUN_KEY, json.dumps(summary))
    except Exception:  # noqa: BLE001 — persisting the summary is best-effort
        logger.warning("Failed to persist db_maintenance last-run summary", exc_info=True)
    return summary


# Read-only table-health query against Postgres' own statistics views. The key
# signal is the dead-tuple ratio (dead / (live + dead)) — that IS bloat; a high
# ratio on a large table means autovacuum isn't keeping up. Timestamps confirm
# maintenance is actually running on each table.
_HEALTH_SQL = text("""
    SELECT
      relname AS table_name,
      n_live_tup AS live,
      n_dead_tup AS dead,
      CASE WHEN n_live_tup + n_dead_tup > 0
           THEN round(100.0 * n_dead_tup / (n_live_tup + n_dead_tup), 1)
           ELSE 0 END AS dead_pct,
      pg_total_relation_size(relid) AS total_bytes,
      n_mod_since_analyze AS mod_since_analyze,
      GREATEST(last_vacuum, last_autovacuum) AS last_vacuum,
      GREATEST(last_analyze, last_autoanalyze) AS last_analyze
    FROM pg_stat_user_tables
    ORDER BY n_dead_tup DESC, total_bytes DESC
""")


async def get_table_health() -> dict:
    """Per-table health from Postgres statistics + the total database size.

    Read-only (system catalogs only). Returns:
        {"db_bytes": int, "tables": [
            {"table": str, "live": int, "dead": int, "dead_pct": float,
             "total_bytes": int, "mod_since_analyze": int,
             "last_vacuum": iso|None, "last_analyze": iso|None}
        ]}
    """
    async with async_session() as session:
        db_bytes = (
            await session.execute(text("SELECT pg_database_size(current_database())"))
        ).scalar() or 0
        rows = (await session.execute(_HEALTH_SQL)).mappings().all()

    tables = [
        {
            "table": r["table_name"],
            "live": int(r["live"] or 0),
            "dead": int(r["dead"] or 0),
            "dead_pct": float(r["dead_pct"] or 0),
            "total_bytes": int(r["total_bytes"] or 0),
            "mod_since_analyze": int(r["mod_since_analyze"] or 0),
            "last_vacuum": iso(r["last_vacuum"]),
            "last_analyze": iso(r["last_analyze"]),
        }
        for r in rows
    ]
    return {"db_bytes": int(db_bytes), "tables": tables}


async def get_last_run() -> dict | None:
    """Return the most recent run summary, or None if it has never run."""
    raw = await get_admin_setting(_LAST_RUN_KEY, "")
    if not raw:
        return None
    try:
        return json.loads(raw)
    except (ValueError, TypeError):
        return None