lct-hack/scripts/load_db.py

#!/usr/bin/env python3
"""Проверка локальной пропускной способности записи PostgreSQL."""

import argparse
import asyncio
import json
import math
import os
import statistics
import time
import urllib.parse
from uuid import uuid4

import asyncpg


def local_database_url(value: str) -> str:
    normalized = value.replace("postgresql+asyncpg://", "postgresql://", 1)
    parsed = urllib.parse.urlsplit(normalized)
    if parsed.scheme not in {"postgres", "postgresql"}:
        raise ValueError("поддерживается только PostgreSQL")
    if parsed.hostname not in {"127.0.0.1", "localhost", "::1"}:
        raise ValueError("предохранитель: benchmark разрешён только для локальной БД")
    return normalized


def percentile(values: list[float], percent: float) -> float:
    ordered = sorted(values)
    return ordered[max(0, math.ceil(percent / 100 * len(ordered)) - 1)]


async def run(args: argparse.Namespace) -> int:
    if not 100 <= args.operations <= 10_000:
        raise ValueError("operations должно быть от 100 до 10000")
    if not 1 <= args.concurrency <= 100:
        raise ValueError("concurrency должно быть от 1 до 100")
    if args.workload == "temporary" and args.concurrency != 1:
        raise ValueError("temporary workload поддерживает только concurrency=1")
    database_url = local_database_url(args.database_url)
    worker_count = min(args.concurrency, args.operations)
    pool = (
        await asyncpg.create_pool(
            database_url, min_size=worker_count, max_size=worker_count,
            command_timeout=30,
        )
        if args.workload == "audit" and worker_count > 1 else None
    )
    connection = await asyncpg.connect(database_url, timeout=10) if pool is None else None
    latencies: list[float] = []
    run_marker = None
    cleanup_rows = 0
    try:
        if args.workload == "temporary":
            assert connection is not None
            async with connection.transaction():
                await connection.execute("""
                    CREATE TEMPORARY TABLE lct_write_benchmark (
                        sequence_no integer PRIMARY KEY,
                        written_at timestamptz NOT NULL DEFAULT now(),
                        payload text NOT NULL
                    ) ON COMMIT DROP
                """)
                started = time.perf_counter()
                for sequence_no in range(args.operations):
                    operation_started = time.perf_counter()
                    await connection.execute(
                        "INSERT INTO lct_write_benchmark(sequence_no, payload) VALUES($1, $2)",
                        sequence_no, f"local-load-{sequence_no:05d}",
                    )
                    latencies.append((time.perf_counter() - operation_started) * 1000)
                elapsed = time.perf_counter() - started
                count = await connection.fetchval("SELECT count(*) FROM lct_write_benchmark")
        else:
            if not args.allow_audit_writes:
                raise ValueError("для workload=audit требуется явный --allow-audit-writes")
            assert pool is not None or connection is not None
            fetchval = pool.fetchval if pool is not None else connection.fetchval
            if await fetchval("SELECT to_regclass('public.audit_log')") is None:
                raise ValueError("таблица public.audit_log отсутствует; сначала примените миграции")
            # This opt-in mode measures durable product-table writes. Every
            # INSERT autocommits and exercises the real audit index; exact
            # run-scoped rows are removed in finally, even after a partial run.
            run_marker = f"load-benchmark:{uuid4().hex}"
            started = time.perf_counter()

            async def write_partition(worker_index: int) -> None:
                if pool is None:
                    assert connection is not None
                    worker_connection = connection
                    for sequence_no in range(worker_index, args.operations, worker_count):
                        operation_started = time.perf_counter()
                        await worker_connection.execute(
                            """INSERT INTO public.audit_log(id, actor, role, action, object_id, detail)
                               VALUES($1, $2, 'system', 'benchmark.write', $3, $4)""",
                            uuid4(), "load-benchmark", run_marker, f"sequence={sequence_no}",
                        )
                        latencies.append((time.perf_counter() - operation_started) * 1000)
                else:
                    async with pool.acquire() as worker_connection:
                        for sequence_no in range(worker_index, args.operations, worker_count):
                            operation_started = time.perf_counter()
                            await worker_connection.execute(
                                """INSERT INTO public.audit_log(id, actor, role, action, object_id, detail)
                                   VALUES($1, $2, 'system', 'benchmark.write', $3, $4)""",
                                uuid4(), "load-benchmark", run_marker, f"sequence={sequence_no}",
                            )
                            latencies.append((time.perf_counter() - operation_started) * 1000)

            await asyncio.gather(*(write_partition(worker) for worker in range(worker_count)))
            elapsed = time.perf_counter() - started
            count = await fetchval(
                "SELECT count(*) FROM public.audit_log WHERE object_id = $1", run_marker
            )
    finally:
        try:
            if run_marker is not None:
                if pool is not None:
                    cleanup_rows = await pool.fetchval(
                        "WITH removed AS (DELETE FROM public.audit_log WHERE object_id = $1 RETURNING 1) "
                        "SELECT count(*) FROM removed", run_marker,
                    )
                elif connection is not None:
                    cleanup_rows = await connection.fetchval(
                        "WITH removed AS (DELETE FROM public.audit_log WHERE object_id = $1 RETURNING 1) "
                        "SELECT count(*) FROM removed", run_marker,
                    )
        finally:
            if pool is not None:
                await pool.close()
            elif connection is not None:
                await connection.close()

    throughput = args.operations / elapsed
    result = {
        "database": "local PostgreSQL",
        "workload": args.workload,
        "operations_requested": args.operations,
        "operations_written": count,
        "cleanup_rows_deleted": cleanup_rows,
        "concurrency": worker_count,
        "elapsed_seconds": elapsed,
        "operations_per_second": throughput,
        "operation_latency_ms": {
            "min": min(latencies),
            "median": statistics.median(latencies),
            "p95": percentile(latencies, 95),
            "max": max(latencies),
        },
        "pass_100_writes_per_second": (
            count == args.operations
            and throughput >= 100
            and (args.workload != "audit" or cleanup_rows == args.operations)
        ),
        "scope": (
            f"autocommitted INSERTs over {worker_count} local asyncpg connection(s) into the "
            "indexed product audit_log table; this run's rows are deleted in finally; "
            "not target 6-core/32-GB hardware"
            if args.workload == "audit" else
                        "sequential INSERT statements over one local asyncpg connection into a temporary "
            "table; table is dropped on commit; not target 6-core/32-GB hardware"
        ),
    }
    rendered = json.dumps(result, ensure_ascii=False, indent=2)
    print(rendered)
    if args.output:
        with open(args.output, "w", encoding="utf-8") as stream:
            stream.write(rendered + "\n")
    return 0 if result["pass_100_writes_per_second"] else 1


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--database-url",
        default=os.environ.get("DATABASE_URL", "postgresql://lct:lct@127.0.0.1:5432/lct"),
    )
    parser.add_argument("--operations", type=int, default=1000)
    parser.add_argument("--concurrency", type=int, default=1,
                        help="number of concurrent PostgreSQL clients (audit workload only)")
    parser.add_argument(
        "--workload", choices=("temporary", "audit"), default="temporary",
        help="temporary benchmark table (default) or real indexed audit_log writes",
    )
    parser.add_argument(
        "--allow-audit-writes", action="store_true",
        help="required opt-in: workload=audit inserts then removes rows in audit_log",
    )
    parser.add_argument("--output", help="сохранить JSON-отчёт")
    try:
        raise SystemExit(asyncio.run(run(parser.parse_args())))
    except Exception as exc:
        print(f"database load failed: {type(exc).__name__}: {exc}")
        raise SystemExit(2)