Copilot commented on code in PR #12960:
URL: https://github.com/apache/gravitino/pull/12960#discussion_r3948414755
##########
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:
The table still lists the default value as `none`, but the CLI now defaults
`--metalake` to an empty string for HTTP transport. Consider updating this row
to reflect the new behavior (e.g., default: empty/\"\" for HTTP; required: yes
for stdio, no for HTTP) so the argument table doesn’t contradict the
implementation.
##########
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:
`_MAX_CACHED_CLIENTS` is now applied independently to two caches
(`_clients_by_auth` and `_service_clients`), meaning the process can hold up to
roughly `2 * _MAX_CACHED_CLIENTS` pooled clients (plus `_default_client`). If
`_MAX_CACHED_CLIENTS` is intended to be a global cap on total open pools,
consider enforcing an overall limit across both caches (or using a single
shared LRU) to keep resource usage bounded as expected.
##########
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:
`_MAX_CACHED_CLIENTS` is now applied independently to two caches
(`_clients_by_auth` and `_service_clients`), meaning the process can hold up to
roughly `2 * _MAX_CACHED_CLIENTS` pooled clients (plus `_default_client`). If
`_MAX_CACHED_CLIENTS` is intended to be a global cap on total open pools,
consider enforcing an overall limit across both caches (or using a single
shared LRU) to keep resource usage bounded as expected.
##########
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()
Review Comment:
The `_clients_by_auth` type annotation uses a bare `tuple`, which loses the
meaning of the key structure. Consider tightening it to something like
`OrderedDict[tuple[str, str], object]` (Authorization, metalake) to make the
cache contract clearer to readers and static analyzers.
##########
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:
This test name/docstring claims concurrency, but the calls are executed
sequentially. Either (a) rename the test to avoid asserting concurrent
behavior, or (b) make it truly concurrent (e.g., run two tasks and verify
isolation) so it actually validates the concurrency guarantee described in the
docstring/PR description.
--
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]
