FabledSteward/steward/core/metrics_retention.py

"""Bound plugin_metrics growth: roll old raw samples up to hourly, prune the rest.

plugin_metrics grows by (sources × resources × sample cadence) — host agents push
host-level + per-core/mount/iface sub-resources every ~30s, so a fleet accrues
millions of rows. We keep raw samples for a short window, aggregate everything
older into hourly averages (plugin_metrics_hourly) and delete the raw rows, then
prune hourly beyond a longer window. Charts read raw for the recent part of a
range and hourly for the older part.

Driven by the core cleanup task (steward.core.cleanup). Runs inside the caller's
open transaction; never opens or commits its own.
"""
from __future__ import annotations

from datetime import datetime, timedelta


def _hour_floor(dt: datetime) -> datetime:
    """Truncate a datetime down to the start of its hour (drops min/sec/µs)."""
    return dt.replace(minute=0, second=0, microsecond=0)


def _rollup_cutoff(now: datetime, raw_days: int) -> datetime:
    """Hour-aligned boundary below which raw metrics get rolled up + deleted.

    Aligning to the hour means we only roll up *whole* elapsed hours — a bucket
    is never split across the keep/roll boundary, so a re-run can't produce a
    partial-then-complete duplicate for the same hour.
    """
    return _hour_floor(now - timedelta(days=raw_days))


async def rollup_plugin_metrics(
    session,
    *,
    raw_days: int,
    rollup_days: int,
    now: datetime | None = None,
) -> dict:
    """Roll up + prune plugin_metrics. Returns a counts dict for logging.

    1. Aggregate plugin_metrics older than the (hour-aligned) raw window into
       plugin_metrics_hourly (avg/max per source/resource/metric/hour), upserting
       so a re-run is idempotent, then delete those raw rows.
    2. Prune rolled-up rows older than the rollup window.
    """
    from datetime import timezone

    from sqlalchemy import delete, func, select
    from sqlalchemy.dialects.postgresql import insert as pg_insert

    from steward.models.metrics import PluginMetric, PluginMetricHourly

    if now is None:
        now = datetime.now(timezone.utc)

    rolled = rolled_rows = rollup_pruned = 0

    # ── 1. Roll up raw metrics older than the raw window into hourly buckets ──
    raw_cutoff = _rollup_cutoff(now, raw_days)
    hour = func.date_trunc("hour", PluginMetric.recorded_at)
    agg = (
        select(
            PluginMetric.source_module,
            PluginMetric.resource_name,
            PluginMetric.metric_name,
            hour.label("bucket"),
            func.avg(PluginMetric.value).label("value_avg"),
            func.max(PluginMetric.value).label("value_max"),
            func.count().label("sample_count"),
        )
        .where(PluginMetric.recorded_at < raw_cutoff)
        .group_by(
            PluginMetric.source_module, PluginMetric.resource_name,
            PluginMetric.metric_name, hour,
        )
    )
    for r in (await session.execute(agg)).all():
        avg_v = float(r.value_avg or 0.0)
        max_v = float(r.value_max or 0.0)
        cnt = int(r.sample_count or 0)
        await session.execute(
            pg_insert(PluginMetricHourly)
            .values(
                source_module=r.source_module, resource_name=r.resource_name,
                metric_name=r.metric_name, bucket=r.bucket,
                value_avg=avg_v, value_max=max_v, sample_count=cnt,
            )
            .on_conflict_do_update(
                constraint="uq_plugin_metrics_hourly_bucket",
                set_={"value_avg": avg_v, "value_max": max_v, "sample_count": cnt},
            )
        )
        rolled += 1
        rolled_rows += cnt
    if rolled:
        await session.execute(
            delete(PluginMetric).where(PluginMetric.recorded_at < raw_cutoff)
        )

    # ── 2. Prune rolled-up rows beyond the rollup window ──
    rollup_cutoff = now - timedelta(days=rollup_days)
    res = await session.execute(
        delete(PluginMetricHourly).where(PluginMetricHourly.bucket < rollup_cutoff)
    )
    rollup_pruned = res.rowcount or 0

    return {
        "buckets_rolled": rolled,
        "raw_rows_rolled": rolled_rows,
        "rollup_pruned": rollup_pruned,
    }