DFiP_Budget_planing/perf/analysis/explain_manual.py

"""EXPLAIN (ANALYZE, BUFFERS) вложенных запросов SQL-функций, которые pg_stat_statements показывает только как nested.

    cd perf && .venv/bin/python analysis/explain_manual.py --scale 1 --out report-data/20261026-s1/explain-manual

Тексты запросов берутся из последних версий функций в api/alembic/versions/sql/*.sql (CREATE OR REPLACE в более
поздней миграции перекрывает раннюю; диапазоны строк и якоря - в QUERIES, при расхождении якоря скрипт падает).
Переменные plpgsql (p_form_id, p_sheet, p_sections, флаги секций s_*) подставляются литералами: p_sections = NULL,
значит все флаги TRUE (как при GET листа без sections). form_id - из манифеста масштаба (самый большой лист rf).

Варианты (SET LOCAL, только в транзакции EXPLAIN): по умолчанию, `jit=off`, `work_mem=64MB`, снова по умолчанию
(`.rerun_default`, чтобы разница не объяснялась прогревом кэша). Затем хронометраж: по 10 вызовов v3.v_form_view
(как из репозитория) на каждую комбинацию jit/work_mem - bench.json; затем pg_stat_statements_reset() и 8 вызовов
v_form_view(AHR) в новом соединении - разбивка по вложенным запросам (pgss_form_view_ahr.json). Сброс статистики
pg_stat_statements: не запускать параллельно с замерами Playwright.
Безопасность: только чтение (v_* функции и их тела - SELECT), каждый запрос в транзакции с ROLLBACK,
statement_timeout 300 с. Вариант `work_mem=64MB` (SET LOCAL) - эксперимент: показывает, уходят ли временные файлы.
Нужна живая БД стенда с данными масштаба (после `perf/stand.sh reset` + gen + grow).

Вывод: <out>/<name>[.<variant>].json (план целиком) и <out>/summary.md (время, буферы, временные файлы, топ узлов).
"""
from __future__ import annotations

import argparse
import asyncio
import json
import os
import re
import statistics
import time
from pathlib import Path

import asyncpg

ROOT = Path(__file__).resolve().parents[2]
SQL_DIR = ROOT / "api" / "alembic" / "versions" / "sql"
BENCH_N = 10
PGSS_N = 8
FLAGS = ["s_plan", "s_seq_d", "s_seq_s", "s_appr", "s_cd", "s_book", "s_cs", "s_al", "s_ca", "s_ckk",
         "s_q1", "s_q2", "s_q3", "s_q4", "s_tot", "s_need_ap"]


def extract(file: str, start: int, end: int, anchor_start: str, anchor_end: str) -> str:
    lines = (SQL_DIR / file).read_text().splitlines()
    first, last = lines[start - 1], lines[end - 1]
    if anchor_start not in first or anchor_end not in last:
        raise SystemExit(f"{file}:{start}-{end}: якоря не совпали ({first!r} / {last!r}) - функция изменилась, обновите QUERIES")
    return "\n".join(lines[start - 1:end]).rstrip().rstrip(";")


def subst(sql: str, form_id: int, sheet: str) -> str:
    sql = re.sub(r"\bp_form_id\b", str(form_id), sql)
    sql = re.sub(r"\bp_sheet\b", f"'{sheet}'::varchar", sql)
    sql = re.sub(r"\bp_sections\b", "NULL::text[]", sql)
    for f in FLAGS:
        sql = re.sub(rf"\b{f}\b", "TRUE", sql)
    return sql


