fix(connect): send OIDC requests through our own httpx client (#2220)

mock_tests/test_auth.py

@@ -6,6 +6,7 @@
 from typing import List, Union
 
 import grpc
+import httpx
 import pytest
 from pytest_httpserver import HTTPServer
 from werkzeug import Request, Response
@@ -85,6 +86,67 @@ def test_client_credentials(weaviate_auth_mock: HTTPServer, start_grpc_server: g
     weaviate_auth_mock.check_assertions()
 
 
+OIDC_CREDENTIALS = {
+    "bearer_token": lambda: weaviate.auth.AuthBearerToken(
+        ACCESS_TOKEN, refresh_token=REFRESH_TOKEN
+    ),
+    "client_password": lambda: weaviate.auth.AuthClientPassword("AUsername", "SomePassWord"),
+    "client_credentials": lambda: weaviate.auth.AuthClientCredentials(
+        client_secret=CLIENT_SECRET, scope=SCOPE
+    ),
+}
+
+
+def _expect_oidc_token(weaviate_auth_mock: HTTPServer) -> None:
+    weaviate_auth_mock.expect_request("/auth").respond_with_json(
+        {"access_token": ACCESS_TOKEN, "expires_in": 500, "refresh_token": REFRESH_TOKEN}
+    )
+    weaviate_auth_mock.expect_request(
+        "/v1/schema", headers={"Authorization": "Bearer " + ACCESS_TOKEN}
+    ).respond_with_json({"classes": []})
+
+
+@pytest.mark.parametrize("credentials", OIDC_CREDENTIALS.keys())
+def test_oidc_requests_use_own_http_client(
+    weaviate_auth_mock: HTTPServer, start_grpc_server: grpc.Server, credentials: str
+) -> None:
+    """Requests must go through the client's own httpx client, never the authlib session.
+
+    Authlib >= 1.8 builds its sessions on httpx2 when it is installed; handing such a
+    session the client's httpx Timeout failed every request with a TypeError (#2220).
+    """
+    _expect_oidc_token(weaviate_auth_mock)
+
+    with weaviate.connect_to_local(
+        host=MOCK_IP,
+        port=MOCK_PORT,
+        grpc_port=MOCK_PORT_GRPC,
+        auth_credentials=OIDC_CREDENTIALS[credentials](),
+    ) as client:
+        assert type(client._connection._client) is httpx.Client
+        client.collections.list_all()
+    assert client._connection._oauth_session is None  # close() closed the session too
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("credentials", OIDC_CREDENTIALS.keys())
+async def test_oidc_requests_use_own_http_client_async(
+    weaviate_auth_mock: HTTPServer, start_grpc_server: grpc.Server, credentials: str
+) -> None:
+    """Async twin of the test above."""
+    _expect_oidc_token(weaviate_auth_mock)
+
+    async with weaviate.use_async_with_local(
+        host=MOCK_IP,
+        port=MOCK_PORT,
+        grpc_port=MOCK_PORT_GRPC,
+        auth_credentials=OIDC_CREDENTIALS[credentials](),
+    ) as client:
+        assert type(client._connection._client) is httpx.AsyncClient
+        await client.collections.list_all()
+    assert client._connection._oauth_session is None  # close() closed the session too
+
+
 @pytest.mark.asyncio
 async def test_client_credentials_refresh_async(
     weaviate_auth_mock: HTTPServer, start_grpc_server: grpc.Server

weaviate/connect/v4.py

@@ -17,6 +17,7 @@
     Literal,
     Optional,
     Tuple,
+    Type,
     TypeVar,
     Union,
     cast,
@@ -105,7 +106,18 @@
 
 Session = Union[Client, OAuth2Client]
 AsyncSession = Union[AsyncClient, AsyncOAuth2Client]
-HttpClient = Union[AsyncClient, AsyncOAuth2Client, Client, OAuth2Client]
+HttpClient = Union[AsyncClient, Client]
+# Holds and refreshes the OIDC token only; requests go through HttpClient. Authlib >= 1.8 builds these
+# sessions on httpx2 when it is installed, so they must never be handed httpx objects (timeouts, transports).
+OAuthSession = Union[AsyncOAuth2Client, OAuth2Client]
+
+try:
+    # an httpx2-based session raises httpx2 errors when fetching a token
+    from httpx2 import HTTPError as _Httpx2HTTPError  # type: ignore
+
+    _OAUTH_HTTP_ERRORS: Tuple[Type[Exception], ...] = (HTTPError, _Httpx2HTTPError)
+except ImportError:
+    _OAUTH_HTTP_ERRORS = (HTTPError,)
 
 PERMISSION_DENIED = "PERMISSION_DENIED"
 
@@ -143,6 +155,7 @@ def __init__(
         self.__additional_headers = {}
         self._auth = auth_client_secret
         self._client: Optional[HttpClient] = None
+        self._oauth_session: Optional[OAuthSession] = None
         self._connection_params = connection_params
         self._grpc_stub: Optional[weaviate_pb2_grpc.WeaviateStub] = None
         self._grpc_channel: Union[AsyncChannel, SyncChannel, None] = None
@@ -274,8 +287,8 @@ def is_connected(self) -> bool:
     def get_current_bearer_token(self) -> str:
         if "authorization" in self._headers:
             return self._headers["authorization"]
-        elif isinstance(self._client, (OAuth2Client, AsyncOAuth2Client)):
-            return f"Bearer {self._client.token['access_token']}"
+        elif self._oauth_session is not None:
+            return f"Bearer {self._oauth_session.token['access_token']}"
         return ""
 
     def _prepare_grpc_headers(self) -> None:
@@ -484,11 +497,12 @@ async def execute() -> None:
                             )
                         )
                         try:
-                            self._client = await _auth.aresult(_auth.get_auth_session())
-                        except HTTPError as e:
+                            self._oauth_session = await _auth.aresult(_auth.get_auth_session())
+                        except _OAUTH_HTTP_ERRORS as e:
                             raise AuthenticationFailedError(
                                 f"Failed to authenticate with OIDC: {repr(e)}"
                             )
+                        self.__make_clients(colour)
 
                         if isinstance(auth_client_secret, AuthClientCredentials):
                             # credentials should only be saved for client credentials, otherwise use refresh token
@@ -507,11 +521,12 @@ async def execute() -> None:
                         )
                     )
                     try:
