DataOps-platform/scripts/validation/rule_execution_capacity.py

"""Recalculable capacity evidence and fail-closed runtime budget gates."""

from __future__ import annotations

import argparse
import gc
import hashlib
import json
import math
import os
import platform
import tempfile
import threading
import time
from collections.abc import Iterable
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any


@dataclass(frozen=True)
class CapacityLimits:
    runner_memory_bytes: int
    datasource_pool_budget: int
    artifact_bytes: int

    def __post_init__(self) -> None:
        for value in (
            self.runner_memory_bytes,
            self.datasource_pool_budget,
            self.artifact_bytes,
        ):
            if isinstance(value, bool) or not isinstance(value, int) or value < 1:
                raise ValueError("capacity limits must be positive integers")


@dataclass(frozen=True)
class CapacityObservation:
    backend: str
    rows: int
    elapsed_seconds: float
    peak_memory_bytes: int
    pool_peak: int
    artifact_bytes: int
    node_durations_seconds: tuple[float, ...]

    def __post_init__(self) -> None:
        if self.backend not in {"sql_pushdown", "polars_batch"}:
            raise ValueError("capacity backend is unsupported")
        if isinstance(self.rows, bool) or not isinstance(self.rows, int):
            raise ValueError("capacity rows must be an integer")
        if self.rows < 1 or self.elapsed_seconds <= 0:
            raise ValueError("capacity observation must contain work")
        if (
            self.peak_memory_bytes < 0
            or self.pool_peak < 0
            or self.artifact_bytes < 0
        ):
            raise ValueError("capacity metrics cannot be negative")
        if not self.node_durations_seconds or any(
            value < 0 for value in self.node_durations_seconds
        ):
            raise ValueError("node durations must be non-negative")


def _p95(values: tuple[float, ...]) -> float:
    ordered = sorted(values)
    index = max(0, math.ceil(len(ordered) * 0.95) - 1)
    return float(ordered[index])


def assert_capacity_within_limits(
    observation: CapacityObservation,
    limits: CapacityLimits,
) -> None:
    exceeded = []
    if observation.peak_memory_bytes > limits.runner_memory_bytes:
        exceeded.append("runner_memory_bytes")
    if observation.pool_peak > limits.datasource_pool_budget:
        exceeded.append("datasource_pool_budget")
    if observation.artifact_bytes > limits.artifact_bytes:
        exceeded.append("artifact_bytes")
    if exceeded:
        raise ValueError(
            "capacity limit exceeded: " + ", ".join(exceeded)
        )


def digest_payload(payload: dict[str, Any]) -> str:
    value = dict(payload)
    value.pop("evidence_sha256", None)
    canonical = json.dumps(
        value,
        ensure_ascii=False,
        sort_keys=True,
        separators=(",", ":"),
    )
    return hashlib.sha256(canonical.encode("utf-8")).hexdigest()


def write_capacity_evidence(
    path: str | Path,
    *,
    observations: Iterable[CapacityObservation],
    limits: CapacityLimits,
    machine: dict[str, Any],
) -> dict[str, Any]:
    if not isinstance(machine, dict) or not machine:
        raise ValueError("capacity machine metadata is required")
    rows = []
    for observation in observations:
        assert_capacity_within_limits(observation, limits)
        item = asdict(observation)
        item["node_durations_seconds"] = list(
            observation.node_durations_seconds
        )
        item["throughput_rows_per_second"] = round(
            observation.rows / observation.elapsed_seconds, 3
        )
        item["p95_node_seconds"] = _p95(
            observation.node_durations_seconds
        )
        rows.append(item)
    if not rows:
        raise ValueError("capacity evidence requires observations")
    payload = {
        "schema_version": "1.0",
        "limits": asdict(limits),
        "machine": dict(machine),
        "observations": rows,
    }
    payload["evidence_sha256"] = digest_payload(payload)
    target = Path(path)
    target.parent.mkdir(parents=True, exist_ok=True)
    target.write_text(
        json.dumps(
            payload,
            ensure_ascii=False,
            sort_keys=True,
            indent=2,
        )
        + "\n",
        encoding="utf-8",
    )
    return payload


write_capacity_evidence.digest_payload = digest_payload


class _PeakRSSSampler:
    def __init__(self) -> None:
        self.peak = 0
        self._stop = threading.Event()
        self._thread: threading.Thread | None = None

    def __enter__(self):
        import psutil

        process = psutil.Process()

        def sample() -> None:
            while not self._stop.wait(0.02):
                self.peak = max(self.peak, process.memory_info().rss)

        self.peak = process.memory_info().rss
        self._thread = threading.Thread(target=sample, daemon=True)
        self._thread.start()
        return self

    def __exit__(self, exc_type, exc, traceback) -> None:
        self._stop.set()
        if self._thread is not None:
            self._thread.join(timeout=1)


