FabledSteward/plugins/docker/retention.py

# plugins/docker/retention.py
"""Bound Docker time-series growth: roll up old metrics, prune old rows.

Published as the "docker.run_retention" capability (see __init__.setup) so the
core cleanup task can drive it WITHOUT importing the docker models (same
opportunistic-coupling pattern as docker.persist_host_samples). Runs inside the
caller's open transaction; never opens or commits its own.

The scaling concern is docker_metrics: ~2880 rows/container/day at a 30s sample.
We keep raw samples for a short window, then aggregate everything older into
hourly averages (docker_metrics_hourly) and delete the raw rows — so multi-day
history stays cheap to store and query. docker_events is light but unbounded
without a cutoff, so it gets a (longer) window too.
"""
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 ever roll up *whole* elapsed hours — a
    bucket is never split across the keep/roll boundary, so re-running can't
    produce a partial-then-complete duplicate for the same hour.
    """
    return _hour_floor(now - timedelta(days=raw_days))


async def run_docker_retention(
    session,
    *,
    events_days: int,
    metrics_raw_days: int,
    metrics_rollup_days: int,
    now: datetime | None = None,
) -> dict:
    """Roll up + prune Docker time-series. Returns a counts dict for logging.

    1. Aggregate docker_metrics older than the (hour-aligned) raw window into
       docker_metrics_hourly (avg cpu/mem per container per hour), upserting so a
       re-run is idempotent, then delete those raw rows.
    2. Prune rolled-up rows older than the rollup window.
    3. Prune docker_events older than the events window.
    """
    from datetime import timezone

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

    from .models import DockerEvent, DockerMetric, DockerMetricHourly

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

    rolled = rolled_rows = events_pruned = rollup_pruned = 0

    # ── 1. Roll up raw metrics older than the raw window into hourly buckets ──
    raw_cutoff = _rollup_cutoff(now, metrics_raw_days)
    hour = func.date_trunc("hour", DockerMetric.scraped_at)
    agg = (
        select(
            DockerMetric.host_id,
            DockerMetric.container_name,
            hour.label("bucket"),
            func.avg(DockerMetric.cpu_pct).label("cpu_pct"),
            func.avg(DockerMetric.mem_pct).label("mem_pct"),
            func.avg(DockerMetric.mem_usage_bytes).label("mem_usage_bytes"),
            func.count().label("sample_count"),
        )
        .where(DockerMetric.scraped_at < raw_cutoff)
        .group_by(DockerMetric.host_id, DockerMetric.container_name, hour)
    )
    for r in (await session.execute(agg)).all():
        stmt = (
            pg_insert(DockerMetricHourly)
            .values(
                host_id=r.host_id,
                container_name=r.container_name,
                bucket=r.bucket,
                cpu_pct=float(r.cpu_pct or 0.0),
                mem_pct=float(r.mem_pct or 0.0),
                mem_usage_bytes=int(r.mem_usage_bytes or 0),
                sample_count=int(r.sample_count or 0),
            )
            .on_conflict_do_update(
                constraint="uq_docker_metrics_hourly_bucket",
                set_={
                    "cpu_pct": float(r.cpu_pct or 0.0),
                    "mem_pct": float(r.mem_pct or 0.0),
                    "mem_usage_bytes": int(r.mem_usage_bytes or 0),
                    "sample_count": int(r.sample_count or 0),
                },
            )
        )
        await session.execute(stmt)
        rolled += 1
        rolled_rows += int(r.sample_count or 0)
    if rolled:
        await session.execute(
            delete(DockerMetric).where(DockerMetric.scraped_at < raw_cutoff)
        )

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

    # ── 3. Prune lifecycle events beyond the events window ──
    events_cutoff = now - timedelta(days=events_days)
    res = await session.execute(
        delete(DockerEvent).where(DockerEvent.at < events_cutoff)
    )
    events_pruned = res.rowcount or 0

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