Routing Smoke Tests Through a Python Health Gateway

#!/usr/bin/env python3
"""Async readiness gateway for a freshly promoted DR sandbox.

Runs a battery of read-only smoke checks against the sandbox over aiohttp,
tagging every request as synthetic so it never contaminates production
observability, aggregates the checks into a single readiness verdict, and exits
with the code the drill runner branches on to gate traffic routing.

Checks: schema reachable, row-count sanity, critical-query latency, replication
caught up. Any failing check holds routing and quarantines the sandbox.

Exit codes:
    0  READY  -> all checks passed; route real traffic to the sandbox
    1  HOLD   -> at least one check failed; hold routing and quarantine
    2  ABORT  -> usage / configuration / transport error, no verdict formed
"""
import argparse
import asyncio
import sys
import time
from dataclasses import dataclass
from enum import IntEnum
from typing import List

import aiohttp

# Every probe carries these headers so downstream observability drops it from
# production SLO windows instead of counting synthetic load as real traffic.
SYNTHETIC_HEADERS = {
    "X-Synthetic-Probe": "true",
    "X-Drill-Traffic": "readiness-gateway",
    "User-Agent": "dr-health-gateway/1.0 (synthetic; read-only)",
}


class Verdict(IntEnum):
    READY = 0
    HOLD = 1
    ABORT = 2


@dataclass
class CheckResult:
    name: str
    ok: bool
    detail: str
    latency_ms: float


@dataclass
class GatewayConfig:
    base_url: str
    expected_tables: List[str]
    min_rows: int
    max_query_ms: float
    max_replication_lag_s: float
    request_timeout_s: float


def _ms(start: float) -> float:
    return (time.perf_counter() - start) * 1000.0


async def _get_json(session: aiohttp.ClientSession, url: str, timeout: float) -> dict:
    async with session.get(
        url, headers=SYNTHETIC_HEADERS, timeout=aiohttp.ClientTimeout(total=timeout)
    ) as resp:
        resp.raise_for_status()
        return await resp.json()


async def check_schema(session: aiohttp.ClientSession, cfg: GatewayConfig) -> CheckResult:
    start = time.perf_counter()
    try:
        data = await _get_json(session, f"{cfg.base_url}/schema", cfg.request_timeout_s)
    except (aiohttp.ClientError, asyncio.TimeoutError) as exc:
        return CheckResult("schema_reachable", False, f"schema endpoint error: {exc}", _ms(start))
    present = set(data.get("tables", []))
    missing = [t for t in cfg.expected_tables if t not in present]
    if missing:
        return CheckResult("schema_reachable", False, f"missing tables: {missing}", _ms(start))
    return CheckResult("schema_reachable", True, f"{len(present)} tables present", _ms(start))


async def check_row_counts(session: aiohttp.ClientSession, cfg: GatewayConfig) -> CheckResult:
    start = time.perf_counter()
    try:
        data = await _get_json(session, f"{cfg.base_url}/metrics/row-counts", cfg.request_timeout_s)
    except (aiohttp.ClientError, asyncio.TimeoutError) as exc:
        return CheckResult("row_count_sanity", False, f"row-count endpoint error: {exc}", _ms(start))
    counts = data.get("counts", {})
    thin = {t: c for t, c in counts.items() if int(c) < cfg.min_rows}
    if thin:
        return CheckResult("row_count_sanity", False, f"tables below floor {cfg.min_rows}: {thin}", _ms(start))
    return CheckResult("row_count_sanity", True, f"{len(counts)} tables above floor", _ms(start))


async def check_query_latency(session: aiohttp.ClientSession, cfg: GatewayConfig) -> CheckResult:
    start = time.perf_counter()
    payload = {"name": "critical_path", "mode": "read-only"}
    try:
        async with session.post(
            f"{cfg.base_url}/query/critical",
            json=payload,
            headers=SYNTHETIC_HEADERS,
            timeout=aiohttp.ClientTimeout(total=cfg.request_timeout_s),
        ) as resp:
            resp.raise_for_status()
            body = await resp.json()
    except (aiohttp.ClientError, asyncio.TimeoutError) as exc:
        return CheckResult("critical_query_latency", False, f"query endpoint error: {exc}", _ms(start))
    server_ms = float(body.get("elapsed_ms", _ms(start)))
    if server_ms > cfg.max_query_ms:
        return CheckResult(
            "critical_query_latency", False,
            f"{server_ms:.1f}ms exceeds {cfg.max_query_ms:.1f}ms budget", server_ms,
        )
    return CheckResult("critical_query_latency", True, f"{server_ms:.1f}ms within budget", server_ms)


