awoooi/ops/agent99-truth-chain-index-executor-overlay.py

#!/usr/bin/env python3
"""Apply the fixed truth-chain index set without accepting raw SQL."""

from __future__ import annotations

import argparse
import asyncio
import hashlib
import io
import json
import os
import re
import socket
import sys
import time
from contextlib import redirect_stderr, redirect_stdout
from typing import Any

from sqlalchemy import text

from src.db.base import get_engine


_SCHEMA_VERSION = "agent99_truth_chain_index_executor_v1"
_COMMAND_ID = "truth_chain_structured_indexes_v1"
_CANONICAL_ASSET = "database:awoooi:prod/truth-chain-structured-indexes"
_MIGRATION_SHA256 = (
    "657089552ceed8de916d87135164b255c85edc154bf437f870b8c2bbdf14c509"
)
_SAFE_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:@+-]{0,127}$")

_EXTENSION_OPERATIONS = (
    ("pg_trgm", "CREATE EXTENSION IF NOT EXISTS pg_trgm"),
    ("btree_gin", "CREATE EXTENSION IF NOT EXISTS btree_gin"),
)

_INDEX_OPERATIONS = (
    (
        "idx_aol_incident_text_created",
        """
        CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_aol_incident_text_created
            ON automation_operation_log ((incident_id::text), created_at DESC)
            WHERE incident_id IS NOT NULL
        """,
    ),
    (
        "idx_aol_input_incident_created",
        """
        CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_aol_input_incident_created
            ON automation_operation_log ((input ->> 'incident_id'), created_at DESC)
            WHERE input ? 'incident_id'
        """,
    ),
    (
        "idx_aol_output_incident_created",
        """
        CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_aol_output_incident_created
            ON automation_operation_log ((output ->> 'incident_id'), created_at DESC)
            WHERE output ? 'incident_id'
        """,
    ),
    (
        "idx_aol_input_source_incidents_gin",
        """
        CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_aol_input_source_incidents_gin
            ON automation_operation_log
            USING GIN ((COALESCE(input #> '{source_refs,incident_ids}', '[]'::jsonb)))
        """,
    ),
    (
        "idx_aol_output_source_incidents_gin",
        """
        CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_aol_output_source_incidents_gin
            ON automation_operation_log
            USING GIN ((COALESCE(output #> '{source_refs,incident_ids}', '[]'::jsonb)))
        """,
    ),
    (
        "idx_outbound_source_incident_ids_gin",
        """
        CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_outbound_source_incident_ids_gin
            ON awooop_outbound_message
            USING GIN ((COALESCE(source_envelope #> '{source_refs,incident_ids}', '[]'::jsonb)))
        """,
    ),
    (
        "idx_outbound_source_code_refs_gin",
        """
        CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_outbound_source_code_refs_gin
            ON awooop_outbound_message
            USING GIN ((COALESCE(source_envelope #> '{source_refs,code_refs}', '[]'::jsonb)))
        """,
    ),
    (
        "idx_outbound_callback_incident_recent",
        """
        CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_outbound_callback_incident_recent
            ON awooop_outbound_message (
                project_id,
                (source_envelope #>> '{callback_reply,incident_id}'),
                queued_at DESC
            )
            WHERE source_envelope #>> '{callback_reply,incident_id}' IS NOT NULL
        """,
    ),
    (
        "idx_outbound_alert_incident_recent",
        """
        CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_outbound_alert_incident_recent
            ON awooop_outbound_message (
                project_id,
                (source_envelope #>> '{alert_notification,incident_id}'),
                queued_at DESC
            )
            WHERE source_envelope #>> '{alert_notification,incident_id}' IS NOT NULL
        """,
    ),
    (
        "idx_outbound_source_incident_recent",
        """
        CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_outbound_source_incident_recent
            ON awooop_outbound_message (
                project_id,
                (source_envelope #>> '{source_refs,incident_id}'),
                queued_at DESC
            )
            WHERE source_envelope #>> '{source_refs,incident_id}' IS NOT NULL
        """,
    ),
    (
        "idx_outbound_top_incident_recent",
        """
        CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_outbound_top_incident_recent
            ON awooop_outbound_message (
                project_id,
                (source_envelope ->> 'incident_id'),
                queued_at DESC
            )
            WHERE source_envelope ->> 'incident_id' IS NOT NULL
        """,
    ),
    (
        "idx_outbound_project_content_preview_trgm",
        """
        CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_outbound_project_content_preview_trgm
            ON awooop_outbound_message
            USING GIN (project_id, content_preview gin_trgm_ops)
            WHERE content_preview IS NOT NULL
        """,
    ),
)

