diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index a0c0a30de088f..93774f1aeed97 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -197,14 +197,27 @@ def get_json_error(response: httpx.Response): raise err +def _raise_for_status_picklable(response: httpx.Response) -> None: + try: + response.raise_for_status() + except httpx.HTTPStatusError as e: + if 400 <= response.status_code < 600: + raise + wrapped = ServerResponseError(str(e), request=e.request, response=e.response) + wrapped.detail = None + raise wrapped from e + + def raise_on_4xx_5xx(response: httpx.Response): - return get_json_error(response) or response.raise_for_status() + get_json_error(response) + _raise_for_status_picklable(response) # Py 3.11+ version def raise_on_4xx_5xx_with_note(response: httpx.Response): try: - return get_json_error(response) or response.raise_for_status() + get_json_error(response) + _raise_for_status_picklable(response) except httpx.HTTPStatusError as e: if TYPE_CHECKING: assert hasattr(e, "add_note") diff --git a/task-sdk/tests/task_sdk/api/test_client.py b/task-sdk/tests/task_sdk/api/test_client.py index 8bd558b461b1b..248d02244a2a1 100644 --- a/task-sdk/tests/task_sdk/api/test_client.py +++ b/task-sdk/tests/task_sdk/api/test_client.py @@ -198,6 +198,22 @@ def test_server_response_error_pickling(self): assert unpickled.response.status_code == 404 assert unpickled.request.url == "http://error" + def test_redirect_response_raises_picklable_error(self): + responses = [httpx.Response(301, headers={"location": "http://proxy-login"})] + client = make_client_w_responses(responses) + + with pytest.raises(ServerResponseError) as exc_info: + client.get("http://error") + + err = exc_info.value + assert err.response.status_code == 301 + assert err.detail is None + + # Check that the error is picklable, unlike a bare httpx.HTTPStatusError. + unpickled = pickle.loads(pickle.dumps(err)) + assert isinstance(unpickled, ServerResponseError) + assert unpickled.response.status_code == 301 + @pytest.mark.skipif(sys.version_info < (3, 11), reason="Exception notes (PEP 678) require Python 3.11") def test_server_error_detail_added_as_note(self): """Notes survive uncaught propagation, handled sites still log detail directly."""