jerryshao commented on code in PR #12960:
URL: https://github.com/apache/gravitino/pull/12960#discussion_r3966172353
##########
docs/gravitino-mcp-server.md:
##########
@@ -144,7 +144,7 @@ You could config Gravitino MCP server by arguments, `uv run
mcp_server -h` shows
| Argument | Description
| Default value | Required |
|----------------------------------|---------------------------------------------------------------------------------------------------------------------------------|-----------------------------|----------|
-| `--metalake` | The Gravitino metalake name.
| none | Yes |
+| `--metalake` | Default Gravitino metalake, used when a
request names none. See Per-request metalake below.
| none | stdio only |
Review Comment:
Reformatted the whole table so every row shares the column widths — the
`--metalake` row's last cell was 12 chars wide against the table's 10, which is
what broke the alignment.
_🤖 Addressed by [Claude Code](https://claude.com/claude-code)_
##########
docs/gravitino-mcp-server.md:
##########
@@ -225,6 +225,14 @@ For exposed or multi-caller HTTP deployments, set
`--no-service-identity-fallbac
Authorization itself is always enforced by Gravitino: the MCP server forwards
the identity but does not make access-control decisions of its own.
+### Per-request metalake (HTTP)
+
+When the server runs with HTTP transport, a request may name the metalake to
operate on with the `X-Gravitino-Metalake` header, taking priority over the
`--metalake` default configured at startup. This lets one server instance serve
more than one metalake: each request independently resolves its own metalake
from its own header, so the server holds no per-connection or per-session
metalake state and stays correct regardless of how many replicas it runs as.
+
+Falls back to `--metalake` when the header is absent. If neither is set, the
call fails with an error naming the missing argument. Authorization is
unchanged — the caller's identity (see above) determines what it may see in the
requested metalake exactly as it would through the REST API.
Review Comment:
Optional on every transport now. It is the default used by any call that
does not name a metalake itself; the docs state that explicitly and spell out
the resolution order.
_🤖 Addressed by [Claude Code](https://claude.com/claude-code)_
##########
docs/gravitino-mcp-server.md:
##########
@@ -144,7 +144,7 @@ You could config Gravitino MCP server by arguments, `uv run
mcp_server -h` shows
| Argument | Description
| Default value | Required |
|----------------------------------|---------------------------------------------------------------------------------------------------------------------------------|-----------------------------|----------|
-| `--metalake` | The Gravitino metalake name.
| none | Yes |
+| `--metalake` | Default Gravitino metalake, used when a
request names none. Required for `stdio`; optional for HTTP, where each request
may instead name a metalake via the `X-Gravitino-Metalake` header. | none
| stdio only |
Review Comment:
Row rewritten: Required is now `No`, and the description points at the
"Selecting a metalake" section. Default stays `none`, matching how the other
optional arguments in this table express "no value set".
_🤖 Addressed by [Claude Code](https://claude.com/claude-code)_
##########
mcp-server/mcp_server/core/context.py:
##########
@@ -147,38 +169,64 @@ def _service_auth(setting: Setting):
class GravitinoContext:
def __init__(self, setting: Setting):
+ # Enforced here (not only in do_main()) so any path that constructs a
+ # GravitinoContext directly - not just the CLI entrypoint - fails
+ # fast on an invalid Setting, matching the pre-per-request-metalake
+ # behavior where a bad Setting couldn't be constructed at all.
+ setting.validate_metalake()
+ setting.validate_oauth()
self._setting = setting
- self._default_client = RESTClientFactory.create_rest_client(
- setting.metalake,
- setting.gravitino_uri,
- startup_authorization(setting),
- auth=_service_auth(setting),
+ # Eagerly built only when a startup default is configured, so the
+ # common single-metalake deployment pays no extra cost. Left unset
+ # (None) when metalake resolution must come from a per-request header
+ # on every call (HTTP transport with no --metalake default).
+ self._default_client = (
+ RESTClientFactory.create_rest_client(
+ setting.metalake,
+ setting.gravitino_uri,
+ startup_authorization(setting),
+ auth=_service_auth(setting),
+ )
+ if setting.metalake
+ else None
)
- # LRU cache of per-principal clients keyed by the raw Authorization
header.
- # Safe without locking: rest_client() runs on the single asyncio event
- # loop and never awaits between lookup and insert.
- self._clients_by_auth: "OrderedDict[str, object]" = OrderedDict()
+ # LRU cache of per-principal clients keyed by (Authorization header,
+ # metalake). Safe without locking: rest_client() runs on the single
+ # asyncio event loop and never awaits between lookup and insert.
+ self._clients_by_auth: "OrderedDict[tuple, object]" = OrderedDict()
+ # LRU cache of service-identity clients (static token / OAuth) keyed by
+ # metalake, for requests that name a non-default metalake but carry no
+ # per-request Authorization header. The startup default metalake is
+ # served by _default_client instead and never enters this cache.
+ self._service_clients: "OrderedDict[str, object]" = OrderedDict()
Review Comment:
Merged the two caches into one keyed by `(authorization, metalake)`, so
`_MAX_CACHED_CLIENTS` bounds the total number of open pools again.
Service-identity clients use an empty authorization, which only the
no-Authorization branch produces, so the keys cannot collide.
_🤖 Addressed by [Claude Code](https://claude.com/claude-code)_
##########
mcp-server/mcp_server/core/context.py:
##########
@@ -190,24 +238,68 @@ def rest_client(self):
"HTTP request omitted Authorization and "
"--no-service-identity-fallback is set"
)
- return self._default_client
+ return self._service_client(metalake)
- cached = self._clients_by_auth.get(authorization)
+ key = (authorization, metalake)
+ cached = self._clients_by_auth.get(key)
if cached is not None:
- self._clients_by_auth.move_to_end(authorization)
+ self._clients_by_auth.move_to_end(key)
return cached
client = RESTClientFactory.create_rest_client(
- self._setting.metalake,
+ metalake,
self._setting.gravitino_uri,
authorization,
)
- self._clients_by_auth[authorization] = client
- if len(self._clients_by_auth) > _MAX_CACHED_CLIENTS:
- _, evicted = self._clients_by_auth.popitem(last=False)
- self._schedule_close(evicted)
+ self._cache_put(self._clients_by_auth, key, client)
return client
+ def _resolve_metalake(self) -> str:
+ """Resolve the metalake for the current call, header first.
+
+ Raises ``ValueError`` (an invalid/missing request parameter, mapped
+ by FastMCP's error middleware to a client-facing "Invalid params"
+ error rather than an internal-error code) when the request names
+ none and no startup default (``--metalake``) is configured.
+ """
+ metalake = _get_request_metalake() or self._setting.metalake
+ if not metalake:
+ raise ValueError(
+ f"No metalake specified: the request omitted the "
+ f"{METALAKE_HEADER!r} header and no --metalake default is "
+ "configured."
+ )
+ return metalake
+
+ def _service_client(self, metalake: str):
+ """Return the service-identity client (static token / OAuth) for
``metalake``."""
+ if (
+ metalake == self._setting.metalake
+ and self._default_client is not None
+ ):
+ return self._default_client
+
+ cached = self._service_clients.get(metalake)
+ if cached is not None:
+ self._service_clients.move_to_end(metalake)
+ return cached
+
+ client = RESTClientFactory.create_rest_client(
+ metalake,
+ self._setting.gravitino_uri,
+ startup_authorization(self._setting),
+ auth=_service_auth(self._setting),
+ )
+ self._cache_put(self._service_clients, metalake, client)
+ return client
+
+ def _cache_put(self, cache: "OrderedDict", key, client) -> None:
+ """Insert into an LRU cache, evicting (and closing) the oldest past
the cap."""
+ cache[key] = client
+ if len(cache) > _MAX_CACHED_CLIENTS:
+ _, evicted = cache.popitem(last=False)
+ self._schedule_close(evicted)
Review Comment:
Same fix as above — one shared LRU, so the cap is a global bound on open
pools rather than per cache.
_🤖 Addressed by [Claude Code](https://claude.com/claude-code)_
##########
mcp-server/tests/unit/test_per_request_metalake.py:
##########
@@ -0,0 +1,392 @@
+# 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 per-request metalake resolution.
+
+GravitinoContext.rest_client() must resolve the metalake from the current HTTP
+request's X-Gravitino-Metalake header, falling back to the configured startup
+default, so a single server instance can serve more than one metalake without
+any per-connection/session state (multi-node safe by construction).
+"""
+
+import asyncio
+import sys
+import unittest
+from unittest import mock
+from unittest.mock import MagicMock, patch
+
+from mcp_server.client.factory import RESTClientFactory
+from mcp_server.client.plain.plain_rest_client_operation import (
+ PlainRESTClientOperation,
+)
+from mcp_server.core import context as context_module
+from mcp_server.core.context import (
+ METALAKE_HEADER,
+ GravitinoContext,
+ ServiceIdentityFallbackDisabled,
+ _get_request_metalake,
+)
+from mcp_server.core.setting import Setting
+from mcp_server.main import _parse_args, do_main
+
+# Tests intentionally exercise context/client internals (e.g. _default_client,
+# _service_clients, _catalog_operation) to assert per-request isolation;
+# protected access is expected.
+# pylint: disable=protected-access
+
+
+def _mock_request(headers: dict) -> MagicMock:
+ """A fake HTTP request whose headers.get() only knows the given keys."""
+ mock_request = MagicMock()
+ mock_request.headers.get.side_effect = lambda key, default="": headers.get(
+ key, default
+ )
+ return mock_request
+
+
+class TestGetRequestMetalake(unittest.TestCase):
+ """Unit tests for _get_request_metalake() (HTTP context extraction)."""
+
+ def test_returns_header_value(self):
+ with patch(
+ "fastmcp.server.dependencies.get_http_request",
+ return_value=_mock_request({METALAKE_HEADER: "ml_a"}),
+ ):
+ self.assertEqual(_get_request_metalake(), "ml_a")
+
+ def test_returns_empty_when_no_http_context(self):
+ """Simulates stdio mode where get_http_request raises RuntimeError."""
+ with patch(
+ "fastmcp.server.dependencies.get_http_request",
+ side_effect=RuntimeError("no request context"),
+ ):
+ self.assertEqual(_get_request_metalake(), "")
+
+ def test_returns_empty_when_header_absent(self):
+ with patch(
+ "fastmcp.server.dependencies.get_http_request",
+ return_value=_mock_request({}),
+ ):
+ self.assertEqual(_get_request_metalake(), "")
+
+ def test_whitespace_only_header_is_treated_as_absent(self):
+ """A whitespace-only header must not be mistaken for a real metalake
name."""
+ with patch(
+ "fastmcp.server.dependencies.get_http_request",
+ return_value=_mock_request({METALAKE_HEADER: " "}),
+ ):
+ self.assertEqual(_get_request_metalake(), "")
+
+
+class TestGravitinoContextPerRequestMetalake(unittest.TestCase):
+ """GravitinoContext.rest_client() isolates per-request metalakes."""
+
+ def setUp(self):
+ RESTClientFactory.set_rest_client(PlainRESTClientOperation)
+
+ def _make_context(self, metalake: str = "ml_default") -> GravitinoContext:
+ return GravitinoContext(
+ Setting(
+ metalake=metalake,
+ gravitino_uri="http://localhost:8090";,
+ transport="http",
+ )
+ )
+
+ def test_header_overrides_startup_default(self):
+ ctx = self._make_context()
+
+ with patch(
+ "fastmcp.server.dependencies.get_http_request",
+ return_value=_mock_request(
+ {
+ "authorization": "Bearer t",
+ METALAKE_HEADER: "ml_other",
+ }
+ ),
+ ):
+ client = ctx.rest_client()
+
+ self.assertEqual(client._catalog_operation.metalake_name, "ml_other")
+
+ def test_falls_back_to_startup_default_when_header_absent(self):
+ ctx = self._make_context()
+
+ with patch(
+ "fastmcp.server.dependencies.get_http_request",
+ side_effect=LookupError,
+ ):
+ client = ctx.rest_client()
+
+ self.assertIs(client, ctx._default_client)
+ self.assertEqual(client._catalog_operation.metalake_name, "ml_default")
+
+ def test_missing_metalake_raises(self):
+ """No startup default and no header -> explicit error, not a silent
guess."""
+ ctx = self._make_context(metalake="")
+
+ with patch(
+ "fastmcp.server.dependencies.get_http_request",
+ side_effect=LookupError,
+ ):
+ with self.assertRaises(ValueError):
+ ctx.rest_client()
+
+ def test_whitespace_only_header_falls_back_to_startup_default(self):
+ """A whitespace-only header must not be used as a literal metalake
name."""
+ ctx = self._make_context()
+
+ with patch(
+ "fastmcp.server.dependencies.get_http_request",
+ return_value=_mock_request(
+ {"authorization": "Bearer t", METALAKE_HEADER: " "}
+ ),
+ ):
+ client = ctx.rest_client()
+
+ self.assertEqual(client._catalog_operation.metalake_name, "ml_default")
+
+ def
test_missing_metalake_error_takes_priority_over_fallback_disabled(self):
+ """When both a missing metalake and a disabled service-identity
fallback
+ apply, the caller must see the fixable "no metalake" error
(ValueError),
+ not ServiceIdentityFallbackDisabled - metalake resolution runs
first."""
+ ctx = GravitinoContext(
+ Setting(
+ metalake="",
+ gravitino_uri="http://localhost:8090";,
+ transport="http",
+ token="static-token",
+ no_service_identity_fallback=True,
+ )
+ )
+
+ with patch(
+ "fastmcp.server.dependencies.get_http_request",
+ return_value=_mock_request({}),
+ ):
+ with self.assertRaises(ValueError) as raised:
+ ctx.rest_client()
+
+ self.assertNotIsInstance(
+ raised.exception, ServiceIdentityFallbackDisabled
+ )
+
+ def test_two_concurrent_requests_get_different_metalake_clients(self):
+ """Different requested metalakes must produce different client
instances,
+ even for the same caller identity. This is the concurrency guarantee
the
+ header-based design relies on instead of any per-connection session
state.
+ """
+ ctx = self._make_context()
+
+ with patch(
+ "fastmcp.server.dependencies.get_http_request",
+ return_value=_mock_request(
+ {"authorization": "Bearer t", METALAKE_HEADER: "ml_a"}
+ ),
+ ):
+ client_a = ctx.rest_client()
+
+ with patch(
+ "fastmcp.server.dependencies.get_http_request",
+ return_value=_mock_request(
+ {"authorization": "Bearer t", METALAKE_HEADER: "ml_b"}
+ ),
+ ):
+ client_b = ctx.rest_client()
Review Comment:
Rewritten as a genuinely concurrent test through the real MCP protocol: two
overlapping calls held in flight by a barrier, so neither can finish until both
have arrived. Mutation-tested — replacing the context variable with
process-global state fails it.
_🤖 Addressed by [Claude Code](https://claude.com/claude-code)_
--
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]
