diff --git a/mock_tests/test_auth.py b/mock_tests/test_auth.py index cc4480c93..5402190d0 100644 --- a/mock_tests/test_auth.py +++ b/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 diff --git a/weaviate/connect/v4.py b/weaviate/connect/v4.py index 5dcadfaeb..21d561b8c 100644 --- a/weaviate/connect/v4.py +++ b/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,12 +561,12 @@ 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 @@ -559,13 +574,13 @@ def _create_background_token_refresh(self, _auth: Optional[_Auth] = None) -> Non # 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)