def _plan_memory_bytes(value: Any) -> int:
    if isinstance(value, dict):
        memory_kib = 0
        for key, item in value.items():
            if key == "Peak Memory Usage" and isinstance(item, (int, float)) or (
                key == "Sort Space Used"
                and value.get("Sort Space Type") == "Memory"
                and isinstance(item, (int, float))
            ):
                memory_kib = max(memory_kib, int(item))
            memory_kib = max(memory_kib, _plan_memory_bytes(item) // 1024)
        return memory_kib * 1024
    if isinstance(value, list):
        return max((_plan_memory_bytes(item) for item in value), default=0)
    return 0


def measure_sql_pushdown(rows: int, database_url: str) -> CapacityObservation:
    """Run real PostgreSQL normalize/assert/deduplicate operators."""
    import psycopg2

    statements = (
        """
        SELECT sum(length(mobile))
        FROM (
          SELECT btrim(
            CASE WHEN gs %% 20 = 0
              THEN ' 1380013800x ' ELSE ' 13800138000 ' END
          ) AS mobile
          FROM generate_series(1, %s) AS gs
        ) normalized
        """,
        """
        SELECT sum(
          (btrim(
            CASE WHEN gs %% 20 = 0
              THEN ' 1380013800x ' ELSE ' 13800138000 ' END
          ) ~ '^[0-9]{11}$')::int
        )
        FROM generate_series(1, %s) AS gs
        """,
        """
        SELECT count(*)
        FROM (
          SELECT (gs / 2)::bigint AS customer_id, max(gs) AS updated_at
          FROM generate_series(1, %s) AS gs
          GROUP BY (gs / 2)::bigint
        ) deduplicated
        """,
    )
    durations = []
    peak_memory = 0
    started = time.monotonic()
    connection = psycopg2.connect(database_url, connect_timeout=10)
    try:
        connection.set_session(readonly=True, autocommit=True)
        with connection.cursor() as cursor:
            for statement in statements:
                cursor.execute(
                    "EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) " + statement,
                    (rows,),
                )
                document = cursor.fetchone()[0][0]
                durations.append(float(document["Execution Time"]) / 1000)
                peak_memory = max(
                    peak_memory, _plan_memory_bytes(document["Plan"])
                )
    finally:
        connection.close()
    return CapacityObservation(
        backend="sql_pushdown",
        rows=rows,
        elapsed_seconds=max(time.monotonic() - started, 0.000001),
        peak_memory_bytes=peak_memory,
        pool_peak=1,
        artifact_bytes=0,
        node_durations_seconds=tuple(durations),
    )


def measure_polars_batch(rows: int) -> CapacityObservation:
    """Run a real streaming Polars rule and materialize its Parquet artifact."""
    import polars as pl

    durations = []
    started = time.monotonic()
    with tempfile.TemporaryDirectory(prefix="data-rule-capacity-") as directory:
        artifact = Path(directory) / "output.parquet"
        with _PeakRSSSampler() as sampler:
            node_started = time.monotonic()
            source = pl.LazyFrame(
                {
                    "row_id": pl.arange(1, rows + 1, eager=True),
                }
            ).with_columns(
                (pl.col("row_id") // 2).alias("customer_id"),
                pl.when(pl.col("row_id") % 20 == 0)
                .then(pl.lit(" 1380013800x "))
                .otherwise(pl.lit(" 13800138000 "))
                .alias("mobile"),
                pl.col("row_id").alias("updated_at"),
            )
            durations.append(time.monotonic() - node_started)

            node_started = time.monotonic()
            normalized = source.with_columns(
                pl.col("mobile").str.strip_chars().alias("mobile")
            )
            durations.append(time.monotonic() - node_started)

            node_started = time.monotonic()
            checked = normalized.with_columns(
                pl.col("mobile")
                .str.contains(r"^[0-9]{11}$")
                .alias("_rule_valid")
            )
            durations.append(time.monotonic() - node_started)

            node_started = time.monotonic()
            output = (
                checked.sort(["customer_id", "updated_at"])
                .unique(subset=["customer_id"], keep="last")
                .sort("customer_id")
            )
            output.sink_parquet(
                artifact,
                compression="zstd",
                engine="streaming",
            )
            durations.append(time.monotonic() - node_started)
            artifact_bytes = artifact.stat().st_size
        peak_memory = sampler.peak
    del source, normalized, checked, output
    gc.collect()
    return CapacityObservation(
        backend="polars_batch",
        rows=rows,
        elapsed_seconds=max(time.monotonic() - started, 0.000001),
        peak_memory_bytes=peak_memory,
        pool_peak=0,
        artifact_bytes=artifact_bytes,
        node_durations_seconds=tuple(durations),
    )


def _write_target_manifest(
    path: Path,
    *,
    requested_rows: list[int],
    observations: list[CapacityObservation],
    failures: list[dict[str, Any]],
) -> None:
    measured = sorted({item.rows for item in observations})
    maximum = max(measured, default=0)
    missing = sorted(set(requested_rows) - set(measured))
    reason = ""
    if missing:
        reason = "; ".join(
            f"{item['backend']}@{item['rows']}: {item['reason']}"
            for item in failures
        ) or "not measured"
    payload = {
        "schema_version": "1.0",
        "requested_rows": requested_rows,
        "measured_rows": measured,
        "backends": ["sql_pushdown", "polars_batch"],
        "hardware_limit": {
            "maximum_measured_rows": maximum,
            "reason": reason,
        },
        "failures": failures,
    }
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(
        json.dumps(payload, ensure_ascii=False, sort_keys=True, indent=2)
        + "\n",
        encoding="utf-8",
    )


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Measure real SQL and Polars data-rule capacity"
    )
    parser.add_argument(
        "--database-url",
        default=os.environ.get("CAPACITY_DATABASE_URL"),
    )
    parser.add_argument(
        "--rows",
        nargs="+",
        type=int,
        default=[100_000, 1_000_000, 10_000_000],
    )
    parser.add_argument(
        "--evidence",
        type=Path,
        default=Path("docs/validation/data-rule-m5-capacity-evidence.json"),
    )
    parser.add_argument(
        "--targets",
        type=Path,
        default=Path("docs/validation/data-rule-m5-capacity-targets.json"),
    )
    parser.add_argument(
        "--runner-memory-bytes",
        type=int,
        default=2 * 1024 * 1024 * 1024,
    )
    parser.add_argument(
        "--artifact-bytes",
        type=int,
        default=512 * 1024 * 1024,
    )
    parser.add_argument(
        "--logical-cpu-count",
        type=int,
        default=os.cpu_count() or 0,
    )
    args = parser.parse_args()
    if not args.database_url:
        parser.error("CAPACITY_DATABASE_URL or --database-url is required")
    requested = sorted(set(args.rows))
    if not requested or requested[0] < 1:
        parser.error("rows must be positive")
    limits = CapacityLimits(
        runner_memory_bytes=args.runner_memory_bytes,
        datasource_pool_budget=2,
        artifact_bytes=args.artifact_bytes,
    )
    observations: list[CapacityObservation] = []
    failures: list[dict[str, Any]] = []
    for row_count in requested:
        for backend, measure in (
            (
                "sql_pushdown",
                lambda rows=row_count: measure_sql_pushdown(
                    rows, args.database_url
                ),
            ),
            (
                "polars_batch",
                lambda rows=row_count: measure_polars_batch(rows),
            ),
        ):
            try:
                observation = measure()
                assert_capacity_within_limits(observation, limits)
                observations.append(observation)
                print(
                    json.dumps(
                        {
                            "backend": backend,
                            "rows": row_count,
                            "elapsed_seconds": observation.elapsed_seconds,
                            "peak_memory_bytes": observation.peak_memory_bytes,
                            "artifact_bytes": observation.artifact_bytes,
                        },
                        sort_keys=True,
                    ),
                    flush=True,
                )
            except Exception as exc:
                failures.append(
                    {
                        "backend": backend,
                        "rows": row_count,
                        "reason": f"{type(exc).__name__}: {str(exc)[:300]}",
                    }
                )
                print(json.dumps(failures[-1], sort_keys=True), flush=True)
    if observations:
        import psutil

        write_capacity_evidence(
            args.evidence,
            observations=observations,
            limits=limits,
            machine={
                "platform": platform.platform(),
                "logical_cpu_count": (
                    args.logical_cpu_count
                    or psutil.cpu_count(logical=True)
                    or os.cpu_count()
                    or "unavailable"
                ),
                "physical_memory_bytes": psutil.virtual_memory().total,
                "python": platform.python_version(),
                "database": "PostgreSQL 16 Docker",
                "measurement_note": (
                    "SQL memory is PostgreSQL plan-reported peak operator "
                    "memory; Polars memory is process RSS."
                ),
            },
        )
    _write_target_manifest(
        args.targets,
        requested_rows=requested,
        observations=observations,
        failures=failures,
    )
    return 0 if observations and not failures else 1


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