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,
    logs_retention_days: int = 3,
    logs_max_bytes_per_container: int = 5_000_000,
    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.
    4. Prune docker_logs with a per-container size+age ring (m79): drop lines
       older than the age window, then keep only the newest ~cap bytes per
       (host, container).
    """
    from datetime import timezone

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

    from .models import DockerEvent, DockerLog, DockerMetric, DockerMetricHourly

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

    rolled = rolled_rows = events_pruned = rollup_pruned = 0
    logs_age_pruned = logs_size_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

    # ── 4. Container-log ring: age cutoff, then per-container byte cap (m79) ──
    logs_cutoff = now - timedelta(days=logs_retention_days)
    res = await session.execute(
        delete(DockerLog).where(DockerLog.ts < logs_cutoff)
    )
    logs_age_pruned = res.rowcount or 0

    # Size ring: per (host, container), sum line bytes newest-first; delete a row
    # once its strictly-newer siblings already fill the cap. Using the EXCLUSIVE
    # prefix (running total minus this row) means the newest row always survives,
    # so a single line larger than the cap is never wiped out.
    running = func.sum(func.length(DockerLog.line)).over(
        partition_by=[DockerLog.host_id, DockerLog.container_name],
        order_by=[DockerLog.ts.desc(), DockerLog.id.desc()],
    )
    prefix_excl = (running - func.length(DockerLog.line)).label("prefix_excl")
    ranked = select(DockerLog.id, prefix_excl).subquery()
    over_cap = select(ranked.c.id).where(
        ranked.c.prefix_excl >= logs_max_bytes_per_container)
    res = await session.execute(
        delete(DockerLog).where(DockerLog.id.in_(over_cap))
    )
    logs_size_pruned = res.rowcount or 0

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