def queries(f2: int, f4: int) -> list[dict]:
    return [
        {"name": "form2_ahr_tw", "what": "FORM_2/AHR: тело v3.v_form2_sheet_sections (запрос «WITH tw», дерево статей + агрегаты)",
         "src": "api/alembic/versions/sql/0033_form2_old_sort.sql:500-1155",
         "sql": subst(extract("0033_form2_old_sort.sql", 500, 1155, "WITH", "ORDER BY sub._sort;"), f2, "AHR"),
         "variants": [("", {}), (".jit_off", {"jit": "off"}), (".work_mem_64MB", {"work_mem": "64MB"}), (".rerun_default", {})]},
        {"name": "form2_ahr_jsonb", "what": "FORM_2/AHR: запрос v3.v_form2_view (сборка jsonb строк поверх v_form2_sheet_jsonb)",
         "src": "api/alembic/versions/sql/0033_form2_old_sort.sql:18-38",
         "sql": subst(extract("0033_form2_old_sort.sql", 18, 38, "SELECT", "ORDER BY j._sort_path;"), f2, "AHR"),
         "variants": [("", {}), (".jit_off", {"jit": "off"}), (".work_mem_64MB", {"work_mem": "64MB"}), (".rerun_default", {})]},
        {"name": "form_view_ahr", "what": "FORM_2/AHR: вызов из репозитория v3.v_form_view (как GET листа)",
         "src": "api/src/repository (sheet_repository: SELECT ... FROM v3.v_form_view)",
         "sql": f"SELECT row_type, depth, sort_order, data FROM v3.v_form_view({f2}, 'AHR', NULL::text[], NULL::varchar)",
         "variants": [("", {}), (".jit_off", {"jit": "off"}), (".work_mem_64MB", {"work_mem": "64MB"}), (".rerun_default", {})], "user": True},
        {"name": "form4_cap_sheet", "what": "FORM_4/CAP: тело v3.v_form4_sheet_jsonb (300 строк после grow)",
         "src": "api/alembic/versions/sql/0012_del_project.sql:1097-…",
         "sql": None, "func_body": ("0012_del_project.sql", 1097, "WITH inp_raw AS ("), "form_id": f4, "sheet": "CAP",
         "variants": [("", {}), (".jit_off", {"jit": "off"})]},
        {"name": "form_view_form4_cap", "what": "FORM_4/CAP: вызов v3.v_form_view (как GET листа)",
         "src": "api/src/repository (sheet_repository)",
         "sql": f"SELECT row_type, depth, sort_order, data FROM v3.v_form_view({f4}, 'CAP', NULL::text[], NULL::varchar)",
         "variants": [("", {}), (".jit_off", {"jit": "off"})], "user": True},
    ]


def func_body(file: str, start: int, anchor: str) -> str:
    lines = (SQL_DIR / file).read_text().splitlines()
    if anchor not in lines[start - 1]:
        raise SystemExit(f"{file}:{start}: якорь {anchor!r} не совпал")
    out = []
    for ln in lines[start - 1:]:
        if "$function$" in ln:
            break
        out.append(ln)
    return "\n".join(out).rstrip().rstrip(";")


def walk(node, acc, depth=0):
    kids = node.get("Plans", [])
    loops = node.get("Actual Loops", 1) or 1
    total = node.get("Actual Total Time", 0) * loops
    child = sum(k.get("Actual Total Time", 0) * (k.get("Actual Loops", 1) or 1) for k in kids)
    acc.append({
        "node": node.get("Node Type"), "rel": node.get("Relation Name") or node.get("CTE Name") or node.get("Function Name") or "",
        "excl_ms": round(max(total - child, 0), 1), "rows": node.get("Actual Rows"), "loops": loops,
        "temp_read": node.get("Temp Read Blocks", 0), "temp_written": node.get("Temp Written Blocks", 0),
        "sort": node.get("Sort Method") or "", "space_kb": node.get("Sort Space Used") or node.get("Peak Memory Usage") or "",
        "depth": depth,
    })
    for k in kids:
        walk(k, acc, depth + 1)


