"""Airflow REST API client with v1/experimental auto-detection.""" import json import logging import time from urllib.parse import quote import requests from .config import AirflowConfig logger = logging.getLogger(__name__) class AirflowAPIError(Exception): """Raised when Airflow API returns an unexpected error.""" class AirflowClient: """Client for Apache Airflow REST API. Supports both stable API v1 (/api/v1/) and experimental API (/api/experimental/). Auto-detects which is available at startup. """ def __init__(self, config: AirflowConfig): self._config = config self._session = requests.Session() self._session.auth = (config.username, config.password) self._session.verify = config.verify_ssl self._session.headers.update({"Content-Type": "application/json"}) self._base_url = config.base_url.rstrip("/") self._api_prefix: str | None = None self._api_version: str | None = None self._last_request_time: float = 0 logger.debug( "AirflowClient initialized: url=%s, user=%s, timeout=%d, ssl=%s", self._base_url, config.username, config.timeout, config.verify_ssl, ) def detect_api_version(self): """Detect available API version. Tries v1 stable API first, falls back to experimental. Raises AirflowAPIError if neither is available. """ if self._config.api_version != "auto": if self._config.api_version == "v1": self._api_prefix = f"{self._base_url}/api/v1" self._api_version = "v1" else: self._api_prefix = f"{self._base_url}/api/experimental" self._api_version = "experimental" logger.info("Using configured API version: %s", self._api_version) return # Try v1 stable API url = f"{self._base_url}/api/v1/health" logger.debug("Probing API v1: GET %s", url) try: resp = self._session.get(url, timeout=self._config.timeout) logger.debug("Probe v1 response: status=%d, body=%s", resp.status_code, resp.text[:200]) if resp.status_code == 200: self._api_prefix = f"{self._base_url}/api/v1" self._api_version = "v1" logger.info("Detected Airflow API v1 (stable)") return except requests.RequestException as e: logger.debug("Probe v1 failed: %s", e) # Try experimental API url = f"{self._base_url}/api/experimental/test" logger.debug("Probing experimental API: GET %s", url) try: resp = self._session.get(url, timeout=self._config.timeout) logger.debug("Probe experimental response: status=%d", resp.status_code) if resp.status_code == 200: self._api_prefix = f"{self._base_url}/api/experimental" self._api_version = "experimental" logger.info("Detected Airflow experimental API") return except requests.RequestException as e: logger.debug("Probe experimental failed: %s", e) raise AirflowAPIError( f"Cannot connect to Airflow API at {self._base_url}. " "Tried /api/v1/health and /api/experimental/test" ) def _rate_limit(self): """Enforce minimum delay between API requests.""" elapsed = time.time() - self._last_request_time if elapsed < self._config.request_delay: wait = self._config.request_delay - elapsed logger.debug("Rate limiting: waiting %.2fs", wait) time.sleep(wait) def _request(self, method: str, path: str, **kwargs) -> requests.Response: """Make an API request with error handling and rate limiting.""" if not self._api_prefix: raise AirflowAPIError("API version not detected. Call detect_api_version() first") self._rate_limit() url = f"{self._api_prefix}{path}" # Log request details params = kwargs.get("params") body = kwargs.get("json") logger.debug( "API request: %s %s params=%s body=%s", method, url, params, json.dumps(body) if body else None, ) t0 = time.monotonic() try: resp = self._session.request( method, url, timeout=self._config.timeout, **kwargs ) self._last_request_time = time.time() elapsed_ms = (time.monotonic() - t0) * 1000 # Log response details logger.debug( "API response: %s %s → %d (%dms) body=%s", method, path, resp.status_code, elapsed_ms, resp.text[:500] if resp.text else "", ) resp.raise_for_status() return resp except requests.ConnectionError as e: raise AirflowAPIError(f"Connection error to {url}: {e}") from e except requests.Timeout as e: elapsed_ms = (time.monotonic() - t0) * 1000 raise AirflowAPIError(f"Timeout after {elapsed_ms:.0f}ms calling {url}: {e}") from e except requests.HTTPError as e: status = resp.status_code body = resp.text[:500] raise AirflowAPIError(f"HTTP {status} from {url}: {body}") from e def get_enabled_dags(self) -> list[dict]: """Get list of enabled (active and not paused) DAGs. Returns list of dicts with at least 'dag_id' key. """ logger.debug("Fetching enabled DAGs (api_version=%s)", self._api_version) if self._api_version == "v1": return self._get_enabled_dags_v1() return self._get_enabled_dags_experimental() def _get_enabled_dags_v1(self) -> list[dict]: """Fetch DAGs via stable v1 API with pagination.""" dags = [] offset = 0 limit = 100 while True: logger.debug("Fetching DAGs page: offset=%d, limit=%d", offset, limit) resp = self._request( "GET", "/dags", params={"limit": limit, "offset": offset, "only_active": True}, ) data = resp.json() page_dags = data.get("dags", []) total = data.get("total_entries", 0) active_count = 0 paused_count = 0 for dag in page_dags: if not dag.get("is_paused", True): dags.append(dag) active_count += 1 else: paused_count += 1 logger.debug( "DAGs page result: total_entries=%d, page_size=%d, " "active=%d, paused=%d", total, len(page_dags), active_count, paused_count, ) offset += limit if offset >= total or not page_dags: break logger.debug( "All enabled DAGs: %s", [d["dag_id"] for d in dags], ) return dags def _get_enabled_dags_experimental(self) -> list[dict]: """Fetch DAGs via experimental API (limited info).""" resp = self._request("GET", "/dags") return resp.json() if resp.status_code == 200 else [] def get_dag_runs(self, dag_id: str, states: list[str] | None = None) -> list[dict]: """Get recent DAG runs, optionally filtered by state.""" encoded_dag_id = quote(dag_id, safe="") logger.debug("Fetching DAG runs: dag_id=%s, states=%s", dag_id, states) if self._api_version == "v1": if states: resp = self._request( "GET", f"/dags/{encoded_dag_id}/dagRuns", params=[("limit", 1), ("order_by", "-start_date")] + [("state", s) for s in states], ) else: resp = self._request( "GET", f"/dags/{encoded_dag_id}/dagRuns", params={"order_by": "-start_date", "limit": 1}, ) runs = resp.json().get("dag_runs", []) else: resp = self._request("GET", f"/dags/{dag_id}/dag_runs") runs = resp.json() if isinstance(resp.json(), list) else [] if states: runs = [r for r in runs if r.get("state") in states] logger.debug( "DAG runs for %s: count=%d, runs=[%s]", dag_id, len(runs), ", ".join( f"{r.get('dag_run_id', '?')}({r.get('state', '?')})" for r in runs ), ) return runs def get_task_instances(self, dag_id: str, dag_run_id: str) -> list[dict]: """Get task instances for a specific DAG run.""" encoded_dag_id = quote(dag_id, safe="") encoded_run_id = quote(dag_run_id, safe="") logger.debug("Fetching task instances: %s/%s", dag_id, dag_run_id) if self._api_version == "v1": resp = self._request( "GET", f"/dags/{encoded_dag_id}/dagRuns/{encoded_run_id}/taskInstances", ) tasks = resp.json().get("task_instances", []) logger.debug( "Task instances for %s/%s: count=%d, states=[%s]", dag_id, dag_run_id, len(tasks), ", ".join( f"{t.get('task_id', '?')}({t.get('state', '?')})" for t in tasks ), ) return tasks else: logger.warning( "Task instances not supported in experimental API for %s/%s", dag_id, dag_run_id, ) return [] def clear_dag_run(self, dag_id: str, dag_run_id: str) -> bool: """Clear (retry) a DAG run by resetting failed task instances. Returns True on success, False on failure. """ encoded_dag_id = quote(dag_id, safe="") encoded_run_id = quote(dag_run_id, safe="") if self._api_version != "v1": logger.warning( "Clearing DAG runs not supported in experimental API for %s/%s", dag_id, dag_run_id, ) return False logger.debug( "Clearing DAG run: %s/%s (only_failed=True)", dag_id, dag_run_id, ) try: self._request( "POST", f"/dags/{encoded_dag_id}/dagRuns/{encoded_run_id}/clear", json={"dry_run": False}, ) logger.info("Cleared DAG run %s/%s for retry", dag_id, dag_run_id) return True except AirflowAPIError as e: logger.error("Failed to clear DAG run %s/%s: %s", dag_id, dag_run_id, e) return False