This is an automated email from the ASF dual-hosted git repository.
jerryshao pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gravitino.git
The following commit(s) were added to refs/heads/main by this push:
new 7694553668 [#12953] feat(mcp-server): select the metalake per tool
call (#12960)
7694553668 is described below
commit 7694553668f0dbef82ef57141accc71372442bac
Author: Jerry Shao <[email protected]>
AuthorDate: Wed Sep 9 23:40:43 2026 +0800
[#12953] feat(mcp-server): select the metalake per tool call (#12960)
### What changes were proposed in this pull request?
The MCP server's metalake was fixed at startup (`--metalake`) and every
tool
call resolved against that one metalake for the life of the process.
This PR
makes it a per-call choice: every tool accepts an optional `metalake`
argument that takes priority over the startup default.
No tool declares the argument. `MetalakeArgumentMiddleware` advertises
it in
each tool's input schema, strips it from the incoming arguments before
the
tool function runs, and publishes it in a context variable scoped to a
single
call, which `GravitinoContext` resolves against: the call's argument,
then the
`--metalake` default, then an error. Nothing is remembered between
calls, so
concurrent callers stay isolated and the server keeps no state that
would have
to be shared between replicas.
A `list_metalakes` tool is added so an agent can discover what that
argument
accepts. It is the one tool that resolves no metalake, so it works on a
server
started with no `--metalake` at all - the deployment that needs it most.
Two smaller changes come along:
- The statistic tools took their own `metalake_name` argument, shadowing
the
value their client was already constructed with. They now use the shared
resolution like every other tool, so there is one way to name a
metalake.
- Audit entries record the resolved metalake, since "which tenant did
this
touch" is no longer answerable from the server configuration alone.
Tools
that are not metalake-scoped record none rather than the default.
An earlier revision of this PR used an `X-Gravitino-Metalake` request
header.
That was dropped: the header only worked over HTTP, was fixed for the
life of
a client connection (so an agent could not switch metalake
mid-conversation),
and depended on MCP clients forwarding custom headers on every request,
which
they do inconsistently. A tool argument travels in the JSON-RPC body,
works
identically on stdio and HTTP, and can differ per call.
### Why are the changes needed?
A metalake is Gravitino's tenant boundary, and every other interface
takes it
per request. The MCP server was the only one that pinned it at startup,
so an
installation with more than one metalake needed one server process per
metalake, or had to accept that agents could reach only one of them.
Fix: #12953
### Does this PR introduce _any_ user-facing change?
- `--metalake` is now optional on every transport. It is the default
used by
any call that does not name one.
- Every metalake-scoped tool accepts an optional `metalake` argument.
- New `list_metalakes` tool, tag `metalake`.
- **Breaking:** `list_statistics_for_metadata` and
`list_statistics_for_partition`
no longer accept `metalake_name`; callers pass `metalake` instead, like
every
other tool.
- A call that names no metalake on a server with no default now fails
with a
message telling the agent to call `list_metalakes` and retry - or, when
a tag
filter hides that tool, to ask the user which metalake to use.
- Audit records gain a `metalake` field.
### How was this patch tested?
257 unit tests pass; `isort`, `black` and `pylint` (10.00/10) are clean.
Coverage of the changed modules is 96%, with the middleware, settings
and new
tool modules at 100%.
The tests that matter most:
- Concurrency, through the real MCP protocol: two overlapping calls
naming
different metalakes each get their own, held simultaneously in flight by
a
barrier so the overlap is real. Mutation-tested - replacing the context
variable with process-global state fails these tests.
- No leakage between sequential calls, which is what would silently turn
per-call plumbing into session state.
- Cold start with no `--metalake`: `list_metalakes` works, other tools
return a
recoverable error, and naming a metalake per call works.
- The new client operation is exercised against a mocked httpx client,
pinning
the `/api/metalakes` endpoint and the response key.
---------
Co-authored-by: Claude Sonnet 5 <[email protected]>
---
docs/gravitino-mcp-server.md | 45 +-
mcp-server/mcp_server/client/__init__.py | 1 +
.../mcp_server/client/gravitino_operation.py | 11 +
.../client/{__init__.py => metalake_operation.py} | 30 +-
.../plain_rest_client_metalake_operation.py} | 28 +-
.../client/plain/plain_rest_client_operation.py | 11 +
.../plain/plain_rest_client_statistic_operation.py | 7 +-
.../mcp_server/client/statistic_operation.py | 5 +-
mcp-server/mcp_server/core/audit.py | 10 +
mcp-server/mcp_server/core/context.py | 280 ++++++-
mcp-server/mcp_server/core/middleware.py | 171 +++++
mcp-server/mcp_server/core/setting.py | 25 +-
mcp-server/mcp_server/main.py | 8 +-
mcp-server/mcp_server/server.py | 64 +-
mcp-server/mcp_server/tools/__init__.py | 3 +
mcp-server/mcp_server/tools/metalake.py | 66 ++
mcp-server/mcp_server/tools/statistic.py | 9 +-
mcp-server/tests/unit/client/test_url_encoding.py | 36 +-
mcp-server/tests/unit/test_audit.py | 31 +
mcp-server/tests/unit/test_per_request_metalake.py | 833 +++++++++++++++++++++
mcp-server/tests/unit/tools/mock_operation.py | 16 +-
mcp-server/tests/unit/tools/test_statistic.py | 6 +-
22 files changed, 1618 insertions(+), 78 deletions(-)
diff --git a/docs/gravitino-mcp-server.md b/docs/gravitino-mcp-server.md
index 304674766b..e97bdc6b69 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 dfddc047dd..3565de2007 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 0943ab0120..763ef02da8 100644
--- a/mcp-server/mcp_server/client/gravitino_operation.py
+++ b/mcp-server/mcp_server/client/gravitino_operation.py
@@ -20,6 +20,7 @@ from abc import ABC, abstractmethod
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 @@ class GravitinoOperation(ABC):
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/__init__.py
b/mcp-server/mcp_server/client/metalake_operation.py
similarity index 55%
copy from mcp-server/mcp_server/client/__init__.py
copy to mcp-server/mcp_server/client/metalake_operation.py
index dfddc047dd..5109ff4661 100644
--- a/mcp-server/mcp_server/client/__init__.py
+++ b/mcp-server/mcp_server/client/metalake_operation.py
@@ -15,12 +15,24 @@
# specific language governing permissions and limitations
# under the License.
-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.model_operation import ModelOperation
-from mcp_server.client.policy_operation import PolicyOperation
-from mcp_server.client.schema_operation import SchemaOperation
-from mcp_server.client.table_operation import TableOperation
-from mcp_server.client.tag_operation import TagOperation
-from mcp_server.client.topic_operation import TopicOperation
+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/__init__.py
b/mcp-server/mcp_server/client/plain/plain_rest_client_metalake_operation.py
similarity index 53%
copy from mcp-server/mcp_server/client/__init__.py
copy to
mcp-server/mcp_server/client/plain/plain_rest_client_metalake_operation.py
index dfddc047dd..714e2712da 100644
--- a/mcp-server/mcp_server/client/__init__.py
+++ b/mcp-server/mcp_server/client/plain/plain_rest_client_metalake_operation.py
@@ -15,12 +15,22 @@
# specific language governing permissions and limitations
# under the License.
-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.model_operation import ModelOperation
-from mcp_server.client.policy_operation import PolicyOperation
-from mcp_server.client.schema_operation import SchemaOperation
-from mcp_server.client.table_operation import TableOperation
-from mcp_server.client.tag_operation import TagOperation
-from mcp_server.client.topic_operation import TopicOperation
+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 0d35525ffb..63187bedea 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 @@ import httpx
from mcp_server.client import (
CatalogOperation,
GravitinoOperation,
+ MetalakeOperation,
ModelOperation,
PolicyOperation,
SchemaOperation,
@@ -38,6 +39,9 @@ from
mcp_server.client.plain.plain_rest_client_fileset_operation import (
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 @@ class PlainRESTClientOperation(GravitinoOperation):
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 49a56036a5..2befb1062d 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 @@ class PlainRESTClientStatisticOperation(StatisticOperation):
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 @@ class PlainRESTClientStatisticOperation(StatisticOperation):
# 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 @@ class PlainRESTClientStatisticOperation(StatisticOperation):
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 b9c2f1f99d..5fd2165e68 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 @@ class StatisticOperation(ABC):
@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 @@ class StatisticOperation(ABC):
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 0d25ef2f24..a8d7d6224e 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 f38831f639..0bb8f7af73 100644
--- a/mcp-server/mcp_server/core/context.py
+++ b/mcp-server/mcp_server/core/context.py
@@ -19,6 +19,8 @@ import asyncio
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 @@ from mcp_server.core.setting import Setting
_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 @@ _CANONICAL_AUTH_SCHEMES = {
"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 @@ class GravitinoContext:
"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 0000000000..54768b399d
--- /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 112528f6d9..bb69175a6b 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 @@ class Setting: # pylint:
disable=too-many-instance-attributes
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 fd172d3fb1..2a50b8f68d 100644
--- a/mcp-server/mcp_server/main.py
+++ b/mcp-server/mcp_server/main.py
@@ -22,6 +22,7 @@ import os
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 56cca377cb..688e0d17d6 100644
--- a/mcp-server/mcp_server/server.py
+++ b/mcp-server/mcp_server/server.py
@@ -37,10 +37,16 @@ from fastmcp.tools.base import ToolResult
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 @@ class AuditMiddleware(Middleware):
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 @@ class AuditMiddleware(Middleware):
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 e15f6ad8d1..4fd524c088 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.catalog import load_catalog_tools
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 @@ SUPPORTED_TOOL_TAGS = frozenset(
"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 0000000000..c2a93cac7b
--- /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 76dd1d9a6f..8e5365fcba 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 @@ def load_statistic_tools(mcp: FastMCP):
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 @@ def load_statistic_tools(mcp: FastMCP):
"""
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 @@ def load_statistic_tools(mcp: FastMCP):
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 @@ def load_statistic_tools(mcp: FastMCP):
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 3e1c9018da..6ff7a182b1 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_fileset_operation import (
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 @@ _ENCODED_SLASH = "cat%2Fschema"
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 @@ class TestStatisticOperationUrlEncoding(unittest.TestCase):
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 e6bdea1364..6071878a84 100644
--- a/mcp-server/tests/unit/test_audit.py
+++ b/mcp-server/tests/unit/test_audit.py
@@ -187,6 +187,37 @@ class TestAuditMiddlewareIntegration(unittest.TestCase):
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 0000000000..68ec723e83
--- /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 1700950217..1a879e93cd 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 import (
)
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 @@ class MockOperation(GravitinoOperation):
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 @@ class MockJobOperation(JobOperation):
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 @@ class MockStatisticOperation(StatisticOperation):
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 1a8527c6bb..1140cd393a 100644
--- a/mcp-server/tests/unit/tools/test_statistic.py
+++ b/mcp-server/tests/unit/tools/test_statistic.py
@@ -38,13 +38,12 @@ class TestStatisticTool(unittest.TestCase):
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 @@ class TestStatisticTool(unittest.TestCase):
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 @@ class TestStatisticTool(unittest.TestCase):
},
)
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,
)
