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.

    The container-log ring is prune_docker_logs, run in its own transaction.
    """
    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,
    }


async def prune_docker_logs(
    session,
    *,
    retention_days: int,
    max_bytes_per_container: int,
    now: datetime | None = None,
) -> dict:
    """The docker_logs ring (m79): drop lines older than the age window, then
    keep only the newest ~cap bytes per (host, container). Runs inside the
    caller's open transaction, like run_docker_retention.

    Only containers at or over the cap are ranked (#5449). Ranking every line
    in the table each hour to find the few containers over it sorted the whole
    table on slow storage.
    """
    from datetime import timezone

    from sqlalchemy import and_, delete, func, select

    from .models import DockerLog

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

    res = await session.execute(
        delete(DockerLog).where(DockerLog.ts < now - timedelta(days=retention_days))
    )
    age_pruned = res.rowcount or 0

    size = func.length(DockerLog.line)
    over = (
        select(DockerLog.host_id, DockerLog.container_name)
        .group_by(DockerLog.host_id, DockerLog.container_name)
        .having(func.sum(size) >= max_bytes_per_container)
        .subquery()
    )
    # Per container, sum line bytes newest-first; delete a row once its
    # strictly-newer siblings already fill the cap. 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(size).over(
        partition_by=[DockerLog.host_id, DockerLog.container_name],
        order_by=[DockerLog.ts.desc(), DockerLog.id.desc()],
    )
    ranked = (
        select(DockerLog.id, (running - size).label("prefix_excl"))
        .join(over, and_(DockerLog.host_id == over.c.host_id,
                         DockerLog.container_name == over.c.container_name))
        .subquery()
    )
    res = await session.execute(
        delete(DockerLog).where(DockerLog.id.in_(
            select(ranked.c.id).where(ranked.c.prefix_excl >= max_bytes_per_container)))
    )
    return {"logs_age_pruned": age_pruned, "logs_size_pruned": res.rowcount or 0}