This is an automated email from the ASF dual-hosted git repository.
jerryshao pushed a commit to branch branch-1.3
in repository https://gitbox.apache.org/repos/asf/gravitino.git
The following commit(s) were added to refs/heads/branch-1.3 by this push:
new 2fdf6ea396 [Cherry-pick to branch-1.3] [#12530] feat(mcp-server):
fetch and refresh OAuth client-credentials tokens (#12531) (#12705)
2fdf6ea396 is described below
commit 2fdf6ea3966008ea5a409abb423fc88744ffb94c
Author: github-actions[bot]
<41898282+github-actions[bot]@users.noreply.github.com>
AuthorDate: Fri Aug 28 15:48:28 2026 +0800
[Cherry-pick to branch-1.3] [#12530] feat(mcp-server): fetch and refresh
OAuth client-credentials tokens (#12531) (#12705)
**Cherry-pick Information:**
- Original commit: `32b72866ddc153a7bae167ea751f8009430e9ba8`
- Target branch: `branch-1.3`
- Status: ✅ Conflicts resolved manually
**Conflict resolution:**
- Kept the branch-1.3 `fastmcp==3.2.0` dependency and added only
`httpx-auth`.
- Regenerated `uv.lock` so it retains FastMCP 3.2.0 and the branch-1.3
package version.
- Added the OAuth fallback-policy server import without bringing in the
unrelated main-branch tool-tag import.
**Validation:**
- `uv lock --check`
- `env -u GRAVITINO_TOKEN uv run --locked python -m unittest
tests.unit.test_oauth tests.unit.test_auth_flow -v` — 61 tests passed.
- `env -u GRAVITINO_TOKEN uv run --locked pytest
tests/integration/test_oauth_refresh_e2e.py -v` — 5 tests passed.
- `uv run --locked pylint ...` — 10.00/10.
- `git diff --check`
---------
Co-authored-by: Nevin Zheng <[email protected]>
Co-authored-by: Cursor <[email protected]>
Co-authored-by: Nevin Zheng <[email protected]>
Co-authored-by: Qi Yu <[email protected]>
---
docs/gravitino-mcp-server.md | 74 +-
mcp-server/dev/run_authz_integration_test.sh | 4 +-
mcp-server/mcp_server/client/factory.py | 12 +-
.../client/plain/plain_rest_client_operation.py | 17 +-
mcp-server/mcp_server/core/context.py | 65 +-
mcp-server/mcp_server/core/oauth.py | 283 +++++++
mcp-server/mcp_server/core/setting.py | 54 +-
mcp-server/mcp_server/main.py | 63 +-
mcp-server/mcp_server/server.py | 43 +-
mcp-server/pyproject.toml | 2 +
.../tests/integration/test_oauth_refresh_e2e.py | 342 +++++++++
mcp-server/tests/unit/test_audit.py | 21 +
mcp-server/tests/unit/test_oauth.py | 843 +++++++++++++++++++++
mcp-server/tests/unit/tools/mock_operation.py | 2 +-
mcp-server/uv.lock | 16 +-
15 files changed, 1808 insertions(+), 33 deletions(-)
diff --git a/docs/gravitino-mcp-server.md b/docs/gravitino-mcp-server.md
index 50053778df..69b7c8c14b 100644
--- a/docs/gravitino-mcp-server.md
+++ b/docs/gravitino-mcp-server.md
@@ -19,7 +19,7 @@ Gravitino MCP server provides the ability to manage Gravitino
metadata for LLM.
1. Clone the code from GitHub, and change to `mcp-server` directory
2. Create virtual environment, `uv venv`
3. Install the required Python packages. `uv pip install -e .`
-4. Add Gravitino MCP server to corresponding LLM tools. Take `cursor` for
example, edit `~/.cursor/mcp.json`, use following configuration for local
Gravitino MCP server:
+4. Add Gravitino MCP server to corresponding LLM tools. Take Cursor for
example, edit `~/.cursor/mcp.json`, use following configuration for local
Gravitino MCP server:
```json
{
@@ -35,12 +35,20 @@ Gravitino MCP server provides the ability to manage
Gravitino metadata for LLM.
"test",
"--gravitino-uri",
"http://127.0.0.1:8090";
- ]
+ ],
+ "env": {
+ "GRAVITINO_OAUTH_TOKEN_ENDPOINT":
"https://idp.example/realms/gravitino/protocol/openid-connect/token";,
+ "GRAVITINO_OAUTH_CLIENT_ID": "mcp-service",
+ "GRAVITINO_OAUTH_CLIENT_SECRET": "<secret>",
+ "GRAVITINO_OAUTH_SCOPE": "gravitino"
+ }
}
}
}
```
+In Cursor stdio mode the MCP process typically receives no `Authorization`
header from the client. Set the `GRAVITINO_OAUTH_*` variables (or the matching
CLI flags) so MCP fetches a service token with the `client_credentials` grant.
Omit `env` to run anonymously, or use `--token` / `GRAVITINO_TOKEN` for a
static credential instead.
+
Or start an HTTP MCP server by `uv run mcp_server --metalake test
--gravitino-uri http://127.0.0.1:8090 --transport http --mcp-url
http://localhost:8000/mcp`, and use the configuration:
```json
@@ -103,19 +111,24 @@ Gravitino MCP server supports the following tools, and
you could export tool by
You could config Gravitino MCP server by arguments, `uv run mcp_server -h`
shows the detailed information.
-| Argument | Description
| Default value | Required |
-|-------------------|----------------------------------------------------------------------------------|-----------------------------|----------|
-| `--metalake` | The Gravitino metalake name.
| none | Yes |
-| `--gravitino-uri` | The URI of Gravitino server.
| `http://127.0.0.1:8090` | No |
-| `--transport` | Transport protocol: stdio (local), http /
streamable-http (Streamable HTTP). | `stdio` | No
|
-| `--mcp-url` | The URL of MCP server if using HTTP transport.
| `http://127.0.0.1:8000/mcp` | No |
-| `--token` | Static credential for Gravitino; or set
`GRAVITINO_TOKEN`. See Authentication. | none (anonymous) | No
|
-| `--tls-cert` | PEM certificate to serve the endpoint over HTTPS.
Requires `--tls-key`. | none | No |
-| `--tls-key` | PEM private key to serve the endpoint over HTTPS.
Requires `--tls-cert`. | none | No |
+| Argument | Description
| Default value | Required |
+|----------------------------------|---------------------------------------------------------------------------------------------------------------------------------|-----------------------------|----------|
+| `--metalake` | The Gravitino metalake name.
| none | Yes |
+| `--gravitino-uri` | The URI of Gravitino server.
| `http://127.0.0.1:8090` | No |
+| `--transport` | Transport protocol: stdio (local), http /
streamable-http (Streamable HTTP).
| `stdio` | No |
+| `--mcp-url` | The URL of MCP server if using HTTP
transport.
| `http://127.0.0.1:8000/mcp` | No |
+| `--token` | Static credential for Gravitino; or set
`GRAVITINO_TOKEN`. See Authentication. Wins over OAuth client-credentials.
| none (anonymous) | No |
+| `--oauth-token-endpoint` | OAuth2 token URL for client-credentials.
Or `GRAVITINO_OAUTH_TOKEN_ENDPOINT`.
| none | No |
+| `--oauth-client-id` | OAuth2 client id. Or
`GRAVITINO_OAUTH_CLIENT_ID`.
| none | No |
+| `--oauth-client-secret` | OAuth2 client secret. Or
`GRAVITINO_OAUTH_CLIENT_SECRET`.
| none | No |
+| `--oauth-scope` | Optional OAuth2 scope. Or
`GRAVITINO_OAUTH_SCOPE`.
| none | No |
+| `--no-service-identity-fallback` | HTTP only: reject requests with no
`Authorization` when OAuth or `--token` is set. Or
`GRAVITINO_NO_SERVICE_IDENTITY_FALLBACK`. | `false` | No
|
+| `--tls-cert` | PEM certificate to serve the endpoint
over HTTPS. Requires `--tls-key`.
| none | No |
+| `--tls-key` | PEM private key to serve the endpoint
over HTTPS. Requires `--tls-cert`.
| none | No |
## Authentication
-By default the MCP server talks to Gravitino anonymously. There are two ways
to attach an identity, depending on the transport.
+By default the MCP server talks to Gravitino anonymously. There are three ways
to authenticate MCP when calling Gravitino.
### Static startup token (stdio and HTTP)
@@ -144,11 +157,40 @@ export GRAVITINO_TOKEN=<your-token>
uv run mcp_server --metalake test --gravitino-uri http://127.0.0.1:8090
```
-In `stdio` mode this token is used for every request. In HTTP mode it is only
the fallback, used when an incoming request does not carry its own
`Authorization` header.
+In `stdio` mode this token is used for every request. In HTTP mode it is only
the fallback, used when an incoming request does not carry its own
`Authorization` header. If both `--token` and OAuth client-credentials are set,
`--token` wins.
+
+### OAuth client credentials (service identity)
+
+When Gravitino uses `gravitino.authenticators = oauth`, a pasted Bearer access
token in `--token` expires and is not refreshed. For the **service** identity
(Cursor stdio, or HTTP when the caller sends no `Authorization` header),
configure MCP as an OAuth client of the same identity provider Gravitino trusts.
+
+Set `--oauth-token-endpoint`, `--oauth-client-id`, and `--oauth-client-secret`
together, plus optional `--oauth-scope` (or the matching `GRAVITINO_OAUTH_*`
environment variables). MCP requests an access token with the
`client_credentials` grant, caches it, refreshes before expiry, and retries
once on HTTP 401.
+
+In Cursor, put the `GRAVITINO_OAUTH_*` values in the `env` block of
`~/.cursor/mcp.json` (see [Usage](#usage)). `--token` / `GRAVITINO_TOKEN`
overrides OAuth client-credentials and stays static (no refresh). An incoming
HTTP `Authorization` header is forwarded as-is and is not refreshed by MCP.
+
+Gravitino maps the JWT to a metalake principal from claims configured in
[`gravitino.authenticator.oauth.principalFields`](./security/how-to-authenticate.md#server-configuration)
(often `sub`); that principal may differ from `--oauth-client-id`. It must
exist as a metalake user with the needed grants, or tool calls fail with 403.
+
+Prefer environment variables (or the `env` block in `~/.cursor/mcp.json`) for
the client secret so it does not appear in `ps` output or shell history:
+
+```shell
+export
GRAVITINO_OAUTH_TOKEN_ENDPOINT=https://idp.example/realms/gravitino/protocol/openid-connect/token
+export GRAVITINO_OAUTH_CLIENT_ID=mcp-service
+export GRAVITINO_OAUTH_CLIENT_SECRET=<secret>
+export GRAVITINO_OAUTH_SCOPE=gravitino
+
+uv run mcp_server --metalake test --gravitino-uri http://127.0.0.1:8090
+```
+
+The matching CLI flags (`--oauth-token-endpoint`, `--oauth-client-id`,
`--oauth-client-secret`, `--oauth-scope`) work the same way, but avoid passing
`--oauth-client-secret` on the command line in production.
+
+This path does not replace per-request user identity in HTTP mode (see below).
### Per-request identity (HTTP)
-When the server runs with HTTP transport, the `Authorization` header of each
incoming MCP request is forwarded verbatim to Gravitino. The scheme is
preserved, so OAuth2 (`Bearer`), Gravitino simple authentication (`Basic
<base64(user:dummy)>`) and others all work. This keeps concurrent sessions from
different principals isolated — one principal's identity never leaks into
another's calls — and lets Gravitino enforce authorization per caller. The
per-request header takes priority over the [...]
+When the server runs with HTTP transport, the `Authorization` header of each
incoming MCP request is forwarded verbatim to Gravitino. The scheme is
preserved, so OAuth2 (`Bearer`), Gravitino simple authentication (`Basic
<base64(user:dummy)>`) and others all work. This keeps concurrent sessions from
different principals isolated — one principal's identity never leaks into
another's calls — and lets Gravitino enforce authorization per caller. The
per-request header takes priority over the [...]
+
+**Security warning:** When OAuth client-credentials or `--token` is
configured, an HTTP request with **no** `Authorization` header is authenticated
as the **service identity**, not as anonymous. With OAuth, that identity
refreshes automatically and stays valid for a long time. If the MCP HTTP
endpoint is reachable by more than one caller, anyone who omits `Authorization`
receives the service principal's full permissions. Use stdio transport for
single-user integrations (for example Curso [...]
+
+For exposed or multi-caller HTTP deployments, set
`--no-service-identity-fallback` (or
`GRAVITINO_NO_SERVICE_IDENTITY_FALLBACK=1`) so requests without `Authorization`
are rejected instead of using the service identity. The flag is ignored for
stdio transport.
Authorization itself is always enforced by Gravitino: the MCP server forwards
the identity but does not make access-control decisions of its own.
@@ -164,12 +206,12 @@ uv run mcp_server --metalake test --gravitino-uri
http://127.0.0.1:8090 \
## Audit Logging
-Every tool invocation is recorded as one structured JSON line in
`gravitino-mcp-audit.log` (written to the server's working directory). Each
record is attributed to the principal derived from the request's
`Authorization` header.
+Every tool invocation is recorded as one structured JSON line in
`gravitino-mcp-audit.log` (written to the server's working directory). Each
record is attributed to the incoming HTTP `Authorization` header when present;
otherwise to the configured service identity (`--token` or OAuth client id).
| Field | Description
|
|--------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| `timestamp` | UTC ISO-8601 time of the call.
|
-| `principal` | Caller identity: username for `Basic` simple auth,
`bearer:<first-8-chars>` for a Bearer token, or `anonymous` when no identity is
present. |
+| `principal` | Caller identity: username for `Basic` simple auth,
`bearer:<first-8-chars>` for a Bearer token,
`oauth:<first-8-chars-of-client-id>` when OAuth client-credentials is the
service identity (stdio or HTTP with no caller header), or `anonymous` when no
identity is present. |
| `tool` | Name of the invoked MCP tool.
|
| `outcome` | `allow` for successful calls, `deny` for failed ones. `deny`
is emitted for any tool-call exception (authorization denial being the common
case); inspect `error_type` to disambiguate. |
| `error_type` | Exception class name, present only when `outcome` is `deny`.
|
diff --git a/mcp-server/dev/run_authz_integration_test.sh
b/mcp-server/dev/run_authz_integration_test.sh
index 82673806b7..b0ca02faf5 100755
--- a/mcp-server/dev/run_authz_integration_test.sh
+++ b/mcp-server/dev/run_authz_integration_test.sh
@@ -24,8 +24,8 @@
# 2. Enables simple authentication + authorization (serviceAdmins=admin).
# 3. Starts the Gravitino server.
# 4. Starts the MCP server in HTTP transport mode.
-# 5. Runs the pytest integration suite (which provisions metadata as admin
and
-# verifies per-user authorization through MCP).
+# 5. Runs the pytest integration suite (authz e2e against the live Gravitino
+# plus the self-contained OAuth refresh e2e that boots mock
IdP/Gravitino).
# 6. Tears everything down and restores the original config.
#
# Usage:
diff --git a/mcp-server/mcp_server/client/factory.py
b/mcp-server/mcp_server/client/factory.py
index 201625540d..068bdf1b74 100644
--- a/mcp-server/mcp_server/client/factory.py
+++ b/mcp-server/mcp_server/client/factory.py
@@ -29,7 +29,12 @@ class RESTClientFactory:
@classmethod
def create_rest_client(
- cls, metalake_name: str, uri: str, authorization: str = ""
+ cls,
+ metalake_name: str,
+ uri: str,
+ authorization: str = "",
+ *,
+ auth=None,
) -> "PlainRESTClientOperation":
"""
Create a new rest client instance with the specified parameters.
@@ -40,11 +45,14 @@ class RESTClientFactory:
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. ``None`` for static or anonymous.
Returns:
New instance of the configured rest client class
"""
- return cls._rest_client_class(metalake_name, uri, authorization)
+ return cls._rest_client_class(
+ metalake_name, uri, authorization, auth=auth
+ )
@classmethod
def set_rest_client(cls, rest_client_class: type) -> None:
diff --git a/mcp-server/mcp_server/client/plain/plain_rest_client_operation.py
b/mcp-server/mcp_server/client/plain/plain_rest_client_operation.py
index a6adc0e099..cf225dfa2b 100644
--- a/mcp-server/mcp_server/client/plain/plain_rest_client_operation.py
+++ b/mcp-server/mcp_server/client/plain/plain_rest_client_operation.py
@@ -15,6 +15,8 @@
# specific language governing permissions and limitations
# under the License.
+from typing import Optional
+
import httpx
from mcp_server.client import (
@@ -62,7 +64,14 @@ from mcp_server.client.topic_operation import TopicOperation
# pylint: disable=too-many-instance-attributes
class PlainRESTClientOperation(GravitinoOperation):
- def __init__(self, metalake_name: str, uri: str, authorization: str = ""):
+ def __init__(
+ self,
+ metalake_name: str,
+ uri: str,
+ authorization: str = "",
+ *,
+ auth: Optional[httpx.Auth] = None,
+ ):
"""Create a REST client for one identity.
Args:
@@ -72,11 +81,15 @@ class PlainRESTClientOperation(GravitinoOperation):
on every request (for example ``"Bearer <token>"`` for OAuth2
or
``"Basic <base64(user:secret)>"`` for simple or Basic auth).
Empty string means anonymous (no header sent).
+ auth: Optional httpx auth hook used instead of a frozen header
+ (OAuth client-credentials refresh).
"""
headers = {}
if authorization:
headers["Authorization"] = authorization
- _rest_client = httpx.AsyncClient(base_url=uri, headers=headers)
+ _rest_client = httpx.AsyncClient(
+ base_url=uri, headers=headers, auth=auth
+ )
# Kept so the shared connection pool can be closed (see close()).
self._rest_client = _rest_client
self._catalog_operation = PlainRESTClientCatalogOperation(
diff --git a/mcp-server/mcp_server/core/context.py
b/mcp-server/mcp_server/core/context.py
index 79cff1bb6d..f38831f639 100644
--- a/mcp-server/mcp_server/core/context.py
+++ b/mcp-server/mcp_server/core/context.py
@@ -21,6 +21,7 @@ import re
from collections import OrderedDict
from mcp_server.client.factory import RESTClientFactory
+from mcp_server.core.oauth import RefreshableBearerAuth
from mcp_server.core.setting import Setting
_LOG = logging.getLogger(__name__)
@@ -50,6 +51,22 @@ _CANONICAL_AUTH_SCHEMES = {
}
+class ServiceIdentityFallbackDisabled(RuntimeError):
+ """HTTP omitted Authorization while service-identity fallback is
disabled."""
+
+
+def _in_http_request() -> bool:
+ """Return True when ``rest_client()`` runs inside an active HTTP
request."""
+ try:
+ # pylint: disable=import-outside-toplevel
+ from fastmcp.server.dependencies import get_http_request
+
+ get_http_request()
+ return True
+ except (LookupError, RuntimeError):
+ return False
+
+
def _get_request_authorization() -> str:
"""Return the raw ``Authorization`` header of the current HTTP request.
@@ -93,6 +110,41 @@ def startup_authorization(setting: Setting) -> str:
return f"Bearer {token}"
+def service_fallback_authorization(setting: Setting) -> str:
+ """Audit / fallback identity when no hop-1 Authorization header is present.
+
+ Prefers the static ``--token``. When only OAuth client-credentials is
+ configured, returns ``OAuth <client_id>`` so audit logs can attribute
+ stdio / no-header calls to the service client.
+ """
+ static = startup_authorization(setting)
+ if static:
+ return static
+ if setting.has_oauth_client():
+ return f"OAuth {setting.oauth_client_id.strip()}"
+ return ""
+
+
+def _service_auth(setting: Setting):
+ """httpx ``auth=`` hook for service OAuth, or None for static/anonymous."""
+ setting.validate_oauth()
+ static = startup_authorization(setting)
+ if static:
+ if setting.has_oauth_client():
+ _LOG.warning(
+ "Ignoring OAuth client credentials because --token is set"
+ )
+ return None
+ if not setting.has_oauth_client():
+ return None
+ return RefreshableBearerAuth(
+ token_endpoint=setting.oauth_token_endpoint.strip(),
+ client_id=setting.oauth_client_id.strip(),
+ client_secret=setting.oauth_client_secret.strip(),
+ scope=setting.oauth_scope.strip(),
+ )
+
+
class GravitinoContext:
def __init__(self, setting: Setting):
self._setting = setting
@@ -100,6 +152,7 @@ class GravitinoContext:
setting.metalake,
setting.gravitino_uri,
startup_authorization(setting),
+ auth=_service_auth(setting),
)
# LRU cache of per-principal clients keyed by the raw Authorization
header.
# Safe without locking: rest_client() runs on the single asyncio event
@@ -118,7 +171,8 @@ class GravitinoContext:
token. This keeps concurrent sessions with different principals fully
isolated — one principal's identity never leaks into another's calls.
- Falls back to the shared default client (static startup token) when:
+ Falls back to the shared default client (static token or OAuth
+ client-credentials) when:
- running in stdio mode (no HTTP request context), or
- the incoming request carries no Authorization header.
@@ -127,6 +181,15 @@ class GravitinoContext:
"""
authorization = _get_request_authorization()
if not authorization:
+ if (
+ self._setting.no_service_identity_fallback
+ and _in_http_request()
+ and self._setting.has_service_identity()
+ ):
+ raise ServiceIdentityFallbackDisabled(
+ "HTTP request omitted Authorization and "
+ "--no-service-identity-fallback is set"
+ )
return self._default_client
cached = self._clients_by_auth.get(authorization)
diff --git a/mcp-server/mcp_server/core/oauth.py
b/mcp-server/mcp_server/core/oauth.py
new file mode 100644
index 0000000000..17e1922a45
--- /dev/null
+++ b/mcp-server/mcp_server/core/oauth.py
@@ -0,0 +1,283 @@
+# 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 asyncio
+import base64
+import json
+import logging
+import threading
+import time
+from collections.abc import AsyncGenerator, Generator
+from typing import Optional, Union
+
+import httpx
+from httpx_auth import AuthenticationFailed, OAuth2, OAuth2ClientCredentials
+
+_LOG = logging.getLogger(__name__)
+
+# Refresh this many seconds before recorded expiry. httpx-auth default is 30.
+DEFAULT_REFRESH_SKEW_SECONDS = 60
+
+_TokenTuple = Union[tuple[str, str], tuple[str, str, Union[int, str]]]
+
+
+class RefreshableBearerAuth(OAuth2ClientCredentials):
+ """httpx-auth client-credentials with form POST and one 401 retry."""
+
+ requires_request_body = True
+ 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)
+ self._token_lock = asyncio.Lock()
+ self._sync_lock = threading.Lock()
+ self._rejected_tokens: set[str] = set()
+ self._retried_tokens: set[str] = set()
+
+ 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 request_new_token(self) -> _TokenTuple:
+ """POST ``client_credentials`` with id/secret in the form body."""
+ data = self._token_form_data()
+ client = self.client or httpx.Client()
+ self._configure_client(client)
+ try:
+ response = client.post(self.token_url, data=data)
+ self._log_token_http_error(response)
+ response.raise_for_status()
+ body = response.json()
+ finally:
+ if self.client is None:
+ client.close()
+ return self._token_tuple(body)
+
+ async def request_new_token_async(self) -> _TokenTuple:
+ """POST ``client_credentials`` without blocking the event loop."""
+ if self.client is not None:
+ return await asyncio.to_thread(self.request_new_token)
+ data = self._token_form_data()
+ async with httpx.AsyncClient() as client:
+ client.timeout = self.timeout
+ response = await client.post(self.token_url, data=data)
+ self._log_token_http_error(response)
+ response.raise_for_status()
+ body = response.json()
+ return self._token_tuple(body)
+
+ 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."""
+ if self.requires_request_body:
+ request.read()
+ token, fetched = self._apply_token(request)
+ response = yield request
+ if response.status_code != 401:
+ return
+ with self._sync_lock:
+ if not self._begin_401_retry(token, fetched):
+ return
+ self._invalidate_if_still_cached(token)
+ retry_token, _ = self._apply_token(request)
+ response = yield request
+ if response.status_code == 401:
+ self._rejected_tokens = {retry_token}
+
+ async def async_auth_flow(
+ self, request: httpx.Request
+ ) -> AsyncGenerator[httpx.Request, httpx.Response]:
+ """Attach a Bearer without a blocking IdP POST on the event loop."""
+ if self.requires_request_body:
+ await request.aread()
+ token, fetched = await self._apply_token_async(request)
+ response = yield request
+ if response.status_code != 401:
+ return
+ async with self._token_lock:
+ if not self._begin_401_retry(token, fetched):
+ return
+ self._invalidate_if_still_cached(token)
+ retry_token, _ = await self._apply_token_async(request)
+ response = yield request
+ if response.status_code == 401:
+ self._rejected_tokens = {retry_token}
+
+ 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 _token_form_data(self) -> dict:
+ data = dict(self.data)
+ data["client_id"] = self.client_id
+ data["client_secret"] = self.client_secret
+ return data
+
+ def _token_tuple(self, body: dict) -> _TokenTuple:
+ 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:
+ return self.state, token, expires_in
+ jwt_expires_in = self._jwt_expires_in(token)
+ if jwt_expires_in is None:
+ raise ValueError(
+ "OAuth token response omitted expires_in and "
+ "access_token is not a JWT with exp"
+ )
+ return self.state, token, jwt_expires_in
+
+ @staticmethod
+ def _jwt_expires_in(token: str) -> Optional[int]:
+ """Return seconds until JWT ``exp``, or None when not cacheable."""
+ parts = token.split(".")
+ if len(parts) != 3:
+ return None
+ try:
+ padded = parts[1] + "=" * (-len(parts[1]) % 4)
+ payload = json.loads(base64.urlsafe_b64decode(padded))
+ exp = payload.get("exp")
+ if not RefreshableBearerAuth._is_numeric_date(exp):
+ return None
+ expires_at = int(float(exp))
+ return max(expires_at - int(time.time()), 1)
+ except (ValueError, json.JSONDecodeError, TypeError):
+ return None
+
+ @staticmethod
+ def _is_numeric_date(value) -> bool:
+ """Return True when value is a JWT NumericDate (RFC 7519)."""
+ if isinstance(value, bool):
+ return False
+ if isinstance(value, (int, float)):
+ return True
+ if isinstance(value, str):
+ try:
+ float(value)
+ return True
+ except ValueError:
+ return False
+ return False
+
+ @staticmethod
+ def _log_token_http_error(response: httpx.Response) -> None:
+ if response.status_code >= 400:
+ _LOG.error(
+ "OAuth token request failed: HTTP %s", response.status_code
+ )
+
+ def _cached_bearer(self) -> Optional[str]:
+ try:
+ return OAuth2.token_cache.get_token(
+ self.state, early_expiry=self.early_expiry
+ )
+ except AuthenticationFailed:
+ return None
+
+ def _store_and_get(self, fetched: _TokenTuple) -> str:
+ return OAuth2.token_cache.get_token(
+ self.state,
+ early_expiry=self.early_expiry,
+ on_missing_token=lambda: fetched,
+ )
+
+ def _invalidate_if_still_cached(self, token: str) -> None:
+ """Drop the cache entry only when it still holds the rejected token."""
+ if self._cached_bearer() == token:
+ self.invalidate()
+
+ def _begin_401_retry(self, token: str, fetched: bool) -> bool:
+ """Reserve a single 401 retry for this token; skip hopeless cases.
+
+ Sets hold only the current token. Historical JWTs are not retained.
+ """
+ if fetched:
+ self._rejected_tokens = {token}
+ return False
+ if token in self._rejected_tokens or token in self._retried_tokens:
+ return False
+ self._retried_tokens = {token}
+ return True
+
+ def _apply_token(self, request: httpx.Request) -> tuple[str, bool]:
+ fetched = False
+ token = self._cached_bearer()
+ if token is None:
+ token = OAuth2.token_cache.get_token(
+ self.state,
+ early_expiry=self.early_expiry,
+ on_missing_token=self.request_new_token,
+ on_expired_token=self.refresh_token,
+ )
+ fetched = True
+ self._update_user_request(request, token)
+ return token, fetched
+
+ async def _apply_token_async(
+ self, request: httpx.Request
+ ) -> tuple[str, bool]:
+ fetched = False
+ token = self._cached_bearer()
+ if token is None:
+ async with self._token_lock:
+ token = self._cached_bearer()
+ if token is None:
+ token = self._store_and_get(
+ await self.request_new_token_async()
+ )
+ fetched = True
+ self._update_user_request(request, token)
+ return token, fetched
diff --git a/mcp-server/mcp_server/core/setting.py
b/mcp-server/mcp_server/core/setting.py
index f0be1b6a6d..112528f6d9 100644
--- a/mcp-server/mcp_server/core/setting.py
+++ b/mcp-server/mcp_server/core/setting.py
@@ -45,14 +45,66 @@ class Setting: # pylint:
disable=too-many-instance-attributes
# Both must be set to enable TLS; empty means plain HTTP.
tls_cert: str = ""
tls_key: str = ""
+ # OAuth2 client-credentials for the service hop to Gravitino. All three of
+ # endpoint / client id / client secret must be set together. Ignored when
+ # ``token`` is also set (--token wins).
+ oauth_token_endpoint: str = ""
+ oauth_client_id: str = ""
+ oauth_client_secret: str = field(default="", repr=False)
+ oauth_scope: str = ""
+ # HTTP only: reject requests with no Authorization when service OAuth or
+ # --token is configured instead of falling back to the service identity.
+ no_service_identity_fallback: bool = False
+
+ def has_oauth_client(self) -> bool:
+ """Return True when client-credentials is fully configured."""
+ return bool(
+ self.oauth_token_endpoint.strip()
+ and self.oauth_client_id.strip()
+ and self.oauth_client_secret.strip()
+ )
+
+ def has_service_identity(self) -> bool:
+ """Return True when a static token or OAuth client-credentials is
set."""
+ return bool(self.token.strip()) or self.has_oauth_client()
+
+ def validate_oauth(self) -> None:
+ """Reject a partial OAuth client-credentials configuration."""
+ filled = [
+ bool(self.oauth_token_endpoint.strip()),
+ bool(self.oauth_client_id.strip()),
+ bool(self.oauth_client_secret.strip()),
+ ]
+ if any(filled) and not all(filled):
+ raise ValueError(
+ "OAuth client credentials requires "
+ "--oauth-token-endpoint, --oauth-client-id, and "
+ "--oauth-client-secret together "
+ "(or GRAVITINO_OAUTH_TOKEN_ENDPOINT / "
+ "GRAVITINO_OAUTH_CLIENT_ID / GRAVITINO_OAUTH_CLIENT_SECRET)."
+ )
+ if self.oauth_scope.strip() and not all(filled):
+ raise ValueError(
+ "OAuth scope (--oauth-scope / GRAVITINO_OAUTH_SCOPE) requires "
+ "a complete client-credentials configuration "
+ "(--oauth-token-endpoint, --oauth-client-id, and "
+ "--oauth-client-secret)."
+ )
def __str__(self) -> str:
# Mirror startup_authorization: a whitespace-only token is anonymous on
# the wire, so it must not be logged as a configured identity.
token_display = "***" if self.token.strip() else ""
+ secret_display = "***" if self.oauth_client_secret.strip() else ""
return (
f"Setting(metalake={self.metalake},
gravitino_uri={self.gravitino_uri}, "
f"tags={self.tags}, transport={self.transport},
mcp_url={self.mcp_url}, "
f"token={token_display}, tls_cert={self.tls_cert}, "
- f"tls_key={self.tls_key})"
+ f"tls_key={self.tls_key}, "
+ f"oauth_token_endpoint={self.oauth_token_endpoint}, "
+ f"oauth_client_id={self.oauth_client_id}, "
+ f"oauth_client_secret={secret_display}, "
+ f"oauth_scope={self.oauth_scope}, "
+ f"no_service_identity_fallback="
+ f"{self.no_service_identity_fallback})"
)
diff --git a/mcp-server/mcp_server/main.py b/mcp-server/mcp_server/main.py
index f3162b802c..5ab8b1c777 100644
--- a/mcp-server/mcp_server/main.py
+++ b/mcp-server/mcp_server/main.py
@@ -20,7 +20,14 @@ import logging
import os
from mcp_server.core.setting import DefaultSetting, Setting
-from mcp_server.server import GravitinoMCPServer
+from mcp_server.server import (
+ GravitinoMCPServer,
+ log_service_identity_fallback_policy,
+)
+
+
+def _env_truthy(name: str) -> bool:
+ return os.environ.get(name, "").strip().lower() in ("1", "true", "yes")
def do_main():
@@ -34,8 +41,19 @@ def do_main():
token=args.token,
tls_cert=args.tls_cert,
tls_key=args.tls_key,
+ oauth_token_endpoint=args.oauth_token_endpoint,
+ oauth_client_id=args.oauth_client_id,
+ oauth_client_secret=args.oauth_client_secret,
+ oauth_scope=args.oauth_scope,
+ no_service_identity_fallback=args.no_service_identity_fallback,
)
_init_logging(setting)
+ try:
+ setting.validate_oauth()
+ except ValueError as exc:
+ logging.error("%s", exc)
+ raise SystemExit(1) from None
+ log_service_identity_fallback_policy(setting)
logging.info("Gravitino MCP server setting: %s", setting)
server = GravitinoMCPServer(setting)
server.run()
@@ -120,7 +138,48 @@ def _parse_args():
"incoming request carries no Authorization header "
"(per-request identity takes priority). "
"Can also be set via the GRAVITINO_TOKEN environment variable. "
- "When omitted, requests are sent without authentication.",
+ "When omitted, requests are sent without authentication. "
+ "Takes precedence over OAuth client-credentials flags.",
+ )
+
+ parser.add_argument(
+ "--oauth-token-endpoint",
+ type=str,
+ default=os.environ.get("GRAVITINO_OAUTH_TOKEN_ENDPOINT", ""),
+ help="OAuth2 token endpoint for client-credentials (service identity).
"
+ "Requires --oauth-client-id and --oauth-client-secret. "
+ "Can also be set via GRAVITINO_OAUTH_TOKEN_ENDPOINT.",
+ )
+ parser.add_argument(
+ "--oauth-client-id",
+ type=str,
+ default=os.environ.get("GRAVITINO_OAUTH_CLIENT_ID", ""),
+ help="OAuth2 client id for client-credentials. "
+ "Can also be set via GRAVITINO_OAUTH_CLIENT_ID.",
+ )
+ parser.add_argument(
+ "--oauth-client-secret",
+ type=str,
+ default=os.environ.get("GRAVITINO_OAUTH_CLIENT_SECRET", ""),
+ help="OAuth2 client secret for client-credentials. "
+ "Can also be set via GRAVITINO_OAUTH_CLIENT_SECRET.",
+ )
+ parser.add_argument(
+ "--oauth-scope",
+ type=str,
+ default=os.environ.get("GRAVITINO_OAUTH_SCOPE", ""),
+ help="Optional OAuth2 scope for client-credentials. "
+ "Can also be set via GRAVITINO_OAUTH_SCOPE.",
+ )
+
+ parser.add_argument(
+ "--no-service-identity-fallback",
+ action="store_true",
+ default=_env_truthy("GRAVITINO_NO_SERVICE_IDENTITY_FALLBACK"),
+ help="HTTP only: reject incoming requests that omit Authorization "
+ "when OAuth client-credentials or --token is configured, instead of "
+ "using the service identity. Ignored for stdio transport. Can also "
+ "be enabled via GRAVITINO_NO_SERVICE_IDENTITY_FALLBACK.",
)
parser.add_argument(
diff --git a/mcp-server/mcp_server/server.py b/mcp-server/mcp_server/server.py
index b8909120da..aaad4c8f44 100644
--- a/mcp-server/mcp_server/server.py
+++ b/mcp-server/mcp_server/server.py
@@ -39,7 +39,7 @@ from mcp_server.core import audit
from mcp_server.core.context import (
GravitinoContext,
_get_request_authorization,
- startup_authorization,
+ service_fallback_authorization,
)
from mcp_server.core.setting import Setting
from mcp_server.tools import load_tools
@@ -49,9 +49,10 @@ def _get_principal_from_request(fallback_authorization: str
= "") -> str:
"""Derive a display principal for audit attribution.
Uses the incoming HTTP request's Authorization header when present;
- otherwise falls back to the static startup identity (``--token``), which is
- what actually authenticates the call in stdio mode or in HTTP requests that
- carry no Authorization header. Returns "anonymous" when neither is set.
+ otherwise falls back to the service identity (``--token`` or OAuth
+ client-credentials), which is what actually authenticates the call in
+ stdio mode or in HTTP requests that carry no Authorization header.
+ Returns "anonymous" when neither is set.
"""
authorization = _get_request_authorization() or fallback_authorization
# pylint: disable=protected-access
@@ -109,7 +110,7 @@ def _create_gravitino_mcp(setting: Setting) -> FastMCP:
lifespan=_create_lifespan_manager(GravitinoContext(setting)),
)
- mcp.add_middleware(AuditMiddleware(startup_authorization(setting)))
+
mcp.add_middleware(AuditMiddleware(service_fallback_authorization(setting)))
mcp.add_middleware(
LoggingMiddleware(include_payloads=True, max_payload_length=1000)
)
@@ -123,6 +124,38 @@ def _create_gravitino_mcp(setting: Setting) -> FastMCP:
return mcp
+def log_service_identity_fallback_policy(setting: Setting) -> None:
+ """Log how HTTP requests without Authorization are handled at startup."""
+ if setting.transport == "stdio":
+ if setting.no_service_identity_fallback:
+ logging.info(
+ "--no-service-identity-fallback is ignored for stdio transport"
+ )
+ return
+ if not setting.has_service_identity():
+ return
+ if setting.no_service_identity_fallback:
+ logging.info(
+ "HTTP requests without an Authorization header will be "
+ "rejected (--no-service-identity-fallback)"
+ )
+ return
+ if setting.token.strip():
+ logging.warning(
+ "HTTP requests without an Authorization header will use the "
+ "configured --token service identity. Do not expose this "
+ "endpoint to untrusted callers."
+ )
+ return
+ if setting.has_oauth_client():
+ logging.warning(
+ "HTTP requests without an Authorization header will use the "
+ "OAuth client-credentials service identity (%s). Do not expose "
+ "this endpoint to untrusted callers.",
+ setting.oauth_client_id.strip(),
+ )
+
+
def _parse_mcp_url(url: str) -> tuple[str, int, str]:
try:
parsed = urlparse(url)
diff --git a/mcp-server/pyproject.toml b/mcp-server/pyproject.toml
index 958bada491..d08e58c37b 100644
--- a/mcp-server/pyproject.toml
+++ b/mcp-server/pyproject.toml
@@ -24,6 +24,8 @@ requires-python = ">=3.10"
dependencies = [
# Pin FastMCP so breaking API changes are handled explicitly during
dependency upgrades.
"fastmcp==3.2.0",
+ # httpx.Auth plugin for hop-2 client_credentials fetch/cache.
+ "httpx-auth>=0.22,<0.24",
# TODO(#10754): pydocket (a transitive dep of fastmcp>=2.14.0)
# requires fakeredis>=2.32.1, but fakeredis 2.35.0 removed FakeConnection
from
# fakeredis.aioredis, which pydocket still uses. Pin fakeredis to <2.35.0
until the
diff --git a/mcp-server/tests/integration/test_oauth_refresh_e2e.py
b/mcp-server/tests/integration/test_oauth_refresh_e2e.py
new file mode 100644
index 0000000000..46ebb6f918
--- /dev/null
+++ b/mcp-server/tests/integration/test_oauth_refresh_e2e.py
@@ -0,0 +1,342 @@
+# 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.
+
+"""Process-level IT: live MCP + mock IdP + mock Gravitino.
+
+Helm does not ship an MCP chart in this repo. This is the integration layer
+that exists today: real HTTP, real MCP process, no frozen ``--token``.
+
+Proves the service hop:
+ MCP --client_credentials--> mock IdP
+ MCP --Authorization Bearer--> mock Gravitino
+without a hop-1 user header.
+"""
+
+# do_GET/do_POST are BaseHTTPRequestHandler names. pytest fixtures share
+# names with test parameters.
+# pylint: disable=invalid-name,redefined-outer-name
+
+import asyncio
+import json
+import os
+import socket
+import subprocess
+import sys
+import threading
+import time
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from typing import List, Optional
+from urllib.parse import parse_qs
+
+import pytest
+from fastmcp import Client
+from fastmcp.client.transports import StreamableHttpTransport
+
+METALAKE = "oauth_it"
+CATALOG = "it_catalog"
+
+
+class _IdPState:
+ def __init__(self, expires_in: int):
+ self.expires_in = expires_in
+ self.hits = 0
+ self.bodies: List[str] = []
+ self.lock = threading.Lock()
+
+
+class _GravitinoState:
+ def __init__(self):
+ self.authorizations: List[str] = []
+ self.fail_first = False
+ self.lock = threading.Lock()
+
+
+def _start_server(handler_cls, state) -> ThreadingHTTPServer:
+ server = ThreadingHTTPServer(("127.0.0.1", 0), handler_cls)
+ server.state = state
+ thread = threading.Thread(target=server.serve_forever, daemon=True)
+ thread.start()
+ return server
+
+
+def _idp_handler(state: _IdPState):
+ class Handler(BaseHTTPRequestHandler):
+ def log_message(self, *_args):
+ return
+
+ def do_POST(self):
+ length = int(self.headers.get("Content-Length", "0"))
+ body = self.rfile.read(length).decode("utf-8")
+ with state.lock:
+ state.hits += 1
+ state.bodies.append(body)
+ token = f"tok-{state.hits}"
+ expires_in = state.expires_in
+ payload = json.dumps(
+ {
+ "access_token": token,
+ "token_type": "bearer",
+ "expires_in": expires_in,
+ }
+ ).encode("utf-8")
+ self.send_response(200)
+ self.send_header("Content-Type", "application/json")
+ self.send_header("Content-Length", str(len(payload)))
+ self.end_headers()
+ self.wfile.write(payload)
+
+ return Handler
+
+
+def _gravitino_handler(state: _GravitinoState):
+ class Handler(BaseHTTPRequestHandler):
+ def log_message(self, *_args):
+ return
+
+ def do_GET(self):
+ auth = self.headers.get("Authorization", "")
+ with state.lock:
+ state.authorizations.append(auth)
+ fail = state.fail_first and auth == "Bearer tok-1"
+ if fail:
+ payload = json.dumps(
+ {
+ "code": 1,
+ "type": "UnauthorizedException",
+ "message": "expired",
+ }
+ ).encode("utf-8")
+ self.send_response(401)
+ else:
+ payload = json.dumps(
+ {
+ "code": 0,
+ "catalogs": [{"name": CATALOG}],
+ }
+ ).encode("utf-8")
+ self.send_response(200)
+ self.send_header("Content-Type", "application/json")
+ self.send_header("Content-Length", str(len(payload)))
+ self.end_headers()
+ self.wfile.write(payload)
+
+ return Handler
+
+
+def _free_port() -> int:
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+ sock.bind(("127.0.0.1", 0))
+ return sock.getsockname()[1]
+
+
+def _wait_for_port(host: str, port: int, timeout: float = 15.0) -> None:
+ deadline = time.time() + timeout
+ while time.time() < deadline:
+ try:
+ with socket.create_connection((host, port), timeout=0.2):
+ return
+ except OSError:
+ time.sleep(0.1)
+ raise TimeoutError(f"{host}:{port} did not open")
+
+
+def _start_mcp(
+ gravitino_uri: str,
+ token_endpoint: str,
+ mcp_port: int,
+ extra_args: Optional[List[str]] = None,
+) -> subprocess.Popen:
+ mcp_url = f"http://127.0.0.1:{mcp_port}/mcp";
+ env = dict(os.environ)
+ # A leftover GRAVITINO_TOKEN would win over client-credentials.
+ env.pop("GRAVITINO_TOKEN", None)
+ env["NO_PROXY"] = "127.0.0.1,localhost"
+ env["no_proxy"] = "127.0.0.1,localhost"
+ command = [
+ sys.executable,
+ "-m",
+ "mcp_server",
+ "--metalake",
+ METALAKE,
+ "--gravitino-uri",
+ gravitino_uri,
+ "--transport",
+ "http",
+ "--mcp-url",
+ mcp_url,
+ "--oauth-token-endpoint",
+ token_endpoint,
+ "--oauth-client-id",
+ "mcp-it",
+ "--oauth-client-secret",
+ "it-secret",
+ "--oauth-scope",
+ "gravitino",
+ ]
+ if extra_args:
+ command.extend(extra_args)
+ return subprocess.Popen(
+ command,
+ env=env,
+ stdout=subprocess.DEVNULL,
+ stderr=subprocess.DEVNULL,
+ )
+
+
[email protected]
+def oauth_stack():
+ """Mock IdP + mock Gravitino + live MCP process (no --token)."""
+ idp_state = _IdPState(expires_in=3600)
+ gravitino_state = _GravitinoState()
+ idp = _start_server(_idp_handler(idp_state), idp_state)
+ gravitino = _start_server(
+ _gravitino_handler(gravitino_state), gravitino_state
+ )
+ mcp_port = _free_port()
+ token_endpoint = f"http://127.0.0.1:{idp.server_address[1]}/token";
+ gravitino_uri = f"http://127.0.0.1:{gravitino.server_address[1]}";
+ proc = _start_mcp(gravitino_uri, token_endpoint, mcp_port)
+ try:
+ _wait_for_port("127.0.0.1", mcp_port)
+ yield {
+ "mcp_url": f"http://127.0.0.1:{mcp_port}/mcp";,
+ "idp": idp_state,
+ "gravitino": gravitino_state,
+ }
+ finally:
+ proc.terminate()
+ try:
+ proc.wait(timeout=5)
+ except subprocess.TimeoutExpired:
+ proc.kill()
+ idp.shutdown()
+ gravitino.shutdown()
+
+
+def _list_catalogs(mcp_url: str) -> list:
+ async def _run():
+ transport = StreamableHttpTransport(url=mcp_url)
+ async with Client(transport) as client:
+ result = await client.call_tool("get_list_of_catalogs")
+ return json.loads(result.content[0].text)
+
+ return asyncio.run(_run())
+
+
+def test_mcp_fetches_token_and_forwards_bearer(oauth_stack):
+ """No hop-1 header: MCP must mint a Bearer and send it to Gravitino."""
+ catalogs = _list_catalogs(oauth_stack["mcp_url"])
+ assert catalogs[0]["name"] == CATALOG
+
+ idp = oauth_stack["idp"]
+ gravitino = oauth_stack["gravitino"]
+ assert idp.hits == 1
+ form = parse_qs(idp.bodies[0])
+ assert form["grant_type"] == ["client_credentials"]
+ assert form["client_id"] == ["mcp-it"]
+ assert form["client_secret"] == ["it-secret"]
+ assert form["scope"] == ["gravitino"]
+ assert gravitino.authorizations == ["Bearer tok-1"]
+
+
+def test_mcp_reuses_cached_token_on_second_call(oauth_stack):
+ """A long-lived token is fetched once for two tool calls."""
+ _list_catalogs(oauth_stack["mcp_url"])
+ _list_catalogs(oauth_stack["mcp_url"])
+ assert oauth_stack["idp"].hits == 1
+ assert oauth_stack["gravitino"].authorizations == [
+ "Bearer tok-1",
+ "Bearer tok-1",
+ ]
+
+
+def test_mcp_refetches_when_token_is_immediately_stale():
+ """expires_in shorter than refresh skew forces a fetch per call."""
+ idp_state = _IdPState(expires_in=1)
+ gravitino_state = _GravitinoState()
+ idp = _start_server(_idp_handler(idp_state), idp_state)
+ gravitino = _start_server(
+ _gravitino_handler(gravitino_state), gravitino_state
+ )
+ mcp_port = _free_port()
+ proc = _start_mcp(
+ f"http://127.0.0.1:{gravitino.server_address[1]}";,
+ f"http://127.0.0.1:{idp.server_address[1]}/token";,
+ mcp_port,
+ )
+ try:
+ _wait_for_port("127.0.0.1", mcp_port)
+ mcp_url = f"http://127.0.0.1:{mcp_port}/mcp";
+ _list_catalogs(mcp_url)
+ _list_catalogs(mcp_url)
+ assert idp_state.hits == 2
+ assert gravitino_state.authorizations == [
+ "Bearer tok-1",
+ "Bearer tok-2",
+ ]
+ finally:
+ proc.terminate()
+ try:
+ proc.wait(timeout=5)
+ except subprocess.TimeoutExpired:
+ proc.kill()
+ idp.shutdown()
+ gravitino.shutdown()
+
+
+def test_mcp_retries_once_after_gravitino_401(oauth_stack):
+ """A 401 from Gravitino invalidates the cache and retries with a new
token."""
+ _list_catalogs(oauth_stack["mcp_url"])
+ oauth_stack["gravitino"].fail_first = True
+ catalogs = _list_catalogs(oauth_stack["mcp_url"])
+ assert catalogs[0]["name"] == CATALOG
+ assert oauth_stack["idp"].hits == 2
+ assert oauth_stack["gravitino"].authorizations[-2:] == [
+ "Bearer tok-1",
+ "Bearer tok-2",
+ ]
+
+
+def test_static_token_overrides_oauth_and_skips_idp():
+ """``--token`` wins: MCP must not call the IdP."""
+ idp_state = _IdPState(expires_in=3600)
+ gravitino_state = _GravitinoState()
+ idp = _start_server(_idp_handler(idp_state), idp_state)
+ gravitino = _start_server(
+ _gravitino_handler(gravitino_state), gravitino_state
+ )
+ mcp_port = _free_port()
+ proc = _start_mcp(
+ f"http://127.0.0.1:{gravitino.server_address[1]}";,
+ f"http://127.0.0.1:{idp.server_address[1]}/token";,
+ mcp_port,
+ extra_args=["--token", "frozen"],
+ )
+ try:
+ _wait_for_port("127.0.0.1", mcp_port)
+ catalogs = _list_catalogs(f"http://127.0.0.1:{mcp_port}/mcp";)
+ assert catalogs[0]["name"] == CATALOG
+ assert idp_state.hits == 0
+ assert gravitino_state.authorizations == ["Bearer frozen"]
+ finally:
+ proc.terminate()
+ try:
+ proc.wait(timeout=5)
+ except subprocess.TimeoutExpired:
+ proc.kill()
+ idp.shutdown()
+ gravitino.shutdown()
diff --git a/mcp-server/tests/unit/test_audit.py
b/mcp-server/tests/unit/test_audit.py
index 29991582b9..e6bdea1364 100644
--- a/mcp-server/tests/unit/test_audit.py
+++ b/mcp-server/tests/unit/test_audit.py
@@ -203,6 +203,27 @@ class TestAuditMiddlewareIntegration(unittest.TestCase):
record = json.loads(self.log_records[0])
self.assertEqual(record["principal"], "bearer:abcdef12")
+ def test_principal_falls_back_to_oauth_client_id(self):
+ """With no request header, OAuth client-credentials is attributable."""
+ RESTClientFactory.set_rest_client(MockOperation)
+ server = GravitinoMCPServer(
+ Setting(
+ "mock_metalake",
+ oauth_token_endpoint="https://idp/token";,
+ oauth_client_id="mcp-service",
+ oauth_client_secret="s",
+ )
+ )
+
+ async def _run():
+ async with Client(server.mcp) as client:
+ await client.call_tool("get_list_of_catalogs")
+
+ asyncio.run(_run())
+
+ record = json.loads(self.log_records[0])
+ self.assertEqual(record["principal"], "oauth:mcp-serv")
+
def test_failed_tool_call_emits_deny_record(self):
"""A tool call that raises an exception produces an audit record with
outcome=deny."""
diff --git a/mcp-server/tests/unit/test_oauth.py
b/mcp-server/tests/unit/test_oauth.py
new file mode 100644
index 0000000000..2c1d5539c4
--- /dev/null
+++ b/mcp-server/tests/unit/test_oauth.py
@@ -0,0 +1,843 @@
+# 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.
+
+"""Tests for OAuth2 client-credentials fetch, cache, refresh, and 401 retry."""
+
+# pylint: disable=protected-access
+
+import asyncio
+import base64
+import json
+import sys
+import threading
+import time
+import unittest
+from unittest import mock
+
+import httpx
+from httpx_auth import OAuth2
+
+from mcp_server.client.factory import RESTClientFactory
+from mcp_server.client.plain.plain_rest_client_operation import (
+ PlainRESTClientOperation,
+)
+from mcp_server.core.context import (
+ GravitinoContext,
+ ServiceIdentityFallbackDisabled,
+ service_fallback_authorization,
+)
+from mcp_server.core.oauth import RefreshableBearerAuth
+from mcp_server.core.setting import Setting
+from mcp_server.main import _parse_args, do_main
+
+
+def _jwt_with_exp(exp) -> str:
+ header = base64.urlsafe_b64encode(b'{"alg":"none"}').rstrip(b"=").decode()
+ payload = (
+ base64.urlsafe_b64encode(json.dumps({"exp": exp}).encode())
+ .rstrip(b"=")
+ .decode()
+ )
+ return f"{header}.{payload}.sig"
+
+
+class _OAuthHttpTestCase(unittest.TestCase):
+ """Drive ``httpx.AsyncClient(auth=...)`` against a shared MockTransport.
+
+ httpx-auth caches tokens in a process-global map, so every test clears it.
+ Tests inject a sync ``httpx.Client`` on the same transport so IdP calls
+ never leave the process. Production token POSTs use ``AsyncClient``.
+ """
+
+ def setUp(self):
+ OAuth2.token_cache.clear()
+
+ def tearDown(self):
+ OAuth2.token_cache.clear()
+
+ def _auth(self, handler, **kwargs) -> RefreshableBearerAuth:
+ transport = httpx.MockTransport(handler)
+ self.addCleanup(OAuth2.token_cache.clear)
+ token_client = httpx.Client(transport=transport)
+ self.addCleanup(token_client.close)
+ auth = RefreshableBearerAuth(
+ token_endpoint="https://idp.example/token";,
+ client_id="mcp",
+ client_secret="s3cret",
+ client=token_client,
+ **kwargs,
+ )
+ auth._test_transport = transport
+ return auth
+
+ def _get(self, auth: RefreshableBearerAuth, path: str = "/api"):
+ async def _run():
+ async with httpx.AsyncClient(
+ auth=auth,
+ transport=auth._test_transport,
+ base_url="https://gravitino.example";,
+ ) as client:
+ return await client.get(path)
+
+ return asyncio.run(_run())
+
+ def _get_twice(self, auth: RefreshableBearerAuth):
+ async def _run():
+ async with httpx.AsyncClient(
+ auth=auth,
+ transport=auth._test_transport,
+ base_url="https://gravitino.example";,
+ ) as client:
+ first = await client.get("/api")
+ second = await client.get("/api")
+ return first, second
+
+ return asyncio.run(_run())
+
+
+class TestRefreshableBearerAuth(_OAuthHttpTestCase):
+ def test_fetches_and_reuses_while_fresh(self):
+ calls = {"token": 0, "api": 0}
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ if request.method == "POST":
+ calls["token"] += 1
+ self.assertEqual(str(request.url), "https://idp.example/token";)
+ body = request.content.decode()
+ self.assertIn("grant_type=client_credentials", body)
+ self.assertIn("client_id=mcp", body)
+ self.assertIn("client_secret=s3cret", body)
+ self.assertIsNone(request.headers.get("authorization"))
+ return httpx.Response(
+ 200, json={"access_token": "tok-1", "expires_in": 3600}
+ )
+ calls["api"] += 1
+ self.assertEqual(
+ request.headers.get("authorization"), "Bearer tok-1"
+ )
+ return httpx.Response(200, json={"ok": True})
+
+ first, second = self._get_twice(self._auth(handler))
+ self.assertEqual(first.status_code, 200)
+ self.assertEqual(second.status_code, 200)
+ self.assertEqual(calls["token"], 1)
+ self.assertEqual(calls["api"], 2)
+
+ def test_parallel_cold_cache_uses_single_token_post(self):
+ """Concurrent async callers must not stampede the IdP on cache miss."""
+ calls = {"token": 0, "api": 0}
+ counter_lock = threading.Lock()
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ if request.method == "POST":
+ with counter_lock:
+ calls["token"] += 1
+ time.sleep(0.05)
+ return httpx.Response(
+ 200, json={"access_token": "tok-1", "expires_in": 3600}
+ )
+ with counter_lock:
+ calls["api"] += 1
+ return httpx.Response(200, json={"ok": True})
+
+ auth = self._auth(handler)
+
+ async def run_parallel() -> None:
+ async with httpx.AsyncClient(
+ auth=auth,
+ transport=auth._test_transport,
+ base_url="https://gravitino.example";,
+ ) as client:
+ await asyncio.gather(*[client.get("/api") for _ in range(10)])
+
+ asyncio.run(run_parallel())
+ self.assertEqual(calls["token"], 1)
+ self.assertEqual(calls["api"], 10)
+
+ def test_fresh_token_401_skips_retry(self):
+ """A brand-new token that is rejected is not retried or refetched."""
+ calls = {"token": 0, "api": 0}
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ if request.method == "POST":
+ calls["token"] += 1
+ return httpx.Response(
+ 200, json={"access_token": "tok-1", "expires_in": 3600}
+ )
+ calls["api"] += 1
+ return httpx.Response(401)
+
+ response = self._get(self._auth(handler))
+ self.assertEqual(response.status_code, 401)
+ self.assertEqual(calls["token"], 1)
+ self.assertEqual(calls["api"], 1)
+
+ def test_parallel_401_does_not_stampede_token_refresh(self):
+ """Concurrent 401s must not each invalidate and refetch the token."""
+ calls = {"token": 0, "api": 0, "phase": "prime"}
+ counter_lock = threading.Lock()
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ if request.method == "POST":
+ with counter_lock:
+ calls["token"] += 1
+ time.sleep(0.02)
+ return httpx.Response(
+ 200, json={"access_token": "tok-1", "expires_in": 3600}
+ )
+ with counter_lock:
+ calls["api"] += 1
+ if calls["phase"] == "prime":
+ return httpx.Response(200, json={"ok": True})
+ return httpx.Response(401)
+
+ auth = self._auth(handler)
+
+ async def prime_and_parallel() -> None:
+ async with httpx.AsyncClient(
+ auth=auth,
+ transport=auth._test_transport,
+ base_url="https://gravitino.example";,
+ ) as client:
+ await client.get("/api")
+ calls["phase"] = "parallel"
+ token_before = calls["token"]
+ await asyncio.gather(*[client.get("/api") for _ in range(10)])
+ self.assertLessEqual(
+ calls["token"] - token_before,
+ 1,
+ "parallel 401 retry must not stampede the IdP",
+ )
+
+ asyncio.run(prime_and_parallel())
+ self.assertGreaterEqual(calls["api"], 11)
+
+ def test_config_401_does_not_refetch_on_every_call(self):
+ """Persistent Gravitino 401 must not POST to the IdP every call."""
+ calls = {"token": 0, "api": 0}
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ if request.method == "POST":
+ calls["token"] += 1
+ return httpx.Response(
+ 200, json={"access_token": "tok-1", "expires_in": 3600}
+ )
+ calls["api"] += 1
+ return httpx.Response(401)
+
+ auth = self._auth(handler)
+
+ async def run_twice() -> None:
+ async with httpx.AsyncClient(
+ auth=auth,
+ transport=auth._test_transport,
+ base_url="https://gravitino.example";,
+ ) as client:
+ first = await client.get("/api")
+ second = await client.get("/api")
+ self.assertEqual(first.status_code, 401)
+ self.assertEqual(second.status_code, 401)
+
+ asyncio.run(run_twice())
+ self.assertEqual(calls["token"], 1)
+ self.assertEqual(calls["api"], 2)
+
+ def test_sends_scope_when_configured(self):
+ seen = {}
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ if request.method == "POST":
+ seen["body"] = request.content.decode()
+ return httpx.Response(
+ 200, json={"access_token": "tok", "expires_in": 3600}
+ )
+ return httpx.Response(200, json={"ok": True})
+
+ self._get(self._auth(handler, scope="gravitino"))
+ self.assertIn("scope=gravitino", seen["body"])
+
+ def test_refetches_when_expired(self):
+ calls = {"token": 0}
+ seen = []
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ if request.method == "POST":
+ calls["token"] += 1
+ return httpx.Response(
+ 200,
+ json={
+ "access_token": f"tok-{calls['token']}",
+ "expires_in": 1,
+ },
+ )
+ seen.append(request.headers.get("authorization"))
+ return httpx.Response(200, json={"ok": True})
+
+ # Skew > expires_in makes the cache stale immediately after fetch.
+ self._get_twice(self._auth(handler, refresh_skew_seconds=60))
+ self.assertEqual(seen, ["Bearer tok-1", "Bearer tok-2"])
+ self.assertEqual(calls["token"], 2)
+
+ def test_accepts_string_expires_in(self):
+ calls = {"token": 0}
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ if request.method == "POST":
+ calls["token"] += 1
+ return httpx.Response(
+ 200,
+ json={"access_token": "tok-1", "expires_in": "3600"},
+ )
+ return httpx.Response(200, json={"ok": True})
+
+ self._get_twice(self._auth(handler))
+ self.assertEqual(calls["token"], 1)
+
+ def test_expires_in_zero_falls_back_to_jwt_exp(self):
+ """expires_in: 0 is falsy; cache from JWT exp instead of refetching."""
+ token = _jwt_with_exp(int(time.time()) + 3600)
+ calls = {"token": 0}
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ if request.method == "POST":
+ calls["token"] += 1
+ return httpx.Response(
+ 200, json={"access_token": token, "expires_in": 0}
+ )
+ return httpx.Response(200, json={"ok": True})
+
+ self._get_twice(self._auth(handler))
+ self.assertEqual(calls["token"], 1)
+
+ def test_expires_in_zero_without_jwt_exp_raises(self):
+ def handler(request: httpx.Request) -> httpx.Response:
+ if request.method == "POST":
+ return httpx.Response(
+ 200, json={"access_token": "opaque-ref", "expires_in": 0}
+ )
+ return httpx.Response(200)
+
+ with self.assertRaises(ValueError) as raised:
+ self._get(self._auth(handler))
+ self.assertIn("expires_in", str(raised.exception))
+
+ def test_uses_jwt_exp_when_expires_in_missing(self):
+ token = _jwt_with_exp(int(time.time()) + 3600)
+ calls = {"token": 0}
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ if request.method == "POST":
+ calls["token"] += 1
+ return httpx.Response(200, json={"access_token": token})
+ self.assertEqual(
+ request.headers.get("authorization"), f"Bearer {token}"
+ )
+ return httpx.Response(200, json={"ok": True})
+
+ self._get_twice(self._auth(handler))
+ self.assertEqual(calls["token"], 1)
+
+ def test_jwt_exp_float_accepted(self):
+ token = _jwt_with_exp(float(int(time.time()) + 3600))
+ calls = {"token": 0}
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ if request.method == "POST":
+ calls["token"] += 1
+ return httpx.Response(200, json={"access_token": token})
+ return httpx.Response(200, json={"ok": True})
+
+ self._get_twice(self._auth(handler))
+ self.assertEqual(calls["token"], 1)
+
+ def test_jwt_exp_string_accepted(self):
+ token = _jwt_with_exp(str(int(time.time()) + 3600))
+ calls = {"token": 0}
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ if request.method == "POST":
+ calls["token"] += 1
+ return httpx.Response(200, json={"access_token": token})
+ return httpx.Response(200, json={"ok": True})
+
+ self._get_twice(self._auth(handler))
+ self.assertEqual(calls["token"], 1)
+
+ def test_opaque_token_without_expires_in_raises(self):
+ def handler(request: httpx.Request) -> httpx.Response:
+ if request.method == "POST":
+ return httpx.Response(200, json={"access_token": "opaque-ref"})
+ return httpx.Response(200)
+
+ with self.assertRaises(ValueError) as raised:
+ self._get(self._auth(handler))
+ self.assertIn("expires_in", str(raised.exception))
+ self.assertIn("JWT", str(raised.exception))
+
+ def test_jwt_without_exp_and_expires_in_raises(self):
+ header = (
+ base64.urlsafe_b64encode(b'{"alg":"none"}').rstrip(b"=").decode()
+ )
+ payload = (
+ base64.urlsafe_b64encode(b'{"sub":"mcp"}').rstrip(b"=").decode()
+ )
+ token = f"{header}.{payload}.sig"
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ if request.method == "POST":
+ return httpx.Response(200, json={"access_token": token})
+ return httpx.Response(200)
+
+ with self.assertRaises(ValueError) as raised:
+ self._get(self._auth(handler))
+ self.assertIn("expires_in", str(raised.exception))
+
+ def test_http_error_is_raised(self):
+ def handler(request: httpx.Request) -> httpx.Response:
+ if request.method == "POST":
+ return httpx.Response(401, json={"error": "invalid_client"})
+ return httpx.Response(200)
+
+ with self.assertRaises(httpx.HTTPStatusError):
+ self._get(self._auth(handler))
+
+ def test_retries_once_after_401(self):
+ calls = {"token": 0, "api": 0, "primed": False}
+ seen = []
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ if request.method == "POST":
+ calls["token"] += 1
+ return httpx.Response(
+ 200,
+ json={
+ "access_token": f"t{calls['token']}",
+ "expires_in": 3600,
+ },
+ )
+ calls["api"] += 1
+ if not calls["primed"]:
+ return httpx.Response(200, json={"ok": True})
+ seen.append(request.headers.get("authorization"))
+ if len(seen) == 1:
+ return httpx.Response(401)
+ return httpx.Response(200, json={"ok": True})
+
+ auth = self._auth(handler)
+
+ async def _run():
+ async with httpx.AsyncClient(
+ auth=auth,
+ transport=auth._test_transport,
+ base_url="https://gravitino.example";,
+ ) as client:
+ await client.get("/api")
+ calls["primed"] = True
+ return await client.get("/api")
+
+ response = asyncio.run(_run())
+ self.assertEqual(response.status_code, 200)
+ self.assertEqual(seen, ["Bearer t1", "Bearer t2"])
+ self.assertEqual(calls["token"], 2)
+
+ def test_401_retry_sets_do_not_exceed_max_limit(self):
+ """After many clock-skew 401 retries, retained tokens stay <= 10."""
+ max_retained = 10
+ cycles = 20
+ calls = {"token": 0, "expect_401": False}
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ if request.method == "POST":
+ calls["token"] += 1
+ return httpx.Response(
+ 200,
+ json={
+ "access_token": f"tok-{calls['token']}",
+ "expires_in": 3600,
+ },
+ )
+ if calls["expect_401"]:
+ calls["expect_401"] = False
+ return httpx.Response(401)
+ return httpx.Response(200, json={"ok": True})
+
+ auth = self._auth(handler)
+
+ async def _run() -> None:
+ async with httpx.AsyncClient(
+ auth=auth,
+ transport=auth._test_transport,
+ base_url="https://gravitino.example";,
+ ) as client:
+ for _ in range(cycles):
+ await client.get("/api")
+ calls["expect_401"] = True
+ response = await client.get("/api")
+ self.assertEqual(response.status_code, 200)
+ auth.invalidate()
+
+ asyncio.run(_run())
+ retained = len(auth._retried_tokens) + len(auth._rejected_tokens)
+ self.assertLessEqual(
+ retained,
+ max_retained,
+ "401 token sets exceeded max limit "
+ f"{max_retained}: retained={retained} "
+ f"retried={auth._retried_tokens!r} "
+ f"rejected={auth._rejected_tokens!r}",
+ )
+
+ def test_401_retry_replays_streaming_request_body(self):
+ seen_bodies = []
+ primed = {"done": False}
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ if request.method == "POST" and request.url.path == "/token":
+ return httpx.Response(
+ 200, json={"access_token": "t1", "expires_in": 3600}
+ )
+ seen_bodies.append(request.content)
+ if not primed["done"]:
+ return httpx.Response(200, json={"ok": True})
+ if len(seen_bodies) == 2:
+ return httpx.Response(401)
+ return httpx.Response(200, json={"ok": True})
+
+ auth = self._auth(handler)
+
+ async def _run():
+ async def body():
+ yield b'{"name":"fileset"}'
+
+ async with httpx.AsyncClient(
+ auth=auth,
+ transport=auth._test_transport,
+ base_url="https://gravitino.example";,
+ ) as client:
+ await client.get("/api")
+ primed["done"] = True
+ return await client.post("/api", content=body())
+
+ response = asyncio.run(_run())
+ self.assertEqual(response.status_code, 200)
+ self.assertEqual(
+ seen_bodies,
+ [b"", b'{"name":"fileset"}', b'{"name":"fileset"}'],
+ )
+
+ def test_async_auth_flow_posts_token_with_async_client(self):
+ calls = {"token": 0}
+
+ def handler(request: httpx.Request) -> httpx.Response:
+ if request.method == "POST":
+ calls["token"] += 1
+ return httpx.Response(
+ 200, json={"access_token": "tok-async", "expires_in": 3600}
+ )
+ self.assertEqual(
+ request.headers.get("authorization"), "Bearer tok-async"
+ )
+ return httpx.Response(200, json={"ok": True})
+
+ transport = httpx.MockTransport(handler)
+ auth = RefreshableBearerAuth(
+ token_endpoint="https://idp.example/token";,
+ client_id="mcp",
+ client_secret="s3cret",
+ )
+ real_async_client = httpx.AsyncClient
+
+ def async_client_factory(*args, **kwargs):
+ kwargs["transport"] = transport
+ return real_async_client(*args, **kwargs)
+
+ with mock.patch(
+ "mcp_server.core.oauth.httpx.AsyncClient",
+ side_effect=async_client_factory,
+ ) as async_cls, mock.patch(
+ "mcp_server.core.oauth.httpx.Client"
+ ) as sync_cls:
+ sync_cls.side_effect = AssertionError(
+ "async_auth_flow must not use httpx.Client"
+ )
+
+ async def _run():
+ async with httpx.AsyncClient(
+ auth=auth,
+ transport=transport,
+ base_url="https://gravitino.example";,
+ ) as client:
+ return await client.get("/api")
+
+ response = asyncio.run(_run())
+
+ self.assertEqual(response.status_code, 200)
+ self.assertEqual(calls["token"], 1)
+ async_cls.assert_called()
+
+
+class TestSettingOAuth(unittest.TestCase):
+ def test_partial_oauth_is_rejected(self):
+ setting = Setting(
+ metalake="ml",
+ oauth_client_id="mcp",
+ oauth_client_secret="s",
+ )
+ with self.assertRaises(ValueError):
+ setting.validate_oauth()
+
+ def test_oauth_scope_without_credentials_is_rejected(self):
+ setting = Setting(metalake="ml", oauth_scope="gravitino")
+ with self.assertRaises(ValueError) as raised:
+ setting.validate_oauth()
+ self.assertIn("scope", str(raised.exception).lower())
+
+ def test_complete_oauth_is_accepted(self):
+ setting = Setting(
+ metalake="ml",
+ oauth_token_endpoint="https://idp/token";,
+ oauth_client_id="mcp",
+ oauth_client_secret="s",
+ )
+ setting.validate_oauth()
+ self.assertTrue(setting.has_oauth_client())
+
+ def test_secret_masked_in_str(self):
+ setting = Setting(
+ metalake="ml",
+ oauth_token_endpoint="https://idp/token";,
+ oauth_client_id="mcp",
+ oauth_client_secret="super-oauth-secret",
+ )
+ self.assertNotIn("super-oauth-secret", str(setting))
+ self.assertNotIn("super-oauth-secret", repr(setting))
+
+
+class TestOAuthArgParsing(unittest.TestCase):
+ def test_env_vars_used_when_flags_omitted(self):
+ env = {
+ "GRAVITINO_OAUTH_TOKEN_ENDPOINT": "https://idp/token";,
+ "GRAVITINO_OAUTH_CLIENT_ID": "env-id",
+ "GRAVITINO_OAUTH_CLIENT_SECRET": "env-secret",
+ "GRAVITINO_OAUTH_SCOPE": "env-scope",
+ }
+ with mock.patch.dict("os.environ", env, clear=True), mock.patch.object(
+ sys, "argv", ["prog", "--metalake", "ml"]
+ ):
+ args = _parse_args()
+ self.assertEqual(args.oauth_token_endpoint, "https://idp/token";)
+ self.assertEqual(args.oauth_client_id, "env-id")
+ self.assertEqual(args.oauth_client_secret, "env-secret")
+ self.assertEqual(args.oauth_scope, "env-scope")
+
+ def test_cli_overrides_env(self):
+ env = {
+ "GRAVITINO_OAUTH_CLIENT_ID": "env-id",
+ "GRAVITINO_OAUTH_CLIENT_SECRET": "env-secret",
+ "GRAVITINO_OAUTH_TOKEN_ENDPOINT": "https://env/token";,
+ }
+ with mock.patch.dict("os.environ", env), mock.patch.object(
+ sys,
+ "argv",
+ [
+ "prog",
+ "--metalake",
+ "ml",
+ "--oauth-client-id",
+ "cli-id",
+ "--oauth-client-secret",
+ "cli-secret",
+ "--oauth-token-endpoint",
+ "https://cli/token";,
+ ],
+ ):
+ args = _parse_args()
+ self.assertEqual(args.oauth_client_id, "cli-id")
+ self.assertEqual(args.oauth_client_secret, "cli-secret")
+ self.assertEqual(args.oauth_token_endpoint, "https://cli/token";)
+
+ def test_no_service_identity_fallback_from_env(self):
+ env = {"GRAVITINO_NO_SERVICE_IDENTITY_FALLBACK": "true"}
+ with mock.patch.dict("os.environ", env, clear=True), mock.patch.object(
+ sys, "argv", ["prog", "--metalake", "ml"]
+ ):
+ args = _parse_args()
+ self.assertTrue(args.no_service_identity_fallback)
+
+
+class TestMainOAuthValidation(unittest.TestCase):
+ def test_partial_oauth_inits_logging_before_exit(self):
+ with mock.patch(
+ "mcp_server.main._init_logging"
+ ) as init_log, mock.patch(
+ "mcp_server.main.GravitinoMCPServer"
+ ), mock.patch.object(
+ sys,
+ "argv",
+ ["mcp_server", "--metalake", "ml", "--oauth-client-id", "x"],
+ ):
+ with self.assertRaises(SystemExit) as raised:
+ do_main()
+ self.assertEqual(raised.exception.code, 1)
+ init_log.assert_called_once()
+
+
+class TestNoServiceIdentityFallback(unittest.TestCase):
+ def setUp(self):
+ RESTClientFactory.set_rest_client(PlainRESTClientOperation)
+
+ def test_http_without_auth_raises_when_flag_set(self):
+ setting = Setting(
+ metalake="ml",
+ gravitino_uri="http://localhost:8090";,
+ transport="http",
+ oauth_token_endpoint="https://idp/token";,
+ oauth_client_id="mcp",
+ oauth_client_secret="s",
+ no_service_identity_fallback=True,
+ )
+ ctx = GravitinoContext(setting)
+ with mock.patch(
+ "mcp_server.core.context._in_http_request", return_value=True
+ ), mock.patch(
+ "mcp_server.core.context._get_request_authorization",
+ return_value="",
+ ):
+ with self.assertRaises(ServiceIdentityFallbackDisabled):
+ ctx.rest_client()
+
+ def test_http_without_auth_allows_fallback_when_flag_false(self):
+ setting = Setting(
+ metalake="ml",
+ gravitino_uri="http://localhost:8090";,
+ transport="http",
+ oauth_token_endpoint="https://idp/token";,
+ oauth_client_id="mcp",
+ oauth_client_secret="s",
+ )
+ ctx = GravitinoContext(setting)
+ with mock.patch(
+ "mcp_server.core.context._in_http_request", return_value=True
+ ), mock.patch(
+ "mcp_server.core.context._get_request_authorization",
+ return_value="",
+ ):
+ client = ctx.rest_client()
+ self.assertIs(client, ctx._default_client)
+
+ def test_http_flag_noop_without_service_identity(self):
+ setting = Setting(
+ metalake="ml",
+ gravitino_uri="http://localhost:8090";,
+ transport="http",
+ no_service_identity_fallback=True,
+ )
+ ctx = GravitinoContext(setting)
+ with mock.patch(
+ "mcp_server.core.context._in_http_request", return_value=True
+ ), mock.patch(
+ "mcp_server.core.context._get_request_authorization",
+ return_value="",
+ ):
+ client = ctx.rest_client()
+ self.assertIs(client, ctx._default_client)
+
+ def test_stdio_uses_default_client_even_when_flag_set(self):
+ setting = Setting(
+ metalake="ml",
+ gravitino_uri="http://localhost:8090";,
+ transport="stdio",
+ oauth_token_endpoint="https://idp/token";,
+ oauth_client_id="mcp",
+ oauth_client_secret="s",
+ no_service_identity_fallback=True,
+ )
+ ctx = GravitinoContext(setting)
+ with mock.patch(
+ "mcp_server.core.context._in_http_request", return_value=False
+ ), mock.patch(
+ "mcp_server.core.context._get_request_authorization",
+ return_value="",
+ ):
+ client = ctx.rest_client()
+ self.assertIs(client, ctx._default_client)
+
+
+class TestServiceFallbackAuthorization(unittest.TestCase):
+ def test_token_wins_over_oauth(self):
+ setting = Setting(
+ metalake="ml",
+ token="static-token",
+ oauth_token_endpoint="https://idp/token";,
+ oauth_client_id="mcp-service",
+ oauth_client_secret="s",
+ )
+ self.assertEqual(
+ service_fallback_authorization(setting), "Bearer static-token"
+ )
+
+ def test_oauth_client_id_used_when_no_token(self):
+ setting = Setting(
+ metalake="ml",
+ oauth_token_endpoint="https://idp/token";,
+ oauth_client_id="mcp-service",
+ oauth_client_secret="s",
+ )
+ self.assertEqual(
+ service_fallback_authorization(setting), "OAuth mcp-service"
+ )
+
+
+class TestGravitinoContextOAuth(unittest.TestCase):
+ def setUp(self):
+ RESTClientFactory.set_rest_client(PlainRESTClientOperation)
+
+ def test_oauth_uses_httpx_auth_hook(self):
+ setting = Setting(
+ metalake="ml",
+ gravitino_uri="http://localhost:8090";,
+ oauth_token_endpoint="https://idp/token";,
+ oauth_client_id="mcp",
+ oauth_client_secret="s",
+ )
+ ctx = GravitinoContext(setting)
+ client = ctx.rest_client()
+ try:
+ rest = client._catalog_operation.rest_client
+ self.assertIsInstance(rest, httpx.AsyncClient)
+ self.assertIsInstance(rest.auth, RefreshableBearerAuth)
+ self.assertIsNone(rest.headers.get("Authorization"))
+ finally:
+ asyncio.run(client.close())
+
+ def test_static_token_overrides_oauth(self):
+ setting = Setting(
+ metalake="ml",
+ gravitino_uri="http://localhost:8090";,
+ token="frozen",
+ oauth_token_endpoint="https://idp/token";,
+ oauth_client_id="mcp",
+ oauth_client_secret="s",
+ )
+ ctx = GravitinoContext(setting)
+ client = ctx.rest_client()
+ try:
+ rest = client._catalog_operation.rest_client
+ self.assertEqual(rest.headers.get("Authorization"), "Bearer
frozen")
+ self.assertIsNone(rest.auth)
+ finally:
+ asyncio.run(client.close())
diff --git a/mcp-server/tests/unit/tools/mock_operation.py
b/mcp-server/tests/unit/tools/mock_operation.py
index e1fc540442..61f0a3e7fe 100644
--- a/mcp-server/tests/unit/tools/mock_operation.py
+++ b/mcp-server/tests/unit/tools/mock_operation.py
@@ -31,7 +31,7 @@ from mcp_server.client.statistic_operation import
StatisticOperation
class MockOperation(GravitinoOperation):
- def __init__(self, metalake, uri, authorization=""):
+ def __init__(self, metalake, uri, authorization="", *, auth=None):
pass
def as_table_operation(self) -> TableOperation:
diff --git a/mcp-server/uv.lock b/mcp-server/uv.lock
index 5ace329a27..f5c9ba8f48 100644
--- a/mcp-server/uv.lock
+++ b/mcp-server/uv.lock
@@ -459,11 +459,12 @@ wheels = [
[[package]]
name = "gravitino-mcp-server"
-version = "1.4.0.dev0"
+version = "1.3.1.dev0"
source = { virtual = "." }
dependencies = [
{ name = "fakeredis" },
{ name = "fastmcp" },
+ { name = "httpx-auth" },
{ name = "parameterized" },
{ name = "pylint" },
{ name = "pytest" },
@@ -473,6 +474,7 @@ dependencies = [
requires-dist = [
{ name = "fakeredis", specifier = "<2.35.0" },
{ name = "fastmcp", specifier = "==3.2.0" },
+ { name = "httpx-auth", specifier = ">=0.22,<0.24" },
{ name = "parameterized", specifier = ">=0.9.0" },
{ name = "pylint", specifier = ">=2.20.0" },
{ name = "pytest", specifier = ">=8.4.1" },
@@ -515,6 +517,18 @@ wheels = [
{ url =
"https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl";,
hash =
"sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad", size
= 73517, upload-time = "2024-12-06T15:37:21.509Z" },
]
+[[package]]
+name = "httpx-auth"
+version = "0.23.1"
+source = { registry = "https://pypi.org/simple"; }
+dependencies = [
+ { name = "httpx" },
+]
+sdist = { url =
"https://files.pythonhosted.org/packages/a8/d4/6bd616f89d1ce43f602b62ec274e33beee6c2bce3d68396e692daafdb57d/httpx_auth-0.23.1.tar.gz";,
hash =
"sha256:27b5a6022ad1b41a303b8737fa2e3e4bce6bbbe7ab67fed0b261359be62e0434", size
= 121418, upload-time = "2025-01-07T18:47:20.05Z" }
+wheels = [
+ { url =
"https://files.pythonhosted.org/packages/2f/23/a72f91bea596b522ac297b948ffee6decdedb535c034fca8062bd72981ce/httpx_auth-0.23.1-py3-none-any.whl";,
hash =
"sha256:04f8bd0824efe3d9fb79690cc670b0da98ea809babb7aea04a72f334d4fd5ec5", size
= 45328, upload-time = "2025-01-07T18:47:18.694Z" },
+]
+
[[package]]
name = "httpx-sse"
version = "0.4.1"
