Migrate Airflow REST client to httpx2 - #578
kstonekuan merged 2 commits into
Conversation
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Existing callers do not close persistent clients, and diagnostic/test issues remain unresolved.
Get a fresh assessment by requesting another Copilot review.
Review effort: Lite
Findings: 2
Open (4)
What changed in this PR
Migrates the Airflow REST client from urllib to synchronous httpx2 while preserving authentication, parsing, errors, and redirect behavior.
Changes:
- Adds
httpx2.Clientlifecycle management and redirect-aware error handling. - Extends client tests for transport errors, redirects, headers, timeouts, and closing.
- Preserves existing Airflow request and response contracts.
| File | Description |
|---|---|
src/hflow/runtime/_client.py |
Replaces urllib transport with httpx2. |
tests/test_runtime_client.py |
Adds migration and compatibility boundary tests. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| self._auth = auth | ||
| self._request_timeout_s = request_timeout_s | ||
| self._token: str | None = None | ||
| self._client = httpx2.Client() |
There was a problem hiding this comment.
Addressed in d1315e5.
Audited all production AirflowClient ownership paths and added explicit lifecycle management, including CLI ingest, runtime/status helpers, and the server-side resolver cache. Clients are now closed on normal, early-return, and exception paths where appropriate.
Added lifecycle coverage for these paths as well.
| body_excerpt = f"redirect location: {location}" + ( | ||
| f"; {body_excerpt}" if body_excerpt else "" | ||
| ) |
There was a problem hiding this comment.
Addressed in d1315e5.
The redirect Location header is now passed through the existing _body_excerpt() bound, and the final combined diagnostic is bounded as well.
Added a regression test covering an oversized Location header
| timeout = kwargs["timeout"] | ||
| assert isinstance(timeout, float) | ||
| observed_timeouts.append(timeout) | ||
| request = httpx2.Request(method, url) | ||
| return httpx2.Response(200, content=b'{"dag_id":"pipeline_ingest"}', request=request) | ||
|
|
||
| monkeypatch.setattr(client_module.httpx2.Client, "request", fake_request) | ||
| client = AirflowClient( | ||
| "http://airflow.example", auth=BearerToken("token"), request_timeout_s=12.5 | ||
| ) | ||
|
|
||
| assert client.dag("pipeline_ingest") == {"dag_id": "pipeline_ingest"} | ||
| assert observed_timeouts == [12.5] |
There was a problem hiding this comment.
Addressed in d1315e5.
Replaced the mock-based timeout plumbing assertion with a behavioral test using a delayed local HTTP response and a short client timeout.
The test now verifies the caller-visible AirflowClientError outcome rather than the internal httpx2 request arguments.
| def close(self: httpx2.Client) -> None: | ||
| nonlocal close_calls | ||
| close_calls += 1 | ||
|
|
||
| monkeypatch.setattr(client_module.httpx2.Client, "close", close) | ||
| client = AirflowClient("http://airflow.example", auth=BearerToken("token")) | ||
|
|
||
| client.close() | ||
| with AirflowClient("http://airflow.example", auth=BearerToken("token")) as context_client: | ||
| assert isinstance(context_client, AirflowClient) | ||
|
|
||
| assert close_calls == 2 |
There was a problem hiding this comment.
Addressed in d1315e5.
Replaced the mock call-count assertion with behavioral tests that verify the client becomes unusable after close() and after exiting the context manager, using the actual behavior of the pinned httpx2 version.
kstonekuan
left a comment
kstonekuan
left a comment
There was a problem hiding this comment.
This is a very good first contribution. The design you posted first is most of why: the review had nothing to discover.
All three transport guards hold under mutation. Following redirects on mutations again reddens 4, dropping the POST/PATCH 3xx refusal reddens 4, and dropping the Location from the message reddens 3. urllib.parse.quote correctly stayed for path escaping, which is not transport.
The part I did not ask for and would have: the server caches a client behind a TTL, so owning one now means eviction can leak a connection pool. _close_resolution runs on refresh, on the exception path, and behind RuntimeResolver.close() with a context manager. That is the failure mode this change introduces, and you closed it before anyone hit it.
Scope is wider than the two files you proposed, but that follows from owning a client that has to be closed rather than from drift, and 545 of the 742 added lines are tests.
Validated on the merge with main: gate clean, 2276 passed.
Summary
Migrate the Airflow REST API v2 client from
urllibto synchronoushttpx2while preserving the existing request, authentication, response parsing, and error-handling behavior.Why
The REST client currently relies on the standard-library
urllibtransport. This change replaces that transport withhttpx2and keeps the migration scoped to the REST client and its tests.GET/HEAD requests continue to follow redirects, while POST/PATCH redirects are explicitly rejected rather than allowing mutation requests to be rewritten. Redirect errors for POST/PATCH also include the
Locationheader for easier diagnosis.The client now owns a synchronous
httpx2.Clientand exposesclose()plus context-manager support for explicit resource lifecycle management.Closes #568
Validation
uv sync --lockeduv run ruff check --fixuv run ruff formatuv run ty checkuv run pytest -qHFLOW_DOCKER_TESTS=1 uv run pytest tests/test_runtime_integration.py -qThe full test suite passed with all tests passing and 8 skipped.
The Docker-backed runtime integration test passed.
Checklist
uv run ruff check --fix,uv run ruff format, anduv run ty check.