_EXPECTED_EXTENSIONS = tuple(name for name, _ in _EXTENSION_OPERATIONS)
_EXPECTED_INDEXES = tuple(name for name, _ in _INDEX_OPERATIONS)

_CATALOG_SQL = text(
    """
    SELECT index_class.relname AS index_name,
           index_state.indisvalid,
           index_state.indisready
    FROM pg_catalog.pg_class AS index_class
    JOIN pg_catalog.pg_index AS index_state
      ON index_state.indexrelid = index_class.oid
    WHERE index_class.relname IN (
        'idx_aol_incident_text_created',
        'idx_aol_input_incident_created',
        'idx_aol_output_incident_created',
        'idx_aol_input_source_incidents_gin',
        'idx_aol_output_source_incidents_gin',
        'idx_outbound_source_incident_ids_gin',
        'idx_outbound_source_code_refs_gin',
        'idx_outbound_callback_incident_recent',
        'idx_outbound_alert_incident_recent',
        'idx_outbound_source_incident_recent',
        'idx_outbound_top_incident_recent',
        'idx_outbound_project_content_preview_trgm'
    )
    """
)


def _normalized_sql(value: str) -> str:
    return re.sub(r"\s+", " ", value.strip()).lower()


def _sha256(value: str) -> str:
    return hashlib.sha256(value.encode("utf-8")).hexdigest()


def _operation_digest() -> str:
    statements = [sql for _, sql in (*_EXTENSION_OPERATIONS, *_INDEX_OPERATIONS)]
    return _sha256("\n".join(_normalized_sql(sql) for sql in statements))


def _safe_error(exc: BaseException) -> dict[str, Any]:
    database_error = getattr(exc, "orig", None)
    sqlstate = str(
        getattr(database_error, "sqlstate", None)
        or getattr(database_error, "pgcode", None)
        or ""
    )
    return {
        "error_class": exc.__class__.__name__,
        "database_error_class": (
            database_error.__class__.__name__ if database_error else None
        ),
        "sqlstate": sqlstate or None,
    }


async def _catalog(connection: Any) -> dict[str, Any]:
    extension_result = await connection.execute(
        text(
            """
            SELECT extname
            FROM pg_catalog.pg_extension
            WHERE extname IN ('pg_trgm', 'btree_gin')
            ORDER BY extname
            """
        )
    )
    extensions = {
        str(row.get("extname") or "")
        for row in extension_result.mappings().all()
    }
    index_result = await connection.execute(_CATALOG_SQL)
    index_rows = list(index_result.mappings().all())
    present = {str(row.get("index_name") or "") for row in index_rows}
    valid_ready = {
        str(row.get("index_name") or "")
        for row in index_rows
        if row.get("indisvalid") and row.get("indisready")
    }
    missing_extensions = sorted(set(_EXPECTED_EXTENSIONS) - extensions)
    missing_indexes = sorted(set(_EXPECTED_INDEXES) - present)
    invalid_or_unready = sorted(present - valid_ready)
    return {
        "expected_extension_count": len(_EXPECTED_EXTENSIONS),
        "extension_count": len(extensions),
        "missing_extensions": missing_extensions,
        "expected_index_count": len(_EXPECTED_INDEXES),
        "present_index_count": len(present),
        "valid_ready_index_count": len(valid_ready),
        "missing_indexes": missing_indexes,
        "invalid_or_unready_indexes": invalid_or_unready,
        "passed": not missing_extensions
        and not missing_indexes
        and not invalid_or_unready,
    }


async def _execute_fixed_operation(
    connection: Any,
    *,
    operation_type: str,
    name: str,
    sql: str,
) -> dict[str, Any]:
    started = time.perf_counter()
    try:
        await connection.execute(text(sql))
    except Exception as exc:
        return {
            "type": operation_type,
            "name": name,
            "statement_sha256": _sha256(_normalized_sql(sql)),
            "duration_ms": round((time.perf_counter() - started) * 1000),
            "success": False,
            **_safe_error(exc),
        }
    return {
        "type": operation_type,
        "name": name,
        "statement_sha256": _sha256(_normalized_sql(sql)),
        "duration_ms": round((time.perf_counter() - started) * 1000),
        "success": True,
        "error_class": None,
        "database_error_class": None,
        "sqlstate": None,
    }


