Repository: encode/httpx

Issue:
Supporting async auth flows.
Hello,

I try to make asynchronous calls to the Microsoft Graph API using `httpx.AsyncClient`. For authentication I use the `azure-identity` package that provides several client classes to easily get an access token:

```python
from azure.identity.aio import EnvironmentCredential

async with EnvironmentCredential() as credential:
    token = await credential.get_token(*scopes)
```

I'd like then to subclass `httpx.Auth` and implement the `auth_flow` method

```python
class AzureAuth(Auth):

    async def auth_flow(self, request):
        token = await self._credential.get_token(*self._scopes) # Get the access token or retrieve it from cache
        request.headers["Authorization"] = "Bearer " + token.token
        yield request
```

so that I can do things like

```python
async with AsyncClient(base_url="https://graph.microsoft.com", auth=AzureAuth(...)) as client:
    response = await client.get("/groups")
```

The problem is that `httpx.Auth.auth_flow` is **synchronous**. How can I make it asynchronous ?



Discussion:
Hi!

This is not supported currently.

After thinking about this with alternatives such as running `await` calls in the event loop (works for asyncio, but trio doesn't support it, and we lose any concurrency benefits), it seems indeed that we might want to consider how to add async support to auth implementations - your use case here is legit.

It's a bit tricky in terms of design and API, but let's consider options…

* Allow `auth_flow` to be either sync, or async (as you propose above), which means `def auth_flow()` _or_ `async def auth_flow()`.
  * Cons:
    * Doesn't allow providing both sync and async implementations, so most likely a no-no.
* Add a new `async def async_auth_flow()` method to the `Auth` interface.
  * On the client side, `AsyncClient` must be able to handle both methods. It would prefer `async_auth_flow()` and use `auth_flow()` as a fallback. `Client` will only call into `auth_flow()`.
  * Pros:
    * Backwards compatible (fully additive).
    * We can implement both sync and async on the same `Auth` subclass, with the choice of implementing them up to the author (perhaps with different libraries).
    * Transparent to the user: authors can add async support without requiring the user to change their code.
  * Cons:
    * Quite verbose in the async case.
* Introduce `AsyncAuth` with an `async def auth_flow()` method.
  * The `AsyncClient` must be able to handle both kinds of auth (`Auth` and `AsyncAuth`) transparently, while `Client` only accepts `Auth`.
  * Pros:
    * Backwards compatible: when the auth flow doesn't do I/O, authors and users can stick to `Auth`.
    * When the auth flow must do I/O, author can provide an async class, and tell users to use the async one with `AsyncClient`.
  * Cons:
    * Instead of a single class, developers now need to deal with "two kinds of classes". But this is only _required_ when the auth class does I/O (if it doesn't, there's no point in having an `AsyncAuth`). So, pretty confusing.

Right now I'm thinking 2/ might be the best option…
Now, another option: "do nothing". If an auth scheme really did require to perform I/O _on each request_, that would be quite wasteful, right? In practice, this doesn't seem like something that would pop up very often, if at all.

Indeed, for your specific use case it looks like the acquisition of a `token` does not depend on the request.

So you could perhaps just write a generic `BearerAuth` class, that takes a static token…

```python
class BearerAuth:
    def __init__(self, token):
        self._token = token

    def auth_flow(self, request):
        request.headers["Authorization"] = f"Bearer {self.token}"
        yield request
```

And then use it like this…

```python
from azure.identity.aio import EnvironmentCredential

scopes = ...

async with EnvironmentCredential() as credential:
    token = await credential.get_token(*scopes)

    async with httpx.AsyncClient(auth=BearerAuth(token.token)):
        ...
```
Closing since from the information provided, I think my suggestion above should solve your use case. If it doesn't, feel free to reach back!
Other options here...

* Create an `AsyncClient` subclass and override `.send(...)` to deal with 403 responses, and renew the auth token when needed.
* Figure out what the actual HTTP request flow is, so that you can implement it as a custom auth flow, rather than by relying on the helper lib.

Of those two, the first would be tied to async and is marginally less convenient if you wanted to provide it as a third party package, but *is* likely to provide the hook you need. The second option gives you a single implementation that works with both, and will end up making *either* sync or async requests as required, but can't give you eg async disk accesses or other async non-HTTP I/O (which for almost all auth schemes shouldn't really matter).
Hi! Sorry if I'm being intrusive on this matter, but as in my company we use httpx against multiple services, I've seen multiple different auth flows on our side, so:

