DFiP_Budget_planing/perf/perfkit/sqlstats.py

"""Статистика SQL из pg_stat_statements и EXPLAIN самых тяжёлых чтений.

CLI (его зовут globalSetup/globalTeardown Playwright):
    python -m perfkit.sqlstats reset
    python -m perfkit.sqlstats dump <run_dir> [--top 15] [--out <dir>]
  --out: каталог для sql_top.json и explain/ (по умолчанию run_dir). Сетевой лог для параметров EXPLAIN читается
  всегда из <run_dir>/network.jsonl. Его задаёт globalTeardown по PERF_SQLSTATS_TAG (<run_dir>/sql/<tag>):
  run-all.sh гоняет несколько вызовов Playwright в одном run_id, и без --out каждый перезатирал бы sql_top.json.
Окружение: PERF_DB_DSN (по умолчанию postgresql://perf:perf@localhost:55432/perf). Если Postgres недоступен,
команды печатают предупреждение в stderr и завершаются с кодом 0 (прогон сценариев не должен падать из-за статистики).

dump(run_dir, top):
  * sql_top.json: {"by_total": [...], "by_mean": [...]}; строка: query, calls, total_exec_time, mean_exec_time,
    rows, shared_blks_hit/read, toplevel.
  * explain/<n>.json для топ-5 по total_exec_time среди ВЫПОЛНЯЕМЫХ кандидатов (toplevel, SELECT/WITH, не писатели:
    v3.add_*/upd_*/del_*, не служебные set_config/pg_*; DML не берётся никогда).

Подстановка параметров. pg_stat_statements хранит нормализованный текст ($1..$n), значения не хранит. Механизм:
  1. Для известных вызовов репозитория (api/src/repository/sheet_repository.py -> v3.v_form_view(form_id, sheet,
     sections[, direction]); summary_repository.py -> v3.v_summary_view(year, sheet, org_ids)) параметры берутся из
     сетевого лога <run_dir>/network.jsonl: подходящий GET (/api/v1/form/{id}/sheet/{sheet}?direction=&sections=,
     /api/v1/svod/1/{sheet}/{year}?org_ids=) разбирается так же, как это делает роутер. Режим "analyze":
     EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) с реальными значениями; перед ним SET LOCAL app.user_id = id роли из лога.
  2. Для остальных кандидатов нет надёжного соответствия лог -> параметры, поэтому режим "generic_plan":
     PREPARE + plan_cache_mode=force_generic_plan + EXPLAIN (FORMAT JSON) EXECUTE с NULL (без выполнения, без ANALYZE) - план показывает только вызов
     функции; помечается "analyze": false. Если и он не строится - кандидат пропускается, причина в explain/skipped.json.
  Всё выполняется в транзакции, которая откатывается (ROLLBACK), statement_timeout 180 с.
"""
from __future__ import annotations

import asyncio
import json
import os
import re
import sys
import urllib.parse
from pathlib import Path

DSN = os.environ.get("PERF_DB_DSN", "postgresql://perf:perf@localhost:55432/perf")
WRITER_PREFIXES = ("add_", "upd_", "del_")
SKIP_FUNCS = ("set_config", "pg_stat_statements", "current_setting", "pg_sleep")
V3_CALL = re.compile(r"\bv3\.(\w+)\s*\(", re.I)
STATEMENT_TIMEOUT_MS = 180_000


def v3_functions(query: str) -> list[str]:
    return V3_CALL.findall(query)


def is_explainable(query: str) -> bool:
    q = query.strip()
    if not re.match(r"(?is)^(select|with)\b", q):
        return False
    low = q.lower()
    if re.search(r"\bfor\s+(update|share|no\s+key\s+update)\b", low):
        return False
    if any(f in low for f in SKIP_FUNCS):
        return False
    if re.search(r"(?is)\b(insert\s+into|update\s+\w|delete\s+from)\b", q) and not v3_functions(q):
        return False
    return not any(f.lower().startswith(WRITER_PREFIXES) for f in v3_functions(q))


def explain_candidates(rows, n: int = 5) -> list[dict]:
    """Топ-n по total_exec_time среди строк, которые можно EXPLAIN-ить."""
    ok = [r for r in rows if is_explainable(r["query"])]
    return sorted(ok, key=lambda r: r["total_exec_time"], reverse=True)[:n]


def pick_params(query: str, samples: dict[str, list]) -> list | None:
    return samples.get(query)


# ---------- параметры из сетевого лога ----------