async def _run(args: argparse.Namespace) -> dict[str, Any]:
    for identity in (args.trace_id, args.run_id, args.work_item_id):
        if not _SAFE_ID.fullmatch(identity):
            raise ValueError("identity_invalid")
    if not re.fullmatch(r"[0-9a-f]{40}", args.source_revision):
        raise ValueError("source_revision_invalid")
    if args.migration_sha256 != _MIGRATION_SHA256:
        raise ValueError("migration_sha256_mismatch")

    started = time.perf_counter()
    operations: list[dict[str, Any]] = []
    engine = get_engine()
    async with engine.connect() as raw_connection:
        connection = await raw_connection.execution_options(
            isolation_level="AUTOCOMMIT"
        )
        before = await _catalog(connection)
        if args.mode == "apply" and before["invalid_or_unready_indexes"]:
            after = before
            terminal = "blocked_invalid_or_unready_index"
        else:
            if args.mode == "apply":
                await connection.execute(text("SET lock_timeout = '15s'"))
                await connection.execute(text("SET statement_timeout = '10min'"))
                for name, sql in _EXTENSION_OPERATIONS:
                    if name not in before["missing_extensions"]:
                        continue
                    operation = await _execute_fixed_operation(
                        connection,
                        operation_type="create_extension",
                        name=name,
                        sql=sql,
                    )
                    operations.append(operation)
                    if not operation["success"]:
                        break
                if not operations or operations[-1]["success"]:
                    for name, sql in _INDEX_OPERATIONS:
                        if name not in before["missing_indexes"]:
                            continue
                        operation = await _execute_fixed_operation(
                            connection,
                            operation_type="create_index_concurrently",
                            name=name,
                            sql=sql,
                        )
                        operations.append(operation)
                        if not operation["success"]:
                            break
            after = await _catalog(connection)
            failed_operation = any(not operation["success"] for operation in operations)
            if args.mode == "check":
                terminal = "check_ready"
            elif args.mode == "verify":
                terminal = "verified_healthy" if after["passed"] else "verification_failed"
            elif failed_operation:
                terminal = "apply_failed"
            else:
                terminal = (
                    "apply_verified_local"
                    if after["passed"]
                    else "verification_failed"
                )

    process_identity = ":".join(
        [socket.gethostname(), str(os.getpid()), args.run_id, args.mode]
    )
    write_attempted = bool(operations)
    writes_performed = any(operation["success"] for operation in operations)
    return {
        "schema_version": _SCHEMA_VERSION,
        "command_id": _COMMAND_ID,
        "canonical_asset": _CANONICAL_ASSET,
        "mode": args.mode,
        "trace_id": args.trace_id,
        "run_id": args.run_id,
        "work_item_id": args.work_item_id,
        "source_revision": args.source_revision,
        "migration_sha256": args.migration_sha256,
        "operation_digest_sha256": _operation_digest(),
        "pod_name": socket.gethostname(),
        "process_identity_sha256": _sha256(process_identity),
        "elapsed_ms": round((time.perf_counter() - started) * 1000),
        "catalog_before": before,
        "catalog_after": after,
        "operations": operations,
        "write_attempted": write_attempted,
        "writes_performed": writes_performed,
        "raw_sql_accepted": False,
        "raw_sql_emitted": False,
        "query_parameters_emitted": False,
        "secret_value_read": False,
        "arbitrary_command_allowed": False,
        "cross_domain_fallback_allowed": False,
        "verifier": {
            "passed": args.mode == "check" or bool(after["passed"]),
            "target_ready": bool(after["passed"]),
        },
        "terminal": terminal,
    }


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--mode", choices=("check", "apply", "verify"), required=True)
    parser.add_argument("--trace-id", required=True)
    parser.add_argument("--run-id", required=True)
    parser.add_argument("--work-item-id", required=True)
    parser.add_argument("--source-revision", required=True)
    parser.add_argument("--migration-sha256", required=True)
    return parser.parse_args()


def main() -> int:
    try:
        application_output = io.StringIO()
        with redirect_stdout(application_output), redirect_stderr(application_output):
            receipt = asyncio.run(_run(_parse_args()))
    except Exception as exc:
        receipt = {
            "schema_version": _SCHEMA_VERSION,
            "command_id": _COMMAND_ID,
            "terminal": "failed",
            **_safe_error(exc),
            "write_attempted": False,
            "writes_performed": False,
            "raw_sql_accepted": False,
            "raw_sql_emitted": False,
            "query_parameters_emitted": False,
            "secret_value_read": False,
            "arbitrary_command_allowed": False,
            "cross_domain_fallback_allowed": False,
        }
    print(json.dumps(receipt, separators=(",", ":"), sort_keys=True))
    return 0 if receipt.get("terminal") in {
        "check_ready",
        "apply_verified_local",
        "verified_healthy",
    } else 3


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