RESTClient: middleware / interceptors

[joscha]

Feature description

I would like to be able to use a middleware on the rest client that uses a Pydantic model on the response.
This has some similarity to using hooks, bit hooks currently don't allow changing the return value, they are only evaluated alongside.

Currently, I am doing something like this:

response = rest_client.get("companies")
companies = CompanyPaged.model_validate_json(json_data=response.text)

for pagination it is a bit more tricky even, as the value returned from .paginate is only the data part of the response, e.g. for something like:

class CompanyPaged(BaseModel):
    data: List[Company]
    """
    A page of Company results
    """
    pagination: Pagination

only List[Company] is returned from .paginate, as the pagination property is evaluated and used internally.

In my ideal world I'd be able to do something like this:

from dlt.sources.helpers.requests import Response

def my_transformer(response: Response):
    return CompanyPaged.model_validate_json(json_data=response.text)

for companies in rest_client.paginate("companies", middleware=[my_transformer]):
    for company in companies:
        print(company)

or even better:

from dlt.sources.helpers.requests import Response

def my_transformer(response: Response):
    return CompanyPaged.model_validate_json(json_data=response.text)

def transform_errors(response: Response):
    error_model = ...   # match response against a list of error models
    raise CustomApiException("Oops", error_model)

try:
    # yields `List[Company]`s
    yield from rest_client.paginate("companies", middleware={
        "success": [my_transformer],
        "error": [transform_errors]
    })
except CustomException as e:
     print(f"Payload: {e.error_model}")
     # handle error (gracefully?) here

Other rest clients have this notion - for example Axios interceptors.

Possibly this can be done via the low-level access to request's Session added in #1844 and mounting a connection adapter, however it doesn't necessarily need to be that low-level, and the involvement is quite substantial.

My current workaround

def raise_if_error(response: Response, *args: Any, **kwargs: Any) -> None:
    if response.status_code < 200 or response.status_code >= 300:
        error_adapter = TypeAdapter(AuthenticationErrors | NotFoundErrors | AuthorizationErrors | ValidationErrors)
        error = error_adapter.validate_json(response.text)
        response.reason = "\n".join([e.message for e in error.errors])
        response.raise_for_status()

hooks = {
    "response": [raise_if_error]
}

AuthenticationErrors | NotFoundErrors | AuthorizationErrors | ValidationErrors are pydantic models looking roughly like this:

class AuthenticationError(BaseModel):
    code: Literal['authentication']
    """
    Error code
    """
    message: str
    """
    Error message
    """


class AuthenticationErrors(BaseModel):
    errors: List[AuthenticationError]
    """
    AuthenticationError errors
    """

Are you a dlt user?

Yes, I'm already a dlt user.

Use case

postprocessing rest client responses.

Proposed solution

Expose/add a middleware / interceptor / connection adapter interface to RESTClient.

If we want to specifically support Pydantic, support for the TypeAdapter might be a good way to do it.

Related issues

#1593, #1844

[joscha]

Real-world example here: https://github.com/joscha/verified-sources/blob/d0b721004ac637f393df5f23174ba4c947441c67/sources/affinity/__init__.py, specifically: https://github.com/joscha/verified-sources/blob/d0b721004ac637f393df5f23174ba4c947441c67/sources/affinity/__init__.py#L40-L85