def _sections(qs: dict) -> list[str] | None:
    v = qs.get("sections", [""])[0]
    return [s for s in v.split(",") if s] or None


def _params_form_view(e: dict, nparams: int) -> list | None:
    m = re.fullmatch(r"/api/v1/form/(\d+)/sheet/(\w+)", e["path"])
    if not m or e["method"] != "GET":
        return None
    qs = urllib.parse.parse_qs(e.get("query") or "")
    direction = qs.get("direction", [None])[0]
    if (nparams == 4) != bool(direction):
        return None
    p = [int(m.group(1)), m.group(2), _sections(qs)]
    return p + [direction] if nparams == 4 else p


def _params_summary(e: dict, nparams: int) -> list | None:
    m = re.fullmatch(r"/api/v1/svod/1/(\w+)/(\d+)", e["path"])
    if not m or e["method"] != "GET" or nparams != 3:
        return None
    qs = urllib.parse.parse_qs(e.get("query") or "")
    ids = [int(x) for x in qs.get("org_ids", [""])[0].split(",") if x.strip()] or None
    return [int(m.group(2)), m.group(1), ids]


BUILDERS = {"v_form_view": _params_form_view, "v_summary_view": _params_summary}


def build_samples(rows, netlog) -> dict[str, tuple[list, dict]]:
    """query -> (params, лог-запись). Берётся лог-запись самого долгого подходящего запроса (худший случай)."""
    out = {}
    for r in rows:
        funcs = v3_functions(r["query"])
        if len(funcs) != 1 or funcs[0] not in BUILDERS:
            continue
        n = len(set(re.findall(r"\$(\d+)", r["query"].split("FROM", 1)[-1])))
        best = None
        for e in netlog:
            p = BUILDERS[funcs[0]](e, n)
            if p is not None and (best is None or (e.get("duration_ms") or 0) > (best[1].get("duration_ms") or 0)):
                best = (p, e)
        if best:
            out[r["query"]] = best
    return out


# ---------- работа с БД ----------

def _warn(msg: str) -> None:
    print(f"WARNING [perf sqlstats]: {msg}", file=sys.stderr)


async def _connect():
    import asyncpg
    return await asyncpg.connect(DSN, timeout=5)


def reset() -> bool:
    async def go():
        c = await _connect()
        try:
            await c.execute("SELECT pg_stat_statements_reset()")
        finally:
            await c.close()
    try:
        asyncio.run(go())
        return True
    except Exception as ex:  # noqa: BLE001
        _warn(f"reset не выполнен, SQL-статистика не будет собрана: {type(ex).__name__}: {ex}")
        return False


STAT_SQL = """
SELECT query, calls, total_exec_time, mean_exec_time, rows, shared_blks_hit, shared_blks_read, toplevel
FROM pg_stat_statements
WHERE dbid = (SELECT oid FROM pg_database WHERE datname = current_database())
  AND query NOT ILIKE '%pg_stat_statements%'
  AND query NOT ILIKE 'EXPLAIN%'
  AND query NOT LIKE '%perf-sqlstats%'
"""


async def _explain_one(conn, query: str, params, user_id: int | None) -> tuple[str, list]:
    tr = conn.transaction()
    await tr.start()
    prepared = False
    try:
        await conn.execute(f"SET LOCAL statement_timeout = {STATEMENT_TIMEOUT_MS}")
        if user_id is not None:
            await conn.execute(f"SET LOCAL app.user_id = {int(user_id)}")
        if params is not None:
            raw = await conn.fetchval("EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) " + query, *params)
            return "analyze", json.loads(raw) if isinstance(raw, str) else raw
        # без значений параметров: PREPARE + force_generic_plan + EXPLAIN EXECUTE с NULL (GENERIC_PLAN не работает
        # по расширенному протоколу asyncpg). План показывает общую форму запроса, без выполнения.
        n = max([int(x) for x in re.findall(r"\$(\d+)", query)] or [0])
        await conn.execute("SET LOCAL plan_cache_mode = force_generic_plan")
        await conn.execute("PREPARE perf_gp AS " + query)
        prepared = True
        args = ", ".join(["NULL"] * n)
        raw = await conn.fetchval("EXPLAIN (FORMAT JSON) EXECUTE perf_gp" + (f"({args})" if n else ""))
        return "generic_plan", json.loads(raw) if isinstance(raw, str) else raw
    finally:
        await tr.rollback()
        if prepared:
            await conn.execute("DEALLOCATE perf_gp")  # PREPARE живёт в сессии, а не в транзакции


