This is an automated email from the ASF dual-hosted git repository.
yuqi1129 pushed a commit to branch cherry-pick-696efe1c-to-branch-1.3
in repository https://gitbox.apache.org/repos/asf/gravitino.git
The following commit(s) were added to
refs/heads/cherry-pick-696efe1c-to-branch-1.3 by this push:
new 8192ff1ea5 fix(mcp): resolve read timeout backport conflicts on
branch-1.3
8192ff1ea5 is described below
commit 8192ff1ea50748aac6a4bdf545269f51d81ae7b8
Author: yuqi <[email protected]>
AuthorDate: Mon Nov 2 20:26:46 2026 +0800
fix(mcp): resolve read timeout backport conflicts on branch-1.3
---
docs/gravitino-mcp-server.md | 19 +-
mcp-server/mcp_server/core/context.py | 94 +--
mcp-server/mcp_server/core/setting.py | 8 +-
mcp-server/tests/unit/test_per_request_metalake.py | 839 ---------------------
mcp-server/tests/unit/test_request_timeout.py | 32 +-
mcp-server/tests/unit/tools/test_semantic_model.py | 643 ----------------
6 files changed, 25 insertions(+), 1610 deletions(-)
diff --git a/docs/gravitino-mcp-server.md b/docs/gravitino-mcp-server.md
index eb8df91a6e..0422a113ba 100644
--- a/docs/gravitino-mcp-server.md
+++ b/docs/gravitino-mcp-server.md
@@ -111,11 +111,11 @@ Gravitino MCP server supports the following tools, and
you could export tool by
You could config Gravitino MCP server by arguments, `uv run mcp_server -h`
shows the detailed information.
-<<<<<<< HEAD
| Argument | Description
| Default value | Required |
|----------------------------------|---------------------------------------------------------------------------------------------------------------------------------|-----------------------------|----------|
| `--metalake` | The Gravitino metalake name.
| none | Yes |
| `--gravitino-uri` | The URI of Gravitino server.
| `http://127.0.0.1:8090` | No |
+| `--gravitino-read-timeout` | Gravitino response read timeout in
seconds; or set `GRAVITINO_READ_TIMEOUT`. Must be finite and positive. Connect,
write and pool timeouts remain 5 seconds. | `60` | 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 |
| `--token` | Static credential for Gravitino; or set
`GRAVITINO_TOKEN`. See Authentication. Wins over OAuth client-credentials.
| none (anonymous) | No |
@@ -126,25 +126,8 @@ You could config Gravitino MCP server by arguments, `uv
run mcp_server -h` shows
| `--no-service-identity-fallback` | HTTP only: reject requests with no
`Authorization` when OAuth or `--token` is set. Or
`GRAVITINO_NO_SERVICE_IDENTITY_FALLBACK`. | `false` | No
|
| `--tls-cert` | PEM certificate to serve the endpoint
over HTTPS. Requires `--tls-key`.
| none | No |
| `--tls-key` | PEM private key to serve the endpoint
over HTTPS. Requires `--tls-cert`.
| none | No |
-=======
-| Argument | Description
| Default value | Required |
-| -------------------------------- |
------------------------------------------------------------------------------------------------------------------------------------------------------------
| --------------------------- | -------- |
-| `--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 |
-| `--gravitino-read-timeout` | Gravitino response read timeout in
seconds; or set `GRAVITINO_READ_TIMEOUT`. Must be finite and positive. Connect,
write and pool timeouts remain 5 seconds. | `60` | 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
|
-| `--token` | Static credential for Gravitino; or set
`GRAVITINO_TOKEN`. See Authentication. Wins over OAuth client-credentials.
| none (anonymous) | No |
-| `--oauth-token-endpoint` | OAuth2 token URL for client-credentials.
Or `GRAVITINO_OAUTH_TOKEN_ENDPOINT`.
| none | No |
-| `--oauth-client-id` | OAuth2 client id. Or
`GRAVITINO_OAUTH_CLIENT_ID`.
| none
| No |
-| `--oauth-client-secret` | OAuth2 client secret. Or
`GRAVITINO_OAUTH_CLIENT_SECRET`.
| none
| No |
-| `--oauth-scope` | Optional OAuth2 scope. Or
`GRAVITINO_OAUTH_SCOPE`.
| none
| No |
-| `--no-service-identity-fallback` | HTTP only: reject requests with no
`Authorization` when OAuth or `--token` is set. Or
`GRAVITINO_NO_SERVICE_IDENTITY_FALLBACK`. |
`false` | No |
-| `--tls-cert` | PEM certificate to serve the endpoint
over HTTPS. Requires `--tls-key`.
| none | No
|
-| `--tls-key` | PEM private key to serve the endpoint
over HTTPS. Requires `--tls-cert`.
| none | No
|
A response read timeout means the MCP server stopped waiting for Gravitino.
The operation may still complete on the server. Check the result before
retrying, especially for writes. This timeout bounds waiting for response data,
not the total duration of the server operation.
->>>>>>> 696efe1c2 ([#14498] fix(mcp): configure Gravitino response read
timeout (#14506))
## Authentication
diff --git a/mcp-server/mcp_server/core/context.py
b/mcp-server/mcp_server/core/context.py
index 34effad2b1..f32aedc26d 100644
--- a/mcp-server/mcp_server/core/context.py
+++ b/mcp-server/mcp_server/core/context.py
@@ -148,33 +148,12 @@ def _service_auth(setting: Setting):
class GravitinoContext:
def __init__(self, setting: Setting):
self._setting = setting
-<<<<<<< HEAD
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,
- read_timeout=setting.gravitino_read_timeout,
- )
- if setting.metalake
- else None
->>>>>>> 696efe1c2 ([#14498] fix(mcp): configure Gravitino response read
timeout (#14506))
+ read_timeout=setting.gravitino_read_timeout,
)
# LRU cache of per-principal clients keyed by the raw Authorization
header.
# Safe without locking: rest_client() runs on the single asyncio event
@@ -212,72 +191,6 @@ class GravitinoContext:
"HTTP request omitted Authorization and "
"--no-service-identity-fallback is set"
)
-<<<<<<< HEAD
-=======
- 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,
- read_timeout=self._setting.gravitino_read_timeout,
- )
- 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
- ):
->>>>>>> 696efe1c2 ([#14498] fix(mcp): configure Gravitino response read
timeout (#14506))
return self._default_client
cached = self._clients_by_auth.get(authorization)
@@ -288,13 +201,8 @@ class GravitinoContext:
client = RESTClientFactory.create_rest_client(
self._setting.metalake,
self._setting.gravitino_uri,
-<<<<<<< HEAD
authorization,
-=======
- startup_authorization(self._setting),
- auth=self._service_auth,
read_timeout=self._setting.gravitino_read_timeout,
->>>>>>> 696efe1c2 ([#14498] fix(mcp): configure Gravitino response read
timeout (#14506))
)
self._clients_by_auth[authorization] = client
if len(self._clients_by_auth) > _MAX_CACHED_CLIENTS:
diff --git a/mcp-server/mcp_server/core/setting.py
b/mcp-server/mcp_server/core/setting.py
index 0555d9fee7..fb223f6ac0 100644
--- a/mcp-server/mcp_server/core/setting.py
+++ b/mcp-server/mcp_server/core/setting.py
@@ -56,15 +56,10 @@ class Setting: # pylint:
disable=too-many-instance-attributes
# HTTP only: reject requests with no Authorization when service OAuth or
# --token is configured instead of falling back to the service identity.
no_service_identity_fallback: bool = False
+
gravitino_read_timeout: float = 60.0
-<<<<<<< HEAD
-=======
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()
if (
not math.isfinite(self.gravitino_read_timeout)
or self.gravitino_read_timeout <= 0
@@ -73,7 +68,6 @@ class Setting: # pylint: disable=too-many-instance-attributes
"Gravitino read timeout must be a finite positive number of
seconds"
)
->>>>>>> 696efe1c2 ([#14498] fix(mcp): configure Gravitino response read
timeout (#14506))
def has_oauth_client(self) -> bool:
"""Return True when client-credentials is fully configured."""
return bool(
diff --git a/mcp-server/tests/unit/test_per_request_metalake.py
b/mcp-server/tests/unit/test_per_request_metalake.py
deleted file mode 100644
index bc4fc190f8..0000000000
--- a/mcp-server/tests/unit/test_per_request_metalake.py
+++ /dev/null
@@ -1,839 +0,0 @@
-# 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, read_timeout=60.0
- ):
- 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, read_timeout=60.0
- ):
- 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, read_timeout=60.0
- ):
- 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_full_name": "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/test_request_timeout.py
b/mcp-server/tests/unit/test_request_timeout.py
index 6cbb979305..aa0cb7587a 100644
--- a/mcp-server/tests/unit/test_request_timeout.py
+++ b/mcp-server/tests/unit/test_request_timeout.py
@@ -41,26 +41,36 @@ from mcp_server.server import GravitinoMCPServer
class TestRequestTimeout(unittest.TestCase):
def test_configuration(self):
- self.assertEqual(60.0, Setting().gravitino_read_timeout)
+ self.assertEqual(60.0, Setting(metalake="m").gravitino_read_timeout)
for value in (0, -1, math.inf, math.nan):
with self.assertRaises(ValueError):
- Setting(gravitino_read_timeout=value)
+ Setting(metalake="m", gravitino_read_timeout=value)
with patch.dict(os.environ, {"GRAVITINO_READ_TIMEOUT": "120"}):
- with patch.object(sys, "argv", ["mcp_server"]):
+ with patch.object(sys, "argv", ["mcp_server", "--metalake", "m"]):
self.assertEqual(120, _parse_args().gravitino_read_timeout)
with patch.object(
- sys, "argv", ["mcp_server", "--gravitino-read-timeout", "90"]
+ sys,
+ "argv",
+ [
+ "mcp_server",
+ "--metalake",
+ "m",
+ "--gravitino-read-timeout",
+ "90",
+ ],
):
self.assertEqual(90, _parse_args().gravitino_read_timeout)
def test_invalid_cli_and_environment_timeout(self):
with patch.dict(os.environ, {"GRAVITINO_READ_TIMEOUT": "nan"}):
- with patch.object(sys, "argv", ["mcp_server"]):
+ with patch.object(sys, "argv", ["mcp_server", "--metalake", "m"]):
with self.assertRaises(SystemExit) as failure:
_parse_args()
self.assertEqual(2, failure.exception.code)
with patch.object(
- sys, "argv", ["mcp_server", "--gravitino-read-timeout", "0"]
+ sys,
+ "argv",
+ ["mcp_server", "--metalake", "m", "--gravitino-read-timeout", "0"],
):
with self.assertRaises(SystemExit):
_parse_args()
@@ -84,10 +94,12 @@ class TestRequestTimeout(unittest.TestCase):
context = GravitinoContext(
Setting(metalake="m", gravitino_read_timeout=120)
)
- clients = [
- context._default_client,
- context._service_client("other"),
- ]
+ clients = [context._default_client]
+ with patch(
+ "mcp_server.core.context._get_request_authorization",
+ return_value="",
+ ):
+ self.assertIs(context._default_client, context.rest_client())
with patch(
"mcp_server.core.context._get_request_authorization",
return_value="Bearer caller",
diff --git a/mcp-server/tests/unit/tools/test_semantic_model.py
b/mcp-server/tests/unit/tools/test_semantic_model.py
deleted file mode 100644
index e5ec52e703..0000000000
--- a/mcp-server/tests/unit/tools/test_semantic_model.py
+++ /dev/null
@@ -1,643 +0,0 @@
-# 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.
-
-"""Semantic Model tools exercised through FastMCP with REST and stub
clients."""
-
-import json
-import unittest
-from decimal import Decimal
-from unittest.mock import patch
-
-import httpx
-import simplejson
-from fastmcp import Client
-from fastmcp.exceptions import ToolError
-
-from mcp_server.client.factory import RESTClientFactory
-from mcp_server.client.plain.plain_rest_client_operation import (
- PlainRESTClientOperation,
-)
-from mcp_server.core import Setting
-from mcp_server.server import GravitinoMCPServer
-from tests.unit.tools import MockOperation
-
-_NAMESPACE = {"catalog_name": "sales", "schema_name": "mart"}
-_MODEL_ARGS = {**_NAMESPACE, "semantic_model_name": "sales_model"}
-_MODELS_PATH = (
- "/api/metalakes/test_metalake/catalogs/sales/schemas/mart/semantic-models"
-)
-_METRICS_PATH =
"/api/metalakes/test_metalake/catalogs/sales/schemas/mart/semantic-model-metrics"
-_METRIC_TOOLS = (
- ("list_metrics", "identifiers", "false"),
- ("list_metric_infos", "metrics", "true"),
-)
-
-
-class TestSemanticModelTool(unittest.IsolatedAsyncioTestCase):
- def setUp(self):
- self.definition = {
- "aiContext": {"instructions": "Use certified sales definitions"},
- "datasets": [
- {
- "name": "orders",
- "source": {
- "namespace": ["sales", "mart"],
- "name": "orders",
- },
- "primaryKey": ["order_id"],
- "fields": [],
- }
- ],
- "relationships": [],
- "metrics": [],
- "customExtensions": [{"vendorName": "example", "data": "{}"}],
- }
- self.model = {
- "name": "sales_model",
- "comment": "Sales definitions",
- "definition": self.definition,
- "properties": {"domain": "sales"},
- "audit": {"creator": "admin", "createTime":
"2026-10-13T00:00:00Z"},
- }
- self.response = {"code": 0, "semanticModel": self.model}
- self.status = 200
- self.requests = []
- self.http_clients = []
- factory_patch = patch.object(
- RESTClientFactory, "_rest_client_class", self._create_client
- )
- factory_patch.start()
- self.addCleanup(factory_patch.stop)
- self.mcp = GravitinoMCPServer(
- Setting("test_metalake", token="semantic-test-token")
- ).mcp
-
- async def asyncTearDown(self):
- for client in self.http_clients:
- await client.aclose()
-
- def _respond(self, request):
- self.requests.append(request)
- return httpx.Response(
- self.status,
- content=(
- self.response
- if isinstance(self.response, bytes)
- else simplejson.dumps(self.response, use_decimal=True)
- ),
- headers={"Content-Type": "application/json"},
- )
-
- def _create_client(self, *args, **kwargs):
- # Keep the production factory/client wiring and inject only transport.
- async_client_type = httpx.AsyncClient
-
- def create_http_client(*client_args, **client_kwargs):
- client = async_client_type(
- *client_args,
- transport=httpx.MockTransport(self._respond),
- **client_kwargs,
- )
- self.http_clients.append(client)
- return client
-
- with patch(
-
"mcp_server.client.plain.plain_rest_client_operation._GravitinoAsyncClient",
- side_effect=create_http_client,
- ):
- return PlainRESTClientOperation(*args, **kwargs)
-
- async def _call(self, name, arguments):
- async with Client(self.mcp) as client:
- result = await client.call_tool(name, arguments)
- return json.loads(result.content[0].text, parse_float=Decimal)
-
- def _assert_request(self, method, path, body=None):
- request = self.requests[-1]
- self.assertEqual(request.method, method)
- self.assertEqual(request.url.raw_path.decode(), path)
- self.assertEqual(
- request.headers["Authorization"], "Bearer semantic-test-token"
- )
- if body is None:
- self.assertEqual(request.content, b"")
- else:
- self.assertEqual(json.loads(request.content), body)
-
- async def test_list_semantic_models(self):
- for identifiers in (
- [
- {
- "namespace": ["test_metalake", "sales", "mart"],
- "name": "sales_model",
- }
- ],
- [],
- ):
- with self.subTest(identifiers=identifiers):
- self.response = {"code": 0, "identifiers": identifiers}
- result = await self._call("list_of_semantic_models",
_NAMESPACE)
- self.assertEqual(result, identifiers)
- self._assert_request("GET", _MODELS_PATH)
-
- async def test_list_metrics_preserves_full_identifiers_and_order(self):
- namespace = ["test_metalake", "sales", "mart", "a.model"]
- identifiers = [
- {"namespace": namespace, "name": name}
- for name in ("Revenue", "revenue", "收入 /é+%")
- ] + [{"namespace": namespace[:-1] + ["b"], "name": "Revenue"}]
- self.response = {"code": 0, "identifiers": identifiers}
- self.assertEqual(
- identifiers, await self._call("list_metrics", _NAMESPACE)
- )
- self.assertEqual(1, len(self.requests))
- self._assert_request("GET", _METRICS_PATH + "?details=false")
-
- async def test_metric_infos_preserve_complete_metrics_and_numbers(self):
- metric = {
- "name": "Revenue",
- "expression": {
- "dialects": [
- {"dialect": "ANSI_SQL", "expression":
"SUM(orders.amount)"},
- {
- "dialect": "SPARK_SQL",
- "expression": "sum(orders.amount)",
- },
- ]
- },
- "description": "Revenue - café 雪",
- "datatype": "Decimal",
- "aiContext": {
- "instructions": "Use certified values",
- "synonyms": ["Sales"],
- "examples": ["Revenue by year"],
- "nested": {
- "precise": Decimal("0.12345678901234567890123456789"),
- "large": 1234567890123456789012345678901234567890,
- "values": [True, None, "extra", Decimal("1.00")],
- },
- },
- "customExtensions": [
- {"vendorName": "vendor", "data": '{"opaque": "雪"}'}
- ],
- }
- namespace = ["test_metalake", "sales", "mart", "a"]
- infos = [
- {"namespace": namespace, "metric": metric},
- {
- "namespace": namespace,
- "metric": {
- **metric,
- "name": "revenue",
- "aiContext": "Certified",
- },
- },
- {
- "namespace": namespace[:-1] + ["b"],
- "metric": {
- "name": "Revenue",
- "expression": metric["expression"],
- },
- },
- ]
- self.response = {"code": 0, "metrics": infos}
- result = await self._call("list_metric_infos", _NAMESPACE)
- self.assertEqual(infos, result)
- self.assertIsInstance(
- result[0]["metric"]["aiContext"]["nested"]["precise"], Decimal
- )
- self.assertEqual(1, len(self.requests))
- self._assert_request("GET", _METRICS_PATH + "?details=true")
-
- self.definition["metrics"] = [metric]
- self.response = {"code": 0, "semanticModel": self.model}
- model = await self._call("load_semantic_model", _MODEL_ARGS)
- self.assertEqual(model, self.model)
- self.assertEqual(result[0]["metric"],
model["definition"]["metrics"][0])
- self._assert_request("GET", f"{_MODELS_PATH}/sales_model")
-
- async def test_metric_discovery_empty_results(self):
- for tool, field, details in _METRIC_TOOLS:
- with self.subTest(tool=tool):
- self.requests.clear()
- self.response = {"code": 0, field: []}
- self.assertEqual([], await self._call(tool, _NAMESPACE))
- self.assertEqual(1, len(self.requests))
- self._assert_request(
- "GET", _METRICS_PATH + "?details=" + details
- )
-
- async def test_metric_discovery_path_encoding_and_request_scope(self):
- self.response = {"code": 0, "identifiers": [], "metrics": []}
- arguments = {
- "metalake": "ml/雪",
- "catalog_name": "cat?",
- "schema_name": "sch #+%",
- }
- path =
"/api/metalakes/ml%2F%E9%9B%AA/catalogs/cat%3F/schemas/sch%20%23%2B%25/semantic-model-metrics"
- for tool, _, details in _METRIC_TOOLS:
- with self.subTest(tool=tool):
- with patch(
- "mcp_server.core.context._get_request_authorization",
- return_value="Bearer request-test-token",
- ):
- await self._call(tool, arguments)
- request = self.requests[-1]
- self.assertEqual(
- path + "?details=" + details, request.url.raw_path.decode()
- )
- self.assertEqual(
- "Bearer request-test-token",
- request.headers["Authorization"],
- )
- await self._call(tool, _NAMESPACE)
- self._assert_request(
- "GET", _METRICS_PATH + "?details=" + details
- )
-
- async def test_metric_discovery_rejects_ambiguous_path_names(self):
- for tool, _, _ in _METRIC_TOOLS:
- for argument in ("metalake", "catalog_name", "schema_name"):
- invalid_names = (
- (".", "..") if argument == "metalake" else ("", ".", "..")
- )
- for name in invalid_names:
- with self.subTest(tool=tool, argument=argument, name=name):
- with self.assertRaisesRegex(
- ToolError, "REST path names"
- ):
- await self._call(
- tool, {**_NAMESPACE, argument: name}
- )
- self.assertEqual([], self.requests)
-
- async def test_metric_discovery_propagates_rest_failures(self):
- for tool, field, _ in _METRIC_TOOLS:
- for status, code, error_type in (
- (401, 1011, "UnauthorizedException"),
- (403, 1008, "ForbiddenException"),
- (404, 1003, "NoSuchSchemaException"),
- (500, 1002, "RuntimeException"),
- (501, 1006, "UnsupportedOperationException"),
- ):
- with self.subTest(tool=tool, status=status):
- self.status = status
- self.response = {
- "code": code,
- "type": error_type,
- "message": "Discovery rejected",
- field: [],
- }
- with self.assertRaisesRegex(ToolError, error_type):
- await self._call(tool, _NAMESPACE)
-
- async def test_metric_discovery_rejects_invalid_responses(self):
- for tool, field, _ in _METRIC_TOOLS:
- for body in (
- [],
- {},
- {field: []},
- {"code": 0},
- {"code": 0, field: None},
- {"code": 0, field: {}},
- {"code": 0, field: "[]"},
- ):
- with self.subTest(tool=tool, body=body):
- self.response = body
- with self.assertRaisesRegex(
- ToolError, "Invalid Gravitino response"
- ):
- await self._call(tool, _NAMESPACE)
- self.status = 502
- self.response = {"code": 0, field: []}
- with self.assertRaisesRegex(ToolError, "502"):
- await self._call(tool, _NAMESPACE)
- for body in (b"Bad Gateway", b'\xff\xfe{"code":0}'):
- with self.subTest(tool=tool, body=body):
- self.response = body
- self.status = 502
- with self.assertRaisesRegex(ToolError, "502"):
- await self._call(tool, _NAMESPACE)
- self.status = 200
- with self.assertRaisesRegex(
- ToolError, "Invalid Gravitino response: invalid JSON"
- ):
- await self._call(tool, _NAMESPACE)
-
- async def test_load_semantic_model_preserves_native_definition(self):
- result = await self._call("load_semantic_model", _MODEL_ARGS)
- self.assertEqual(result, self.model)
- self._assert_request("GET", f"{_MODELS_PATH}/sales_model")
-
- async def test_create_semantic_model_optional_fields(self):
- required = {"name": "sales_model", "definition": self.definition}
- for optional, expected in (
- ({}, {}),
- ({"comment": None, "properties": None}, {}),
- (
- {"comment": "", "properties": {}},
- {"comment": "", "properties": {}},
- ),
- (
- {
- "comment": "Sales definitions",
- "properties": {"domain": "sales"},
- },
- {
- "comment": "Sales definitions",
- "properties": {"domain": "sales"},
- },
- ),
- ):
- with self.subTest(optional=optional):
- result = await self._call(
- "create_semantic_model",
- {**_NAMESPACE, **required, **optional},
- )
- self.assertEqual(result, self.model)
- self._assert_request(
- "POST",
- _MODELS_PATH,
- {**required, **expected},
- )
-
- async def test_alter_semantic_model_preserves_order_and_null_comment(self):
- updates = [
- {"@type": "rename", "newName": "sales_model_v2"},
- {"@type": "updateComment", "newComment": None},
- {"@type": "setProperty", "property": "domain", "value": "sales"},
- {"@type": "removeProperty", "property": "legacy"},
- {"@type": "replaceDefinition", "definition": self.definition},
- ]
- self.model["name"] = "sales_model_v2"
- self.model["comment"] = None
- result = await self._call(
- "alter_semantic_model", {**_MODEL_ARGS, "updates": updates}
- )
- self.assertEqual(result, self.model)
- self._assert_request(
- "PUT", f"{_MODELS_PATH}/sales_model", {"updates": updates}
- )
-
- async def test_drop_semantic_model(self):
- for dropped in (True, False):
- with self.subTest(dropped=dropped):
- self.response = {"code": 0, "dropped": dropped}
- result = await self._call("drop_semantic_model", _MODEL_ARGS)
- self.assertIs(result, dropped)
- self._assert_request("DELETE", f"{_MODELS_PATH}/sales_model")
-
- async def test_all_operations_encode_each_path_segment(self):
- namespace = {
- "metalake": "ml/a",
- "catalog_name": "cat?",
- "schema_name": "sch#",
- }
- model_name = "sales %"
- models_path =
"/api/metalakes/ml%2Fa/catalogs/cat%3F/schemas/sch%23/semantic-models"
- for tool, method, extra, suffix in (
- ("list_of_semantic_models", "GET", {}, ""),
- (
- "create_semantic_model",
- "POST",
- {"name": model_name, "definition": self.definition},
- "",
- ),
- (
- "load_semantic_model",
- "GET",
- {"semantic_model_name": model_name},
- "/sales%20%25",
- ),
- (
- "alter_semantic_model",
- "PUT",
- {
- "semantic_model_name": model_name,
- "updates": [
- {"@type": "updateComment", "newComment": "new"}
- ],
- },
- "/sales%20%25",
- ),
- (
- "drop_semantic_model",
- "DELETE",
- {"semantic_model_name": model_name},
- "/sales%20%25",
- ),
- ):
- with self.subTest(tool=tool):
- await self._call(tool, {**namespace, **extra})
- request = self.requests[-1]
- self.assertEqual(request.method, method)
- self.assertEqual(
- request.url.raw_path.decode(), models_path + suffix
- )
- if tool == "create_semantic_model":
- self.assertEqual(
- json.loads(request.content)["name"], model_name
- )
-
- async def test_service_and_request_identity_and_metalake(self):
- async with Client(self.mcp) as client:
- with patch(
- "mcp_server.core.context._get_request_authorization",
- return_value="Bearer request-test-token",
- ):
- await client.call_tool(
- "load_semantic_model",
- {**_MODEL_ARGS, "metalake": "other_ml"},
- )
- await client.call_tool("load_semantic_model", _MODEL_ARGS)
-
- self.assertEqual(
- self.requests[0].url.path,
- _MODELS_PATH.replace("test_metalake", "other_ml") + "/sales_model",
- )
- self.assertEqual(
- self.requests[0].headers["Authorization"],
- "Bearer request-test-token",
- )
- self._assert_request("GET", f"{_MODELS_PATH}/sales_model")
-
- async def test_ambiguous_path_names_never_send_a_request(self):
- for argument in (
- "metalake",
- "catalog_name",
- "schema_name",
- "semantic_model_name",
- ):
- # An empty per-call metalake selects the configured default.
- invalid_names = (
- (".", "..") if argument == "metalake" else ("", ".", "..")
- )
- for name in invalid_names:
- with self.subTest(argument=argument, name=name):
- with self.assertRaisesRegex(ToolError, "REST path names"):
- await self._call(
- "drop_semantic_model",
- {**_MODEL_ARGS, argument: name},
- )
- self.assertEqual(self.requests, [])
-
- async def test_dots_within_names_are_preserved(self):
- await self._call(
- "load_semantic_model",
- {**_MODEL_ARGS, "semantic_model_name": "sales.v2"},
- )
- self._assert_request("GET", f"{_MODELS_PATH}/sales.v2")
-
- async def test_rest_errors_reach_the_mcp_caller(self):
- for tool, arguments, status, error_type in (
- (
- "list_of_semantic_models",
- _NAMESPACE,
- 403,
- "ForbiddenException",
- ),
- (
- "load_semantic_model",
- _MODEL_ARGS,
- 404,
- "NoSuchSemanticModelException",
- ),
- (
- "create_semantic_model",
- {
- **_NAMESPACE,
- "name": "sales_model",
- "definition": self.definition,
- },
- 409,
- "SemanticModelAlreadyExistsException",
- ),
- (
- "alter_semantic_model",
- {**_MODEL_ARGS, "updates": []},
- 400,
- "IllegalArgumentException",
- ),
- ("drop_semantic_model", _MODEL_ARGS, 409, "NotInUseException"),
- ("drop_semantic_model", _MODEL_ARGS, 403, "ForbiddenException"),
- ):
- with self.subTest(tool=tool, status=status):
- self.status = status
- self.response = {
- "code": status,
- "type": error_type,
- "message": "Semantic Model operation rejected",
- }
- with self.assertRaisesRegex(ToolError, error_type):
- await self._call(tool, arguments)
-
- async def test_create_requires_definition_before_sending_request(self):
- with self.assertRaisesRegex(ToolError, "definition"):
- await self._call(
- "create_semantic_model",
- {**_NAMESPACE, "name": "sales_model"},
- )
- self.assertEqual(self.requests, [])
-
- async def test_semantic_model_tag_exposes_all_tools(self):
- self.mcp = GravitinoMCPServer(
- Setting("test_metalake", tags={"semantic_model"})
- ).mcp
- async with Client(self.mcp) as client:
- tools = {tool.name: tool for tool in await client.list_tools()}
- self.assertEqual(
- set(tools),
- {
- "list_of_semantic_models",
- "list_metrics",
- "list_metric_infos",
- "load_semantic_model",
- "create_semantic_model",
- "alter_semantic_model",
- "drop_semantic_model",
- },
- )
- for tool in tools.values():
- self.assertIn("metalake", tool.inputSchema["properties"])
- for name, _, _ in _METRIC_TOOLS:
- tool = tools[name]
- self.assertEqual(
- {"catalog_name", "schema_name"},
- set(tool.inputSchema["required"]),
- )
- for name in (
- "list_of_semantic_models",
- "list_metrics",
- "list_metric_infos",
- "load_semantic_model",
- ):
- tool = tools[name]
- self.assertTrue(tool.annotations.readOnlyHint)
- self.assertTrue(tool.annotations.idempotentHint)
- self.assertFalse(tool.annotations.destructiveHint)
- self.assertEqual(
- set(tools["create_semantic_model"].inputSchema["required"]),
- {"catalog_name", "schema_name", "name", "definition"},
- )
- self.assertEqual(
- set(tools["drop_semantic_model"].inputSchema["required"]),
- {"catalog_name", "schema_name", "semantic_model_name"},
- )
-
-
-class TestSemanticModelToolWithMockOperation(unittest.IsolatedAsyncioTestCase):
- async def test_all_tools_return_json_from_mock_operation(self):
- model_args = {
- **_NAMESPACE,
- "semantic_model_name": "mock_semantic_model",
- }
- model = {"name": "mock_semantic_model"}
- cases = (
- ("list_of_semantic_models", _NAMESPACE, []),
- ("list_metrics", _NAMESPACE, []),
- ("list_metric_infos", _NAMESPACE, []),
- ("load_semantic_model", model_args, model),
- (
- "create_semantic_model",
- {**_NAMESPACE, "name": "mock_semantic_model", "definition":
{}},
- model,
- ),
- (
- "alter_semantic_model",
- {
- **model_args,
- "updates": [
- {"@type": "updateComment", "newComment": "Updated"}
- ],
- },
- model,
- ),
- ("drop_semantic_model", model_args, True),
- )
- with patch.object(
- RESTClientFactory, "_rest_client_class", MockOperation
- ):
- mcp = GravitinoMCPServer(
- Setting("mock_metalake", tags={"semantic_model"})
- ).mcp
- async with Client(mcp) as client:
- for tool, arguments, expected in cases:
- with self.subTest(tool=tool):
- result = await client.call_tool(tool, arguments)
- self.assertEqual(
- json.loads(result.content[0].text), expected
- )
