jerryshao commented on code in PR #12960:
URL: https://github.com/apache/gravitino/pull/12960#discussion_r3966176441
##########
mcp-server/mcp_server/core/context.py:
##########
@@ -190,23 +255,77 @@ 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
+ self._cache_put(key, client)
+ return client
+
+ def _resolve_metalake(self) -> str:
+ """Resolve the metalake for the current call, tool argument 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 call names none and
+ no startup default (``--metalake``) is configured. The message is
+ written for the agent that will read it: it names the recovery path so
+ a model can correct itself instead of just reporting the failure.
+ """
+ metalake = get_request_metalake() or self._setting.metalake
+ if not metalake:
+ # Only point at the discovery tool when this deployment actually
+ # exposes it; a tag filter can hide it, and naming a tool the
+ # agent cannot call leaves it with no way forward.
+ recovery = (
+ "Call 'list_metalakes' to see the metalakes you can access, "
+ f"then retry this call with the '{METALAKE_ARGUMENT}' "
+ "argument set to one of them."
+ if self._setting.exposes_metalake_discovery()
+ else f"Retry this call with the '{METALAKE_ARGUMENT}' argument
"
+ "set to the metalake to use, or ask the user which one to use."
+ )
+ raise ValueError(f"No metalake specified. {recovery}")
+ 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
+
+ key = ("", metalake)
+ cached = self._clients_by_auth.get(key)
+ if cached is not None:
+ self._clients_by_auth.move_to_end(key)
+ return cached
+
+ client = RESTClientFactory.create_rest_client(
+ metalake,
+ self._setting.gravitino_uri,
+ startup_authorization(self._setting),
+ auth=_service_auth(self._setting),
+ )
+ self._cache_put(key, client)
+ return client
+
+ def _cache_put(self, key: "tuple[str, str]", client) -> None:
+ """Cache a client, evicting (and closing) the oldest past the cap."""
+ self._clients_by_auth[key] = client
if len(self._clients_by_auth) > _MAX_CACHED_CLIENTS:
_, evicted = self._clients_by_auth.popitem(last=False)
self._schedule_close(evicted)
Review Comment:
Fixed. Borrows are now tracked per tool call: eviction defers close() for a
client that is still serving a call, and the last borrower closes it.
Regression test added (`TestEvictionDoesNotCloseAClientInUse`) — it parks a
call, fills the cache past the cap, asserts the client is still open, then
asserts it closes once the call finishes. Mutation-tested: restoring the
immediate close fails it with exactly this symptom.
_🤖 Addressed by [Claude Code](https://claude.com/claude-code)_
##########
mcp-server/mcp_server/core/middleware.py:
##########
@@ -0,0 +1,121 @@
+# 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.
+
+from typing import Any, Dict, Sequence
+
+import mcp.types as mt
+from fastmcp.server.middleware.middleware import (
+ CallNext,
+ Middleware,
+ MiddlewareContext,
+)
+from fastmcp.tools.base import Tool, ToolResult
+
+from mcp_server.core.context import (
+ METALAKE_ARGUMENT,
+ reset_request_metalake,
+ set_request_metalake,
+)
+
+# Tools that never resolve a metalake, so advertising the argument on them
+# would offer the agent a knob that does nothing. `list_metalakes` is the
+# discovery tool itself (its whole point is working without a metalake) and
+# `metadata_type_to_fullname_formats` is pure computation that never calls
+# Gravitino. A tool missing from this set only gets a harmless no-op argument.
+TOOLS_WITHOUT_METALAKE = frozenset(
+ {"list_metalakes", "metadata_type_to_fullname_formats"}
+)
+
+_METALAKE_ARGUMENT_DESCRIPTION = (
+ "Metalake to operate on. Omit to use the server's configured default "
+ "metalake. Call 'list_metalakes' to discover which metalakes are "
+ "available to you."
+)
+
+
+def _schema_with_metalake(parameters: Dict[str, Any]) -> Dict[str, Any]:
+ """Return ``parameters`` with an optional ``metalake`` property added.
+
+ Copied rather than mutated so the registered Tool objects keep the schema
+ their functions actually declare; the argument exists only on the wire.
+ ``required`` is deliberately left alone - omitting the argument is what
+ every single-metalake deployment does.
+ """
+ schema = dict(parameters)
+ properties = dict(schema.get("properties") or {})
+ # Never shadow a parameter a tool declares itself.
+ if METALAKE_ARGUMENT in properties:
+ return parameters
+ properties[METALAKE_ARGUMENT] = {
+ "type": "string",
+ "description": _METALAKE_ARGUMENT_DESCRIPTION,
+ }
+ schema["properties"] = properties
+ return schema
+
+
+class MetalakeArgumentMiddleware(Middleware):
+ """Lets any tool call name the metalake it operates on.
+
+ Every tool gains an optional ``metalake`` argument without declaring it:
+ this middleware advertises it in each tool's input schema, strips it from
+ the incoming arguments before the tool function runs, and publishes it for
+ ``GravitinoContext.rest_client()`` to resolve against.
+
+ The value lives in a context variable for the duration of one tool call
+ only, so no metalake state is carried between calls or shared between
+ server replicas.
+ """
+
+ async def on_list_tools(
+ self,
+ context: MiddlewareContext[mt.ListToolsRequest],
+ call_next: CallNext[mt.ListToolsRequest, Sequence[Tool]],
+ ) -> Sequence[Tool]:
+ tools = await call_next(context)
+ return [
+ (
+ tool
+ if tool.name in TOOLS_WITHOUT_METALAKE
+ else tool.model_copy(
+ update={
+ "parameters": _schema_with_metalake(tool.parameters)
+ }
+ )
+ )
+ for tool in tools
+ ]
+
+ async def on_call_tool(
+ self,
+ context: MiddlewareContext[mt.CallToolRequestParams],
+ call_next: CallNext[mt.CallToolRequestParams, ToolResult],
+ ) -> ToolResult:
+ arguments = context.message.arguments
+ # Popped so the tool function never sees an argument it cannot accept.
+ metalake = (
+ arguments.pop(METALAKE_ARGUMENT, "")
+ if isinstance(arguments, dict)
+ else ""
+ )
+ token = set_request_metalake(metalake)
Review Comment:
Fixed, and thanks for the precise repro. The raw argument is now published
verbatim and validated in `_resolve_metalake()` rather than in the middleware,
so the rejection happens inside the error-handling and audit middleware instead
of outside them. Verified through the protocol: `false`, `0`, `[]` and `42` all
return `Invalid params` with no REST client constructed; `null` and an omitted
argument still mean the default. Tests cover each value and assert no client is
built.
_🤖 Addressed by [Claude Code](https://claude.com/claude-code)_
##########
mcp-server/mcp_server/core/context.py:
##########
@@ -190,23 +255,77 @@ 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
+ self._cache_put(key, client)
+ return client
+
+ def _resolve_metalake(self) -> str:
+ """Resolve the metalake for the current call, tool argument 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 call names none and
+ no startup default (``--metalake``) is configured. The message is
+ written for the agent that will read it: it names the recovery path so
+ a model can correct itself instead of just reporting the failure.
+ """
+ metalake = get_request_metalake() or self._setting.metalake
+ if not metalake:
+ # Only point at the discovery tool when this deployment actually
+ # exposes it; a tag filter can hide it, and naming a tool the
+ # agent cannot call leaves it with no way forward.
+ recovery = (
+ "Call 'list_metalakes' to see the metalakes you can access, "
+ f"then retry this call with the '{METALAKE_ARGUMENT}' "
+ "argument set to one of them."
+ if self._setting.exposes_metalake_discovery()
+ else f"Retry this call with the '{METALAKE_ARGUMENT}' argument
"
+ "set to the metalake to use, or ask the user which one to use."
+ )
+ raise ValueError(f"No metalake specified. {recovery}")
+ 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
+
+ key = ("", metalake)
+ cached = self._clients_by_auth.get(key)
+ if cached is not None:
+ self._clients_by_auth.move_to_end(key)
+ return cached
+
+ client = RESTClientFactory.create_rest_client(
+ metalake,
+ self._setting.gravitino_uri,
+ startup_authorization(self._setting),
+ auth=_service_auth(self._setting),
Review Comment:
Fixed — the service auth object is built once in `GravitinoContext.__init__`
and shared by the default, discovery and per-metalake clients. Test asserts
every metalake client receives the same instance.
_🤖 Addressed by [Claude Code](https://claude.com/claude-code)_
##########
mcp-server/mcp_server/tools/statistic.py:
##########
@@ -22,7 +22,6 @@ def load_statistic_tools(mcp: FastMCP):
@mcp.tool(tags={"statistic"})
async def list_statistics_for_metadata(
ctx: Context,
- metalake_name: str,
metadata_type: str,
Review Comment:
Agreed, and you are right that this needed compatibility: `metalake_name`
has shipped since v1.0.0. It is now accepted on both statistic tools as a
deprecated alias, folded into the shared metalake in the middleware, and no
longer advertised so new callers see only `metalake`. Supplying both with
different values is rejected before any REST call. Tests cover old, new,
agreeing, conflicting, not-advertised, and that the alias does not leak to
other tools.
_🤖 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]