-                        self._client = _auth.result(_auth.get_auth_session())
-                    except HTTPError as e:
+                        self._oauth_session = _auth.result(_auth.get_auth_session())
+                    except _OAUTH_HTTP_ERRORS as e:
                         raise AuthenticationFailedError(
                             f"Failed to authenticate with OIDC: {repr(e)}"
                         )
+                    self.__make_clients(colour)
 
                     if isinstance(auth_client_secret, AuthClientCredentials):
                         # credentials should only be saved for client credentials, otherwise use refresh token
@@ -546,26 +561,26 @@ async def execute() -> None:
     def _create_background_token_refresh(self, _auth: Optional[_Auth] = None) -> None:
         """Create a background thread that periodically refreshes access and refresh tokens.
 
-        While the underlying library refreshes tokens, it does not have an internal cronjob that checks every
-        X-seconds if a token has expired. If there is no activity for longer than the refresh tokens lifetime, it will
-        expire. Therefore, refresh manually shortly before expiration time is up.
+        Requests go through our own HTTP client with the session's current token, so the session never refreshes
+        on its own. Refresh shortly before the token expires, and keep refreshing so the refresh token does not
+        expire either.
         """
-        assert isinstance(self._client, (OAuth2Client, AsyncOAuth2Client))
-        if "refresh_token" not in self._client.token and _auth is None:
+        assert isinstance(self._oauth_session, (OAuth2Client, AsyncOAuth2Client))
+        if "refresh_token" not in self._oauth_session.token and _auth is None:
             return
 
         # stop the refresher from an earlier connect(), if any
         self._cancel_background_token_refresh()
 
         # refresh 30s before the token expires (assume 1 minute if the token does not say);
         # the loops below always wait at least 1s
-        refresh_in: int = self._client.token.get("expires_in", 60) - 30
+        refresh_in: int = self._oauth_session.token.get("expires_in", 60) - 30
         # the refresher keeps its own event: after close() + connect() the old refresher
         # must stop on this one instead of running on with the new one
         shutdown = Event()
         self._shutdown_background_event = shutdown
 
-        if isinstance(self._client, AsyncOAuth2Client):
+        if isinstance(self._oauth_session, AsyncOAuth2Client):
             # async client: refresh in an asyncio task on the current loop, not in a thread
             # (threads cannot start under WASM/Pyodide)
             self.__token_refresh_task = asyncio.get_running_loop().create_task(
@@ -578,8 +593,8 @@ def periodic_refresh_token(refresh_time: int, _auth: Optional[_Auth]) -> None:
             while not shutdown.wait(timeout=max(refresh_time, 1)):
                 try:
                     # use one client for the whole round: close()/connect() may replace
-                    # self._client while a refresh is running
-                    client = self._client
+                    # self._oauth_session while a refresh is running
+                    client = self._oauth_session
                     if not isinstance(client, OAuth2Client):
                         continue
                     if "refresh_token" in client.token:
@@ -628,7 +643,7 @@ async def __periodic_token_refresh_async(
         while not shutdown.is_set():
             await asyncio.sleep(max(refresh_time, 1))
             try:
-                client = self._client
+                client = self._oauth_session
                 if not isinstance(client, AsyncOAuth2Client):
                     continue
                 if "refresh_token" in client.token:
@@ -769,6 +784,10 @@ async def execute() -> None:
                     assert isinstance(self._client, AsyncClient)
                     await self._client.aclose()
                     self._client = None
+                if self._oauth_session is not None:
+                    assert isinstance(self._oauth_session, AsyncOAuth2Client)
+                    await self._oauth_session.aclose()
+                    self._oauth_session = None
                 if self._grpc_stub is not None:
                     assert self._grpc_channel is not None
                     assert isinstance(self._grpc_channel, AsyncChannel)
@@ -782,6 +801,10 @@ async def execute() -> None:
             assert isinstance(self._client, Client)
             self._client.close()
             self._client = None
+        if self._oauth_session is not None:
+            assert isinstance(self._oauth_session, OAuth2Client)
+            self._oauth_session.close()
+            self._oauth_session = None
         if self._grpc_stub is not None:
             assert self._grpc_channel is not None
             assert isinstance(self._grpc_channel, SyncChannel)