async def check_replication(session: aiohttp.ClientSession, cfg: GatewayConfig) -> CheckResult:
    start = time.perf_counter()
    try:
        data = await _get_json(session, f"{cfg.base_url}/replication/status", cfg.request_timeout_s)
    except (aiohttp.ClientError, asyncio.TimeoutError) as exc:
        return CheckResult("replication_caught_up", False, f"replication endpoint error: {exc}", _ms(start))
    lag = float(data.get("lag_seconds", 1e9))
    if lag > cfg.max_replication_lag_s:
        return CheckResult(
            "replication_caught_up", False,
            f"lag {lag:.1f}s exceeds {cfg.max_replication_lag_s:.1f}s tolerance", _ms(start),
        )
    return CheckResult("replication_caught_up", True, f"lag {lag:.1f}s within tolerance", _ms(start))


async def run_checks(cfg: GatewayConfig) -> List[CheckResult]:
    checks = [check_schema, check_row_counts, check_query_latency, check_replication]
    async with aiohttp.ClientSession() as session:
        results = await asyncio.gather(*(check(session, cfg) for check in checks))
    return list(results)


def build_config(args: argparse.Namespace) -> GatewayConfig:
    return GatewayConfig(
        base_url=args.base_url.rstrip("/"),
        expected_tables=[t.strip() for t in args.expected_tables.split(",") if t.strip()],
        min_rows=args.min_rows,
        max_query_ms=args.max_query_ms,
        max_replication_lag_s=args.max_lag,
        request_timeout_s=args.timeout,
    )


def main(argv=None) -> int:
    parser = argparse.ArgumentParser(description="Readiness gateway for a promoted DR sandbox")
    parser.add_argument("base_url", help="base URL of the sandbox data-access API")
    parser.add_argument("--expected-tables", dest="expected_tables", required=True,
                        help="comma-separated tables that must be present")
    parser.add_argument("--min-rows", type=int, default=1)
    parser.add_argument("--max-query-ms", type=float, default=250.0)
    parser.add_argument("--max-lag", type=float, default=5.0)
    parser.add_argument("--timeout", type=float, default=10.0)
    args = parser.parse_args(argv)

    if not args.base_url.startswith(("http://", "https://")):
        print("base_url must be an http(s) URL", file=sys.stderr)
        return int(Verdict.ABORT)
    cfg = build_config(args)
    if not cfg.expected_tables:
        print("--expected-tables must list at least one table", file=sys.stderr)
        return int(Verdict.ABORT)

    try:
        results = asyncio.run(run_checks(cfg))
    except Exception as exc:  # setup/transport failure short-circuits before any verdict
        print(f"gateway failed to execute checks: {exc}", file=sys.stderr)
        return int(Verdict.ABORT)

    for r in results:
        print(f"[{'PASS' if r.ok else 'FAIL'}] {r.name} ({r.latency_ms:.1f}ms): {r.detail}")

    if all(r.ok for r in results):
        print("verdict READY: gating open, routing real traffic to the sandbox")
        return int(Verdict.READY)
    failed = [r.name for r in results if not r.ok]
    print(f"verdict HOLD: gating closed, quarantine sandbox; failed: {failed}", file=sys.stderr)
    return int(Verdict.HOLD)


if __name__ == "__main__":
    sys.exit(main())

[PASS] schema_reachable (18.4ms): 3 tables present
[PASS] row_count_sanity (22.1ms): 3 tables above floor
[PASS] critical_query_latency (86.7ms): 86.7ms within budget
[PASS] replication_caught_up (15.9ms): lag 1.2s within tolerance
verdict READY: gating open, routing real traffic to the sandbox

[PASS] schema_reachable (17.9ms): 3 tables present
[PASS] row_count_sanity (24.0ms): 3 tables above floor
[PASS] critical_query_latency (91.3ms): 91.3ms within budget
[FAIL] replication_caught_up (16.2ms): lag 42.0s exceeds 5.0s tolerance
verdict HOLD: gating closed, quarantine sandbox; failed: ['replication_caught_up']