cc4dfc0692
- only_failed появился в Airflow 2.10+, на 2.9.1 возвращает 400 - limit=1 — берём только последний run для каждого DAG
292 lines
11 KiB
Python
292 lines
11 KiB
Python
"""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 "<empty>",
|
|
)
|
|
|
|
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
|