Skip to content

Fix Dag processor API client not catching Task SDK client errors #74728

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
ephraimbuddy wants to merge 3 commits into from
Closed

Fix Dag processor API client not catching Task SDK client errors #74728

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
9749e84
Fix Dag processor API client not catching Task SDK client errors
ephraimbuddy Oct 14, 2026
cf6c70f
Parse Execution API JSON errors that carry media type parameters
ephraimbuddy Oct 14, 2026
e6532c2
Ban new httpx imports in airflow-core
ephraimbuddy Oct 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 18 additions & 2 deletions airflow-core/pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -103,8 +103,9 @@ dependencies = [
# parsed with no field-count or field-size bound at all.
"starlette>=1.3.1",
"httpx>=0.25.0",
# Only for InProcessExecutionAPI's transports, which feed httpx2-based Task SDK clients.
# The rest of airflow-core is still on httpx; tracked at https://github.com/apache/airflow/issues/70522
# Only for code that talks to the httpx2-based Task SDK client: InProcessExecutionAPI's transports
# and the Dag processor's API client. The rest of airflow-core is still on httpx; tracked at
# https://github.com/apache/airflow/issues/70522
"httpx2>=2.0.0",
'importlib_metadata>=6.5;python_version<"3.12"',
'importlib_metadata>=7.0;python_version>="3.12"',
Expand Down Expand Up @@ -332,6 +333,21 @@ required-version = ">=0.11.8"
apache-airflow-core = {workspace = true}
apache-airflow-devel-common = { workspace = true }

[tool.ruff]
extend = "../pyproject.toml"

# Not yet migrated to httpx2; tracked at https://github.com/apache/airflow/issues/70522
[tool.ruff.lint.extend-per-file-ignores]
"src/airflow/api/client/local_client.py" = ["TID251"]
"src/airflow/api_fastapi/auth/tokens.py" = ["TID251"]
"src/airflow/example_dags/example_dag_decorator.py" = ["TID251"]
"tests/unit/api_fastapi/auth/test_tokens.py" = ["TID251"]
"tests/unit/api_fastapi/core_api/routes/public/test_dag_sources.py" = ["TID251"]
"tests/unit/api_fastapi/execution_api/versions/head/test_hitl.py" = ["TID251"]

[tool.ruff.lint.flake8-tidy-imports.extend-banned-api]
"httpx".msg = "Use httpx2; httpx objects do not interoperate with the httpx2-based Task SDK client. See https://github.com/apache/airflow/issues/70522"

[tool.airflow]
shared_distributions = [
"apache-airflow-shared-configuration",
Expand Down
34 changes: 17 additions & 17 deletions airflow-core/src/airflow/dag_processing/api_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@
from typing import Any
from uuid import UUID

import httpx
import httpx2

@kaxil kaxil Oct 15, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This file and both test modules are already on main through #74731, which merged after this branch was last pushed. git diff between this head and main is empty for all three. After a rebase, what's left is the from_response media-type change, its test, and the ruff ban. Could you retitle around those? As it stands the squash commit would claim the #74731 fix a second time in the changelog, and the description's "fails 86 of 98 tests on main" no longer describes main.

import structlog
from uuid6 import uuid7

Expand Down Expand Up @@ -73,10 +73,10 @@ class DagProcessorJobAlreadyRunning(RuntimeError):
"""The session's previous Job has not completed or stopped heartbeating yet."""


def get_error_reason(error: httpx.HTTPStatusError) -> str | None:
def get_error_reason(error: httpx2.HTTPStatusError) -> str | None:
try:
payload = error.response.json()
except (ValueError, httpx.ResponseNotRead):
except (ValueError, httpx2.ResponseNotRead):
return None
detail = payload.get("detail") if isinstance(payload, dict) else None
return detail.get("reason") if isinstance(detail, dict) else None
Expand Down Expand Up @@ -174,8 +174,8 @@ def _get_parse_token(self, context: DagParseContext, *, retry: bool) -> str:
return context.token
try:
self._exchange_parse_token(context, retry=False, timeout=self._get_bounded_timeout())
except (httpx.HTTPError, ValueError) as error:
if isinstance(error, httpx.HTTPStatusError) and error.response.status_code < 500:
except (httpx2.HTTPError, ValueError) as error:
if isinstance(error, httpx2.HTTPStatusError) and error.response.status_code < 500:
raise
context.renew_at = monotonic() + 30
log.warning(
Expand All @@ -189,7 +189,7 @@ def _get_parse_token(self, context: DagParseContext, *, retry: bool) -> str:
return self._exchange_parse_token(context, retry=retry)

def _exchange_parse_token(
self, context: DagParseContext, *, retry: bool, timeout: httpx.Timeout | None = None
self, context: DagParseContext, *, retry: bool, timeout: httpx2.Timeout | None = None
) -> str:
self._ensure_job_token(retry=retry)
started_at = monotonic()
Expand All @@ -208,7 +208,7 @@ def _exchange_parse_token(
context.renew_at = started_at + lifetime * 0.8
return parsed.token

def _update_auth(self, response: httpx.Response) -> None:
def _update_auth(self, response: httpx2.Response) -> None:
# Task-token refresh headers cannot replace a provisioned or Job-bound credential.
pass

Expand Down Expand Up @@ -241,14 +241,14 @@ def register_job(self, *, retry: bool = True) -> int:
self._check_can_run()
try:
return self._register_job(retry=retry)
except httpx.HTTPStatusError as error:
except httpx2.HTTPStatusError as error:
if error.response.status_code == 409 and get_error_reason(error) == "job_running":
raise DagProcessorJobAlreadyRunning(
"The previous Dag processor Job is still alive"
) from error
raise

def _register_job(self, *, retry: bool, timeout: httpx.Timeout | None = None) -> int:
def _register_job(self, *, retry: bool, timeout: httpx2.Timeout | None = None) -> int:
if self.restart_required:
raise DagProcessorRegistrationRetired(
"The Dag processor registration has ended; restart required"
Expand All @@ -269,7 +269,7 @@ def _register_job(self, *, retry: bool, timeout: httpx.Timeout | None = None) ->
timeout=timeout or self.timeout,
)
break
except httpx.HTTPStatusError as error:
except httpx2.HTTPStatusError as error:
if error.response.status_code == 409 and get_error_reason(error) == "registration_retired":
self._restart_required = True
raise DagProcessorRegistrationRetired(
Expand All @@ -296,8 +296,8 @@ def _register_job(self, *, retry: bool, timeout: httpx.Timeout | None = None) ->
self._retry_renewal_at = 0.0
return registered.job_id

def _get_bounded_timeout(self) -> httpx.Timeout:
return httpx.Timeout(
def _get_bounded_timeout(self) -> httpx2.Timeout:
return httpx2.Timeout(
**{
key: min(value if value is not None else 1.0, 1.0)
for key, value in self.timeout.as_dict().items()
Expand All @@ -315,7 +315,7 @@ def _ensure_job_token(self, *, retry: bool) -> None:
if self._read_session_token() != self._registered_with or monotonic() >= self._renew_at:
# Do not spend the runtime request's retry budget on an optional early renewal.
self._register_job(retry=False, timeout=self._get_bounded_timeout())
except (httpx.HTTPError, OSError, ValueError, DagProcessorRegistrationRetired) as error:
except (httpx2.HTTPError, OSError, ValueError, DagProcessorRegistrationRetired) as error:
self._retry_renewal_at = monotonic() + 30
log.warning(
"Unable to renew Dag processor Job token",
Expand All @@ -326,11 +326,11 @@ def _ensure_job_token(self, *, retry: bool) -> None:
if monotonic() >= self._expires_at:
self._register_job(retry=retry, timeout=None if retry else self._get_bounded_timeout())

def request(self, *args, retry: bool = True, **kwargs) -> httpx.Response:
def request(self, *args, retry: bool = True, **kwargs) -> httpx2.Response:
"""Use a parsing credential for subprocess requests, and the Job credential for manager work."""
self._check_can_run()
try:
headers = httpx.Headers(kwargs.get("headers"))
headers = httpx2.Headers(kwargs.get("headers"))
if context := self._parse_context.get():
kwargs["auth"] = BearerAuth(self._get_parse_token(context, retry=retry))
headers["Airflow-Dag-Bundle"] = context.request.bundle_name
Expand All @@ -342,7 +342,7 @@ def request(self, *args, retry: bool = True, **kwargs) -> httpx.Response:
headers.setdefault("Content-Type", "application/json")
kwargs["headers"] = headers
return super().request(*args, retry=retry, **kwargs)
except httpx.HTTPStatusError as error:
except httpx2.HTTPStatusError as error:
if error.response.status_code == 403 and get_error_reason(error) == "job_closed":
self._restart_required = True
raise DagProcessorRegistrationRetired(
Expand Down Expand Up @@ -377,7 +377,7 @@ def complete_job(self, state: TerminalJobState) -> None:
json=body.model_dump(mode="json"),
headers=_JOB_API_HEADERS,
)
except httpx.HTTPStatusError as error:
except httpx2.HTTPStatusError as error:
if (
retried_expiry
or error.response.status_code not in (401, 403)
Expand Down
18 changes: 9 additions & 9 deletions airflow-core/tests/unit/api_fastapi/execution_api/test_dag_processor_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
from datetime import UTC, datetime
from unittest import mock

import httpx
import httpx2
import jwt
import pytest
import uvicorn
Expand Down Expand Up @@ -210,9 +210,9 @@ def test_rotation_stop_and_completion_over_http(
}


class LostAcknowledgmentTransport(httpx.BaseTransport):
class LostAcknowledgmentTransport(httpx2.BaseTransport):
def __init__(self):
self.transport = httpx.HTTPTransport()
self.transport = httpx2.HTTPTransport()
self.lost = set()

def handle_request(self, request):
Expand All @@ -225,7 +225,7 @@ def handle_request(self, request):
self.lost.add(request.url.path)
response.read()
response.close()
raise httpx.ReadError("Lost acknowledgment after commit", request=request)
raise httpx2.ReadError("Lost acknowledgment after commit", request=request)
return response

def close(self):
Expand All @@ -248,7 +248,7 @@ def test_lost_registration_and_completion_responses_over_http(api_url, provision
)
with processor.use_parse(context):
assert processor.variables.get("parse-token-key").value == "value"
with pytest.raises(httpx.HTTPStatusError) as error:
with pytest.raises(httpx2.HTTPStatusError) as error:
processor.post(f"jobs/{job_id}/heartbeat")
assert error.value.response.status_code == 403
processor.complete_job(TerminalJobState.SUCCESS)
Expand All @@ -269,7 +269,7 @@ def test_expired_credentials_fail_until_provisioning_resumes(clock, api_url, pro
time_machine.shift(700)
clock.return_value = 700

with pytest.raises(httpx.HTTPStatusError) as error:
with pytest.raises(httpx2.HTTPStatusError) as error:
processor.heartbeat()
assert error.value.response.status_code == 403

Expand Down Expand Up @@ -301,7 +301,7 @@ def test_changed_session_signals_restart_but_allows_drain_and_completion(
assert session.get(Job, job_id).state == JobState.SUCCESS.value


class CompletionFailureTransport(httpx.HTTPTransport):
class CompletionFailureTransport(httpx2.HTTPTransport):
def __init__(self, *, commit):
super().__init__()
self.commit = commit
Expand All @@ -313,7 +313,7 @@ def handle_request(self, request):
response = super().handle_request(request)
response.read()
response.close()
raise httpx.ReadError("Completion acknowledgment unavailable", request=request)
raise httpx2.ReadError("Completion acknowledgment unavailable", request=request)
return super().handle_request(request)


Expand All @@ -329,7 +329,7 @@ def test_completion_recovery_after_token_expiry(
base_url=api_url, token_file=token_file, hostname="processor-1", transport=transport
) as processor:
job_id = processor.register_job()
with pytest.raises(httpx.ReadError):
with pytest.raises(httpx2.ReadError):
processor.complete_job(TerminalJobState.SUCCESS)
time_machine.shift(350)
clock.return_value = 350
Expand Down
Loading
Loading