diff --git a/.env.example b/.env.example index 8b3a359..35ef616 100644 --- a/.env.example +++ b/.env.example @@ -1,3 +1,5 @@ AIRFLOW_BASE_URL=http://localhost:8080 AIRFLOW_USERNAME=your_username AIRFLOW_PASSWORD=your_password +AIRFLOW_API_VERSION=v2 +AIRFLOW_AUTH_MODE=token diff --git a/README.md b/README.md index c14cd54..dca3c2c 100644 --- a/README.md +++ b/README.md @@ -121,8 +121,15 @@ Configure: AIRFLOW_BASE_URL=http://localhost:8080 AIRFLOW_USERNAME=your_username AIRFLOW_PASSWORD=your_password +AIRFLOW_API_VERSION=v2 +AIRFLOW_AUTH_MODE=token ``` +Use `AIRFLOW_API_VERSION=v1` with `AIRFLOW_AUTH_MODE=basic` for Airflow 2.x +Stable REST API deployments. Airflow 3.x uses the `v2` API and typically uses +token authentication. Authentication still depends on the API auth backend +configured in the Airflow deployment. + Load the environment variables: ```bash @@ -288,6 +295,7 @@ The current implementation should be considered experimental and is not yet inte Planned areas include: - DAG-level analysis models +- richer CLI reporting - broader Airflow compatibility testing ## License diff --git a/src/flowsense/config.py b/src/flowsense/config.py index ea0c32e..c3d6947 100644 --- a/src/flowsense/config.py +++ b/src/flowsense/config.py @@ -9,6 +9,8 @@ class AirflowConfig: base_url: str username: str password: str + api_version: str = "v2" + auth_mode: str = "token" def get_airflow_config() -> AirflowConfig: @@ -19,6 +21,8 @@ def get_airflow_config() -> AirflowConfig: username = os.getenv("AIRFLOW_USERNAME") password = os.getenv("AIRFLOW_PASSWORD") + api_version = os.getenv("AIRFLOW_API_VERSION", "v2") + auth_mode = os.getenv("AIRFLOW_AUTH_MODE", "token") if not username: raise RuntimeError("AIRFLOW_USERNAME environment variable is required.") @@ -30,4 +34,6 @@ def get_airflow_config() -> AirflowConfig: base_url=base_url, username=username, password=password, + api_version=api_version, + auth_mode=auth_mode, ) diff --git a/src/flowsense/infrastructure/airflow/client.py b/src/flowsense/infrastructure/airflow/client.py index 18a3948..e76d81f 100644 --- a/src/flowsense/infrastructure/airflow/client.py +++ b/src/flowsense/infrastructure/airflow/client.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Any, Self +from typing import Any, Literal, Self import httpx @@ -18,6 +18,8 @@ ) PAGE_SIZE = 100 +AirflowApiVersion = Literal["v1", "v2"] +AirflowAuthMode = Literal["basic", "token"] class AirflowClient: @@ -26,6 +28,8 @@ def __init__( base_url: str | None = None, username: str | None = None, password: str | None = None, + api_version: AirflowApiVersion | None = None, + auth_mode: AirflowAuthMode | None = None, http_client: httpx.Client | None = None, ): config = get_airflow_config() @@ -33,6 +37,15 @@ def __init__( self.base_url = (base_url or config.base_url).rstrip("/") self.username = username or config.username self.password = password or config.password + self.api_version = api_version or config.api_version + self.auth_mode = auth_mode or config.auth_mode + + if self.api_version not in {"v1", "v2"}: + raise ValueError("api_version must be 'v1' or 'v2'") + + if self.auth_mode not in {"basic", "token"}: + raise ValueError("auth_mode must be 'basic' or 'token'") + self._owns_http_client = http_client is None self._http_client = http_client or httpx.Client(timeout=10.0) self._token: str | None = None @@ -94,10 +107,21 @@ def _get_token(self) -> str: return token def _headers(self) -> dict[str, str]: - return { - "Authorization": f"Bearer {self._get_token()}", - "Accept": "application/json", - } + headers = {"Accept": "application/json"} + + if self.auth_mode == "token": + headers["Authorization"] = f"Bearer {self._get_token()}" + + return headers + + def _authentication(self) -> httpx.BasicAuth | None: + if self.auth_mode == "basic": + return httpx.BasicAuth(self.username, self.password) + + return None + + def _api_url(self, path: str) -> str: + return f"{self.base_url}/api/{self.api_version}{path}" def _get_paginated( self, @@ -109,11 +133,18 @@ def _get_paginated( last_page: dict = {} while True: + request_kwargs: dict[str, object] = { + "headers": self._headers(), + "params": {"limit": PAGE_SIZE, "offset": offset}, + } + authentication = self._authentication() + if authentication is not None: + request_kwargs["auth"] = authentication + response = self._request( method="GET", url=url, - headers=self._headers(), - params={"limit": PAGE_SIZE, "offset": offset}, + **request_kwargs, ) last_page = response.json() @@ -140,7 +171,7 @@ def _get_paginated( def get_dag_runs(self, dag_id: str) -> dict: return self._get_paginated( - url=f"{self.base_url}/api/v2/dags/{dag_id}/dagRuns", + url=self._api_url(f"/dags/{dag_id}/dagRuns"), collection_key="dag_runs", ) @@ -150,10 +181,7 @@ def get_task_instances( dag_run_id: str, ) -> dict: return self._get_paginated( - url=( - f"{self.base_url}/api/v2/dags/{dag_id}/dagRuns/" - f"{dag_run_id}/taskInstances" - ), + url=self._api_url(f"/dags/{dag_id}/dagRuns/{dag_run_id}/taskInstances"), collection_key="task_instances", ) @@ -164,7 +192,7 @@ def collect_task_runs(self, dag_id: str) -> list[TaskRun]: ] def run_timestamp(run: AirflowDagRunDTO) -> float: - timestamp = run.run_after or run.queued_at + timestamp = run.run_after or run.logical_date or run.queued_at return timestamp.timestamp() if timestamp is not None else float("-inf") dag_runs.sort(key=run_timestamp) @@ -198,7 +226,7 @@ def run_timestamp(run: AirflowDagRunDTO) -> float: def get_dag_tasks(self, dag_id: str) -> dict: return self._get_paginated( - url=f"{self.base_url}/api/v2/dags/{dag_id}/tasks", + url=self._api_url(f"/dags/{dag_id}/tasks"), collection_key="tasks", ) diff --git a/src/flowsense/infrastructure/airflow/dto.py b/src/flowsense/infrastructure/airflow/dto.py index 043b4de..bc5395d 100644 --- a/src/flowsense/infrastructure/airflow/dto.py +++ b/src/flowsense/infrastructure/airflow/dto.py @@ -1,6 +1,6 @@ from datetime import datetime -from pydantic import BaseModel, ConfigDict, Field +from pydantic import AliasChoices, BaseModel, ConfigDict, Field class AirflowDTO(BaseModel): @@ -10,6 +10,10 @@ class AirflowDTO(BaseModel): class AirflowDagRunDTO(AirflowDTO): dag_run_id: str state: str | None = None + logical_date: datetime | None = Field( + default=None, + validation_alias=AliasChoices("logical_date", "execution_date"), + ) run_after: datetime | None = None queued_at: datetime | None = None diff --git a/tests/test_airflow_compatibility.py b/tests/test_airflow_compatibility.py new file mode 100644 index 0000000..89ab41a --- /dev/null +++ b/tests/test_airflow_compatibility.py @@ -0,0 +1,203 @@ +from unittest.mock import MagicMock, patch + +import httpx +import pytest + +from flowsense.config import AirflowConfig, get_airflow_config +from flowsense.infrastructure.airflow.client import AirflowClient +from flowsense.infrastructure.airflow.dto import ( + AirflowDagRunDTO, + AirflowTaskDTO, + AirflowTaskInstanceDTO, +) + + +def _client( + http_client: MagicMock, + *, + api_version: str, + auth_mode: str, +) -> AirflowClient: + with patch( + "flowsense.infrastructure.airflow.client.get_airflow_config", + return_value=AirflowConfig( + base_url="http://airflow.test", + username="airflow", + password="airflow", + api_version=api_version, + auth_mode=auth_mode, + ), + ): + return AirflowClient(http_client=http_client) + + +def test_reads_api_compatibility_settings_from_environment( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("AIRFLOW_BASE_URL", "http://airflow.test") + monkeypatch.setenv("AIRFLOW_USERNAME", "airflow") + monkeypatch.setenv("AIRFLOW_PASSWORD", "airflow") + monkeypatch.setenv("AIRFLOW_API_VERSION", "v1") + monkeypatch.setenv("AIRFLOW_AUTH_MODE", "basic") + + config = get_airflow_config() + + assert config.api_version == "v1" + assert config.auth_mode == "basic" + + +@pytest.mark.parametrize("api_version", ["v1", "v2"]) +def test_uses_configured_stable_api_version(api_version: str) -> None: + http_client = MagicMock(spec=httpx.Client) + response = MagicMock() + response.json.return_value = {"dag_runs": [], "total_entries": 0} + http_client.request.return_value = response + client = _client(http_client, api_version=api_version, auth_mode="token") + client._token = "token" + + client.get_dag_runs("demo") + + assert http_client.request.call_args.kwargs["url"] == ( + f"http://airflow.test/api/{api_version}/dags/demo/dagRuns" + ) + + +def test_uses_basic_auth_for_airflow_2_api() -> None: + http_client = MagicMock(spec=httpx.Client) + response = MagicMock() + response.json.return_value = {"dag_runs": [], "total_entries": 0} + http_client.request.return_value = response + client = _client(http_client, api_version="v1", auth_mode="basic") + + client.get_dag_runs("demo") + + request_kwargs = http_client.request.call_args.kwargs + assert request_kwargs["headers"] == {"Accept": "application/json"} + assert isinstance(request_kwargs["auth"], httpx.BasicAuth) + + +def test_uses_bearer_token_for_airflow_3_api() -> None: + http_client = MagicMock(spec=httpx.Client) + response = MagicMock() + response.json.return_value = {"dag_runs": [], "total_entries": 0} + http_client.request.return_value = response + client = _client(http_client, api_version="v2", auth_mode="token") + client._token = "token" + + client.get_dag_runs("demo") + + request_kwargs = http_client.request.call_args.kwargs + assert request_kwargs["headers"] == { + "Accept": "application/json", + "Authorization": "Bearer token", + } + assert "auth" not in request_kwargs + + +@pytest.mark.parametrize( + ("payload", "expected_timestamp"), + [ + ( + { + "dag_run_id": "airflow_2", + "state": "success", + "execution_date": "2026-01-01T10:00:00Z", + }, + "2026-01-01T10:00:00+00:00", + ), + ( + { + "dag_run_id": "airflow_2_modern", + "state": "success", + "logical_date": "2026-01-02T10:00:00Z", + }, + "2026-01-02T10:00:00+00:00", + ), + ( + { + "dag_run_id": "airflow_3", + "state": "success", + "run_after": "2026-01-03T10:00:00Z", + }, + "2026-01-03T10:00:00+00:00", + ), + ], +) +def test_accepts_airflow_dag_run_timestamp_variants( + payload: dict[str, object], + expected_timestamp: str, +) -> None: + dag_run = AirflowDagRunDTO.model_validate(payload) + timestamp = dag_run.run_after or dag_run.logical_date or dag_run.queued_at + + assert timestamp is not None + assert timestamp.isoformat() == expected_timestamp + + +@pytest.mark.parametrize( + "payload", + [ + { + "task_id": "transform", + "state": "success", + "start_date": "2026-01-01T10:00:00Z", + "end_date": "2026-01-01T10:00:03Z", + "duration": 3.0, + "try_number": 1, + }, + { + "id": "task-instance-id", + "task_id": "transform", + "state": "success", + "start_date": "2026-01-01T10:00:00Z", + "end_date": "2026-01-01T10:00:03Z", + "duration": 3.0, + "try_number": 1, + "map_index": 2, + "task_display_name": "Transform data", + "dag_version": {"version_number": 4}, + }, + ], +) +def test_accepts_airflow_2_and_3_task_instance_payloads( + payload: dict[str, object], +) -> None: + task = AirflowTaskInstanceDTO.model_validate(payload) + + assert task.task_id == "transform" + assert task.duration == 3.0 + + +def test_accepts_extra_fields_in_airflow_task_payload() -> None: + task = AirflowTaskDTO.model_validate( + { + "task_id": "extract", + "downstream_task_ids": ["transform"], + "operator_name": "PythonOperator", + "is_mapped": False, + } + ) + + assert task.downstream_task_ids == ["transform"] + + +@pytest.mark.parametrize( + ("api_version", "auth_mode", "message"), + [ + ("v3", "token", "api_version"), + ("v2", "oauth", "auth_mode"), + ], +) +def test_rejects_unsupported_client_configuration( + api_version: str, + auth_mode: str, + message: str, +) -> None: + http_client = MagicMock(spec=httpx.Client) + + with pytest.raises(ValueError, match=message): + _client( + http_client, + api_version=api_version, + auth_mode=auth_mode, + )