Can a Custom Transport class method access the client session that called it?

[mijolabs]

Is there a clever way for the Custom Transport handle_request() method to access the client instance that called it, so that the same client session can be re-used for any outgoing requests?

To prevent the outgoing request from calling handle_request() again, extensions would be used.

[tomchristie]

Good question. Could give an example of the kind of functionality you're trying to achieve here, just for a bit more context?

[mijolabs]

@tomchristie

The use case is baking-in an API token flow into the client object and fetching new tokens from the token endpoint automatically when needed. It would be great reusing the same client instance so that the connection can be reused, and also because it has already been configured with e.g. the base_url, mTLS certs, timeouts, etc.

Further context: So the use case is basically what is shown in the docs for customizing authentication, which I was initially going to use instead of a custom transport since a token flow seems like one of the primary intended use cases for a custom auth class. The reason for going with a custom transport instead is that I just haven't been able to think of a clean way to prevent repeatedly triggering the token fetch in auth_flow() since extensions isn't available to custom auth. It's not useful checking the response status code (as illustrated in the docs example), because the token endpoint could return the same status code that would be used to trigger the token fetch in the first place. I suppose one could place a check on e.g. the requested URL path, but it just doesn't feel as clean as using extensions like we can for custom transports. Another thing to note is that the example in the docs doesn't show how the request to the token endpoint is being issued. This is where it would have been nice to re-use the same client session.

Nonetheless, my use case applies to both custom auth and custom transport classes; having access to their own client session would be useful to both of them.

[tomchristie]

Can you break that out into a simple pseudo-code?
We're typically suggesting a pattern where transports are composed, for example...

• Possibility to add new transports at runtime #2947 (reply in thread)

It'd be helpful to clearly understand if that's actually sufficient here or not.

class CustomTransport(httpx.BaseTransport):
    def __init__(self, transport=None, **kwargs):
        self.transport = transport or httpx.HTTPTransport(**kwargs)

    def handle_request(self, request):
        ...  # flow
        response = self.transport.handle_request(request)
        ...
        return response

You've given three examples of state that you'd like to be available, so I'll touch on these briefly...

• base_url - This isn't available anywhere. I don't fully understand the use-case?
• timeout - This is passed through as part of the request throughout.
• ssl_context - This is configured on the transport, not the client. Am doing some working on clarifying and simplifying this at the moment.

[mijolabs]

(got this to work but it's a bit clunky, see reply below)

[mijolabs]

I had some time to test a bit more and the below seems to work:

class CustomTransport(AsyncBaseTransport):
    def __init__(self):
        self.transport = AsyncHTTPTransport(
            cert=("public.pem", "private.pem"),
        )
        self.basic_auth = self.get_basic_auth_b64(
            os.getenv("CLIENT_ID"), os.getenv("CLIENT_SECRET")
        )
        self.token = None

    async def handle_async_request(self, request: Request) -> Response:
        if (
            not request.extensions.get("disable_hook", False)
            and (self.token is None or self.token_has_expired())
        ):
            self.token = await self.fetch_token(request)

        request.headers["authorization"] = f"Bearer {self.token}"

        return await self.transport.handle_async_request(request)
    
    async def fetch_token(self, request: Request) -> str:
        endpoint = "/token_endpoint"
        headers = {
            "authorization": f"Basic {self.basic_auth}",
        }
        data = {
            "grant_type": "client_credentials",
        }

        token_request = Request(
            "POST",
            f"{request.url.scheme}://{request.url.host}{endpoint}",
            headers=headers,
            data=data,
            extensions={"disable_hook": True},
        )

        _response = await self.transport.handle_async_request(token_request)
        await _response.aread()

        return _response.json()["access_token"]

    def token_has_expired(self) -> bool:
        ...

    def get_basic_auth_b64(self) -> str:
        ...

We are able to pass the client certificates to AsyncHTTPTransport (as pointed out by @tomchristie earlier in the thread). But for the token-fetching request the base_url is extracted from the original Request, the base64 Basic auth header has to be added manually, and a new Request object is created to pass to handle_async_request().

So this appears to work but isn't too elegant. If we could e.g. access the original client instance, it could already be instantiated with the mTLS certs, the basic auth, and the base_url, and we could just perform a client.post() to fetch the token.

Hopefully this does a better job at clarifying what I'm trying to achieve. I noted that #3012 mentions enabling per-request event hooks, which could probably be used for this case instead of a custom transport.