AirflowDAGMonitor/airflow_monitor/analyzer.py

"""DAG run analyzer - classifies runs and detects issues."""

import dataclasses
import logging
from datetime import datetime, timezone
from typing import Callable

logger = logging.getLogger(__name__)


@dataclasses.dataclass
class DagIssue:
    """Represents a detected problem with a DAG run."""

    dag_id: str
    dag_run_id: str
    issue_type: str  # "failed" or "long_running"
    state: str  # airflow state string
    start_time: str  # ISO 8601 UTC
    duration_seconds: float
    error_info: str = ""

    def to_dict(self) -> dict:
        return dataclasses.asdict(self)


class DagAnalyzer:
    """Analyzes DAG runs and produces a list of issues."""

    def __init__(self, long_running_threshold: int, max_failed_age: int = 86400):
        self._threshold = long_running_threshold
        self._max_failed_age = max_failed_age  # ignore failed runs older than this (seconds)
        logger.debug(
            "DagAnalyzer initialized: long_running_threshold=%ds (%.1f min), max_failed_age=%ds (%.1f hours)",
            long_running_threshold, long_running_threshold / 60,
            max_failed_age, max_failed_age / 3600,
        )

    def analyze_dag_runs(
        self,
        dag_id: str,
        runs: list[dict],
        task_instances_fn: Callable[[str, str], list[dict]],
    ) -> list[DagIssue]:
        """Analyze DAG runs for a single DAG.

        Args:
            dag_id: DAG identifier.
            runs: List of DAG run dicts from Airflow API.
            task_instances_fn: Callable(dag_id, dag_run_id) -> list of task instances.
                Called lazily only for failed runs to extract error details.

        Returns:
            List of DagIssue objects for problematic runs.
        """
        issues = []
        now = datetime.now(timezone.utc)
        logger.debug("Analyzing %d runs for DAG '%s' (now=%s)", len(runs), dag_id, now.isoformat())

        for run in runs:
            run_id = run.get("dag_run_id", "")
            state = run.get("state", "")
            start_date_str = run.get("start_date") or run.get("execution_date", "")

            if not start_date_str:
                logger.warning("DAG run %s/%s has no start_date, skipping", dag_id, run_id)
                continue

            start_date = self._parse_datetime(start_date_str)
            if start_date is None:
                logger.warning(
                    "Cannot parse start_date '%s' for %s/%s",
                    start_date_str, dag_id, run_id,
                )
                continue

            duration = (now - start_date).total_seconds()

            logger.debug(
                "  Run %s/%s: state=%s, start=%s, duration=%.0fs (%.1f min), threshold=%ds",
                dag_id, run_id, state, start_date_str,
                duration, duration / 60, self._threshold,
            )

            if state == "failed":
                # Skip old failed runs — only alert on recent failures
                if duration > self._max_failed_age:
                    logger.debug(
                        "  → SKIP: failed run too old (%.0fs > %ds max_failed_age)",
                        duration, self._max_failed_age,
                    )
                    continue

                error_info = self._extract_error(dag_id, run_id, task_instances_fn)
                issue = DagIssue(
                    dag_id=dag_id,
                    dag_run_id=run_id,
                    issue_type="failed",
                    state=state,
                    start_time=start_date.isoformat(),
                    duration_seconds=round(duration, 1),
                    error_info=error_info,
                )
                issues.append(issue)
                logger.debug("  → ISSUE DETECTED: %s (error: %s)", issue.issue_type, error_info or "n/a")

            elif state == "running" and duration > self._threshold:
                issue = DagIssue(
                    dag_id=dag_id,
                    dag_run_id=run_id,
                    issue_type="long_running",
                    state=state,
                    start_time=start_date.isoformat(),
                    duration_seconds=round(duration, 1),
                )
                issues.append(issue)
                logger.debug(
                    "  → ISSUE DETECTED: long_running (%.0fs > %ds threshold)",
                    duration, self._threshold,
                )

            elif state == "running":
                logger.debug("  → OK: running within threshold (%.0fs <= %ds)", duration, self._threshold)

            else:
                logger.debug("  → SKIP: state=%s (not actionable)", state)

        logger.debug("Analysis complete for DAG '%s': %d issues found", dag_id, len(issues))
        return issues

    def _extract_error(
        self,
        dag_id: str,
        dag_run_id: str,
        task_instances_fn: Callable[[str, str], list[dict]],
    ) -> str:
        """Extract error info from failed task instances."""
        logger.debug("Extracting error info for %s/%s", dag_id, dag_run_id)
        try:
            tasks = task_instances_fn(dag_id, dag_run_id)
        except Exception as e:
            logger.warning("Failed to fetch task instances for %s/%s: %s", dag_id, dag_run_id, e)
            return ""

        failed_tasks = [t for t in tasks if t.get("state") == "failed"]
        logger.debug(
            "Task instances for %s/%s: total=%d, failed=%d",
            dag_id, dag_run_id, len(tasks), len(failed_tasks),
        )

        if not failed_tasks:
            return ""

        # Return info about the first failed task
        task = failed_tasks[0]
        task_id = task.get("task_id", "unknown")
        # Try to get the error message from different possible fields
        error = (
            task.get("rendered_fields", {}).get("error", "")
            or task.get("note", "")
            or ""
        )
        # Truncate long error messages
        if len(error) > 500:
            error = error[:500] + "..."

        result = f"Task '{task_id}' failed" + (f": {error}" if error else "")
        logger.debug("Error extracted for %s/%s: %s", dag_id, dag_run_id, result)
        return result

    @staticmethod
    def _parse_datetime(dt_string: str) -> datetime | None:
        """Parse ISO 8601 datetime string from Airflow API.

        Handles both 'Z' suffix and '+00:00' timezone offset.
        """
        if not dt_string:
            return None

        # Normalize 'Z' to '+00:00' for fromisoformat
        cleaned = dt_string.replace("Z", "+00:00")

        try:
            dt = datetime.fromisoformat(cleaned)
            # Ensure timezone-aware (UTC)
            if dt.tzinfo is None:
                dt = dt.replace(tzinfo=timezone.utc)
            return dt
        except ValueError:
            return None