AirflowDAGMonitor/airflow_monitor/monitor.py

"""Main monitoring loop - orchestrates the full monitoring cycle."""

import logging
import threading
import time

from .actions import ActionHandler, DataExporter
from .analyzer import DagAnalyzer
from .client import AirflowAPIError, AirflowClient
from .config import AppConfig
from .state import StateManager

logger = logging.getLogger(__name__)


class Monitor:
    """Orchestrates the Airflow DAG monitoring cycle.

    Cycle flow:
    1. Load state
    2. Fetch enabled DAGs
    3. Get active runs for each DAG
    4. Analyze for issues (failed, long-running)
    5. Retry where possible
    6. Wait and re-check retried runs
    7. Alert via Zabbix for unresolved issues
    8. Send heartbeat, cleanup, save state
    """

    def __init__(self, config: AppConfig, shutdown_event: threading.Event):
        self._config = config
        self._shutdown = shutdown_event
        self._client = AirflowClient(config.airflow)
        self._state = StateManager(config.monitor.state_file, config.monitor.state_max_age)
        self._analyzer = DagAnalyzer(
            config.monitor.long_running_threshold,
            config.monitor.max_failed_age,
        )
        self._exporter = DataExporter(config.zabbix)
        self._actions = ActionHandler(
            self._client, self._state, self._exporter, config.monitor,
        )
        self._cycle_count = 0

    def run(self):
        """Main loop. Runs until shutdown_event is set."""
        logger.info(
            "Starting Airflow DAG monitor (cycle=%ds, threshold=%ds, retries=%d, retry_wait=%ds)",
            self._config.monitor.cycle_interval,
            self._config.monitor.long_running_threshold,
            self._config.monitor.max_retries,
            self._config.monitor.retry_wait,
        )

        try:
            self._client.detect_api_version()
        except AirflowAPIError as e:
            logger.critical("Cannot connect to Airflow API: %s", e)
            return

        while not self._shutdown.is_set():
            try:
                self._cycle()
            except Exception:
                logger.exception("Unhandled error in monitoring cycle")

            # Interruptible sleep between cycles
            logger.debug(
                "Sleeping %ds until next cycle",
                self._config.monitor.cycle_interval,
            )
            self._shutdown.wait(timeout=self._config.monitor.cycle_interval)

        logger.info("Shutting down gracefully")
        self._state.save()

    def _cycle(self):
        """Execute one complete monitoring cycle."""
        self._cycle_count += 1
        cycle_start = time.monotonic()
        logger.info("=== Monitoring cycle #%d started ===", self._cycle_count)
        self._state.load()

        # --- Phase 1: Collect issues ---
        phase_start = time.monotonic()
        all_issues = []
        try:
            dags = self._client.get_enabled_dags()
        except AirflowAPIError as e:
            logger.error("Failed to fetch DAG list: %s", e)
            self._exporter.export_heartbeat()
            return

        logger.info(
            "Phase 1/5 [Collect]: found %d enabled DAGs (%.1fs)",
            len(dags), time.monotonic() - phase_start,
        )

        for i, dag in enumerate(dags, 1):
            if self._shutdown.is_set():
                return

            dag_id = dag.get("dag_id", "")
            if not dag_id:
                continue

            logger.debug("Processing DAG %d/%d: %s", i, len(dags), dag_id)

            try:
                runs = self._client.get_dag_runs(dag_id, states=["running", "failed"])
            except AirflowAPIError as e:
                logger.warning("Failed to fetch runs for DAG '%s': %s", dag_id, e)
                continue

            if not runs:
                logger.debug("  No active/failed runs for %s", dag_id)
                continue

            logger.debug("  Found %d active/failed runs for %s", len(runs), dag_id)

            issues = self._analyzer.analyze_dag_runs(
                dag_id, runs,
                task_instances_fn=self._client.get_task_instances,
            )
            all_issues.extend(issues)

        if all_issues:
            logger.warning(
                "Phase 1 result: %d issues found: %s",
                len(all_issues),
                ", ".join(
                    f"{i.dag_id}/{i.dag_run_id}({i.issue_type}, {i.duration_seconds:.0f}s)"
                    for i in all_issues
                ),
            )
        else:
            logger.info("Phase 1 result: no issues found across %d DAGs", len(dags))

        # --- Phase 2: Handle issues (retry or mark for alert) ---
        phase_start = time.monotonic()
        needs_recheck = []
        for issue in all_issues:
            action = self._actions.handle_issue(issue, all_issues)
            logger.info(
                "Phase 2/5 [Handle]: DAG %s run %s → %s (type=%s, duration=%.0fs)",
                issue.dag_id, issue.dag_run_id, action,
                issue.issue_type, issue.duration_seconds,
            )
            if action == "retried":
                needs_recheck.append(issue)

        logger.info(
            "Phase 2/5 [Handle]: processed %d issues, %d retried, (%.1fs)",
            len(all_issues), len(needs_recheck),
            time.monotonic() - phase_start,
        )

        # --- Phase 3: Wait and re-check retried runs ---
        if needs_recheck and not self._shutdown.is_set():
            wait_time = self._config.monitor.retry_wait
            logger.info(
                "Phase 3/5 [Wait]: waiting %ds to re-check %d retried runs",
                wait_time, len(needs_recheck),
            )
            self._shutdown.wait(timeout=wait_time)

            if not self._shutdown.is_set():
                phase_start = time.monotonic()
                still_failing = self._recheck_retried(needs_recheck)
                # Replace all_issues with only the still-failing ones for alerting
                # Keep non-retried issues that need alerting
                non_retried_issues = [
                    i for i in all_issues if i not in needs_recheck
                ]
                all_issues = non_retried_issues + still_failing
                logger.info(
                    "Phase 3/5 [Recheck]: %d/%d still failing after retry (%.1fs)",
                    len(still_failing), len(needs_recheck),
                    time.monotonic() - phase_start,
                )
        else:
            logger.debug("Phase 3/5 [Wait]: skipped (no retried runs)")

        # --- Phase 4: Alert via Zabbix ---
        phase_start = time.monotonic()
        self._actions.collect_and_alert(all_issues)
        logger.info(
            "Phase 4/5 [Alert]: alert phase complete (%.1fs)",
            time.monotonic() - phase_start,
        )

        # --- Phase 5: Housekeeping ---
        phase_start = time.monotonic()
        self._exporter.export_heartbeat()
        self._state.purge_old_entries()
        self._state.save()

        cycle_elapsed = time.monotonic() - cycle_start
        self._exporter.export_status(
            self._cycle_count, len(dags), len(all_issues), cycle_elapsed,
        )
        logger.info(
            "=== Monitoring cycle #%d complete: %d DAGs checked, "
            "%d issues, cycle_time=%.1fs ===",
            self._cycle_count, len(dags), len(all_issues), cycle_elapsed,
        )

    def _recheck_retried(self, retried_issues: list) -> list:
        """Re-check DAG runs that were retried.

        Returns list of DagIssue objects that are still failing.
        """
        still_failing = []

        for issue in retried_issues:
            if self._shutdown.is_set():
                break

            logger.debug("Re-checking retried run: %s/%s", issue.dag_id, issue.dag_run_id)

            try:
                runs = self._client.get_dag_runs(
                    issue.dag_id, states=["running", "failed"],
                )
            except AirflowAPIError as e:
                logger.warning(
                    "Failed to re-check DAG '%s': %s. Treating as still failing.",
                    issue.dag_id, e,
                )
                still_failing.append(issue)
                continue

            recheck_issues = self._analyzer.analyze_dag_runs(
                issue.dag_id, runs,
                task_instances_fn=self._client.get_task_instances,
            )

            # Check if the same dag_run_id still has problems
            for ri in recheck_issues:
                if ri.dag_run_id == issue.dag_run_id:
                    logger.warning(
                        "DAG %s/%s still failing after retry: type=%s, duration=%.0fs",
                        ri.dag_id, ri.dag_run_id, ri.issue_type, ri.duration_seconds,
                    )
                    still_failing.append(ri)
                    break
            else:
                logger.info(
                    "DAG %s/%s recovered after retry",
                    issue.dag_id, issue.dag_run_id,
                )

        return still_failing