diff --git a/docs/gravitino-mcp-server.md b/docs/gravitino-mcp-server.md index 304674766b6..e97bdc6b69f 100644 --- a/docs/gravitino-mcp-server.md +++ b/docs/gravitino-mcp-server.md @@ -71,6 +71,7 @@ Gravitino MCP server supports the following tools, and you could export tool by | Tool name | Description | Tag | |-------------------------------------|--------------------------------------------------------------------------------|--------------| +| `list_metalakes` | Retrieve the metalakes the caller can access. | `metalake` | | `get_list_of_catalogs` | Retrieve a list of all catalogs in the system. | `catalog` | | `create_catalog` | Create a new catalog. | `catalog` | | `alter_catalog` | Alter an existing catalog. | `catalog` | @@ -144,7 +145,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 metalake, used by any tool call that does not name one. See Selecting a metalake. | none | No | | `--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 | @@ -235,6 +236,48 @@ uv run mcp_server --metalake test --gravitino-uri http://127.0.0.1:8090 \ --tls-cert /path/to/cert.pem --tls-key /path/to/key.pem ``` +## Selecting a metalake + +A metalake is Gravitino's top-level tenant boundary, and every tool operates inside one. `--metalake` sets the **default**: the metalake used by any tool call that does not name one itself. It is optional on every transport. + +Any tool call may name a different metalake with a `metalake` argument, which takes priority over the default. The argument is optional on every tool, so a server configured with `--metalake` behaves exactly as it always has for callers that ignore it. + +The metalake for a call is resolved in this order: + +1. The call's own `metalake` argument, when it passes one. +2. The `--metalake` startup default, when it is configured. +3. Otherwise the call fails, telling the agent to call `list_metalakes` and retry. + +Because each call carries its own metalake, one server instance can serve several metalakes at once: nothing is remembered between calls, so concurrent callers never see each other's metalake and the server stays correct however many replicas it runs as. This works identically over stdio and HTTP. + +Use the `list_metalakes` tool to discover which metalakes a caller may use. It is the one tool that does not need a metalake, so it works on a server started with no `--metalake` at all. + +The statistic tools (`list_statistics_for_metadata`, `list_statistics_for_partition`) shipped their own `metalake_name` argument before metalake selection was unified. It is still accepted as a deprecated alias for `metalake`, so existing callers keep working; passing both with different values is rejected. New callers should use `metalake`. + +Authorization is unchanged — the caller's identity (see above) determines what it may see in the named metalake exactly as it would through the REST API. Note that a caller can now reach any metalake its credentials permit, so scope the credentials accordingly when that matters. + +### Examples + +Single metalake, agents never think about it — the common case, and unchanged: + +```bash +uv run mcp_server --metalake test --gravitino-uri http://127.0.0.1:8090 +``` + +Several metalakes behind one server, with `prod` as the default: + +```bash +uv run mcp_server --metalake prod --transport http --mcp-url http://0.0.0.0:8000/mcp +``` + +An agent then works in `prod` by default and switches per request when asked — "which catalogs are in the staging metalake?" sends `metalake=staging` on that call alone, without restarting or reconfiguring anything. + +No default at all, every call chooses: + +```bash +uv run mcp_server --transport http --mcp-url http://0.0.0.0:8000/mcp +``` + ## 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 incoming HTTP `Authorization` header when present; otherwise to the configured service identity (`--token` or OAuth client id). diff --git a/mcp-server/mcp_server/client/__init__.py b/mcp-server/mcp_server/client/__init__.py index dfddc047dd3..3565de20078 100644 --- a/mcp-server/mcp_server/client/__init__.py +++ b/mcp-server/mcp_server/client/__init__.py @@ -18,6 +18,7 @@ from mcp_server.client.catalog_operation import CatalogOperation from mcp_server.client.gravitino_operation import GravitinoOperation from mcp_server.client.job_operation import JobOperation +from mcp_server.client.metalake_operation import MetalakeOperation from mcp_server.client.model_operation import ModelOperation from mcp_server.client.policy_operation import PolicyOperation from mcp_server.client.schema_operation import SchemaOperation diff --git a/mcp-server/mcp_server/client/gravitino_operation.py b/mcp-server/mcp_server/client/gravitino_operation.py index 0943ab0120a..763ef02da88 100644 --- a/mcp-server/mcp_server/client/gravitino_operation.py +++ b/mcp-server/mcp_server/client/gravitino_operation.py @@ -20,6 +20,7 @@ from mcp_server.client.catalog_operation import CatalogOperation from mcp_server.client.fileset_operation import FilesetOperation from mcp_server.client.job_operation import JobOperation +from mcp_server.client.metalake_operation import MetalakeOperation from mcp_server.client.model_operation import ModelOperation from mcp_server.client.partition_operation import PartitionOperation from mcp_server.client.policy_operation import PolicyOperation @@ -154,3 +155,13 @@ def as_view_operation(self) -> ViewOperation: ViewOperation: Interface for performing view-level operations """ pass + + @abstractmethod + def as_metalake_operation(self) -> MetalakeOperation: + """ + Access the metalake operation interface of this Gravitino operation. + + Returns: + MetalakeOperation: Interface for metalake-level operations + """ + pass diff --git a/mcp-server/mcp_server/client/metalake_operation.py b/mcp-server/mcp_server/client/metalake_operation.py new file mode 100644 index 00000000000..5109ff4661e --- /dev/null +++ b/mcp-server/mcp_server/client/metalake_operation.py @@ -0,0 +1,38 @@ +# 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 abc import ABC, abstractmethod + + +class MetalakeOperation(ABC): + """ + Abstract base class for Gravitino metalake operations. + + Unlike every other operation, these are not scoped to a single metalake: + they address the server's top-level ``/api/metalakes`` endpoint, so an + agent can discover which metalakes it may operate on before naming one. + """ + + @abstractmethod + async def get_list_of_metalakes(self) -> str: + """ + Retrieve the list of metalakes the caller is allowed to see. + + Returns: + str: JSON-formatted string containing metalake information. + """ + pass diff --git a/mcp-server/mcp_server/client/plain/plain_rest_client_metalake_operation.py b/mcp-server/mcp_server/client/plain/plain_rest_client_metalake_operation.py new file mode 100644 index 00000000000..714e2712dae --- /dev/null +++ b/mcp-server/mcp_server/client/plain/plain_rest_client_metalake_operation.py @@ -0,0 +1,36 @@ +# 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 httpx import AsyncClient + +from mcp_server.client.metalake_operation import MetalakeOperation +from mcp_server.client.plain.utils import extract_content_from_response + + +class PlainRESTClientMetalakeOperation(MetalakeOperation): + """Metalake operations against the server's top-level endpoint. + + Takes no metalake name: this is the one operation that must work before a + metalake has been chosen. + """ + + def __init__(self, rest_client: AsyncClient): + self.rest_client = rest_client + + async def get_list_of_metalakes(self) -> str: + response = await self.rest_client.get("/api/metalakes") + return extract_content_from_response(response, "metalakes", []) 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 0d35525ffbe..63187bedea2 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 @@ -22,6 +22,7 @@ from mcp_server.client import ( CatalogOperation, GravitinoOperation, + MetalakeOperation, ModelOperation, PolicyOperation, SchemaOperation, @@ -38,6 +39,9 @@ from mcp_server.client.plain.plain_rest_client_job_operation import ( PlainRESTClientJobOperation, ) +from mcp_server.client.plain.plain_rest_client_metalake_operation import ( + PlainRESTClientMetalakeOperation, +) from mcp_server.client.plain.plain_rest_client_model_operation import ( PlainRESTClientModelOperation, ) @@ -134,11 +138,18 @@ def __init__( self._view_operation = PlainRESTClientViewOperation( metalake_name, _rest_client ) + # Not metalake-scoped: addresses the server's top-level endpoint. + self._metalake_operation = PlainRESTClientMetalakeOperation( + _rest_client + ) async def close(self) -> None: """Close the shared httpx client and release its connection pool.""" await self._rest_client.aclose() + def as_metalake_operation(self) -> MetalakeOperation: + return self._metalake_operation + def as_catalog_operation(self) -> CatalogOperation: return self._catalog_operation diff --git a/mcp-server/mcp_server/client/plain/plain_rest_client_statistic_operation.py b/mcp-server/mcp_server/client/plain/plain_rest_client_statistic_operation.py index 49a56036a55..2befb1062d0 100644 --- a/mcp-server/mcp_server/client/plain/plain_rest_client_statistic_operation.py +++ b/mcp-server/mcp_server/client/plain/plain_rest_client_statistic_operation.py @@ -28,10 +28,10 @@ def __init__(self, metalake_name: str, rest_client): self.rest_client = rest_client async def list_of_statistics( - self, metalake_name: str, metadata_type: str, metadata_fullname: str + self, metadata_type: str, metadata_fullname: str ) -> str: response = await self.rest_client.get( - f"/api/metalakes/{encode_path_segment(metalake_name)}" + f"/api/metalakes/{encode_path_segment(self.metalake_name)}" f"/objects/{encode_path_segment(metadata_type)}" f"/{encode_path_segment(metadata_fullname)}/statistics" ) @@ -40,7 +40,6 @@ async def list_of_statistics( # pylint: disable=R0917 async def list_statistic_for_partition( self, - metalake_name: str, metadata_type: str, metadata_fullname: str, from_partition_name: str, @@ -49,7 +48,7 @@ async def list_statistic_for_partition( to_inclusive: bool = False, ) -> str: response = await self.rest_client.get( - f"/api/metalakes/{encode_path_segment(metalake_name)}" + f"/api/metalakes/{encode_path_segment(self.metalake_name)}" f"/objects/{encode_path_segment(metadata_type)}" f"/{encode_path_segment(metadata_fullname)}/statistics/partitions", params={ diff --git a/mcp-server/mcp_server/client/statistic_operation.py b/mcp-server/mcp_server/client/statistic_operation.py index b9c2f1f99da..5fd2165e680 100644 --- a/mcp-server/mcp_server/client/statistic_operation.py +++ b/mcp-server/mcp_server/client/statistic_operation.py @@ -25,12 +25,11 @@ class StatisticOperation(ABC): @abstractmethod async def list_of_statistics( - self, metalake_name: str, metadata_type: str, metadata_fullname: str + self, metadata_type: str, metadata_fullname: str ) -> str: """ Retrieve the list of statistics for a specific metadata type and fullname within a metalake. Args: - metalake_name: Name of the metalake metadata_type: Type of metadata (e.g., table, column) metadata_fullname: Full name of the metadata item @@ -43,7 +42,6 @@ async def list_of_statistics( @abstractmethod async def list_statistic_for_partition( self, - metalake_name: str, metadata_type: str, metadata_fullname: str, from_partition_name: str, @@ -57,7 +55,6 @@ async def list_statistic_for_partition( So `metadata_type` should always be "table". Args: - metalake_name: Name of the metalake metadata_type: Type of metadata, should be "table" for partition statistics metadata_fullname: Full name of the metadata item, the format should be "{catalog}.{schema}.{table}". diff --git a/mcp-server/mcp_server/core/audit.py b/mcp-server/mcp_server/core/audit.py index 0d25ef2f249..a8d7d6224eb 100644 --- a/mcp-server/mcp_server/core/audit.py +++ b/mcp-server/mcp_server/core/audit.py @@ -63,6 +63,7 @@ def emit( tool: str, outcome: str, error_type: str = "", + metalake: str = "", ) -> None: """Write one structured JSON audit record to the audit logger. @@ -74,6 +75,13 @@ def emit( authorization denial being the common case), not only authorization failures; inspect error_type to disambiguate. error_type: Exception class name when outcome is "deny", empty otherwise. + metalake: Metalake the call operated on, resolved - so a call that + relied on the server default records that default. Empty + only for tools that are not metalake-scoped, such as the + metalake listing, which spans every tenant the caller can + see. Recorded because one server can now serve several + metalakes, so "which tenant did this touch" is no longer + answerable from the server config alone. """ record = { "timestamp": datetime.now(timezone.utc).isoformat(), @@ -81,6 +89,8 @@ def emit( "tool": tool, "outcome": outcome, } + if metalake: + record["metalake"] = metalake if error_type: record["error_type"] = error_type diff --git a/mcp-server/mcp_server/core/context.py b/mcp-server/mcp_server/core/context.py index f38831f6394..0bb8f7af733 100644 --- a/mcp-server/mcp_server/core/context.py +++ b/mcp-server/mcp_server/core/context.py @@ -19,6 +19,8 @@ import logging import re from collections import OrderedDict +from contextvars import ContextVar +from typing import Any from mcp_server.client.factory import RESTClientFactory from mcp_server.core.oauth import RefreshableBearerAuth @@ -26,11 +28,12 @@ _LOG = logging.getLogger(__name__) -# Upper bound on the number of per-principal REST clients kept alive at once. -# Each client owns an httpx connection pool; caching by Authorization header lets -# repeated calls from the same principal reuse a pool instead of opening a new one -# per tool call, while the LRU bound keeps memory/sockets in check as principals -# (e.g. rotating tokens) come and go. +# Upper bound on the number of cached REST clients kept alive at once, across +# every (principal, metalake) combination. Each client owns an httpx connection +# pool; caching lets repeated calls from the same principal against the same +# metalake reuse a pool instead of opening a new one per tool call, while the +# LRU bound keeps memory/sockets in check as principals (e.g. rotating tokens) +# and metalakes come and go. _MAX_CACHED_CLIENTS = 128 # An RFC 9110 auth-scheme uses the HTTP token syntax. Here it must be followed by @@ -50,6 +53,65 @@ "negotiate": "Negotiate", } +# Name of the optional argument that every tool accepts to name the metalake +# it should operate on. The argument is not declared on any tool function: +# MetalakeArgumentMiddleware advertises it in each tool's input schema, strips +# it from the incoming arguments, and publishes it on _REQUEST_METALAKE below. +METALAKE_ARGUMENT = "metalake" + +# Sentinel for "the call did not pass the argument at all", so an explicitly +# supplied bad value is never mistaken for an omitted one. +MISSING_METALAKE = object() + + +class InvalidMetalakeArgument: + """An argument the middleware rejected, carrying the reason to report. + + The middleware runs outside the error-handling and audit middleware, so it + cannot raise directly without bypassing both. It publishes this instead and + _resolve_metalake() raises inside them, before any REST call is made. + """ + + def __init__(self, reason: str): + self.reason = reason + + +# The metalake argument of the tool call currently being served, exactly as the +# client sent it, or MISSING_METALAKE when it sent none. Held raw rather than +# validated so that _resolve_metalake() - which runs inside the error-handling +# and audit middleware - is what rejects a bad value; validating in the +# outermost middleware would bypass both. Scoped to a single tool invocation +# (the middleware resets it in a finally block), so this is request plumbing, +# not session state: nothing is remembered between calls and no state is +# shared between server replicas. +_REQUEST_METALAKE: ContextVar[Any] = ContextVar( + "request_metalake", default=MISSING_METALAKE +) + +# Clients handed out during the tool call currently being served, as +# (context, client) pairs. A client evicted from the cache while it is still +# serving a call must not have its connection pool closed underneath that call, +# so eviction defers the close until the last borrower releases it. None when +# no call is in flight (lifespan setup, direct unit tests), in which case +# nothing is tracked and eviction closes immediately as before. +_BORROWED_CLIENTS: ContextVar[Any] = ContextVar( + "borrowed_clients", default=None +) + + +def begin_request_clients(): + """Start tracking the clients this tool call borrows.""" + return _BORROWED_CLIENTS.set([]) + + +def release_request_clients(token) -> None: + """Release every client this tool call borrowed, closing evicted ones.""" + borrowed = _BORROWED_CLIENTS.get() or [] + _BORROWED_CLIENTS.reset(token) + for owner, client in borrowed: + # pylint: disable=protected-access + owner._release_client(client) + class ServiceIdentityFallbackDisabled(RuntimeError): """HTTP omitted Authorization while service-identity fallback is disabled.""" @@ -87,6 +149,32 @@ def _get_request_authorization() -> str: return "" +def set_request_metalake(metalake): + """Publish the raw ``metalake`` argument of the current tool call. + + Takes the value verbatim - validation happens in + :meth:`GravitinoContext._resolve_metalake`. Returns the token the caller + must pass to :func:`reset_request_metalake` once the call finishes, so + nothing leaks into the next one. + """ + return _REQUEST_METALAKE.set(metalake) + + +def reset_request_metalake(token) -> None: + """Undo :func:`set_request_metalake` at the end of a tool call.""" + _REQUEST_METALAKE.reset(token) + + +def get_request_metalake() -> str: + """The metalake this call named, or "" when it named none or named it badly. + + Never raises: audit logging calls this on the failure path too, where the + value may be exactly the malformed input that caused the failure. + """ + metalake = _REQUEST_METALAKE.get() + return metalake.strip() if isinstance(metalake, str) else "" + + def startup_authorization(setting: Setting) -> str: """The static --token rendered as an ``Authorization`` header value. @@ -147,38 +235,79 @@ 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. + 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), + # Built once and shared by every service-identity client. Separate + # instances share the token cache key but each owns its own refresh + # lock and 401 retry state, so per-metalake instances would bypass + # refresh coalescing and hit the IdP once per metalake. + self._service_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=self._service_auth, + ) + if setting.metalake + else None ) - # LRU cache of per-principal clients keyed by the raw Authorization header. + # One LRU cache for every cached client, keyed by (Authorization + # header, metalake), so _MAX_CACHED_CLIENTS bounds the total number of + # open connection pools rather than being applied per cache. An empty + # Authorization means the service identity (static token / OAuth), + # which is only ever produced by the no-Authorization branch of + # rest_client(), so it can never collide with a real principal's key. # 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() + self._clients_by_auth: "OrderedDict[tuple[str, str], object]" = ( + OrderedDict() + ) + # How many in-flight calls are using each handed-out client, and the + # clients evicted while still in use, to be closed once idle. + self._borrows: "dict" = {} + self._close_when_idle: "set" = set() # Strong references to in-flight background close tasks; the event loop # only keeps weak references, so without this they could be GC'd before # running. Entries are discarded when each task completes. self._pending_closes: "set[asyncio.Task]" = set() - def rest_client(self): - """Return a REST client carrying the correct identity for this request. + def rest_client(self, *, require_metalake: bool = True): + """Return a REST client carrying the correct identity and metalake. + + The metalake is resolved per call: the ``metalake`` argument of the + current tool call takes priority, falling back to the configured + startup default (``--metalake``). Raises ``ValueError`` when neither is + available. + + ``require_metalake=False`` skips metalake resolution entirely, for the + metalake-listing tool: it is what an agent calls when it does not know + a metalake yet, so it must work on a server with no default configured. + The returned client can only be used for operations that are not + metalake-scoped. - In HTTP transport mode the incoming request's ``Authorization`` header is - forwarded verbatim to Gravitino, taking priority over the static startup - token. This keeps concurrent sessions with different principals fully - isolated — one principal's identity never leaks into another's calls. + Identity resolution is unchanged: in HTTP transport mode the incoming + request's ``Authorization`` header is forwarded verbatim to Gravitino, + taking priority over the static startup token. This keeps concurrent + sessions with different principals and/or metalakes fully isolated — + one caller's identity or metalake never leaks into another's calls. - Falls back to the shared default client (static token or OAuth + Falls back to a service-identity client (static token or OAuth client-credentials) when: - running in stdio mode (no HTTP request context), or - the incoming request carries no Authorization header. - Per-principal clients are cached (and their connection pools reused) so a - new pool is not opened on every tool call. + Clients are cached per (identity, metalake) combination (and their + connection pools reused) so a new pool is not opened on every call. """ + metalake = self._resolve_metalake() if require_metalake else "" authorization = _get_request_authorization() if not authorization: if ( @@ -190,23 +319,118 @@ def rest_client(self): "HTTP request omitted Authorization and " "--no-service-identity-fallback is set" ) + return self._borrow(self._service_client(metalake)) + + key = (authorization, metalake) + cached = self._clients_by_auth.get(key) + if cached is not None: + self._clients_by_auth.move_to_end(key) + return self._borrow(cached) + + client = RESTClientFactory.create_rest_client( + metalake, + self._setting.gravitino_uri, + authorization, + ) + self._cache_put(key, client) + return self._borrow(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. + """ + requested = _REQUEST_METALAKE.get() + if isinstance(requested, InvalidMetalakeArgument): + raise ValueError(requested.reason) + if requested is not MISSING_METALAKE and not isinstance( + requested, (str, type(None)) + ): + # An explicitly supplied non-string must not be silently treated as + # an omitted argument: `false`, `0` and `[]` would otherwise route + # a call - including a mutation - to the default metalake. + raise ValueError( + f"The '{METALAKE_ARGUMENT}' argument must be a string naming a " + f"metalake, but got {type(requested).__name__}." + ) + + 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 - cached = self._clients_by_auth.get(authorization) + key = ("", 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, + startup_authorization(self._setting), + auth=self._service_auth, ) - self._clients_by_auth[authorization] = client + self._cache_put(key, client) + return client + + def _borrow(self, client): + """Mark ``client`` as in use for the duration of the current call.""" + borrowed = _BORROWED_CLIENTS.get() + if borrowed is None: + # Not inside a tool call: nothing will release the borrow, so + # tracking it would pin the client forever. + return client + self._borrows[client] = self._borrows.get(client, 0) + 1 + borrowed.append((self, client)) + return client + + def _release_client(self, client) -> None: + """Drop one borrow, closing the client if it was evicted while in use.""" + remaining = self._borrows.get(client, 0) - 1 + if remaining > 0: + self._borrows[client] = remaining + return + self._borrows.pop(client, None) + if client in self._close_when_idle: + self._close_when_idle.discard(client) + self._schedule_close(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) - return client + if self._borrows.get(evicted): + # Still serving a call: closing now would drop that call's + # connection mid-request. The last borrower closes it instead. + self._close_when_idle.add(evicted) + else: + self._schedule_close(evicted) def _schedule_close(self, client) -> None: """Best-effort close of an evicted client's connection pool. diff --git a/mcp-server/mcp_server/core/middleware.py b/mcp-server/mcp_server/core/middleware.py new file mode 100644 index 00000000000..54768b399d3 --- /dev/null +++ b/mcp-server/mcp_server/core/middleware.py @@ -0,0 +1,171 @@ +# 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, + MISSING_METALAKE, + InvalidMetalakeArgument, + begin_request_clients, + release_request_clients, + 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"} +) + +# Tools that shipped their own `metalake_name` argument before metalake +# selection was unified (v1.0.0). It is still accepted as a deprecated alias so +# existing callers and clients holding a cached schema keep working, but it is +# no longer advertised: new callers see only `metalake`. +DEPRECATED_METALAKE_NAME_ARGUMENT = "metalake_name" +_TOOLS_WITH_METALAKE_NAME_ALIAS = frozenset( + {"list_statistics_for_metadata", "list_statistics_for_partition"} +) + +_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 + + +def _apply_metalake_name_alias(tool_name, arguments, metalake): + """Fold a legacy ``metalake_name`` argument into the shared metalake. + + Returns the metalake to publish: the alias when only it was given, the + canonical argument otherwise, or a rejection when a call supplies both + with different values - silently picking one could send a write to the + wrong tenant. + """ + if tool_name not in _TOOLS_WITH_METALAKE_NAME_ALIAS or not isinstance( + arguments, dict + ): + return metalake + alias = arguments.pop(DEPRECATED_METALAKE_NAME_ARGUMENT, MISSING_METALAKE) + if alias is MISSING_METALAKE: + return metalake + if metalake is not MISSING_METALAKE and metalake != alias: + return InvalidMetalakeArgument( + f"'{METALAKE_ARGUMENT}' and the deprecated " + f"'{DEPRECATED_METALAKE_NAME_ARGUMENT}' were both given with " + f"different values ({metalake!r} and {alias!r}). Pass only " + f"'{METALAKE_ARGUMENT}'." + ) + return alias + + +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, + # and passed on verbatim: popping it here removes it from FastMCP's own + # schema validation, so _resolve_metalake() has to be able to tell an + # omitted argument from an explicitly supplied bad one. + metalake = ( + arguments.pop(METALAKE_ARGUMENT, MISSING_METALAKE) + if isinstance(arguments, dict) + else MISSING_METALAKE + ) + metalake = _apply_metalake_name_alias( + context.message.name if context.message else "", + arguments, + metalake, + ) + token = set_request_metalake(metalake) + clients_token = begin_request_clients() + try: + return await call_next(context) + finally: + # Release before the metalake reset so a client evicted while this + # call was running is closed now rather than leaking. + release_request_clients(clients_token) + # Without this the metalake would leak into the next call served on + # this context, turning per-call plumbing into implicit state. + reset_request_metalake(token) diff --git a/mcp-server/mcp_server/core/setting.py b/mcp-server/mcp_server/core/setting.py index 112528f6d98..bb69175a6b2 100644 --- a/mcp-server/mcp_server/core/setting.py +++ b/mcp-server/mcp_server/core/setting.py @@ -18,6 +18,11 @@ from dataclasses import dataclass, field from typing import Set +# Tag carried by the metalake discovery tools. Defined here because +# Setting is what interprets --include-tool-tags; tools/metalake.py +# imports it so the tag and the check below cannot drift apart. +METALAKE_TOOL_TAG = "metalake" + @dataclass class DefaultSetting: @@ -28,7 +33,10 @@ class DefaultSetting: @dataclass class Setting: # pylint: disable=too-many-instance-attributes - metalake: str + # Default metalake, used by any tool call that does not name one itself + # via the `metalake` argument. Optional on every transport: a deployment + # serving several metalakes can leave it unset and let each call choose. + metalake: str = "" gravitino_uri: str = DefaultSetting.default_gravitino_uri tags: Set[str] = field(default_factory=set) transport: str = DefaultSetting.default_transport @@ -56,6 +64,12 @@ class Setting: # pylint: disable=too-many-instance-attributes # --token is configured instead of falling back to the service identity. no_service_identity_fallback: bool = False + def __post_init__(self) -> None: + # A whitespace-only --metalake (e.g. a shell-quoting mistake) must be + # treated as "no default configured", the same as an empty string, + # rather than silently used as a nonsensical metalake name. + self.metalake = self.metalake.strip() + def has_oauth_client(self) -> bool: """Return True when client-credentials is fully configured.""" return bool( @@ -64,6 +78,15 @@ def has_oauth_client(self) -> bool: and self.oauth_client_secret.strip() ) + def exposes_metalake_discovery(self) -> bool: + """Whether the `list_metalakes` tool is reachable in this deployment. + + --include-tool-tags is an allowlist, so a tag filter that omits + "metalake" hides the discovery tool. Callers use this to avoid telling + an agent to call a tool it cannot see. + """ + return not self.tags or METALAKE_TOOL_TAG in self.tags + 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() diff --git a/mcp-server/mcp_server/main.py b/mcp-server/mcp_server/main.py index fd172d3fb1f..2a50b8f68dc 100644 --- a/mcp-server/mcp_server/main.py +++ b/mcp-server/mcp_server/main.py @@ -22,6 +22,7 @@ from mcp_server.core.setting import DefaultSetting, Setting from mcp_server.server import ( GravitinoMCPServer, + log_metalake_policy, log_service_identity_fallback_policy, ) from mcp_server.tools import SUPPORTED_TOOL_TAGS @@ -54,6 +55,7 @@ def do_main(): except ValueError as exc: logging.error("%s", exc) raise SystemExit(1) from None + log_metalake_policy(setting) log_service_identity_fallback_policy(setting) logging.info("Gravitino MCP server setting: %s", setting) server = GravitinoMCPServer(setting) @@ -90,8 +92,10 @@ def _parse_args(): parser.add_argument( "--metalake", type=str, - required=True, - help="Gravitino metalake name.", + default="", + help="Default Gravitino metalake name, used by any tool call that " + "does not name one itself via its 'metalake' argument. Optional: a " + "server with no default serves whichever metalake each call names.", ) parser.add_argument( "--gravitino-uri", diff --git a/mcp-server/mcp_server/server.py b/mcp-server/mcp_server/server.py index 56cca377cb3..688e0d17d62 100644 --- a/mcp-server/mcp_server/server.py +++ b/mcp-server/mcp_server/server.py @@ -37,10 +37,16 @@ from mcp_server.core import audit from mcp_server.core.context import ( + METALAKE_ARGUMENT, GravitinoContext, _get_request_authorization, + get_request_metalake, service_fallback_authorization, ) +from mcp_server.core.middleware import ( + TOOLS_WITHOUT_METALAKE, + MetalakeArgumentMiddleware, +) from mcp_server.core.setting import Setting from mcp_server.tools import load_tools @@ -62,9 +68,25 @@ def _get_principal_from_request(fallback_authorization: str = "") -> str: class AuditMiddleware(Middleware): """Emit a structured audit record for every tool invocation.""" - def __init__(self, fallback_authorization: str = ""): + def __init__( + self, fallback_authorization: str = "", default_metalake: str = "" + ): super().__init__() self._fallback_authorization = fallback_authorization + self._default_metalake = default_metalake + + def _metalake(self, tool_name: str) -> str: + """The metalake the call actually used, named or defaulted. + + Recording the resolved value keeps each record self-contained: an + auditor can tell which tenant was touched without joining against the + server's startup configuration. Tools that are not metalake-scoped + record nothing rather than the default, which they never touch - a + metalake listing spans every tenant the caller can see. + """ + if tool_name in TOOLS_WITHOUT_METALAKE: + return "" + return get_request_metalake() or self._default_metalake async def on_call_tool( self, @@ -75,7 +97,12 @@ async def on_call_tool( principal = _get_principal_from_request(self._fallback_authorization) try: result = await call_next(context) - audit.emit(principal=principal, tool=tool_name, outcome="allow") + audit.emit( + principal=principal, + tool=tool_name, + outcome="allow", + metalake=self._metalake(tool_name), + ) return result except Exception as exc: audit.emit( @@ -83,6 +110,7 @@ async def on_call_tool( tool=tool_name, outcome="deny", error_type=type(exc).__name__, + metalake=self._metalake(tool_name), ) raise @@ -106,7 +134,14 @@ def _create_gravitino_mcp(setting: Setting) -> FastMCP: # Allowlist mode: disable everything, then re-enable the wanted tags. mcp.enable(tags=setting.tags, only=True) - mcp.add_middleware(AuditMiddleware(service_fallback_authorization(setting))) + # Added first so it wraps the others: it must publish the call's metalake + # before AuditMiddleware records it. + mcp.add_middleware(MetalakeArgumentMiddleware()) + mcp.add_middleware( + AuditMiddleware( + service_fallback_authorization(setting), setting.metalake + ) + ) mcp.add_middleware( LoggingMiddleware(include_payloads=True, max_payload_length=1000) ) @@ -152,6 +187,29 @@ def log_service_identity_fallback_policy(setting: Setting) -> None: ) +def log_metalake_policy(setting: Setting) -> None: + """Log how the metalake is resolved for tool calls at startup.""" + if setting.metalake: + logging.info( + "Default metalake '%s' configured; a tool call may override it " + "with the '%s' argument.", + setting.metalake, + METALAKE_ARGUMENT, + ) + return + logging.info( + "No default --metalake configured; every tool call must name one " + "with the '%s' argument%s.", + METALAKE_ARGUMENT, + ( + " (see the 'list_metalakes' tool)" + if setting.exposes_metalake_discovery() + else ", and --include-tool-tags hides the 'list_metalakes' tool " + "that would let an agent discover them" + ), + ) + + def _parse_mcp_url(url: str) -> tuple[str, int, str]: try: parsed = urlparse(url) diff --git a/mcp-server/mcp_server/tools/__init__.py b/mcp-server/mcp_server/tools/__init__.py index e15f6ad8d12..4fd524c088a 100644 --- a/mcp-server/mcp_server/tools/__init__.py +++ b/mcp-server/mcp_server/tools/__init__.py @@ -21,6 +21,7 @@ from mcp_server.tools.fileset import load_fileset_tools from mcp_server.tools.job import load_job_tool from mcp_server.tools.metadata import load_metadata_tool +from mcp_server.tools.metalake import load_metalake_tools from mcp_server.tools.model import load_model_tools from mcp_server.tools.partition import load_partition_tools from mcp_server.tools.policy import load_policy_tools @@ -37,6 +38,7 @@ "catalog", "fileset", "job", + "metalake", "model", "partition", "policy", @@ -60,6 +62,7 @@ def load_tools(mcp: FastMCP): load_fileset_tools(mcp) load_tag_tool(mcp) load_metadata_tool(mcp) + load_metalake_tools(mcp) load_statistic_tools(mcp) load_policy_tools(mcp) load_partition_tools(mcp) diff --git a/mcp-server/mcp_server/tools/metalake.py b/mcp-server/mcp_server/tools/metalake.py new file mode 100644 index 00000000000..c2a93cac7b2 --- /dev/null +++ b/mcp-server/mcp_server/tools/metalake.py @@ -0,0 +1,66 @@ +# 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 fastmcp import Context, FastMCP + +from mcp_server.core.setting import METALAKE_TOOL_TAG + + +def load_metalake_tools(mcp: FastMCP): + @mcp.tool(tags={METALAKE_TOOL_TAG}) + async def list_metalakes(ctx: Context) -> str: + """ + Retrieve the list of metalakes the caller is allowed to access. + + A metalake is the top-level tenant boundary in Gravitino. Every other + tool operates inside one: it uses the server's configured default + metalake unless the call passes a `metalake` argument. Use this tool to + discover which values that argument accepts - for example when a call + failed because no metalake was specified, or when the user asks about a + metalake other than the default. + + Args: + ctx (Context): The request context. + + Returns: + str: A JSON string containing the list of metalakes. + + Example Return Value: + [ + { + "name": "metalake_a", + "comment": "Production metadata", + "properties": {}, + "audit": { + "creator": "anonymous", + "createTime": "2025-08-20T07:33:41.233089Z" + } + } + ] + + name: The name of the metalake, i.e. the value to pass as the + `metalake` argument of other tools. + comment: A human-readable description of the metalake. + properties: Metalake properties. + audit: Metadata about the metalake's creation and modification. + """ + # require_metalake=False: this is the tool an agent calls when it does + # not know a metalake yet, so it must work with none configured. + client = ctx.request_context.lifespan_context.rest_client( + require_metalake=False + ) + return await client.as_metalake_operation().get_list_of_metalakes() diff --git a/mcp-server/mcp_server/tools/statistic.py b/mcp-server/mcp_server/tools/statistic.py index 76dd1d9a6fc..8e5365fcba1 100644 --- a/mcp-server/mcp_server/tools/statistic.py +++ b/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, metadata_fullname: str, ) -> str: @@ -36,9 +35,8 @@ async def list_statistics_for_metadata( Args: ctx (Context): The request context. - metalake_name (str): The name of the metalake. metadata_type (str): The type of metadata (e.g., table, column). For - more, please refer to too 'metadata_type_to_fullname_formats' + more, please refer to the tool 'metadata_type_to_fullname_formats' metadata_fullname (str): The full name of the metadata object. For more, please refer to tool 'metadata_type_to_fullname_formats'. @@ -70,14 +68,13 @@ async def list_statistics_for_metadata( """ client = ctx.request_context.lifespan_context.rest_client() return await client.as_statistic_operation().list_of_statistics( - metalake_name, metadata_type, metadata_fullname + metadata_type, metadata_fullname ) # pylint: disable=R0917 @mcp.tool(tags={"statistic"}) async def list_statistics_for_partition( ctx: Context, - metalake_name: str, metadata_type: str, metadata_fullname: str, from_partition_name: str, @@ -92,7 +89,6 @@ async def list_statistics_for_partition( Args: ctx (Context): The request context. - metalake_name (str): The name of the metalake. metadata_type (str): The type of metadata, should be "table" for partition statistics. metadata_fullname (str): The full name of the metadata item, the format should be "{catalog}.{schema}.{table}". @@ -136,7 +132,6 @@ async def list_statistics_for_partition( client = ctx.request_context.lifespan_context.rest_client() return ( await client.as_statistic_operation().list_statistic_for_partition( - metalake_name, metadata_type, metadata_fullname, from_partition_name, diff --git a/mcp-server/tests/unit/client/test_url_encoding.py b/mcp-server/tests/unit/client/test_url_encoding.py index 3e1c9018dad..6ff7a182b1e 100644 --- a/mcp-server/tests/unit/client/test_url_encoding.py +++ b/mcp-server/tests/unit/client/test_url_encoding.py @@ -31,6 +31,9 @@ from mcp_server.client.plain.plain_rest_client_job_operation import ( PlainRESTClientJobOperation, ) +from mcp_server.client.plain.plain_rest_client_metalake_operation import ( + PlainRESTClientMetalakeOperation, +) from mcp_server.client.plain.plain_rest_client_model_operation import ( PlainRESTClientModelOperation, ) @@ -96,6 +99,36 @@ def _called_params(mock_method): METALAKE = "my_metalake" +class TestMetalakeOperation(unittest.TestCase): + """The one operation that is not scoped to a metalake. + + Everything else in the test suite reaches list_metalakes through + MockOperation, so without this the real endpoint path and the response key + the server actually returns are never executed. + """ + + def test_lists_metalakes_from_the_top_level_endpoint(self): + client = _make_mock_client( + {"metalakes": [{"name": "ml_a"}, {"name": "ml_b"}]} + ) + op = PlainRESTClientMetalakeOperation(client) + + result = asyncio.run(op.get_list_of_metalakes()) + + # Not under /api/metalakes/{metalake}/... - it must not be scoped. + self.assertEqual(_called_url(client.get), "/api/metalakes") + self.assertIn("ml_a", result) + self.assertIn("ml_b", result) + + def test_returns_the_default_when_the_response_has_no_metalakes_key(self): + """Guards the response key: a typo here would silently return nothing + rather than failing, and every mock-based test would still pass.""" + client = _make_mock_client({"code": 0}) + op = PlainRESTClientMetalakeOperation(client) + + self.assertEqual(asyncio.run(op.get_list_of_metalakes()), "[]") + + class TestCatalogOperationUrlEncoding(unittest.TestCase): def test_get_list_of_catalogs_encodes_metalake(self): client = _make_mock_client({"catalogs": []}) @@ -499,7 +532,7 @@ class TestStatisticOperationUrlEncoding(unittest.TestCase): def test_list_of_statistics_encodes_metadata_fullname(self): client = _make_mock_client({"statistics": []}) op = PlainRESTClientStatisticOperation(METALAKE, client) - asyncio.run(op.list_of_statistics(METALAKE, "table", _PATH_TRAVERSAL)) + asyncio.run(op.list_of_statistics("table", _PATH_TRAVERSAL)) url = _called_url(client.get) self.assertIn(_ENCODED_PATH_TRAVERSAL, url) self.assertNotIn("../../", url) @@ -510,7 +543,6 @@ def test_list_statistic_for_partition_uses_params_for_partition_names(self): op = PlainRESTClientStatisticOperation(METALAKE, client) asyncio.run( op.list_statistic_for_partition( - METALAKE, "table", "catalog.schema.table", from_partition_name=_QUERY_INJECTION, diff --git a/mcp-server/tests/unit/test_audit.py b/mcp-server/tests/unit/test_audit.py index e6bdea13644..6071878a84f 100644 --- a/mcp-server/tests/unit/test_audit.py +++ b/mcp-server/tests/unit/test_audit.py @@ -187,6 +187,37 @@ async def _run(): self.assertEqual(record["outcome"], "allow") self.assertEqual(record["principal"], "anonymous") + def test_record_carries_the_resolved_metalake(self): + """Defaulted calls must record the metalake too, not just named ones - + otherwise the log cannot say which tenant was touched without also + knowing the server's startup configuration.""" + + async def _run(): + async with Client(self.mcp) as client: + await client.call_tool("get_list_of_catalogs") + await client.call_tool( + "get_list_of_catalogs", {"metalake": "named_ml"} + ) + + asyncio.run(_run()) + + metalakes = [json.loads(r)["metalake"] for r in self.log_records] + self.assertEqual(metalakes, ["mock_metalake", "named_ml"]) + + def test_non_metalake_scoped_tool_records_no_metalake(self): + """list_metalakes spans every metalake the caller can see, so tagging + it with the server default would claim a tenant it never touched.""" + + async def _run(): + async with Client(self.mcp) as client: + await client.call_tool("list_metalakes") + + asyncio.run(_run()) + + record = json.loads(self.log_records[0]) + self.assertEqual(record["tool"], "list_metalakes") + self.assertNotIn("metalake", record) + def test_principal_falls_back_to_startup_token(self): """With no request header, the audit principal uses the startup --token.""" RESTClientFactory.set_rest_client(MockOperation) diff --git a/mcp-server/tests/unit/test_per_request_metalake.py b/mcp-server/tests/unit/test_per_request_metalake.py new file mode 100644 index 00000000000..68ec723e838 --- /dev/null +++ b/mcp-server/tests/unit/test_per_request_metalake.py @@ -0,0 +1,833 @@ +# 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-call metalake selection. + +Any tool call may name the metalake it operates on with a `metalake` argument, +falling back to the server's `--metalake` default. The argument is advertised +and consumed by MetalakeArgumentMiddleware, so no tool declares it, and it is +carried in a context variable scoped to one call - never across calls, never +shared between server replicas. +""" + +import asyncio +import sys +import unittest +from unittest import mock +from unittest.mock import MagicMock, patch + +from fastmcp import Client + +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_ARGUMENT, + GravitinoContext, + ServiceIdentityFallbackDisabled, + get_request_metalake, + reset_request_metalake, + set_request_metalake, +) +from mcp_server.core.middleware import ( + TOOLS_WITHOUT_METALAKE, + _schema_with_metalake, +) +from mcp_server.core.setting import Setting +from mcp_server.main import _parse_args, do_main +from mcp_server.server import GravitinoMCPServer +from tests.unit.tools import MockOperation + +# Tests intentionally exercise context internals (_default_client, +# _clients_by_auth, _catalog_operation) to assert per-call isolation; +# protected access is expected. +# pylint: disable=protected-access + + +class TestSchemaInjection(unittest.TestCase): + """_schema_with_metalake() advertises the argument without breaking tools.""" + + def test_adds_optional_metalake_property(self): + schema = _schema_with_metalake( + {"type": "object", "properties": {"name": {"type": "string"}}} + ) + self.assertIn(METALAKE_ARGUMENT, schema["properties"]) + self.assertEqual( + schema["properties"][METALAKE_ARGUMENT]["type"], "string" + ) + + def test_does_not_make_metalake_required(self): + """Omitting it is what every single-metalake deployment does.""" + schema = _schema_with_metalake( + { + "type": "object", + "properties": {"name": {"type": "string"}}, + "required": ["name"], + } + ) + self.assertEqual(schema["required"], ["name"]) + + def test_does_not_mutate_the_original_schema(self): + original = { + "type": "object", + "properties": {"name": {"type": "string"}}, + } + _schema_with_metalake(original) + self.assertNotIn(METALAKE_ARGUMENT, original["properties"]) + + def test_never_shadows_a_parameter_the_tool_declares(self): + original = { + "type": "object", + "properties": {METALAKE_ARGUMENT: {"type": "integer"}}, + } + schema = _schema_with_metalake(original) + self.assertEqual( + schema["properties"][METALAKE_ARGUMENT]["type"], "integer" + ) + + +class TestMiddlewareOverTheProtocol(unittest.TestCase): + """End-to-end through a real FastMCP client, not a stubbed context.""" + + def setUp(self): + RESTClientFactory.set_rest_client(MockOperation) + self.mcp = GravitinoMCPServer(Setting("mock_metalake")).mcp + + def tearDown(self): + RESTClientFactory.set_rest_client(PlainRESTClientOperation) + + def test_every_metalake_scoped_tool_advertises_the_argument(self): + async def _run(): + async with Client(self.mcp) as client: + return await client.list_tools() + + tools = asyncio.run(_run()) + self.assertTrue(tools) + missing = [ + t.name + for t in tools + if t.name not in TOOLS_WITHOUT_METALAKE + and METALAKE_ARGUMENT not in (t.inputSchema.get("properties") or {}) + ] + self.assertEqual(missing, []) + + def test_tool_call_accepts_and_consumes_the_argument(self): + """The tool function must not receive an argument it cannot accept.""" + + async def _run(): + async with Client(self.mcp) as client: + return await client.call_tool( + "get_list_of_catalogs", {METALAKE_ARGUMENT: "other_ml"} + ) + + # Would raise if the argument reached the tool function. + result = asyncio.run(_run()) + self.assertIsNotNone(result) + + def test_metalake_does_not_leak_into_the_next_call(self): + """The middleware must reset the context variable after every call. + + Without the reset this per-call plumbing would silently become session + state - exactly what this design exists to avoid. + """ + seen = [] + + async def _run(): + async with Client(self.mcp) as client: + await client.call_tool( + "get_list_of_catalogs", {METALAKE_ARGUMENT: "first_ml"} + ) + seen.append(get_request_metalake()) + await client.call_tool("get_list_of_catalogs") + seen.append(get_request_metalake()) + + asyncio.run(_run()) + self.assertEqual(seen, ["", ""]) + + +class _EchoCatalogOperation: + """A catalog listing that reports which metalake its client was built for, + and parks until every in-flight call has arrived.""" + + barrier = None + + def __init__(self, metalake_name): + self._metalake_name = metalake_name + + async def get_list_of_catalogs(self) -> str: + if _EchoCatalogOperation.barrier is not None: + await _EchoCatalogOperation.barrier() + return self._metalake_name + + +class _EchoMetalakeClient(MockOperation): + """MockOperation that remembers the metalake it was constructed with.""" + + def __init__(self, metalake, uri, authorization="", *, auth=None): + super().__init__(metalake, uri, authorization, auth=auth) + self._metalake = metalake + + def as_catalog_operation(self): + return _EchoCatalogOperation(self._metalake) + + +class TestConcurrentToolCallsOverTheProtocol(unittest.TestCase): + """The isolation guarantee, proven on the real path. + + The context-level concurrency test drives set_request_metalake() directly. + This one goes through the middleware and the MCP protocol, which is what + actually sets and resets the context variable per call - line coverage of + the middleware does not prove two overlapping calls stay separate. + """ + + def setUp(self): + RESTClientFactory.set_rest_client(_EchoMetalakeClient) + self.mcp = GravitinoMCPServer(Setting("ml_default")).mcp + + def tearDown(self): + _EchoCatalogOperation.barrier = None + RESTClientFactory.set_rest_client(PlainRESTClientOperation) + + def test_overlapping_calls_each_use_their_own_metalake(self): + arrived = asyncio.Event() + counter = {"n": 0} + + async def _barrier(): + # Neither call may finish until both are inside the tool, so the + # two requests are genuinely in flight at the same time. + counter["n"] += 1 + if counter["n"] == 2: + arrived.set() + await arrived.wait() + + _EchoCatalogOperation.barrier = _barrier + + async def _run(): + async with Client(self.mcp) as client: + return await asyncio.gather( + client.call_tool( + "get_list_of_catalogs", {METALAKE_ARGUMENT: "ml_a"} + ), + client.call_tool( + "get_list_of_catalogs", {METALAKE_ARGUMENT: "ml_b"} + ), + ) + + first, second = asyncio.run(asyncio.wait_for(_run(), timeout=10)) + + self.assertEqual( + [first.content[0].text, second.content[0].text], ["ml_a", "ml_b"] + ) + + def test_overlapping_calls_do_not_poison_the_default(self): + """A call that names no metalake must still get the default even while + another call naming one is in flight.""" + arrived = asyncio.Event() + counter = {"n": 0} + + async def _barrier(): + counter["n"] += 1 + if counter["n"] == 2: + arrived.set() + await arrived.wait() + + _EchoCatalogOperation.barrier = _barrier + + async def _run(): + async with Client(self.mcp) as client: + return await asyncio.gather( + client.call_tool( + "get_list_of_catalogs", {METALAKE_ARGUMENT: "ml_named"} + ), + client.call_tool("get_list_of_catalogs"), + ) + + named, defaulted = asyncio.run(asyncio.wait_for(_run(), timeout=10)) + + self.assertEqual(named.content[0].text, "ml_named") + self.assertEqual(defaulted.content[0].text, "ml_default") + + +class _RecordingClient(MockOperation): + """Records the metalake and auth object each client was built with.""" + + built = [] + auths = [] + + def __init__(self, metalake, uri, authorization="", *, auth=None): + super().__init__(metalake, uri, authorization, auth=auth) + _RecordingClient.built.append(metalake) + _RecordingClient.auths.append(auth) + + +class _ClosableClient(MockOperation): + """Records close() so a test can tell when a pool was actually torn down.""" + + # Reassigned per test in setUp; a list here so the type is unambiguous. + closed: list = [] + park_on = None + parked = None + released = None + + def __init__(self, metalake, uri, authorization="", *, auth=None): + super().__init__(metalake, uri, authorization, auth=auth) + self._metalake = metalake + + def as_catalog_operation(self): + return _ClosableCatalogOperation(self._metalake, self) + + async def close(self): + _ClosableClient.closed.append(self._metalake) + + +class _ClosableCatalogOperation: + def __init__(self, metalake, owner): + self._metalake = metalake + self._owner = owner + + async def get_list_of_catalogs(self) -> str: + if self._metalake == _ClosableClient.park_on: + _ClosableClient.parked.set() + await _ClosableClient.released.wait() + return self._metalake + + +class TestEvictionDoesNotCloseAClientInUse(unittest.TestCase): + """A call in flight must keep its connection pool. + + Eviction used to close immediately, so a slow call lost its connection as + soon as other calls filled the cache - and for a write, the backend may + already have committed, leaving an ambiguous outcome. + """ + + def setUp(self): + _ClosableClient.closed = [] + _ClosableClient.park_on = "victim" + _ClosableClient.parked = asyncio.Event() + _ClosableClient.released = asyncio.Event() + RESTClientFactory.set_rest_client(_ClosableClient) + self.mcp = GravitinoMCPServer(Setting("ml_default")).mcp + + def tearDown(self): + _ClosableClient.park_on = None + RESTClientFactory.set_rest_client(PlainRESTClientOperation) + + def test_evicted_client_is_closed_only_after_its_call_finishes(self): + cap = context_module._MAX_CACHED_CLIENTS + + async def _run(): + async with Client(self.mcp) as client: + slow = asyncio.ensure_future( + client.call_tool( + "get_list_of_catalogs", {METALAKE_ARGUMENT: "victim"} + ) + ) + await _ClosableClient.parked.wait() + + # Fill the cache past its bound while that call is parked. + for i in range(cap + 1): + await client.call_tool( + "get_list_of_catalogs", {METALAKE_ARGUMENT: f"ml_{i}"} + ) + + evicted_while_in_flight = "victim" not in _ClosableClient.closed + + _ClosableClient.released.set() + result = await slow + # Let the deferred close task run. + await asyncio.sleep(0) + return evicted_while_in_flight, result + + still_open, result = asyncio.run(asyncio.wait_for(_run(), timeout=30)) + + self.assertTrue( + still_open, + "the client serving an in-flight call was closed by eviction", + ) + # The call completed against its own metalake, not a recycled client. + self.assertEqual(result.content[0].text, "victim") + self.assertIn( + "victim", + _ClosableClient.closed, + "the evicted client should be closed once its call finished", + ) + + +class TestDiscoveryWithoutADefaultMetalake(unittest.TestCase): + """A server with no --metalake must still be usable from a cold start. + + list_metalakes is the tool an agent reaches for when it does not know a + metalake yet, so it must not be gated behind having one - otherwise it is + unusable on exactly the deployment that needs it. + """ + + def setUp(self): + RESTClientFactory.set_rest_client(MockOperation) + self.mcp = GravitinoMCPServer(Setting(metalake="")).mcp + + def tearDown(self): + RESTClientFactory.set_rest_client(PlainRESTClientOperation) + + def test_list_metalakes_works_with_no_metalake_configured(self): + async def _run(): + async with Client(self.mcp) as client: + return await client.call_tool("list_metalakes") + + result = asyncio.run(_run()) + self.assertEqual(result.content[0].text, "mock_metalakes") + + def test_other_tools_report_how_to_recover(self): + """The agent should be told to call list_metalakes, not just fail.""" + + async def _run(): + async with Client(self.mcp) as client: + return await client.call_tool("get_list_of_catalogs") + + with self.assertRaises(Exception) as raised: + asyncio.run(_run()) + self.assertIn("list_metalakes", str(raised.exception)) + + def test_naming_a_metalake_per_call_works_with_no_default(self): + async def _run(): + async with Client(self.mcp) as client: + return await client.call_tool( + "get_list_of_catalogs", {METALAKE_ARGUMENT: "ml_a"} + ) + + self.assertIsNotNone(asyncio.run(_run())) + + +class TestToolsThatNeverResolveAMetalake(unittest.TestCase): + """Tools that ignore the metalake must not advertise the argument.""" + + def setUp(self): + RESTClientFactory.set_rest_client(MockOperation) + self.mcp = GravitinoMCPServer(Setting("mock_metalake")).mcp + + def tearDown(self): + RESTClientFactory.set_rest_client(PlainRESTClientOperation) + + def test_discovery_and_pure_computation_tools_skip_the_argument(self): + """Offering a knob that does nothing invites the model to misuse it - + on list_metalakes it would otherwise be the only parameter, reading + like a filter for the listing.""" + + async def _run(): + async with Client(self.mcp) as client: + return { + t.name: t.inputSchema for t in await client.list_tools() + } + + schemas = asyncio.run(_run()) + for name in ("list_metalakes", "metadata_type_to_fullname_formats"): + self.assertNotIn( + METALAKE_ARGUMENT, + schemas[name].get("properties") or {}, + f"{name} should not advertise the metalake argument", + ) + + +class TestRecoveryHintMatchesTheDeployment(unittest.TestCase): + """--include-tool-tags is an allowlist and can hide list_metalakes.""" + + def setUp(self): + RESTClientFactory.set_rest_client(MockOperation) + + def tearDown(self): + RESTClientFactory.set_rest_client(PlainRESTClientOperation) + + def _call_with_tags(self, tags): + mcp = GravitinoMCPServer(Setting("", tags=tags)).mcp + + async def _run(): + async with Client(mcp) as client: + await client.call_tool("get_list_of_catalogs") + + with self.assertRaises(Exception) as raised: + asyncio.run(_run()) + return str(raised.exception) + + def test_hint_points_at_discovery_when_it_is_exposed(self): + self.assertIn("list_metalakes", self._call_with_tags(set())) + + def test_hint_omits_discovery_when_a_tag_filter_hides_it(self): + """Naming a tool the agent cannot call leaves it with no way forward.""" + message = self._call_with_tags({"catalog"}) + self.assertNotIn("list_metalakes", message) + self.assertIn(f"'{METALAKE_ARGUMENT}' argument", message) + + +class TestMetalakeResolution(unittest.TestCase): + """GravitinoContext resolves the metalake per call.""" + + 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_call_argument_overrides_startup_default(self): + ctx = self._make_context() + token = set_request_metalake("ml_other") + try: + client = ctx.rest_client() + finally: + reset_request_metalake(token) + + self.assertEqual(client._catalog_operation.metalake_name, "ml_other") + + def test_falls_back_to_startup_default(self): + ctx = self._make_context() + client = ctx.rest_client() + + self.assertIs(client, ctx._default_client) + self.assertEqual(client._catalog_operation.metalake_name, "ml_default") + + def test_whitespace_only_argument_is_treated_as_absent(self): + ctx = self._make_context() + token = set_request_metalake(" ") + try: + client = ctx.rest_client() + finally: + reset_request_metalake(token) + + self.assertEqual(client._catalog_operation.metalake_name, "ml_default") + + def test_missing_metalake_raises_with_a_recoverable_message(self): + """The message is read by an agent, so it must name the way out.""" + ctx = self._make_context(metalake="") + + with self.assertRaises(ValueError) as raised: + ctx.rest_client() + + message = str(raised.exception) + self.assertIn("list_metalakes", message) + self.assertIn(f"'{METALAKE_ARGUMENT}' argument", message) + + def test_missing_metalake_takes_priority_over_fallback_disabled(self): + 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", + side_effect=LookupError, + ): + with self.assertRaises(ValueError) as raised: + ctx.rest_client() + + self.assertNotIsInstance( + raised.exception, ServiceIdentityFallbackDisabled + ) + + def test_discovery_works_with_no_metalake_anywhere(self): + """list_metalakes is what an agent calls before it knows a metalake, + so it must not require one - otherwise it is unusable on exactly the + server that needs it.""" + ctx = self._make_context(metalake="") + + client = ctx.rest_client(require_metalake=False) + + self.assertIsNotNone(client.as_metalake_operation()) + + def test_two_concurrent_calls_get_different_metalake_clients(self): + """Two calls naming different metalakes, in flight at the same time, + must each get their own client. Both tasks publish their metalake and + then park until the other has too, so the calls genuinely overlap.""" + ctx = self._make_context() + + async def _call(metalake, ready, go): + token = set_request_metalake(metalake) + try: + ready.set() + await go.wait() + return ctx.rest_client() + finally: + reset_request_metalake(token) + + async def _drive(): + ready_a, ready_b, go = ( + asyncio.Event(), + asyncio.Event(), + asyncio.Event(), + ) + task_a = asyncio.ensure_future(_call("ml_a", ready_a, go)) + task_b = asyncio.ensure_future(_call("ml_b", ready_b, go)) + await ready_a.wait() + await ready_b.wait() + go.set() + return await asyncio.gather(task_a, task_b) + + client_a, client_b = asyncio.run(_drive()) + + self.assertEqual(client_a._catalog_operation.metalake_name, "ml_a") + self.assertEqual(client_b._catalog_operation.metalake_name, "ml_b") + self.assertIsNot(client_a, client_b) + + def test_same_identity_and_metalake_reuses_cached_client(self): + ctx = self._make_context() + token = set_request_metalake("ml_a") + try: + first = ctx.rest_client() + second = ctx.rest_client() + finally: + reset_request_metalake(token) + + self.assertIs(first, second) + + def test_one_bound_covers_principal_and_service_clients_together(self): + """_MAX_CACHED_CLIENTS bounds the total number of open connection + pools, not each cache separately.""" + ctx = self._make_context() + cap = context_module._MAX_CACHED_CLIENTS + + def _mock_request(authorization): + request = MagicMock() + request.headers.get.side_effect = ( + lambda key, default="": authorization + ) + return request + + for i in range(cap): + with patch( + "fastmcp.server.dependencies.get_http_request", + return_value=_mock_request(f"Bearer t{i}"), + ): + ctx.rest_client() + + for i in range(10): + token = set_request_metalake(f"ml_{i}") + try: + ctx.rest_client() + finally: + reset_request_metalake(token) + + self.assertLessEqual(len(ctx._clients_by_auth), cap) + + +class TestMetalakeArgumentValidation(unittest.TestCase): + """Popping the argument removes it from FastMCP's schema validation, so + this middleware has to reject bad values itself - and must do it where the + error handling and audit middleware still see the failure.""" + + def setUp(self): + RESTClientFactory.set_rest_client(_RecordingClient) + _RecordingClient.built = [] + self.mcp = GravitinoMCPServer(Setting("prod")).mcp + + def tearDown(self): + RESTClientFactory.set_rest_client(PlainRESTClientOperation) + + def _call(self, arguments): + # The server built its own default client at construction; only count + # the clients this call causes. + _RecordingClient.built = [] + + async def _run(): + async with Client(self.mcp) as client: + return await client.call_tool("get_list_of_catalogs", arguments) + + return asyncio.run(_run()) + + def test_non_string_values_are_rejected_before_any_request(self): + """`false`, `0` and `[]` are falsy: without an explicit check they + would silently route the call to the default metalake.""" + for value in (False, 0, [], 42): + with self.subTest(value=value): + with self.assertRaises(Exception) as raised: + self._call({METALAKE_ARGUMENT: value}) + self.assertIn("must be a string", str(raised.exception)) + self.assertEqual( + _RecordingClient.built, + [], + "no REST client should be built for a rejected argument", + ) + + def test_rejection_is_reported_as_a_client_error(self): + """It is the caller's argument that is wrong, not the server.""" + with self.assertRaises(Exception) as raised: + self._call({METALAKE_ARGUMENT: 42}) + self.assertIn("Invalid params", str(raised.exception)) + + def test_null_and_omitted_both_mean_the_default(self): + for arguments in ({METALAKE_ARGUMENT: None}, {}): + with self.subTest(arguments=arguments): + self.assertIsNotNone(self._call(arguments)) + + +class TestDeprecatedMetalakeNameAlias(unittest.TestCase): + """`metalake_name` shipped on the statistic tools in v1.0.0.""" + + def setUp(self): + RESTClientFactory.set_rest_client(_RecordingClient) + _RecordingClient.built = [] + self.mcp = GravitinoMCPServer(Setting("prod")).mcp + + def tearDown(self): + RESTClientFactory.set_rest_client(PlainRESTClientOperation) + + def _call(self, arguments): + # The server built its own default client at construction; only count + # the clients this call causes. + _RecordingClient.built = [] + + async def _run(): + async with Client(self.mcp) as client: + return await client.call_tool( + "list_statistics_for_metadata", arguments + ) + + return asyncio.run(_run()) + + def _base(self): + return {"metadata_type": "table", "metadata_fullname": "c.s.t"} + + def test_legacy_argument_still_selects_the_metalake(self): + self._call({**self._base(), "metalake_name": "legacy_ml"}) + self.assertEqual(_RecordingClient.built, ["legacy_ml"]) + + def test_new_argument_works_the_same(self): + self._call({**self._base(), METALAKE_ARGUMENT: "new_ml"}) + self.assertEqual(_RecordingClient.built, ["new_ml"]) + + def test_agreeing_values_are_accepted(self): + self._call( + {**self._base(), METALAKE_ARGUMENT: "ml", "metalake_name": "ml"} + ) + self.assertEqual(_RecordingClient.built, ["ml"]) + + def test_conflicting_values_are_rejected_before_any_request(self): + """Silently picking one could send a write to the wrong tenant.""" + with self.assertRaises(Exception) as raised: + self._call( + { + **self._base(), + METALAKE_ARGUMENT: "ml_a", + "metalake_name": "ml_b", + } + ) + self.assertIn("different values", str(raised.exception)) + self.assertEqual(_RecordingClient.built, []) + + def test_alias_is_not_advertised(self): + """Accepted for compatibility, but new callers should see one way.""" + + async def _run(): + async with Client(self.mcp) as client: + return { + t.name: t.inputSchema for t in await client.list_tools() + } + + schema = asyncio.run(_run())["list_statistics_for_metadata"] + self.assertNotIn("metalake_name", schema.get("properties") or {}) + self.assertIn(METALAKE_ARGUMENT, schema["properties"]) + + def test_alias_is_scoped_to_the_tools_that_shipped_it(self): + """Other tools must not silently accept an unknown argument.""" + + async def _run(): + async with Client(self.mcp) as client: + return await client.call_tool( + "get_list_of_catalogs", {"metalake_name": "ml"} + ) + + with self.assertRaises(Exception): + asyncio.run(_run()) + + +class TestServiceAuthIsSharedAcrossMetalakes(unittest.TestCase): + """Separate RefreshableBearerAuth instances share a token cache key but + not the refresh lock, so per-metalake instances would hit the IdP once per + metalake on a cold cache.""" + + def setUp(self): + RESTClientFactory.set_rest_client(_RecordingClient) + _RecordingClient.auths = [] + _RecordingClient.built = [] + + def tearDown(self): + RESTClientFactory.set_rest_client(PlainRESTClientOperation) + + def test_every_metalake_client_shares_one_auth_object(self): + ctx = GravitinoContext( + Setting( + metalake="ml_default", + gravitino_uri="http://localhost:8090", + transport="http", + oauth_token_endpoint="https://idp/token", + oauth_client_id="mcp", + oauth_client_secret="s", + ) + ) + for metalake in ("ml_a", "ml_b", "ml_c"): + token = set_request_metalake(metalake) + try: + ctx.rest_client() + finally: + reset_request_metalake(token) + + auths = [a for a in _RecordingClient.auths if a is not None] + self.assertTrue(auths) + self.assertEqual( + len(set(id(a) for a in auths)), + 1, + "each metalake built its own auth object", + ) + + +class TestSettingMetalake(unittest.TestCase): + def test_whitespace_only_metalake_is_stripped_to_empty(self): + """A shell-quoting mistake like --metalake " " must not read as a + configured default.""" + self.assertEqual(Setting(metalake=" ").metalake, "") + + def test_stdio_without_metalake_is_allowed(self): + """stdio can now name a metalake per call like any other transport, so + --metalake is no longer required there.""" + GravitinoContext(Setting(metalake="", transport="stdio")) + + +class TestMetalakeArgParsing(unittest.TestCase): + def test_metalake_is_optional_and_defaults_to_empty(self): + with mock.patch.object(sys, "argv", ["mcp_server"]): + args = _parse_args() + self.assertEqual(args.metalake, "") + + def test_server_starts_without_a_default_metalake(self): + with mock.patch("mcp_server.main._init_logging"), mock.patch( + "mcp_server.main.GravitinoMCPServer" + ) as server, mock.patch.object(sys, "argv", ["mcp_server"]): + do_main() + server.assert_called_once() + + +if __name__ == "__main__": + unittest.main() diff --git a/mcp-server/tests/unit/tools/mock_operation.py b/mcp-server/tests/unit/tools/mock_operation.py index 17009502170..1a879e93cd2 100644 --- a/mcp-server/tests/unit/tools/mock_operation.py +++ b/mcp-server/tests/unit/tools/mock_operation.py @@ -27,6 +27,7 @@ ) from mcp_server.client.fileset_operation import FilesetOperation from mcp_server.client.job_operation import JobOperation +from mcp_server.client.metalake_operation import MetalakeOperation from mcp_server.client.partition_operation import PartitionOperation from mcp_server.client.statistic_operation import StatisticOperation from mcp_server.client.view_operation import ViewOperation @@ -72,6 +73,14 @@ def as_partition_operation(self) -> PartitionOperation: def as_view_operation(self) -> ViewOperation: return MockViewOperation() + def as_metalake_operation(self) -> MetalakeOperation: + return MockMetalakeOperation() + + +class MockMetalakeOperation(MetalakeOperation): + async def get_list_of_metalakes(self) -> str: + return "mock_metalakes" + class MockCatalogOperation(CatalogOperation): async def get_list_of_catalogs(self) -> str: @@ -429,14 +438,13 @@ async def cancel_job(self, job_id: str) -> str: class MockStatisticOperation(StatisticOperation): async def list_of_statistics( - self, metalake_name: str, metadata_type: str, metadata_fullname: str + self, metadata_type: str, metadata_fullname: str ) -> str: - return f"mock_statistics: {metalake_name}, {metadata_type}, {metadata_fullname}" + return f"mock_statistics: {metadata_type}, {metadata_fullname}" # pylint: disable=R0917 async def list_statistic_for_partition( self, - metalake_name: str, metadata_type: str, metadata_fullname: str, from_partition_name: str, @@ -445,7 +453,7 @@ async def list_statistic_for_partition( to_inclusive: bool = False, ) -> str: return ( - f"mock_statistics_for_partition: {metalake_name}, {metadata_type}, {metadata_fullname}," + f"mock_statistics_for_partition: {metadata_type}, {metadata_fullname}," f" {from_partition_name}, {to_partition_name}, {from_inclusive}, {to_inclusive}" ) diff --git a/mcp-server/tests/unit/tools/test_statistic.py b/mcp-server/tests/unit/tools/test_statistic.py index 1a8527c6bba..1140cd393af 100644 --- a/mcp-server/tests/unit/tools/test_statistic.py +++ b/mcp-server/tests/unit/tools/test_statistic.py @@ -38,13 +38,12 @@ async def _test_list_of_statistics(mcp_server): result = await client.call_tool( "list_statistics_for_metadata", { - "metalake_name": "mock_metalake", "metadata_type": "mock_type", "metadata_fullname": "mock_fullname", }, ) self.assertEqual( - "mock_statistics: mock_metalake, mock_type, mock_fullname", + "mock_statistics: mock_type, mock_fullname", result.content[0].text, ) @@ -56,7 +55,6 @@ async def _test_list_statistics_for_partition(mcp_server): result = await client.call_tool( "list_statistics_for_partition", { - "metalake_name": "mock_metalake", "metadata_type": "mock_type", "metadata_fullname": "mock_fullname", "from_partition_name": "from_partition", @@ -64,7 +62,7 @@ async def _test_list_statistics_for_partition(mcp_server): }, ) self.assertEqual( - "mock_statistics_for_partition: mock_metalake, mock_type, mock_fullname, " + "mock_statistics_for_partition: mock_type, mock_fullname, " "from_partition, to_partition, True, False", result.content[0].text, )