This is an automated email from the ASF dual-hosted git repository. github-actions[bot] pushed a commit to branch cherry-pick-696efe1c-to-branch-1.3 in repository https://gitbox.apache.org/repos/asf/gravitino.git
commit 9dd8278ede200ae87efb3b46dd35c3a0a4f07c22 Author: Qi Yu <[email protected]> AuthorDate: Mon Nov 2 19:33:59 2026 +0800 [#14498] fix(mcp): configure Gravitino response read timeout (#14506) ### What changes were proposed in this pull request? Configure the Gravitino response read timeout with `--gravitino-read-timeout` / `GRAVITINO_READ_TIMEOUT` (default 60 seconds), validate finite positive values, and apply it to default, service-identity and forwarded-identity clients. Preserve actionable timeout details through FastMCP, including the method, encoded resource path and configured budget. Exclude queries, userinfo, headers and body from the new message. Keep the original cause and do not replay requests. Authentication subrequests retain their own timeout exception. ### Why are the changes needed? HTTPX's implicit five-second read timeout can expire during otherwise valid metadata operations, and FastMCP can hide the relevant Gravitino timeout context. Fixes #14498 ### Does this PR introduce _any_ user-facing change? The default read timeout increases from 5 to 60 seconds; CLI overrides the environment variable. Connect, write and pool timeouts remain five seconds. This bounds waiting for response data, not total server-operation duration. A timed-out mutation may still complete: verify its result before retrying. ### How was this patch tested? - All 311 MCP unit tests passed, including a real registered table tool through GravitinoMCPServer and its production middleware, all identity paths, configuration precedence, no replay, sanitized diagnostics, cause preservation and OAuth timeout isolation. - Reinstating the five-second client timeout makes both configuration-wiring regressions fail; restoring the implementation passes the suite. - Black/isort, full MCP pylint and RAT passed. Independent fresh-context review found no outstanding findings. - No live slow-network, external IdP or third-party MCP client budget test was run; MockTransport injects timeout errors and separate assertions verify client budgets. # Conflicts: # docs/gravitino-mcp-server.md # mcp-server/mcp_server/core/context.py # mcp-server/mcp_server/core/setting.py # mcp-server/tests/unit/test_per_request_metalake.py # mcp-server/tests/unit/tools/test_semantic_model.py --- docs/gravitino-mcp-server.md | 20 + mcp-server/mcp_server/client/factory.py | 8 +- .../client/plain/plain_rest_client_operation.py | 36 +- mcp-server/mcp_server/core/context.py | 94 +++ mcp-server/mcp_server/core/setting.py | 19 + mcp-server/mcp_server/main.py | 24 + mcp-server/tests/unit/test_per_request_metalake.py | 839 +++++++++++++++++++++ mcp-server/tests/unit/test_request_timeout.py | 252 +++++++ mcp-server/tests/unit/tools/mock_operation.py | 4 +- mcp-server/tests/unit/tools/test_semantic_model.py | 643 ++++++++++++++++ 10 files changed, 1935 insertions(+), 4 deletions(-) diff --git a/docs/gravitino-mcp-server.md b/docs/gravitino-mcp-server.md index c196a8f03f..eb8df91a6e 100644 --- a/docs/gravitino-mcp-server.md +++ b/docs/gravitino-mcp-server.md @@ -111,6 +111,7 @@ 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 | @@ -125,6 +126,25 @@ 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/client/factory.py b/mcp-server/mcp_server/client/factory.py index 068bdf1b74..1ed6e13490 100644 --- a/mcp-server/mcp_server/client/factory.py +++ b/mcp-server/mcp_server/client/factory.py @@ -35,6 +35,7 @@ class RESTClientFactory: authorization: str = "", *, auth=None, + read_timeout: float = 60.0, ) -> "PlainRESTClientOperation": """ Create a new rest client instance with the specified parameters. @@ -45,13 +46,18 @@ class RESTClientFactory: authorization: Full Authorization header value forwarded verbatim (e.g. "Bearer <token>" or "Basic <base64(user:dummy)>"). Empty string means anonymous. + read_timeout: Response read timeout in seconds. auth: Optional httpx auth hook. ``None`` for static or anonymous. Returns: New instance of the configured rest client class """ return cls._rest_client_class( - metalake_name, uri, authorization, auth=auth + metalake_name, + uri, + authorization, + auth=auth, + read_timeout=read_timeout, ) @classmethod 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 cf225dfa2b..cdad58ce18 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 @@ -29,6 +29,7 @@ from mcp_server.client import ( TagOperation, ) from mcp_server.client.job_operation import JobOperation +from mcp_server.client.plain.exception import GravitinoException from mcp_server.client.plain.plain_rest_client_catalog_operation import ( PlainRESTClientCatalogOperation, ) @@ -62,6 +63,30 @@ from mcp_server.client.plain.plain_rest_client_topic_operation import ( from mcp_server.client.topic_operation import TopicOperation +class _GravitinoAsyncClient(httpx.AsyncClient): + """Preserve actionable timeout details instead of FastMCP's generic timeout.""" + + async def send(self, request, **kwargs): + try: + return await super().send(request, **kwargs) + except httpx.ReadTimeout as exc: + # OAuth hooks can time out on their own token request, with a different budget. + # Preserve that exception instead of misreporting a Gravitino timeout. + if exc.request.url != request.url: + raise + # Include only the encoded path: queries, userinfo and credentials are excluded. + resource = request.url.raw_path.split(b"?", 1)[0].decode( + "ascii", errors="backslashreplace" + ) + # Do not retry writes: the upstream operation may already have committed. + raise GravitinoException( + f"The MCP server stopped waiting for Gravitino ({request.method} " + f"{resource}) after a response read " + f"timeout of {self.timeout.read:g} seconds. The operation may still " + "complete on the server; check its result before retrying." + ) from exc + + # pylint: disable=too-many-instance-attributes class PlainRESTClientOperation(GravitinoOperation): def __init__( @@ -71,6 +96,7 @@ class PlainRESTClientOperation(GravitinoOperation): authorization: str = "", *, auth: Optional[httpx.Auth] = None, + read_timeout: float = 60.0, ): """Create a REST client for one identity. @@ -81,14 +107,20 @@ class PlainRESTClientOperation(GravitinoOperation): on every request (for example ``"Bearer <token>"`` for OAuth2 or ``"Basic <base64(user:secret)>"`` for simple or Basic auth). Empty string means anonymous (no header sent). + read_timeout: Response read timeout in seconds (default 60). auth: Optional httpx auth hook used instead of a frozen header (OAuth client-credentials refresh). """ headers = {} if authorization: headers["Authorization"] = authorization - _rest_client = httpx.AsyncClient( - base_url=uri, headers=headers, auth=auth + _rest_client = _GravitinoAsyncClient( + base_url=uri, + headers=headers, + auth=auth, + timeout=httpx.Timeout( + read_timeout, connect=5.0, pool=5.0, write=5.0 + ), ) # Kept so the shared connection pool can be closed (see close()). self._rest_client = _rest_client diff --git a/mcp-server/mcp_server/core/context.py b/mcp-server/mcp_server/core/context.py index f38831f639..34effad2b1 100644 --- a/mcp-server/mcp_server/core/context.py +++ b/mcp-server/mcp_server/core/context.py @@ -148,11 +148,33 @@ 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)) ) # LRU cache of per-principal clients keyed by the raw Authorization header. # Safe without locking: rest_client() runs on the single asyncio event @@ -190,6 +212,72 @@ 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) @@ -200,7 +288,13 @@ 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 112528f6d9..0555d9fee7 100644 --- a/mcp-server/mcp_server/core/setting.py +++ b/mcp-server/mcp_server/core/setting.py @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. +import math from dataclasses import dataclass, field from typing import Set @@ -55,7 +56,24 @@ 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 + ): + raise ValueError( + "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( @@ -99,6 +117,7 @@ class Setting: # pylint: disable=too-many-instance-attributes return ( f"Setting(metalake={self.metalake}, gravitino_uri={self.gravitino_uri}, " f"tags={self.tags}, transport={self.transport}, mcp_url={self.mcp_url}, " + f"gravitino_read_timeout={self.gravitino_read_timeout}, " f"token={token_display}, tls_cert={self.tls_cert}, " f"tls_key={self.tls_key}, " f"oauth_token_endpoint={self.oauth_token_endpoint}, " diff --git a/mcp-server/mcp_server/main.py b/mcp-server/mcp_server/main.py index 5ab8b1c777..912cdbaac5 100644 --- a/mcp-server/mcp_server/main.py +++ b/mcp-server/mcp_server/main.py @@ -17,6 +17,7 @@ import argparse import logging +import math import os from mcp_server.core.setting import DefaultSetting, Setting @@ -34,6 +35,7 @@ def do_main(): args = _parse_args() setting = Setting( metalake=args.metalake, + gravitino_read_timeout=args.gravitino_read_timeout, gravitino_uri=args.gravitino_uri, tags=args.include_tool_tags, transport=args.transport, @@ -80,6 +82,20 @@ def _comma_separated_set(value) -> set: return set(item.strip() for item in value.split(",") if item.strip()) +def _positive_timeout(value): + try: + timeout = float(value) + except ValueError as exc: + raise argparse.ArgumentTypeError( + "timeout must be a finite positive number" + ) from exc + if not math.isfinite(timeout) or timeout <= 0: + raise argparse.ArgumentTypeError( + "timeout must be a finite positive number" + ) + return timeout + + def _parse_args(): parser = argparse.ArgumentParser( description="Gravitino Mcp server", @@ -99,6 +115,14 @@ def _parse_args(): help=f"The uri of Gravitino server. (default: {DefaultSetting.default_gravitino_uri})", ) + parser.add_argument( + "--gravitino-read-timeout", + type=_positive_timeout, + default=os.environ.get("GRAVITINO_READ_TIMEOUT", "60"), + help="Gravitino response read timeout in seconds (default: 60). " + "Can also be set with GRAVITINO_READ_TIMEOUT. Must be finite and positive.", + ) + parser.add_argument( "--include-tool-tags", type=_comma_separated_set, 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..bc4fc190f8 --- /dev/null +++ b/mcp-server/tests/unit/test_per_request_metalake.py @@ -0,0 +1,839 @@ +# 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 new file mode 100644 index 0000000000..6cbb979305 --- /dev/null +++ b/mcp-server/tests/unit/test_request_timeout.py @@ -0,0 +1,252 @@ +# 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 inspect client configuration and identity caches deliberately. +# pylint: disable=protected-access + +import asyncio +import math +import os +import sys +import unittest +from unittest.mock import patch + +import httpx +from fastmcp import Client, FastMCP + +from mcp_server.client.factory import RESTClientFactory +from mcp_server.client.plain.exception import GravitinoException +from mcp_server.client.plain.plain_rest_client_operation import ( + _GravitinoAsyncClient, +) +from mcp_server.core.context import GravitinoContext +from mcp_server.core.setting import Setting +from mcp_server.main import _parse_args +from mcp_server.server import GravitinoMCPServer + + +class TestRequestTimeout(unittest.TestCase): + def test_configuration(self): + self.assertEqual(60.0, Setting().gravitino_read_timeout) + for value in (0, -1, math.inf, math.nan): + with self.assertRaises(ValueError): + Setting(gravitino_read_timeout=value) + with patch.dict(os.environ, {"GRAVITINO_READ_TIMEOUT": "120"}): + with patch.object(sys, "argv", ["mcp_server"]): + self.assertEqual(120, _parse_args().gravitino_read_timeout) + with patch.object( + sys, "argv", ["mcp_server", "--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 self.assertRaises(SystemExit) as failure: + _parse_args() + self.assertEqual(2, failure.exception.code) + with patch.object( + sys, "argv", ["mcp_server", "--gravitino-read-timeout", "0"] + ): + with self.assertRaises(SystemExit): + _parse_args() + + def test_client_timeout(self): + async def check(): + client = RESTClientFactory.create_rest_client( + "m", "http://localhost";, read_timeout=90 + ) + try: + self.assertEqual(90, client._rest_client.timeout.read) + self.assertEqual(5, client._rest_client.timeout.connect) + self.assertEqual(5, client._rest_client.timeout.pool) + finally: + await client._rest_client.aclose() + + asyncio.run(check()) + + def test_all_identity_paths_use_configured_timeout(self): + async def check(): + context = GravitinoContext( + Setting(metalake="m", gravitino_read_timeout=120) + ) + clients = [ + context._default_client, + context._service_client("other"), + ] + with patch( + "mcp_server.core.context._get_request_authorization", + return_value="Bearer caller", + ): + clients.append(context.rest_client()) + try: + for client in clients: + self.assertEqual(120, client._rest_client.timeout.read) + finally: + for client in clients: + await client.close() + + asyncio.run(check()) + + def test_read_timeout_is_actionable_and_not_retried(self): + async def check(): + calls = [] + + def handler(request): + calls.append(request) + raise httpx.ReadTimeout( + "secret backend detail", request=request + ) + + async with _GravitinoAsyncClient( + transport=httpx.MockTransport(handler), + timeout=httpx.Timeout(60), + ) as client: + with self.assertRaises(GravitinoException) as failure: + await client.post( + "http://user:password@localhost/tables?token=query_secret";, + headers={"Authorization": "Bearer auth_secret"}, + json={}, + ) + self.assertIn("POST /tables", str(failure.exception)) + self.assertNotIn("password", str(failure.exception)) + self.assertNotIn("query_secret", str(failure.exception)) + self.assertIn("60 seconds", str(failure.exception)) + self.assertIn("may still complete", str(failure.exception)) + self.assertNotIn("secret", str(failure.exception)) + self.assertIsInstance( + failure.exception.__cause__, httpx.ReadTimeout + ) + self.assertEqual(1, len(calls)) + + asyncio.run(check()) + + def test_mcp_tool_preserves_timeout_details(self): + async def check(): + def handler(request): + raise httpx.ReadTimeout("slow", request=request) + + server = FastMCP("timeout-test") + async with _GravitinoAsyncClient( + transport=httpx.MockTransport(handler), + timeout=httpx.Timeout(90), + ) as rest: + + @server.tool + async def load_table() -> str: + """Load metadata from Gravitino.""" + await rest.get("http://localhost/table";) + return "loaded" + + async with Client(server) as client: + result = await client.call_tool( + "load_table", {}, raise_on_error=False + ) + self.assertTrue(result.is_error) + message = " ".join(item.text for item in result.content) + self.assertIn("90 seconds", message) + self.assertIn("may still complete", message) + + asyncio.run(check()) + + def test_registered_tool_preserves_timeout_through_production_server(self): + async def check(): + calls = [] + + def handler(request): + calls.append(request) + raise httpx.ReadTimeout("slow", request=request) + + original_init = _GravitinoAsyncClient.__init__ + + def init_with_transport(client, *args, **kwargs): + kwargs["transport"] = httpx.MockTransport(handler) + original_init(client, *args, **kwargs) + + with patch.object( + _GravitinoAsyncClient, "__init__", init_with_transport + ): + server = GravitinoMCPServer( + Setting(metalake="m", gravitino_read_timeout=75) + ) + async with Client(server.mcp) as client: + result = await client.call_tool( + "get_table_metadata_details", + { + "catalog_name": "c", + "schema_name": "s", + "table_name": "t", + }, + raise_on_error=False, + ) + self.assertTrue(result.is_error) + message = " ".join(item.text for item in result.content) + self.assertIn("75 seconds", message) + self.assertIn( + "GET /api/metalakes/m/catalogs/c/schemas/s/tables/t", + message, + ) + self.assertIn("may still complete", message) + self.assertEqual(1, len(calls)) + + asyncio.run(check()) + + def test_authentication_timeout_is_not_reported_as_gravitino(self): + async def check(): + calls = [] + + def token_handler(request): + raise httpx.ReadTimeout("IdP timeout", request=request) + + async with httpx.AsyncClient( + transport=httpx.MockTransport(token_handler) + ) as token_client: + + class TokenAuth(httpx.Auth): + async def async_auth_flow(self, request): + await token_client.get("http://identity/token";) + yield request + + def gravitino_handler(request): + calls.append(request) + return httpx.Response(200) + + async with _GravitinoAsyncClient( + auth=TokenAuth(), + transport=httpx.MockTransport(gravitino_handler), + timeout=httpx.Timeout(120), + ) as client: + with self.assertRaises(httpx.ReadTimeout) as failure: + await client.get("http://gravitino/table";) + self.assertEqual("IdP timeout", str(failure.exception)) + self.assertEqual([], calls) + + asyncio.run(check()) + + def test_other_results_are_preserved(self): + async def check(): + def handler(request): + return httpx.Response(403, json={"message": "denied"}) + + async with _GravitinoAsyncClient( + transport=httpx.MockTransport(handler) + ) as client: + response = await client.get("http://localhost/tables";) + self.assertEqual(403, response.status_code) + self.assertEqual("denied", response.json()["message"]) + + asyncio.run(check()) diff --git a/mcp-server/tests/unit/tools/mock_operation.py b/mcp-server/tests/unit/tools/mock_operation.py index 61f0a3e7fe..0817b91eed 100644 --- a/mcp-server/tests/unit/tools/mock_operation.py +++ b/mcp-server/tests/unit/tools/mock_operation.py @@ -31,7 +31,9 @@ from mcp_server.client.statistic_operation import StatisticOperation class MockOperation(GravitinoOperation): - def __init__(self, metalake, uri, authorization="", *, auth=None): + def __init__( + self, metalake, uri, authorization="", *, auth=None, read_timeout=60.0 + ): pass def as_table_operation(self) -> TableOperation: diff --git a/mcp-server/tests/unit/tools/test_semantic_model.py b/mcp-server/tests/unit/tools/test_semantic_model.py new file mode 100644 index 0000000000..e5ec52e703 --- /dev/null +++ b/mcp-server/tests/unit/tools/test_semantic_model.py @@ -0,0 +1,643 @@ +# 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 + )
