vignette/apps/api/app/services/admin_health_maintenance.py

"""Maintenance helpers for persisted admin health samples."""

from __future__ import annotations

from dataclasses import dataclass

from ..db import acquire


@dataclass(frozen=True)
class AdminHealthMaintenanceResult:
    applied: bool
    rollup_days: int
    retention_days: int
    rollup_event_count: int
    rollup_bucket_count: int
    upserted_rollups: int
    prunable_event_count: int
    deleted_events: int


def _validate_windows(*, rollup_days: int, retention_days: int) -> None:
    if rollup_days < 1:
        raise ValueError("rollup_days must be 1 or greater")
    if retention_days < rollup_days:
        raise ValueError("retention_days must be greater than or equal to rollup_days")


def _deleted_count(command_tag: str) -> int:
    parts = (command_tag or "").split()
    if len(parts) >= 2 and parts[0].upper() == "DELETE":
        try:
            return int(parts[-1])
        except ValueError:
            return 0
    return 0


def _row_int(row, key: str) -> int:
    if row is None:
        return 0
    try:
        return int(row[key] or 0)
    except (KeyError, TypeError, ValueError):
        return 0


async def maintain_admin_health_events(
    *,
    rollup_days: int,
    retention_days: int,
    apply: bool = False,
) -> AdminHealthMaintenanceResult:
    """Roll up old health samples and optionally prune retained raw rows.

    Dry-run mode returns the affected event and bucket counts without mutating
    data.  Apply mode upserts daily service rollups first, then deletes only raw
    events older than the retention window.
    """
    _validate_windows(rollup_days=rollup_days, retention_days=retention_days)
    async with acquire(role="admin") as conn:
        rollup_stats = await conn.fetchrow(
            """
            SELECT
              COUNT(*)::int AS event_count,
              COUNT(DISTINCT (observed_at::date, environment, engine_mode, service_key))::int
                AS bucket_count
            FROM app.admin_health_event
            WHERE observed_at::date < current_date - $1::int
            """,
            rollup_days,
        )
        prune_stats = await conn.fetchrow(
            """
            SELECT COUNT(*)::int AS event_count
            FROM app.admin_health_event
            WHERE observed_at::date < current_date - $1::int
            """,
            retention_days,
        )
        upserted_rollups = 0
        deleted_events = 0
        if apply:
            rows = await conn.fetch(
                """
                WITH rolled AS (
                  SELECT
                    observed_at::date AS rollup_date,
                    environment,
                    engine_mode,
                    service_key,
                    (array_agg(service_name ORDER BY observed_at DESC, id DESC))[1] AS service_name,
                    COUNT(*)::int AS sample_count,
                    COUNT(*) FILTER (WHERE service_status = 'ok')::int AS ok_samples,
                    COUNT(*) FILTER (WHERE service_status = 'degraded')::int AS degraded_samples,
                    COUNT(*) FILTER (WHERE service_status = 'down')::int AS down_samples,
                    MIN(observed_at) AS first_observed_at,
                    MAX(observed_at) AS last_observed_at,
                    (array_agg(service_status ORDER BY observed_at DESC, id DESC))[1] AS latest_status,
                    MAX(observed_at) FILTER (WHERE service_status = 'down') AS last_down_at,
                    MAX(load)::real AS max_load
                  FROM app.admin_health_event
                  WHERE observed_at::date < current_date - $1::int
                  GROUP BY observed_at::date, environment, engine_mode, service_key
                )
                INSERT INTO app.admin_health_daily_rollup (
                  rollup_date,
                  environment,
                  engine_mode,
                  service_key,
                  service_name,
                  sample_count,
                  ok_samples,
                  degraded_samples,
                  down_samples,
                  first_observed_at,
                  last_observed_at,
                  latest_status,
                  last_down_at,
                  max_load,
                  generated_at
                )
                SELECT
                  rollup_date,
                  environment,
                  engine_mode,
                  service_key,
                  service_name,
                  sample_count,
                  ok_samples,
                  degraded_samples,
                  down_samples,
                  first_observed_at,
                  last_observed_at,
                  latest_status,
                  last_down_at,
                  max_load,
                  now()
                FROM rolled
                ON CONFLICT (rollup_date, environment, engine_mode, service_key)
                DO UPDATE SET
                  service_name = EXCLUDED.service_name,
                  sample_count = EXCLUDED.sample_count,
                  ok_samples = EXCLUDED.ok_samples,
                  degraded_samples = EXCLUDED.degraded_samples,
                  down_samples = EXCLUDED.down_samples,
                  first_observed_at = EXCLUDED.first_observed_at,
                  last_observed_at = EXCLUDED.last_observed_at,
                  latest_status = EXCLUDED.latest_status,
                  last_down_at = EXCLUDED.last_down_at,
                  max_load = EXCLUDED.max_load,
                  generated_at = now()
                RETURNING 1
                """,
                rollup_days,
            )
            upserted_rollups = len(rows)
            deleted_events = _deleted_count(
                await conn.execute(
                    """
                    DELETE FROM app.admin_health_event
                    WHERE observed_at::date < current_date - $1::int
                    """,
                    retention_days,
                )
            )

    return AdminHealthMaintenanceResult(
        applied=apply,
        rollup_days=rollup_days,
        retention_days=retention_days,
        rollup_event_count=_row_int(rollup_stats, "event_count"),
        rollup_bucket_count=_row_int(rollup_stats, "bucket_count"),
        upserted_rollups=upserted_rollups,
        prunable_event_count=_row_int(prune_stats, "event_count"),
        deleted_events=deleted_events,
    )