async def run(dsn: str, scale: int, out: Path):
    man = json.loads((ROOT / "perf" / "results" / f"manifest-s{scale}.json").read_text())
    rf = [s for s in man["sheets"] if s.get("owner", "rf") == "rf"]
    f2 = next(s["form_id"] for s in rf if s["form_type"] == "FORM_2" and s["sheet"] == "AHR")
    f4 = next(s["form_id"] for s in rf if s["form_type"] == "FORM_4" and s["sheet"] == "CAP")
    con = await asyncpg.connect(dsn)
    uid = await con.fetchval("SELECT id FROM v3.app_user WHERE username = $1", man["users"]["rf"]["username"])
    wm = await con.fetchval("SHOW work_mem")
    jit = await con.fetchval("SHOW jit")
    pg = await con.fetchval("SHOW server_version")
    md = ["# EXPLAIN (ANALYZE, BUFFERS) вложенных запросов (explain_manual.py)", "",
          f"Postgres {pg}, work_mem по умолчанию {wm}; масштаб {scale}; FORM_2/AHR form_id={f2}, FORM_4/CAP form_id={f4}; "
          f"app.user_id={uid} (perf_rf). Каждый запрос - в транзакции с ROLLBACK.", ""]
    for q in queries(f2, f4):
        sql = q["sql"]
        if sql is None:
            sql = subst(func_body(*q["func_body"]), q["form_id"], q["sheet"])
        for suffix, settings in q["variants"]:
            tr = con.transaction()
            await tr.start()
            try:
                await con.execute("SET LOCAL statement_timeout = '300s'")
                if q.get("user"):
                    await con.execute(f"SET LOCAL app.user_id = '{uid}'")
                for k, val in settings.items():
                    await con.execute(f"SET LOCAL {k} = '{val}'")
                plan = json.loads(await con.fetchval("EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) " + sql))[0]
            finally:
                await tr.rollback()
            name = q["name"] + suffix
            sett = ", ".join(f"{k}={val}" for k, val in settings.items()) or f"по умолчанию (work_mem {wm}, jit {jit})"
            (out / f"{name}.json").write_text(json.dumps({"what": q["what"], "src": q["src"], "settings": settings,
                                                          "sql": sql, "plan": plan}, ensure_ascii=False, indent=1))
            top = plan["Plan"]
            nodes: list = []
            walk(top, nodes)
            tr_ = sum(n["temp_read"] for n in nodes)
            tw_ = sum(n["temp_written"] for n in nodes)
            jt = (plan.get("JIT") or {}).get("Timing") or {}
            jtot = jt.get("Total") if isinstance(jt.get("Total"), (int, float)) else None
            md += [f"## {name}", "", f"{q['what']}. Источник: `{q['src']}`. Настройки: {sett}.", "",
                   f"- Execution Time: **{plan['Execution Time']:.1f} мс**, Planning Time {plan['Planning Time']:.1f} мс; строк {top.get('Actual Rows')}; "
                   f"стоимость плана {top.get('Total Cost')}",
                   f"- JIT: " + (f"{(plan.get('JIT') or {}).get('Functions')} функций, {jtot:.1f} мс (emission {jt.get('Emission', 0):.1f})" if jtot is not None else "нет"),
                   f"- Shared Hit {top.get('Shared Hit Blocks', 0)}, Shared Read {top.get('Shared Read Blocks', 0)}; "
                   f"Temp Read/Written (узел верхнего уровня) {top.get('Temp Read Blocks', 0)}/{top.get('Temp Written Blocks', 0)}, "
                   f"сумма по узлам {tr_}/{tw_} блоков по 8 КБ", "",
                   "| узел | объект | excl мс | строк | loops | temp r/w | sort |", "|---|---|---|---|---|---|---|"]
            for n in sorted(nodes, key=lambda n: -n["excl_ms"])[:10]:
                md.append(f"| {n['node']} | {n['rel']} | {n['excl_ms']} | {n['rows']} | {n['loops']} | {n['temp_read']}/{n['temp_written']} | {n['sort']} {n['space_kb']} |")
            md.append("")
            print(f"{name}: {plan['Execution Time']:.0f} ms, jit {jtot}, temp {tr_}/{tw_}", flush=True)
    # хронометраж вызовов как из репозитория (без EXPLAIN, его инструментирование добавляет накладные расходы)
    md += ["## Хронометраж вызовов v3.v_form_view (как GET листа), мс", "",
           f"По {BENCH_N} вызовов на вариант после 1 прогрева, клиент asyncpg; медиана / мин / макс.", "",
           "| вызов | настройки | медиана | мин | макс |", "|---|---|---|---|---|"]
    bench = []
    for label, call in ((f"FORM_2/AHR form {f2}", f"SELECT count(*) FROM v3.v_form_view({f2}, 'AHR', NULL::text[], NULL::varchar)"),
                        (f"FORM_4/CAP form {f4}", f"SELECT count(*) FROM v3.v_form_view({f4}, 'CAP', NULL::text[], NULL::varchar)")):
        for settings in ({}, {"jit": "off"}, {"work_mem": "64MB"}, {"jit": "off", "work_mem": "64MB"}):
            ts = []
            for i in range(BENCH_N + 1):
                tr = con.transaction()
                await tr.start()
                try:
                    await con.execute(f"SET LOCAL app.user_id = '{uid}'")
                    for k, val in settings.items():
                        await con.execute(f"SET LOCAL {k} = '{val}'")
                    t0 = time.perf_counter()
                    await con.fetchval(call)
                    dt = (time.perf_counter() - t0) * 1000
                finally:
                    await tr.rollback()
                if i:
                    ts.append(dt)
            sett = ", ".join(f"{k}={val}" for k, val in settings.items()) or "по умолчанию"
            row = {"call": label, "settings": settings, "median_ms": round(statistics.median(ts), 1), "min_ms": round(min(ts), 1), "max_ms": round(max(ts), 1)}
            bench.append(row)
            md.append(f"| {label} | {sett} | {row['median_ms']} | {row['min_ms']} | {row['max_ms']} |")
            print(row, flush=True)
    (out / "bench.json").write_text(json.dumps(bench, ensure_ascii=False, indent=1))
    md.append("")
    # что внутри вызова в установившемся режиме: pg_stat_statements по вложенным запросам, новое соединение,
    # PGSS_N вызовов v_form_view(AHR) по умолчанию (первые 5 вызовов plpgsql - custom-план, дальше может быть generic)
    await con.execute("SELECT pg_stat_statements_reset()")
    c2 = await asyncpg.connect(dsn)
    per_call = []
    for i in range(PGSS_N):
        async with c2.transaction():
            await c2.execute(f"SET LOCAL app.user_id = '{uid}'")
            t0 = time.perf_counter()
            await c2.fetchval(f"SELECT count(*) FROM v3.v_form_view({f2}, 'AHR', NULL::text[], NULL::varchar)")
            per_call.append(round((time.perf_counter() - t0) * 1000, 1))
    await c2.close()
    rows = await con.fetch("""SELECT left(regexp_replace(query, '\\s+', ' ', 'g'), 90) AS q, calls, round(mean_exec_time::numeric, 1) AS mean_ms,
                                     round(total_exec_time::numeric, 1) AS total_ms, jit_functions, round(jit_emission_time::numeric, 1) AS jit_emission_ms,
                                     round((jit_generation_time + jit_inlining_time + jit_optimization_time + jit_emission_time)::numeric, 1) AS jit_total_ms,
                                     temp_blks_written, toplevel
                                FROM pg_stat_statements
                               WHERE query NOT ILIKE '%pg_stat_statements%' AND query NOT ILIKE 'SET %' AND total_exec_time > 5
                               ORDER BY total_exec_time DESC LIMIT 12""")
    pgss = [dict(r) for r in rows]
    for r in pgss:
        for k, v in r.items():
            if hasattr(v, "is_finite"):
                r[k] = float(v)
    (out / "pgss_form_view_ahr.json").write_text(json.dumps({"per_call_ms": per_call, "statements": pgss}, ensure_ascii=False, indent=1))
    md += [f"## Внутри v3.v_form_view(FORM_2/AHR): pg_stat_statements за {PGSS_N} вызовов в новом соединении", "",
           f"Время вызовов по порядку, мс: {per_call}", "",
           "| запрос | calls | mean мс | total мс | jit функций | jit мс | temp written | toplevel |", "|---|---|---|---|---|---|---|---|"]
    md += [f"| `{r['q']}` | {r['calls']} | {r['mean_ms']} | {r['total_ms']} | {r['jit_functions']} | {r['jit_total_ms']} | {r['temp_blks_written']} | {r['toplevel']} |" for r in pgss]
    md.append("")
    print(per_call, flush=True)
    await con.close()
    (out / "summary.md").write_text("\n".join(md) + "\n")


def main(argv=None):
    ap = argparse.ArgumentParser()
    ap.add_argument("--scale", type=int, default=1)
    ap.add_argument("--out", required=True)
    ap.add_argument("--dsn", default=os.environ.get("PERF_DB_DSN", "postgresql://perf:perf@localhost:55432/perf"))
    a = ap.parse_args(argv)
    out = Path(a.out)
    out.mkdir(parents=True, exist_ok=True)
    asyncio.run(run(a.dsn, a.scale, out))


if __name__ == "__main__":
    main()