AirflowDAGMonitor/airflow_monitor/state.py

"""Persistent state manager for tracking DAG run retries and alerts."""

import json
import logging
import os
import pathlib
import tempfile
import time
from datetime import datetime, timezone

logger = logging.getLogger(__name__)


class StateManager:
    """Manages persistent state in a JSON file.

    State tracks retry counts and alert flags per DAG run to ensure
    idempotent behavior across monitoring cycles.
    """

    def __init__(self, state_file: str, max_age: int = 86400):
        self._path = pathlib.Path(state_file)
        self._max_age = max_age
        self._data: dict = {"version": 1, "entries": {}}
        logger.debug(
            "StateManager initialized: file=%s, max_age=%ds",
            state_file, max_age,
        )

    def load(self):
        """Load state from disk. Start fresh if missing or corrupt."""
        if not self._path.exists():
            logger.debug("State file not found at %s, starting with empty state", self._path)
            self._data = {"version": 1, "entries": {}}
            return

        try:
            with open(self._path, "r", encoding="utf-8") as f:
                raw = f.read()
            self._data = json.loads(raw)
            if not isinstance(self._data.get("entries"), dict):
                raise ValueError("Invalid state structure: 'entries' is not a dict")
            entry_count = len(self._data["entries"])
            logger.debug(
                "State loaded: %d entries, file_size=%d bytes",
                entry_count, len(raw),
            )
            if entry_count > 0:
                logger.debug(
                    "State entries: %s",
                    ", ".join(
                        f"{k}(retries={v.get('retry_count', 0)}, alerted={v.get('alerted', False)})"
                        for k, v in self._data["entries"].items()
                    ),
                )
        except (json.JSONDecodeError, ValueError, KeyError) as e:
            logger.warning("Corrupt state file %s, starting fresh: %s", self._path, e)
            self._data = {"version": 1, "entries": {}}

    def save(self):
        """Save state to disk atomically (write tmp -> rename)."""
        self._path.parent.mkdir(parents=True, exist_ok=True)
        try:
            content = json.dumps(self._data, indent=2, ensure_ascii=False)
            fd, tmp_path = tempfile.mkstemp(
                dir=str(self._path.parent), suffix=".tmp"
            )
            with os.fdopen(fd, "w", encoding="utf-8") as f:
                f.write(content)
            os.replace(tmp_path, str(self._path))
            logger.debug(
                "State saved: %d entries, %d bytes → %s",
                len(self._data["entries"]), len(content), self._path,
            )
        except OSError as e:
            logger.error("Failed to save state to %s: %s", self._path, e)
            # Clean up temp file if it exists
            try:
                os.unlink(tmp_path)
            except OSError:
                pass

    @staticmethod
    def _key(dag_id: str, dag_run_id: str) -> str:
        return f"{dag_id}::{dag_run_id}"

    @staticmethod
    def _now_iso() -> str:
        return datetime.now(timezone.utc).isoformat()

    def get_entry(self, dag_id: str, dag_run_id: str) -> dict | None:
        """Get stored entry for a DAG run, or None."""
        return self._data["entries"].get(self._key(dag_id, dag_run_id))

    def ensure_entry(self, dag_id: str, dag_run_id: str) -> dict:
        """Get or create an entry for a DAG run."""
        key = self._key(dag_id, dag_run_id)
        if key not in self._data["entries"]:
            self._data["entries"][key] = {
                "dag_id": dag_id,
                "dag_run_id": dag_run_id,
                "retry_count": 0,
                "last_retry_time": None,
                "alerted": False,
                "first_seen": self._now_iso(),
                "last_seen": self._now_iso(),
            }
            logger.debug("State: new entry created for %s", key)
        else:
            self._data["entries"][key]["last_seen"] = self._now_iso()
            logger.debug("State: updated last_seen for %s", key)
        return self._data["entries"][key]

    def get_retry_count(self, dag_id: str, dag_run_id: str) -> int:
        """Get current retry count for a DAG run."""
        entry = self.get_entry(dag_id, dag_run_id)
        count = entry["retry_count"] if entry else 0
        logger.debug("State: retry_count for %s/%s = %d", dag_id, dag_run_id, count)
        return count

    def increment_retry(self, dag_id: str, dag_run_id: str):
        """Increment retry counter and record time."""
        entry = self.ensure_entry(dag_id, dag_run_id)
        old_count = entry["retry_count"]
        entry["retry_count"] += 1
        entry["last_retry_time"] = self._now_iso()
        logger.debug(
            "State: retry_count for %s/%s incremented %d → %d",
            dag_id, dag_run_id, old_count, entry["retry_count"],
        )

    def mark_alerted(self, dag_id: str, dag_run_id: str):
        """Mark a DAG run as alerted (avoid duplicate alerts)."""
        entry = self.ensure_entry(dag_id, dag_run_id)
        entry["alerted"] = True
        logger.debug("State: marked alerted for %s/%s", dag_id, dag_run_id)

    def is_alerted(self, dag_id: str, dag_run_id: str) -> bool:
        """Check if alert was already sent for this DAG run."""
        entry = self.get_entry(dag_id, dag_run_id)
        return entry["alerted"] if entry else False

    def purge_old_entries(self):
        """Remove entries older than max_age seconds."""
        cutoff = time.time() - self._max_age
        keys_to_remove = []

        for key, entry in self._data["entries"].items():
            last_seen = entry.get("last_seen", entry.get("first_seen", ""))
            try:
                ts = datetime.fromisoformat(last_seen).timestamp()
                if ts < cutoff:
                    keys_to_remove.append(key)
            except (ValueError, TypeError):
                keys_to_remove.append(key)

        for key in keys_to_remove:
            logger.debug("State: purging old entry %s", key)
            del self._data["entries"][key]

        if keys_to_remove:
            logger.info("Purged %d old state entries (max_age=%ds)", len(keys_to_remove), self._max_age)
        else:
            logger.debug("State: no entries to purge")