jerryshao commented on code in PR #12531: URL: https://github.com/apache/gravitino/pull/12531#discussion_r3820357277
########## mcp-server/mcp_server/core/oauth.py: ########## @@ -0,0 +1,131 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""httpx ``auth=`` hook for MCP → Gravitino OAuth2 client-credentials. + +Uses ``httpx-auth`` for fetch and cache on the existing ``httpx.AsyncClient``. +Credentials go in the form body (``client_secret_post``), matching the +Java/Python Gravitino clients. httpx-auth defaults to HTTP Basic. This class +retries once after Gravitino HTTP 401. +""" + +import logging +from collections.abc import Generator +from typing import Optional, Union + +import httpx +from httpx_auth import OAuth2, OAuth2ClientCredentials + +_LOG = logging.getLogger(__name__) + +# Refresh this many seconds before recorded expiry. httpx-auth default is 30. +DEFAULT_REFRESH_SKEW_SECONDS = 60 + + +class RefreshableBearerAuth(OAuth2ClientCredentials): + """httpx-auth client-credentials with form POST and one 401 retry.""" + + requires_response_body = True + + def __init__( + self, + *, + token_endpoint: str, + client_id: str, + client_secret: str, + scope: str = "", + refresh_skew_seconds: int = DEFAULT_REFRESH_SKEW_SECONDS, + client: Optional[httpx.Client] = None, + ): + """Build an ``auth=`` hook for the service hop. + + Args: + token_endpoint: Identity-provider token URL. + client_id: OAuth2 client id. + client_secret: OAuth2 client secret. + scope: Optional OAuth2 scope. + refresh_skew_seconds: httpx-auth ``early_expiry``. + client: Optional sync httpx client used only for token POSTs + (tests inject ``MockTransport`` here). + """ + kwargs = {"early_expiry": float(refresh_skew_seconds)} + if scope: + kwargs["scope"] = scope + if client is not None: + kwargs["client"] = client + super().__init__(token_endpoint, client_id, client_secret, **kwargs) + + def invalidate(self) -> None: + """Drop the cached token so the next call fetches a new one.""" + cache = OAuth2.token_cache + # TokenMemoryCache.clear() wipes every client; only drop ours. + with cache._forbid_concurrent_cache_access: # pylint: disable=protected-access + cache.tokens.pop(self.state, None) + + def _configure_client(self, client: httpx.Client) -> None: + """Do not send HTTP Basic; id and secret go in the form body.""" + client.timeout = self.timeout + + def request_new_token( + self, + ) -> Union[tuple[str, str], tuple[str, str, Union[int, str]]]: + """POST ``client_credentials`` with id/secret in the form body.""" + data = dict(self.data) + data["client_id"] = self.client_id + data["client_secret"] = self.client_secret + client = self.client or httpx.Client() Review Comment: `RefreshableBearerAuth` only defines the sync `auth_flow` generator (no `async_auth_flow` override), and `request_new_token()` does a blocking `httpx.Client().post()` here. httpx's default `Auth.async_auth_flow` just drives the sync `auth_flow` generator directly when a class doesn't override it — so under `AsyncClient`, every token fetch/refresh blocks the whole asyncio event loop for the IdP round-trip (including its full timeout on IdP slowness). In HTTP transport mode with multiple concurrent MCP sessions sharing this auth hook, one session's token refresh would stall every other session's in-flight tool calls. Worth overriding `async_auth_flow` to do the token POST with an async client, or at minimum documenting the tradeoff if this is intentional for the current deployment model. ########## mcp-server/mcp_server/core/oauth.py: ########## @@ -0,0 +1,131 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""httpx ``auth=`` hook for MCP → Gravitino OAuth2 client-credentials. + +Uses ``httpx-auth`` for fetch and cache on the existing ``httpx.AsyncClient``. +Credentials go in the form body (``client_secret_post``), matching the +Java/Python Gravitino clients. httpx-auth defaults to HTTP Basic. This class +retries once after Gravitino HTTP 401. +""" + +import logging +from collections.abc import Generator +from typing import Optional, Union + +import httpx +from httpx_auth import OAuth2, OAuth2ClientCredentials + +_LOG = logging.getLogger(__name__) + +# Refresh this many seconds before recorded expiry. httpx-auth default is 30. +DEFAULT_REFRESH_SKEW_SECONDS = 60 + + +class RefreshableBearerAuth(OAuth2ClientCredentials): + """httpx-auth client-credentials with form POST and one 401 retry.""" + + requires_response_body = True + + def __init__( + self, + *, + token_endpoint: str, + client_id: str, + client_secret: str, + scope: str = "", + refresh_skew_seconds: int = DEFAULT_REFRESH_SKEW_SECONDS, + client: Optional[httpx.Client] = None, + ): + """Build an ``auth=`` hook for the service hop. + + Args: + token_endpoint: Identity-provider token URL. + client_id: OAuth2 client id. + client_secret: OAuth2 client secret. + scope: Optional OAuth2 scope. + refresh_skew_seconds: httpx-auth ``early_expiry``. + client: Optional sync httpx client used only for token POSTs + (tests inject ``MockTransport`` here). + """ + kwargs = {"early_expiry": float(refresh_skew_seconds)} + if scope: + kwargs["scope"] = scope + if client is not None: + kwargs["client"] = client + super().__init__(token_endpoint, client_id, client_secret, **kwargs) + + def invalidate(self) -> None: + """Drop the cached token so the next call fetches a new one.""" + cache = OAuth2.token_cache + # TokenMemoryCache.clear() wipes every client; only drop ours. + with cache._forbid_concurrent_cache_access: # pylint: disable=protected-access + cache.tokens.pop(self.state, None) + + def _configure_client(self, client: httpx.Client) -> None: + """Do not send HTTP Basic; id and secret go in the form body.""" + client.timeout = self.timeout + + def request_new_token( + self, + ) -> Union[tuple[str, str], tuple[str, str, Union[int, str]]]: + """POST ``client_credentials`` with id/secret in the form body.""" + data = dict(self.data) + data["client_id"] = self.client_id + data["client_secret"] = self.client_secret + client = self.client or httpx.Client() + self._configure_client(client) + try: + response = client.post(self.token_url, data=data) + if response.status_code >= 400: + _LOG.error( + "OAuth token request failed: HTTP %s", response.status_code + ) + response.raise_for_status() + body = response.json() + finally: + if self.client is None: + client.close() + token = body.get(self.token_field_name) + if not token or not isinstance(token, str): + raise ValueError("OAuth token response missing access_token") + expires_in = body.get("expires_in") + _LOG.info("Fetched OAuth access token") + if expires_in in (None, ""): + return self.state, token Review Comment: When the IdP omits `expires_in`, this returns a 2-tuple `(state, token)`. IIRC httpx-auth's `TokenMemoryCache` treats a 2-tuple as "derive expiry from the JWT `exp` claim", which means it does `header, body, other = token.split(".")` on the token — an unhandled `ValueError` for any opaque/reference (non-JWT) access token, and `TokenExpiryNotProvided` if it happens to have 3 dot-separated parts but no `exp` claim. That would crash every tool call rather than surfacing a clean auth error. Notably, `clients/client-python/gravitino/auth/default_oauth2_token_provider.py`'s `DefaultOAuth2TokenProvider` already guards this exact case (checks `len(parts) != 3` and returns `None` gracefully) — might be worth mirroring that guard here, or explicitly validating/raising with a clearer message when the IdP's token isn't a JWT and no `expires_in` was given. ########## mcp-server/mcp_server/core/oauth.py: ########## @@ -0,0 +1,131 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""httpx ``auth=`` hook for MCP → Gravitino OAuth2 client-credentials. + +Uses ``httpx-auth`` for fetch and cache on the existing ``httpx.AsyncClient``. +Credentials go in the form body (``client_secret_post``), matching the +Java/Python Gravitino clients. httpx-auth defaults to HTTP Basic. This class +retries once after Gravitino HTTP 401. +""" + +import logging +from collections.abc import Generator +from typing import Optional, Union + +import httpx +from httpx_auth import OAuth2, OAuth2ClientCredentials + +_LOG = logging.getLogger(__name__) + +# Refresh this many seconds before recorded expiry. httpx-auth default is 30. +DEFAULT_REFRESH_SKEW_SECONDS = 60 + + +class RefreshableBearerAuth(OAuth2ClientCredentials): + """httpx-auth client-credentials with form POST and one 401 retry.""" + + requires_response_body = True + + def __init__( + self, + *, + token_endpoint: str, + client_id: str, + client_secret: str, + scope: str = "", + refresh_skew_seconds: int = DEFAULT_REFRESH_SKEW_SECONDS, + client: Optional[httpx.Client] = None, + ): + """Build an ``auth=`` hook for the service hop. + + Args: + token_endpoint: Identity-provider token URL. + client_id: OAuth2 client id. + client_secret: OAuth2 client secret. + scope: Optional OAuth2 scope. + refresh_skew_seconds: httpx-auth ``early_expiry``. + client: Optional sync httpx client used only for token POSTs + (tests inject ``MockTransport`` here). + """ + kwargs = {"early_expiry": float(refresh_skew_seconds)} + if scope: + kwargs["scope"] = scope + if client is not None: + kwargs["client"] = client + super().__init__(token_endpoint, client_id, client_secret, **kwargs) + + def invalidate(self) -> None: + """Drop the cached token so the next call fetches a new one.""" + cache = OAuth2.token_cache + # TokenMemoryCache.clear() wipes every client; only drop ours. + with cache._forbid_concurrent_cache_access: # pylint: disable=protected-access + cache.tokens.pop(self.state, None) + + def _configure_client(self, client: httpx.Client) -> None: + """Do not send HTTP Basic; id and secret go in the form body.""" + client.timeout = self.timeout + + def request_new_token( + self, + ) -> Union[tuple[str, str], tuple[str, str, Union[int, str]]]: + """POST ``client_credentials`` with id/secret in the form body.""" + data = dict(self.data) + data["client_id"] = self.client_id + data["client_secret"] = self.client_secret + client = self.client or httpx.Client() + self._configure_client(client) + try: + response = client.post(self.token_url, data=data) + if response.status_code >= 400: + _LOG.error( + "OAuth token request failed: HTTP %s", response.status_code + ) + response.raise_for_status() + body = response.json() + finally: + if self.client is None: + client.close() + token = body.get(self.token_field_name) + if not token or not isinstance(token, str): + raise ValueError("OAuth token response missing access_token") + expires_in = body.get("expires_in") + _LOG.info("Fetched OAuth access token") + if expires_in in (None, ""): + return self.state, token + return self.state, token, expires_in + + def auth_flow( + self, request: httpx.Request + ) -> Generator[httpx.Request, httpx.Response, None]: + """Attach a cached or freshly fetched Bearer; retry once on HTTP 401.""" + self._apply_token(request) + response = yield request + if response.status_code != 401: + return + self.invalidate() + self._apply_token(request) + yield request Review Comment: On a 401, this re-yields the same `httpx.Request` object for the retry, but the class only sets `requires_response_body = True` (line 42) — not `requires_request_body`. httpx only pre-reads/materializes the request body for auth flows when `requires_request_body` is `True`; otherwise a streaming body may already be exhausted by the first send. Not currently triggered since all Gravitino REST calls here use JSON-serializable (replayable) bodies, but it'd silently send an empty/truncated body on retry the first time this auth hook is used with a streamed/multipart body (e.g. a fileset upload). Might be worth setting `requires_request_body = True` alongside `requires_response_body` for correctness, even if it's a no-op for current callers. ########## mcp-server/mcp_server/client/factory.py: ########## @@ -40,11 +45,17 @@ def create_rest_client( authorization: Full Authorization header value forwarded verbatim (e.g. "Bearer <token>" or "Basic <base64(user:dummy)>"). Empty string means anonymous. + auth: Optional httpx auth hook. Only passed through when set so + test doubles that omit the argument keep working. Returns: Review Comment: This branch to omit `auth` looks like it can be collapsed — both `PlainRESTClientOperation.__init__` and the test double `MockOperation.__init__` already accept `auth: ... = None` as a keyword-only default, so `cls._rest_client_class(metalake_name, uri, authorization, auth=auth)` alone should behave identically whether `auth` is `None` or set. The docstring's justification ("test doubles that omit the argument keep working") doesn't hold anymore since `MockOperation` was updated in this same PR to accept `auth=None`. Minor, but it's a branch that has to be kept in sync manually as new rest-client classes are added. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
