Skip to content

Fix Dag processor client compatibility with httpx2 #74731

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

Merged
pierrejeambrun merged 2 commits into from
Oct 14, 2026
Merged

Fix Dag processor client compatibility with httpx2 #74731

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3a66368
Fix Dag processor client compatibility with httpx2
hussein-awala Oct 14, 2026
deeb7dd
Fix Dag processor HTTP integration test compatibility
hussein-awala 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
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
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