FabledSteward/steward/core/cleanup.py

from __future__ import annotations
import asyncio
import logging
import os
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING

from sqlalchemy import delete

from steward.models.alerts import AlertEvent
from steward.models.ansible import AnsibleRun
from steward.models.audit import AuditEvent
from steward.models.monitors import MonitorResult

if TYPE_CHECKING:
    from quart import Quart

logger = logging.getLogger(__name__)


# Time-series columns pruned at data.retention_days (#5155). A plugin adds its
# own with prune_by_age() in setup(), so the core never imports plugin models.
# Tables kept another way are not listed: plugin_metrics and the docker tables
# are rolled up (below), and the audit log has its own, longer window.
_AGED: list = [MonitorResult.checked_at, AnsibleRun.started_at, AlertEvent.transitioned_at]


def prune_by_age(column) -> None:
    """Have the hourly cleanup delete rows whose `column` (a mapped DateTime
    attribute, e.g. TraefikMetric.scraped_at) is older than data.retention_days."""
    if not any(c is column for c in _AGED):
        _AGED.append(column)


async def run_cleanup(app: "Quart") -> None:
    """Delete rows older than their retention window from every growing table,
    and run logs older than the data window, then run the metric and Docker
    rollups. Windows are read from Settings each run (rule 25: no restart)."""
    from steward.core.settings import get_setting
    from steward.ansible.executor import ARTIFACT_DIR

    now = datetime.now(timezone.utc)
    async with app.db_sessionmaker() as session:
        async with session.begin():
            data_days = int(await get_setting(session, "data.retention_days") or 90)
            audit_days = int(await get_setting(session, "audit.retention_days") or 365)
            cutoff = now - timedelta(days=data_days)
            aged = [(c, cutoff) for c in _AGED]
            aged.append((AuditEvent.timestamp, now - timedelta(days=audit_days)))
            for column, before in aged:
                result = await session.execute(
                    delete(column.class_).where(column < before)
                )
                if result.rowcount:
                    logger.info("Pruned %d rows from %s", result.rowcount,
                                column.class_.__tablename__)

        removed = await asyncio.to_thread(prune_run_logs, ARTIFACT_DIR, cutoff)
        if removed:
            logger.info("Pruned %d Ansible run log files from %s", removed, ARTIFACT_DIR)

        # plugin_metrics is NOT blanket-deleted here — it's rolled up to hourly
        # then pruned, so multi-week host history stays cheap.
        await _run_metrics_retention(session, now)
        await _run_docker_retention(session, now)


def prune_run_logs(directory: str, before: datetime) -> int:
    """Delete the run log files in `directory` last written before `before`,
    the same window as the runs they belong to. Returns how many went."""
    removed = 0
    try:
        entries = list(os.scandir(directory))
    except OSError:          # no run has written one yet, or not readable
        return 0
    limit = before.timestamp()
    for entry in entries:
        try:
            if (entry.name.endswith(".log") and entry.is_file()
                    and entry.stat().st_mtime < limit):
                os.unlink(entry.path)
                removed += 1
        except OSError:
            logger.warning("Could not prune run log %s", entry.path)
    return removed


async def _run_metrics_retention(session, now: datetime) -> None:
    """Roll up + prune plugin_metrics (raw → hourly → gone). Windows read fresh
    from settings each run (rule 25 — UI change takes effect next cleanup, no
    restart). get_setting's SELECT autobegins, so read inside the begin block."""
    from steward.core.metrics_retention import rollup_plugin_metrics
    from steward.core.settings import get_setting

    async with session.begin():
        raw_days = int(await get_setting(session, "metrics.retention.raw_days") or 7)
        rollup_days = int(await get_setting(session, "metrics.retention.rollup_days") or 90)
        counts = await rollup_plugin_metrics(
            session, raw_days=raw_days, rollup_days=rollup_days, now=now,
        )
    if counts and any(counts.values()):
        logger.info("Metrics retention: %s", counts)


async def _run_docker_retention(session, now: datetime) -> None:
    """Drive the docker plugin's rollup + prune via its capability, if loaded.

    Windows are read fresh from settings each run (rule 25 — a change in the
    Settings UI takes effect on the next hourly cleanup, no restart). Kept in its
    own transaction so a docker-side failure can't roll back the generic prune
    above. No-op when the docker plugin is disabled (capability absent).
    """
    from steward.core.capabilities import has_capability, invoke_capability
    if not has_capability("docker.run_retention"):
        return
    from steward.core.settings import get_setting
    from steward.models.users import UserRole

    # Reads + rollup/prune share one transaction — get_setting's SELECT would
    # otherwise autobegin one, making a later session.begin() raise.
    async with session.begin():
        raw_days = int(await get_setting(session, "docker.retention.metrics_raw_days") or 7)
        rollup_days = int(await get_setting(session, "docker.retention.metrics_rollup_days") or 90)
        events_days = int(await get_setting(session, "docker.retention.events_days") or 30)
        logs_days = int(await get_setting(session, "docker.logs.retention_days") or 3)
        logs_cap = int(
            await get_setting(session, "docker.logs.max_bytes_per_container") or 5_000_000)
        counts = await invoke_capability(
            "docker.run_retention", UserRole.viewer, session,
            events_days=events_days, metrics_raw_days=raw_days,
            metrics_rollup_days=rollup_days, logs_retention_days=logs_days,
            logs_max_bytes_per_container=logs_cap, now=now,
        )
    if counts and any(counts.values()):
        logger.info("Docker retention: %s", counts)