"""Shared Apache Airflow integration helpers.""" from __future__ import annotations import logging import os from dataclasses import dataclass from typing import Any from urllib.parse import quote import httpx from pydantic import Field, field_validator from config.strict_config import StrictConfigModel from integrations._validation_helpers import report_classify_failure, report_validation_failure logger = logging.getLogger(__name__) DEFAULT_AIRFLOW_BASE_URL = "http://localhost:8080/api/v1" DEFAULT_AIRFLOW_TIMEOUT_SECONDS = 15.0 DEFAULT_AIRFLOW_MAX_RESULTS = 50 class AirflowConfig(StrictConfigModel): """Normalized Airflow connection settings.""" base_url: str = DEFAULT_AIRFLOW_BASE_URL username: str = "" password: str = "" auth_token: str = "" timeout_seconds: float = Field(default=DEFAULT_AIRFLOW_TIMEOUT_SECONDS, gt=0) verify_ssl: bool = True max_results: int = Field(default=DEFAULT_AIRFLOW_MAX_RESULTS, gt=0, le=200) integration_id: str = "" @field_validator("base_url", mode="before") @classmethod def _normalize_base_url(cls, value: Any) -> str: normalized = str(value or DEFAULT_AIRFLOW_BASE_URL).strip().rstrip("/") return normalized or DEFAULT_AIRFLOW_BASE_URL @field_validator("username", "password", "auth_token", mode="before") @classmethod def _normalize_str(cls, value: Any) -> str: return str(value or "").strip() @property def headers(self) -> dict[str, str]: headers = {"Accept": "application/json"} if self.auth_token: headers["Authorization"] = f"Bearer {self.auth_token}" return headers @property def auth(self) -> tuple[str, str] | None: if self.username and not self.auth_token: return (self.username, self.password) return None @property def is_configured(self) -> bool: return bool(self.base_url and (self.auth_token or self.username)) @dataclass(frozen=True) class AirflowValidationResult: """Result of validating an Airflow integration.""" ok: bool detail: str def build_airflow_config(raw: dict[str, Any] | None) -> AirflowConfig: """Build a normalized Airflow config object from env/store data.""" return AirflowConfig.model_validate(raw or {}) def airflow_config_from_env() -> AirflowConfig | None: """Load an Airflow config from env vars.""" username = os.getenv("AIRFLOW_USERNAME", "").strip() auth_token = os.getenv("AIRFLOW_AUTH_TOKEN", "").strip() if not username and not auth_token: return None return build_airflow_config( { "base_url": os.getenv("AIRFLOW_BASE_URL", DEFAULT_AIRFLOW_BASE_URL).strip() or DEFAULT_AIRFLOW_BASE_URL, "username": username, "password": os.getenv("AIRFLOW_PASSWORD", "").strip(), "auth_token": auth_token, "timeout_seconds": os.getenv( "AIRFLOW_TIMEOUT_SECONDS", str(DEFAULT_AIRFLOW_TIMEOUT_SECONDS), ), "verify_ssl": os.getenv("AIRFLOW_VERIFY_SSL", "true").strip().lower() in ("true", "1", "yes"), "max_results": os.getenv( "AIRFLOW_MAX_RESULTS", str(DEFAULT_AIRFLOW_MAX_RESULTS) ).strip(), } ) def _request_json( config: AirflowConfig, method: str, path: str, *, params: list[tuple[str, str | int | float | bool | None]] | None = None, json: dict[str, Any] | None = None, ) -> Any: """Make an Airflow API request and return parsed JSON.""" url = f"{config.base_url}{path}" response = httpx.request( method, url, headers=config.headers, auth=config.auth, params=params, json=json, timeout=config.timeout_seconds, verify=config.verify_ssl, ) response.raise_for_status() return response.json() def validate_airflow_config(config: AirflowConfig) -> AirflowValidationResult: """Validate Airflow connectivity with a lightweight DAG query.""" if not config.is_configured: return AirflowValidationResult( ok=False, detail="Airflow auth is required. Provide AIRFLOW_AUTH_TOKEN or AIRFLOW_USERNAME/AIRFLOW_PASSWORD.", ) try: payload = validate_airflow_connection(config=config) dags = payload.get("dags", []) if isinstance(payload, dict) else [] total_entries = ( payload.get("total_entries", len(dags)) if isinstance(payload, dict) else len(dags) ) return AirflowValidationResult( ok=True, detail=f"Airflow connectivity successful. Reachable DAG API; total visible DAGs: {total_entries}.", ) except httpx.HTTPStatusError as err: detail = err.response.text.strip() or str(err) return AirflowValidationResult(ok=False, detail=f"Airflow validation failed: {detail}") except Exception as err: report_validation_failure( err, logger=logger, integration="airflow", method="validate_airflow_config", ) return AirflowValidationResult(ok=False, detail=f"Airflow validation failed: {err}") def validate_airflow_connection( *, config: AirflowConfig, ) -> dict[str, Any]: """Validate Airflow connection.""" payload = _request_json( config, "GET", "/dags", params=[("limit", 1)], ) return payload if isinstance(payload, dict) else {} def get_airflow_dag_runs( *, config: AirflowConfig, dag_id: str, limit: int = 10, state: str | None = None, order_by: str = "-start_date", ) -> list[dict[str, Any]]: """Fetch DAG runs for a given DAG.""" effective_limit = min(limit, config.max_results) encoded_dag_id = quote(dag_id, safe="") params: list[tuple[str, str | int | float | bool | None]] = [ ("limit", effective_limit), ("order_by", order_by), ] if state: params.append(("state", state)) payload = _request_json( config, "GET", f"/dags/{encoded_dag_id}/dagRuns", params=params, ) if not isinstance(payload, dict): return [] dag_runs = payload.get("dag_runs", []) return dag_runs if isinstance(dag_runs, list) else [] def get_airflow_task_instances( *, config: AirflowConfig, dag_id: str, dag_run_id: str, ) -> list[dict[str, Any]]: """Fetch task instances for a given DAG run.""" encoded_dag_id = quote(dag_id, safe="") encoded_dag_run_id = quote(dag_run_id, safe="") payload = _request_json( config, "GET", f"/dags/{encoded_dag_id}/dagRuns/{encoded_dag_run_id}/taskInstances", ) if not isinstance(payload, dict): return [] task_instances = payload.get("task_instances", []) return task_instances if isinstance(task_instances, list) else [] def _to_failure_evidence( *, dag_id: str, dag_run: dict[str, Any], task_instance: dict[str, Any], ) -> dict[str, Any]: """Normalize a failed or retrying task instance into investigation-friendly evidence.""" start_date = task_instance.get("start_date") or dag_run.get("start_date") end_date = task_instance.get("end_date") or dag_run.get("end_date") state = task_instance.get("state", "") try_number = task_instance.get("try_number") max_tries = task_instance.get("max_tries") duration = task_instance.get("duration") return { "source": "airflow", "dag_id": dag_id, "dag_run_id": dag_run.get("dag_run_id", ""), "logical_date": dag_run.get("logical_date", ""), "run_type": dag_run.get("run_type", ""), "dag_run_state": dag_run.get("state", ""), "task_id": task_instance.get("task_id", ""), "task_state": state, "operator": task_instance.get("operator", ""), "try_number": try_number, "max_tries": max_tries, "queued_dttm": task_instance.get("queued_dttm", ""), "start_date": start_date, "end_date": end_date, "duration": duration, "hostname": task_instance.get("hostname", ""), "unixname": task_instance.get("unixname", ""), "pool": task_instance.get("pool", ""), "queue": task_instance.get("queue", ""), "priority_weight": task_instance.get("priority_weight"), } def get_recent_airflow_failures( *, config: AirflowConfig, dag_id: str, limit: int = 5, ) -> list[dict[str, Any]]: """Fetch recent failed or retrying task evidence for a DAG. Strategy: - fetch recent DAG runs - fetch task instances for each run - return failed/up_for_retry/upstream_failed task evidence """ dag_runs = get_airflow_dag_runs( config=config, dag_id=dag_id, limit=limit, ) evidence: list[dict[str, Any]] = [] interesting_states = {"failed", "up_for_retry", "upstream_failed"} for dag_run in dag_runs: dag_run_id = str(dag_run.get("dag_run_id", "")).strip() if not dag_run_id: continue try: task_instances = get_airflow_task_instances( config=config, dag_id=dag_id, dag_run_id=dag_run_id, ) except Exception as err: report_validation_failure( err, logger=logger, integration="airflow", method="get_recent_airflow_failures.task_instances", extras={"dag_id": dag_id, "dag_run_id": dag_run_id}, ) continue for task_instance in task_instances: state = str(task_instance.get("state", "")).strip().lower() if state not in interesting_states: continue evidence.append( _to_failure_evidence( dag_id=dag_id, dag_run=dag_run, task_instance=task_instance, ) ) return evidence def classify( credentials: dict[str, Any], record_id: str ) -> tuple[AirflowConfig | None, str | None]: try: cfg = build_airflow_config( { "base_url": credentials.get("base_url", DEFAULT_AIRFLOW_BASE_URL), "username": credentials.get("username", ""), "password": credentials.get("password", ""), "auth_token": credentials.get("auth_token", ""), "timeout_seconds": credentials.get("timeout_seconds", 15.0), "verify_ssl": credentials.get("verify_ssl", True), "max_results": credentials.get("max_results", 50), "integration_id": record_id, } ) except Exception as exc: report_classify_failure(exc, logger=logger, integration="airflow", record_id=record_id) return None, None if cfg.is_configured: return cfg, "airflow" return None, None