This is an automated email from the ASF dual-hosted git repository.
yuqi1129 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/gravitino.git
The following commit(s) were added to refs/heads/main by this push:
new 696efe1c2c [#14498] fix(mcp): configure Gravitino response read
timeout (#14506)
696efe1c2c is described below
commit 696efe1c2ccfe202f1b72fda63839e538c941268
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.
---
docs/gravitino-mcp-server.md | 31 +--
mcp-server/mcp_server/client/factory.py | 8 +-
.../client/plain/plain_rest_client_operation.py | 36 ++-
mcp-server/mcp_server/core/context.py | 3 +
mcp-server/mcp_server/core/setting.py | 10 +
mcp-server/mcp_server/main.py | 24 ++
mcp-server/tests/unit/test_per_request_metalake.py | 12 +-
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 | 2 +-
10 files changed, 360 insertions(+), 22 deletions(-)
diff --git a/docs/gravitino-mcp-server.md b/docs/gravitino-mcp-server.md
index c07bee1615..81087397e9 100644
--- a/docs/gravitino-mcp-server.md
+++ b/docs/gravitino-mcp-server.md
@@ -217,20 +217,23 @@ indicating whether it was dropped. It removes semantic
metadata without deleting
You could config Gravitino MCP server by arguments, `uv run mcp_server -h`
shows the detailed information.
-| 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 |
-| `--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 |
+| 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.
## 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 74afefdac2..5aed84699c 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
@@ -30,6 +30,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,
)
@@ -76,6 +77,30 @@ from mcp_server.client.semantic_model_operation import
SemanticModelOperation
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__(
@@ -85,6 +110,7 @@ class PlainRESTClientOperation(GravitinoOperation):
authorization: str = "",
*,
auth: Optional[httpx.Auth] = None,
+ read_timeout: float = 60.0,
):
"""Create a REST client for one identity.
@@ -95,14 +121,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 0bb8f7af73..8620bd00e8 100644
--- a/mcp-server/mcp_server/core/context.py
+++ b/mcp-server/mcp_server/core/context.py
@@ -255,6 +255,7 @@ class GravitinoContext:
setting.gravitino_uri,
startup_authorization(setting),
auth=self._service_auth,
+ read_timeout=setting.gravitino_read_timeout,
)
if setting.metalake
else None
@@ -331,6 +332,7 @@ class GravitinoContext:
metalake,
self._setting.gravitino_uri,
authorization,
+ read_timeout=self._setting.gravitino_read_timeout,
)
self._cache_put(key, client)
return self._borrow(client)
@@ -394,6 +396,7 @@ class GravitinoContext:
self._setting.gravitino_uri,
startup_authorization(self._setting),
auth=self._service_auth,
+ read_timeout=self._setting.gravitino_read_timeout,
)
self._cache_put(key, client)
return client
diff --git a/mcp-server/mcp_server/core/setting.py
b/mcp-server/mcp_server/core/setting.py
index bb69175a6b..73c5a7d17a 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
@@ -63,12 +64,20 @@ 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
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"
+ )
def has_oauth_client(self) -> bool:
"""Return True when client-credentials is fully configured."""
@@ -122,6 +131,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 2a50b8f68d..9f16f59453 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
@@ -36,6 +37,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,
@@ -83,6 +85,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",
@@ -104,6 +120,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
index 95faea1abb..bc4fc190f8 100644
--- a/mcp-server/tests/unit/test_per_request_metalake.py
+++ b/mcp-server/tests/unit/test_per_request_metalake.py
@@ -179,7 +179,9 @@ class _EchoCatalogOperation:
class _EchoMetalakeClient(MockOperation):
"""MockOperation that remembers the metalake it was constructed with."""
- def __init__(self, metalake, uri, authorization="", *, auth=None):
+ def __init__(
+ self, metalake, uri, authorization="", *, auth=None, read_timeout=60.0
+ ):
super().__init__(metalake, uri, authorization, auth=auth)
self._metalake = metalake
@@ -270,7 +272,9 @@ class _RecordingClient(MockOperation):
built = []
auths = []
- def __init__(self, metalake, uri, authorization="", *, auth=None):
+ 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)
@@ -285,7 +289,9 @@ class _ClosableClient(MockOperation):
parked = None
released = None
- def __init__(self, metalake, uri, authorization="", *, auth=None):
+ def __init__(
+ self, metalake, uri, authorization="", *, auth=None, read_timeout=60.0
+ ):
super().__init__(metalake, uri, authorization, auth=auth)
self._metalake = metalake
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 481bdd650b..8c0c6fe6fa 100644
--- a/mcp-server/tests/unit/tools/mock_operation.py
+++ b/mcp-server/tests/unit/tools/mock_operation.py
@@ -37,7 +37,9 @@ from mcp_server.client.view_operation import ViewOperation
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
index c54d5537f2..e5ec52e703 100644
--- a/mcp-server/tests/unit/tools/test_semantic_model.py
+++ b/mcp-server/tests/unit/tools/test_semantic_model.py
@@ -116,7 +116,7 @@ class
TestSemanticModelTool(unittest.IsolatedAsyncioTestCase):
return client
with patch(
-
"mcp_server.client.plain.plain_rest_client_operation.httpx.AsyncClient",
+
"mcp_server.client.plain.plain_rest_client_operation._GravitinoAsyncClient",
side_effect=create_http_client,
):
return PlainRESTClientOperation(*args, **kwargs)