> If an auth scheme really did require to perform I/O on each request, that would be quite wasteful, right? In practice, this doesn't seem like something that would pop up very often, if at all.

That's true, but some requests (against the same service) might need I/O and some not. I think I can give two flow examples:

1st:
- Fetch a token from a remote API which is valid for only 30 minutes
- Save it in local cache (expired with TTL)
- While it is valid, no I/O is performed
- When it is invalid, repeat the process

2nd:
- Fetch a token from a remote API, which you don't know for how long it is going to be valid
- Use it until you receive a 401 response from some API
- Refetch the token

Those two flows wouldn't make an I/O for each request, but might need to.

I do have these flow on my end, and many more (and having I/O or not in them is around 50/50), and would be great if we could use the `Auth` class for this.
I solved those, and others, on my end bu having a custom `HTTPService` class with a custom flow in each.
Just one snippet example:

```python
    async def get_token(self):
        token = await self._cache.get("TOKEN")
        if token is not None:
            return token

        await self._lock.acquire()
        try:
            token = await self._cache.get("TOKEN")
            if token is not None:
                return token

            token = await self._request_token()
            await self._cache.set("TOKEN", token, 27 * HALF_MINUTE)
            return token
        finally:
            self._lock.release()
```

`await self._request_token()` is actually a pretty complex endpoint request, but just wanted to give an example here.
The locking around token acquiry is a really good example case, yup.
@tomchristie  So should this issue be reopened? Or maybe open another one to track the need for an async authentication flow?
Seconding this request. We're dealing with a system that's very similar to the one described earlier, and having async auth would be really helpful.
Opening this up with a 1.0 milestone so that we can make sure we're absolutely happy with what we've got here.

Have been thinking about this a bit recently, and one thing that jumps out is that there *are* other ways besides the generator approach that would still give us essentially the same auth flow, but just by slightly different means, that would *also* allow us to support async-native auth flows for cases where they really are required.

Here's a bit of an example...

```python
class Auth:
    """
    The base class for all authentication classes.

    Subclasses must implement the `authenticate()` method.
    """

    def authenticate(self, request: Request) -> typing.Union[Request, AuthFlow, AsyncAuthFlow]:
        """
        Perform the authentication, may return one of three things:

        * A `Request` instance, used to make the request. Good for simple authentication cases.
        * An `AuthFlow`, which is a stateful instance that provides an API for making multiple requests.
        * An `AsyncAuthFlow`, which is a stateful instance that provides an async API for making multiple requests.
        """
        ...

class AuthFlow:
    """
    A concrete subclass of AuthFlow may be returned by `authenticate()` in order
    to implement an authentication flow that may make multiple requests.

    Authentication classes using this style may be used by either `Client` or `AsyncClient`,
    unless the implementation uses sync I/O such as disk accesses, in which case the auth
    should be documented as only supporting usage with `Client`. Usage of sync I/O within
    an async codebase will block the event loop, and could cause performance issues.
    """

    def on_request(request) -> Request:
        """
        Called once. Return the initial request that should be used.
        """
        raise NotImplementedError()

    def on_response(response) -> Optional[Request]:
        """
        Called once or more. Optionally return another request to issue.
        """
        raise NotImplementedError()


class AsyncAuthFlow:
    """
    A concrete subclass of AsyncAuthFlow may be returned by `authenticate()` in order
    to implement an async authentication flow that may make multiple requests.

    Authentication classes using this style may only be used by `AsyncClient`.
    """

    async def on_request(request) -> Request:
        """
        Called once. Return the initial request that should be used.
        """
        raise NotImplementedError()

    async def on_response(response) -> Optional[Request]:
        """
        Called once or more. Optionally return another request to issue.
        """
        raise NotImplementedError()
```

We'd still end up with essentially the same implementation within the client codebase, we'd just be interfacing with the auth flow using a different API. We'd also be able to continue to support the existing `auth_flow` with a backwards compatible shim, for however long we'd want to support that.

🤔
@tomchristie Pretty interesting!

So just to be sure I'm getting this right... For cases where library authors want to provide both a sync and async implementation, they would provide two classes:

- `MyAuth`, whose authenticate method returns an `AuthFlow`
- `MyAsyncAuth`, whose authenticate method returns an `AsyncAuthFlow`