async def _dump(run_dir: Path, top: int, out_dir: Path) -> dict:
    conn = await _connect()
    try:
        rows = [dict(r) for r in await conn.fetch(STAT_SQL)]
        for r in rows:
            for k in ("total_exec_time", "mean_exec_time"):
                r[k] = round(float(r[k]), 3)
        by_total = sorted(rows, key=lambda r: r["total_exec_time"], reverse=True)[:top]
        by_mean = sorted(rows, key=lambda r: r["mean_exec_time"], reverse=True)[:top]
        (out_dir / "sql_top.json").write_text(
            json.dumps({"by_total": by_total, "by_mean": by_mean}, ensure_ascii=False, indent=2, default=str), encoding="utf-8")

        netlog = []
        nf = run_dir / "network.jsonl"
        if nf.exists():
            netlog = [json.loads(l) for l in nf.read_text(encoding="utf-8").splitlines() if l.strip()]
        users = {r["username"]: r["id"] for r in await conn.fetch("SELECT /* perf-sqlstats */ id, username FROM v3.app_user")}
        manifest_users = {}
        try:
            from .manifest import load, manifest_path
            manifest_users = {k: u.username for k, u in load(manifest_path(int(os.environ.get("PERF_SCALE", "1")))).users.items()}
        except Exception:  # noqa: BLE001
            pass

        cands = [r for r in rows if r["toplevel"]]
        cands = sorted([r for r in cands if is_explainable(r["query"])], key=lambda r: r["total_exec_time"], reverse=True)
        samples = build_samples(cands, netlog)
        ex_dir = out_dir / "explain"
        ex_dir.mkdir(exist_ok=True)
        for old in ex_dir.glob("*.json"):  # результаты прошлого dump в том же каталоге
            old.unlink()
        done, skipped = 0, []
        for r in cands:
            if done >= 5:
                break
            smp = pick_params(r["query"], samples)
            params, entry = smp if smp else (None, None)
            uid = users.get(manifest_users.get(entry["role"], "")) if entry else None
            try:
                mode, plan = await _explain_one(conn, r["query"], params, uid)
            except Exception as ex:  # noqa: BLE001
                skipped.append({"query": r["query"], "error": f"{type(ex).__name__}: {str(ex)[:300]}"})
                continue
            done += 1
            (ex_dir / f"{done}.json").write_text(json.dumps({
                "query": r["query"], "calls": r["calls"], "total_exec_time": r["total_exec_time"],
                "mean_exec_time": r["mean_exec_time"], "mode": mode, "analyze": mode == "analyze",
                "params": params, "param_source": ({k: entry.get(k) for k in ("method", "url", "role")} if entry else None),
                "plan": plan}, ensure_ascii=False, indent=2, default=str), encoding="utf-8")
        if skipped:
            (ex_dir / "skipped.json").write_text(json.dumps(skipped, ensure_ascii=False, indent=2), encoding="utf-8")
        return {"explained": done, "skipped": len(skipped)}
    finally:
        await conn.close()


def dump(run_dir, top: int = 15, out_dir=None) -> bool:
    run_dir = Path(run_dir)
    out_dir = Path(out_dir) if out_dir is not None else run_dir
    run_dir.mkdir(parents=True, exist_ok=True)
    out_dir.mkdir(parents=True, exist_ok=True)
    try:
        info = asyncio.run(_dump(run_dir, top, out_dir))
        print(f"sqlstats: sql_top.json, explain: {info['explained']} (пропущено {info['skipped']}) -> {out_dir}")
        return True
    except Exception as ex:  # noqa: BLE001
        _warn(f"dump не выполнен, sql_top.json/explain не созданы: {type(ex).__name__}: {ex}")
        return False


def main(argv=None) -> int:
    argv = sys.argv[1:] if argv is None else argv
    if argv and argv[0] == "reset":
        reset()
        return 0
    if len(argv) >= 2 and argv[0] == "dump":
        top = int(argv[argv.index("--top") + 1]) if "--top" in argv else 15
        out = argv[argv.index("--out") + 1] if "--out" in argv else None
        dump(argv[1], top, out)
        return 0
    print("usage: python -m perfkit.sqlstats reset | dump <run_dir> [--top N] [--out DIR]", file=sys.stderr)
    return 2


if __name__ == "__main__":
    raise SystemExit(main())