netx/netx_api/biz_state/retention.py

"""Biz-state batch retention: age/day policies + reference protection."""

from __future__ import annotations

from collections import defaultdict
from datetime import datetime, timedelta
from typing import Any

from sqlalchemy.orm import Session

from ..models import (
    BizCompareJob,
    BizCompareRun,
    BizMigrationProject,
    BizMigrationRun,
    BizStateBatch,
    BizStateBatchCommand,
    BizStateLldpNeighbor,
    BizStateMetricRow,
    BizStateTask,
    BizStateVrfRouteSummary,
)


def _utcnow() -> datetime:
    return datetime.utcnow()


def protected_batch_map(db: Session, *, task_id: str = "") -> dict[str, list[str]]:
    """batch_id → reason codes. Empty task_id = all tasks."""
    out: dict[str, list[str]] = defaultdict(list)

    def add(bid: str, reason: str) -> None:
        b = str(bid or "").strip()
        if not b:
            return
        if reason not in out[b]:
            out[b].append(reason)

    q = db.query(BizStateBatch)
    if task_id:
        q = q.filter(BizStateBatch.task_id == task_id)
    for b in q.filter(BizStateBatch.is_baseline.is_(True)).all():
        add(b.id, "manual_baseline")

    for j in db.query(BizCompareJob).all():
        add(j.before_batch_id, "compare_job_before")
        add(j.after_batch_id, "compare_job_after")

    for r in db.query(BizCompareRun).all():
        add(r.before_batch_id, "compare_run_before")
        add(r.after_batch_id, "compare_run_after")

    for p in db.query(BizMigrationProject).all():
        add(p.old_baseline_batch_id, "migration_old_baseline")
        add(p.new_baseline_batch_id, "migration_new_baseline")

    for r in db.query(BizMigrationRun).all():
        add(r.old_batch_id, "migration_run_old")
        add(r.new_batch_id, "migration_run_new")

    return dict(out)


def protected_batch_ids(db: Session, *, task_id: str = "") -> set[str]:
    return set(protected_batch_map(db, task_id=task_id).keys())


def batch_protect_info(db: Session, batch_id: str) -> dict[str, Any]:
    reasons = protected_batch_map(db).get(batch_id, [])
    b = db.get(BizStateBatch, batch_id)
    if b and bool(getattr(b, "is_baseline", False)) and "manual_baseline" not in reasons:
        reasons = ["manual_baseline", *reasons]
    return {
        "protected": bool(reasons),
        "reasons": reasons,
        "is_baseline": bool(b and getattr(b, "is_baseline", False)),
    }


def delete_batch_data(db: Session, batch_id: str) -> None:
    """Hard-delete one batch and child rows (caller must check protection)."""
    bid = str(batch_id or "").strip()
    if not bid:
        return
    db.query(BizStateLldpNeighbor).filter(BizStateLldpNeighbor.batch_id == bid).delete()
    db.query(BizStateVrfRouteSummary).filter(BizStateVrfRouteSummary.batch_id == bid).delete()
    db.query(BizStateMetricRow).filter(BizStateMetricRow.batch_id == bid).delete()
    db.query(BizStateBatchCommand).filter(BizStateBatchCommand.batch_id == bid).delete()
    b = db.get(BizStateBatch, bid)
    if b:
        db.delete(b)


def purge_task_batches(db: Session, task: BizStateTask) -> dict[str, Any]:
    """Apply retention_days + optional daily_keep policy. Never deletes protected batches.

    Returns counts for logging/UI.
    """
    task_id = str(task.id)
    retention_days = max(1, int(getattr(task, "retention_days", None) or 30))
    daily_on = bool(getattr(task, "daily_keep_enabled", False))
    daily_n = max(1, int(getattr(task, "daily_keep_count", None) or 10))

    protected = protected_batch_ids(db, task_id=task_id)
    rows = (
        db.query(BizStateBatch)
        .filter(BizStateBatch.task_id == task_id)
        .order_by(BizStateBatch.started_at.desc())
        .all()
    )
    now = _utcnow()
    cutoff = now - timedelta(days=retention_days)
    today = now.date()

    to_drop: set[str] = set()

    # 1) Age: older than retention_days
    for b in rows:
        if b.id in protected:
            continue
        started = b.started_at or now
        if started < cutoff:
            to_drop.add(b.id)

    # 2) Daily keep (default off): for each past calendar day, keep N newest
    if daily_on:
        by_day: dict[Any, list[BizStateBatch]] = defaultdict(list)
        for b in rows:
            started = b.started_at or now
            d = started.date()
            if d >= today:
                continue  # 当天的多余留到第二天再清
            by_day[d].append(b)
        for _day, day_rows in by_day.items():
            # already ordered desc globally; re-sort
            day_rows.sort(key=lambda x: x.started_at or now, reverse=True)
            for b in day_rows[daily_n:]:
                if b.id not in protected:
                    to_drop.add(b.id)

    dropped = 0
    for bid in to_drop:
        delete_batch_data(db, bid)
        dropped += 1
    if dropped:
        db.commit()
    return {
        "task_id": task_id,
        "dropped": dropped,
        "protected": len(protected),
        "retention_days": retention_days,
        "daily_keep_enabled": daily_on,
        "daily_keep_count": daily_n,
    }


def purge_all_tasks(db: Session) -> dict[str, Any]:
    """Periodic sweep for all tasks (HF idle tasks included)."""
    tasks = db.query(BizStateTask).all()
    total = 0
    details: list[dict[str, Any]] = []
    for t in tasks:
        info = purge_task_batches(db, t)
        total += int(info.get("dropped") or 0)
        if info.get("dropped"):
            details.append(info)
    return {"dropped": total, "tasks": details}