Is this correct? If so, this doesn't seem to satisfy the  "users shouldn't need to care about sync vs async Auth" property (which I think would be nice have) since they'd need to deal with two possible Auth classes.

Resolving patch:
diff --git a/httpx/_auth.py b/httpx/_auth.py
--- a/httpx/_auth.py
+++ b/httpx/_auth.py
@@ -17,6 +17,11 @@ class Auth:
 
     To implement a custom authentication scheme, subclass `Auth` and override
     the `.auth_flow()` method.
+
+    If the authentication scheme does I/O such as disk access or network calls, or uses
+    synchronization primitives such as locks, you should override `.sync_auth_flow()`
+    and/or `.async_auth_flow()` instead of `.auth_flow()` to provide specialized
+    implementations that will be used by `Client` and `AsyncClient` respectively.
     """
 
     requires_request_body = False
@@ -46,6 +51,56 @@ def auth_flow(self, request: Request) -> typing.Generator[Request, Response, Non
         """
         yield request
 
+    def sync_auth_flow(
+        self, request: Request
+    ) -> typing.Generator[Request, Response, None]:
+        """
+        Execute the authentication flow synchronously.
+
+        By default, this defers to `.auth_flow()`. You should override this method
+        when the authentication scheme does I/O and/or uses concurrency primitives.
+        """
+        if self.requires_request_body:
+            request.read()
+
+        flow = self.auth_flow(request)
+        request = next(flow)
+
+        while True:
+            response = yield request
+            if self.requires_response_body:
+                response.read()
+
+            try:
+                request = flow.send(response)
+            except StopIteration:
+                break
+
+    async def async_auth_flow(
+        self, request: Request
+    ) -> typing.AsyncGenerator[Request, Response]:
+        """
+        Execute the authentication flow asynchronously.
+
+        By default, this defers to `.auth_flow()`. You should override this method
+        when the authentication scheme does I/O and/or uses concurrency primitives.
+        """
+        if self.requires_request_body:
+            await request.aread()
+
+        flow = self.auth_flow(request)
+        request = next(flow)
+
+        while True:
+            response = yield request
+            if self.requires_response_body:
+                await response.aread()
+
+            try:
+                request = flow.send(response)
+            except StopIteration:
+                break
+
 
 class FunctionAuth(Auth):
     """
diff --git a/httpx/_client.py b/httpx/_client.py
--- a/httpx/_client.py
+++ b/httpx/_client.py
@@ -785,15 +785,12 @@ def _send_handling_auth(
         auth: Auth,
         timeout: Timeout,
     ) -> Response:
-        if auth.requires_request_body:
-            request.read()
-
-        auth_flow = auth.auth_flow(request)
+        auth_flow = auth.sync_auth_flow(request)
         request = next(auth_flow)
+
         while True:
             response = self._send_single_request(request, timeout)
-            if auth.requires_response_body:
-                response.read()
+
             try:
                 next_request = auth_flow.send(response)
             except StopIteration:
@@ -1409,18 +1406,15 @@ async def _send_handling_auth(
         auth: Auth,
         timeout: Timeout,
     ) -> Response:
-        if auth.requires_request_body:
-            await request.aread()
+        auth_flow = auth.async_auth_flow(request)
+        request = await auth_flow.__anext__()
 
-        auth_flow = auth.auth_flow(request)
-        request = next(auth_flow)
         while True:
             response = await self._send_single_request(request, timeout)
-            if auth.requires_response_body:
-                await response.aread()
+
             try:
-                next_request = auth_flow.send(response)
-            except StopIteration:
+                next_request = await auth_flow.asend(response)
+            except StopAsyncIteration:
                 return response
             except BaseException as exc:
                 await response.aclose()


Test patch:
diff --git a/tests/client/test_auth.py b/tests/client/test_auth.py
--- a/tests/client/test_auth.py
+++ b/tests/client/test_auth.py
@@ -1,5 +1,12 @@
+"""
+Integration tests for authentication.
+
+Unit tests for auth classes also exist in tests/test_auth.py
+"""
+import asyncio
 import hashlib
 import os
+import threading
 import typing
 
 import httpcore
@@ -183,6 +190,31 @@ def auth_flow(self, request: Request) -> typing.Generator[Request, Response, Non
         yield request
 
 
+class SyncOrAsyncAuth(Auth):
+    """
+    A mock authentication scheme that uses a different implementation for the
+    sync and async cases.
+    """
+
+    def __init__(self) -> None:
+        self._lock = threading.Lock()
+        self._async_lock = asyncio.Lock()
+
+    def sync_auth_flow(
+        self, request: Request
+    ) -> typing.Generator[Request, Response, None]:
+        with self._lock:
+            request.headers["Authorization"] = "sync-auth"
+        yield request
+
+    async def async_auth_flow(
+        self, request: Request
+    ) -> typing.AsyncGenerator[Request, Response]:
+        async with self._async_lock:
+            request.headers["Authorization"] = "async-auth"
+        yield request
+
+
 @pytest.mark.asyncio
 async def test_basic_auth() -> None:
     url = "https://example.org/"
@@ -664,3 +696,34 @@ def test_sync_auth_reads_response_body() -> None:
 
     assert response.status_code == 200
     assert response.json() == {"auth": '{"auth": "xyz"}'}
+
+
+@pytest.mark.asyncio
+async def test_async_auth() -> None:
+    """
+    Test that we can use an auth implementation specific to the async case, to
+    support cases that require performing I/O or using concurrency primitives (such
+    as checking a disk-based cache or fetching a token from a remote auth server).
+    """
+    url = "https://example.org/"
+    auth = SyncOrAsyncAuth()
+
+    async with httpx.AsyncClient(transport=AsyncMockTransport()) as client:
+        response = await client.get(url, auth=auth)
+
+    assert response.status_code == 200
+    assert response.json() == {"auth": "async-auth"}
+
+
+def test_sync_auth() -> None:
+    """
+    Test that we can use an auth implementation specific to the sync case.
+    """
+    url = "https://example.org/"
+    auth = SyncOrAsyncAuth()
+
+    with httpx.Client(transport=SyncMockTransport()) as client:
+        response = client.get(url, auth=auth)
+
+    assert response.status_code == 200
+    assert response.json() == {"auth": "sync-auth"}
diff --git a/tests/test_auth.py b/tests/test_auth.py
new file mode 100644
--- /dev/null
+++ b/tests/test_auth.py
@@ -0,0 +1,63 @@
+"""
+Unit tests for auth classes.
+
+Integration tests also exist in tests/client/test_auth.py
+"""
+import pytest
+
+import httpx
+
+
+def test_basic_auth():
+    auth = httpx.BasicAuth(username="user", password="pass")
+    request = httpx.Request("GET", "https://www.example.com")
+
+    # The initial request should include a basic auth header.
+    flow = auth.sync_auth_flow(request)
+    request = next(flow)
+    assert request.headers["Authorization"].startswith("Basic")
+
+    # No other requests are made.
+    response = httpx.Response(content=b"Hello, world!", status_code=200)
+    with pytest.raises(StopIteration):
+        flow.send(response)
+
+
+def test_digest_auth_with_200():
+    auth = httpx.DigestAuth(username="user", password="pass")
+    request = httpx.Request("GET", "https://www.example.com")
+
+    # The initial request should not include an auth header.
+    flow = auth.sync_auth_flow(request)
+    request = next(flow)
+    assert "Authorization" not in request.headers
+
+    # If a 200 response is returned, then no other requests are made.
+    response = httpx.Response(content=b"Hello, world!", status_code=200)
+    with pytest.raises(StopIteration):
+        flow.send(response)
+
+
+def test_digest_auth_with_401():
+    auth = httpx.DigestAuth(username="user", password="pass")
+    request = httpx.Request("GET", "https://www.example.com")
+
+    # The initial request should not include an auth header.
+    flow = auth.sync_auth_flow(request)
+    request = next(flow)
+    assert "Authorization" not in request.headers
+
+    # If a 401 response is returned, then a digest auth request is made.
+    headers = {
+        "WWW-Authenticate": 'Digest realm="...", qop="auth", nonce="...", opaque="..."'
+    }
+    response = httpx.Response(
+        content=b"Auth required", status_code=401, headers=headers
+    )
+    request = flow.send(response)
+    assert request.headers["Authorization"].startswith("Digest")
+
+    # No other requests are made.
+    response = httpx.Response(content=b"Hello, world!", status_code=200)
+    with pytest.raises(StopIteration):
+        flow.send(response)