Skip to content

Commit

Permalink
Improve error handling in Task SDK api client (apache#44044)
Browse files Browse the repository at this point in the history
There was a logic bug where a non-detail error was not parsing correctly
causing a _different_ error (pydantic ValidationError, rather than a 404 error)
  • Loading branch information
ashb authored Nov 14, 2024
1 parent cc7b7a2 commit fe4aafe
Show file tree
Hide file tree
Showing 2 changed files with 34 additions and 12 deletions.
26 changes: 17 additions & 9 deletions task_sdk/src/airflow/sdk/api/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@

import httpx
import methodtools
import msgspec
import structlog
from pydantic import BaseModel
from uuid6 import uuid7
Expand All @@ -48,7 +49,6 @@
__all__ = [
"Client",
"ConnectionOperations",
"ErrorBody",
"ServerResponseError",
"TaskInstanceOperations",
]
Expand Down Expand Up @@ -177,8 +177,9 @@ def connections(self) -> ConnectionOperations:
return ConnectionOperations(self)


class ErrorBody(BaseModel):
detail: list[RemoteValidationError] | dict[str, Any]
# This is only used for parsing. ServerResponseError is raised instead
class _ErrorBody(BaseModel):
detail: list[RemoteValidationError] | str

def __repr__(self):
return repr(self.detail)
Expand All @@ -188,7 +189,7 @@ class ServerResponseError(httpx.HTTPStatusError):
def __init__(self, message: str, *, request: httpx.Request, response: httpx.Response):
super().__init__(message, request=request, response=response)

detail: ErrorBody
detail: list[RemoteValidationError] | str | dict[str, Any] | None

@classmethod
def from_response(cls, response: httpx.Response) -> ServerResponseError | None:
Expand All @@ -201,16 +202,23 @@ def from_response(cls, response: httpx.Response) -> ServerResponseError | None:
if response.headers.get("content-type") != "application/json":
return None

detail: list[RemoteValidationError] | dict[str, Any] | None = None
try:
err = ErrorBody.model_validate_json(response.read())
if isinstance(err.detail, list):
body = _ErrorBody.model_validate_json(response.read())

if isinstance(body.detail, list):
detail = body.detail
msg = "Remote server returned validation error"
else:
msg = err.detail.get("message", "") or "Un-parseable error"
msg = body.detail or "Un-parseable error"
except Exception:
err = ErrorBody.model_validate_json(response.content)
try:
detail = msgspec.json.decode(response.content)
except Exception:
# Fallback to a normal httpx error
return None
msg = "Server returned error"

self = cls(msg, request=response.request, response=response)
self.detail = err
self.detail = detail
return self
20 changes: 17 additions & 3 deletions task_sdk/tests/api/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
import httpx
import pytest

from airflow.sdk.api.client import Client, ErrorBody, RemoteValidationError, ServerResponseError
from airflow.sdk.api.client import Client, RemoteValidationError, ServerResponseError


class TestClient:
Expand All @@ -40,8 +40,8 @@ def handle_request(request: httpx.Request) -> httpx.Response:
client.get("http://error")

assert isinstance(err.value, ServerResponseError)
assert isinstance(err.value.detail, ErrorBody)
assert err.value.detail.detail == [
assert isinstance(err.value.detail, list)
assert err.value.detail == [
RemoteValidationError(loc=["#0"], msg="err", type="required"),
]

Expand All @@ -60,3 +60,17 @@ def handle_request(request: httpx.Request) -> httpx.Response:
with pytest.raises(httpx.HTTPStatusError) as err:
client.get("http://error")
assert not isinstance(err.value, ServerResponseError)

def test_error_parsing_other_json(self):
def handle_request(request: httpx.Request) -> httpx.Response:
# Some other json than an error body.
return httpx.Response(404, json={"detail": "Not found"})

client = Client(
base_url=None, dry_run=True, token="", mounts={"'http://": httpx.MockTransport(handle_request)}
)

with pytest.raises(ServerResponseError) as err:
client.get("http://error")
assert err.value.args == ("Not found",)
assert err.value.detail is None

0 comments on commit fe4aafe

Please sign in to comment.