This is an automated email from the ASF dual-hosted git repository. yuqi1129 pushed a commit to branch feat/mcp-governance-task3-6 in repository https://gitbox.apache.org/repos/asf/gravitino.git
commit 4c3570b6c0ebae3effb0618ae8c9d1ecbeb25e07 Author: yuqi <[email protected]> AuthorDate: Wed Jun 10 20:11:48 2026 +0800 [#11565][#11566] feat(mcp-server): add --token CLI param and inject Bearer token into Gravitino REST calls - Setting: add `token` field; mask value in __str__ to avoid log leakage - main.py: add `--token` CLI arg (fallback to GRAVITINO_TOKEN env var) - PlainRESTClientOperation: accept token; set Authorization: Bearer header on httpx client when non-empty - RESTClientFactory.create_rest_client: pass token through - GravitinoContext: propagate setting.token to factory - MockOperation: accept token param (test compatibility) - tests/unit/test_auth_flow.py: 7 new tests covering token injection, masking, and context propagation --- mcp-server/mcp_server/client/factory.py | 5 +- .../client/plain/plain_rest_client_operation.py | 7 +- mcp-server/mcp_server/core/context.py | 2 +- mcp-server/mcp_server/core/setting.py | 11 +++ mcp-server/mcp_server/main.py | 11 +++ mcp-server/tests/unit/test_auth_flow.py | 93 ++++++++++++++++++++++ mcp-server/tests/unit/tools/mock_operation.py | 2 +- 7 files changed, 125 insertions(+), 6 deletions(-) diff --git a/mcp-server/mcp_server/client/factory.py b/mcp-server/mcp_server/client/factory.py index d990284d28..2e5489d053 100644 --- a/mcp-server/mcp_server/client/factory.py +++ b/mcp-server/mcp_server/client/factory.py @@ -29,7 +29,7 @@ class RESTClientFactory: @classmethod def create_rest_client( - cls, metalake_name: str, uri: str + cls, metalake_name: str, uri: str, token: str = "" ) -> "PlainRESTClientOperation": """ Create a new rest client instance with the specified parameters. @@ -37,11 +37,12 @@ class RESTClientFactory: Args: metalake_name: Name of the metalake uri: URI of the Gravitino server endpoint + token: Bearer token for Authorization header (empty = anonymous) Returns: New instance of the configured rest client class """ - return cls._rest_client_class(metalake_name, uri) + return cls._rest_client_class(metalake_name, uri, token) @classmethod def set_rest_client(cls, rest_client_class: type) -> None: 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 205562b6d8..8bb8b51aa3 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 @@ -62,8 +62,11 @@ from mcp_server.client.topic_operation import TopicOperation # pylint: disable=too-many-instance-attributes class PlainRESTClientOperation(GravitinoOperation): - def __init__(self, metalake_name: str, uri: str): - _rest_client = httpx.AsyncClient(base_url=uri) + def __init__(self, metalake_name: str, uri: str, token: str = ""): + headers = {} + if token: + headers["Authorization"] = f"Bearer {token}" + _rest_client = httpx.AsyncClient(base_url=uri, headers=headers) self._catalog_operation = PlainRESTClientCatalogOperation( metalake_name, _rest_client ) diff --git a/mcp-server/mcp_server/core/context.py b/mcp-server/mcp_server/core/context.py index 6e540a0b51..83350a7d10 100644 --- a/mcp-server/mcp_server/core/context.py +++ b/mcp-server/mcp_server/core/context.py @@ -22,7 +22,7 @@ from mcp_server.core.setting import Setting class GravitinoContext: def __init__(self, setting: Setting): self.gravitino_client = RESTClientFactory.create_rest_client( - setting.metalake, setting.gravitino_uri + setting.metalake, setting.gravitino_uri, setting.token ) def rest_client(self): diff --git a/mcp-server/mcp_server/core/setting.py b/mcp-server/mcp_server/core/setting.py index 474fdbc6b4..4d159d6c90 100644 --- a/mcp-server/mcp_server/core/setting.py +++ b/mcp-server/mcp_server/core/setting.py @@ -33,3 +33,14 @@ class Setting: tags: Set[str] = field(default_factory=set) transport: str = DefaultSetting.default_transport mcp_url: str = DefaultSetting.default_mcp_url + # Bearer token forwarded to Gravitino on every request. + # Empty string means anonymous (no Authorization header sent). + token: str = "" + + def __str__(self) -> str: + token_display = "***" if self.token else "" + return ( + f"Setting(metalake={self.metalake}, gravitino_uri={self.gravitino_uri}, " + f"tags={self.tags}, transport={self.transport}, mcp_url={self.mcp_url}, " + f"token={token_display})" + ) diff --git a/mcp-server/mcp_server/main.py b/mcp-server/mcp_server/main.py index 73a5aa265f..b1200b3180 100644 --- a/mcp-server/mcp_server/main.py +++ b/mcp-server/mcp_server/main.py @@ -17,6 +17,7 @@ import argparse import logging +import os from mcp_server.core.setting import DefaultSetting, Setting from mcp_server.server import GravitinoMCPServer @@ -30,6 +31,7 @@ def do_main(): tags=args.include_tool_tags, transport=args.transport, mcp_url=args.mcp_url, + token=args.token, ) _init_logging(setting) logging.info("Gravitino MCP server setting: %s", setting) @@ -96,6 +98,15 @@ def _parse_args(): help=f"The url of MCP server if using http transport. (default: {DefaultSetting.default_mcp_url})", ) + parser.add_argument( + "--token", + type=str, + default=os.environ.get("GRAVITINO_TOKEN", ""), + help="Bearer token forwarded as Authorization header to Gravitino on every request. " + "Can also be set via the GRAVITINO_TOKEN environment variable. " + "When omitted, requests are sent without authentication.", + ) + args = parser.parse_args() return args diff --git a/mcp-server/tests/unit/test_auth_flow.py b/mcp-server/tests/unit/test_auth_flow.py new file mode 100644 index 0000000000..a413515bf9 --- /dev/null +++ b/mcp-server/tests/unit/test_auth_flow.py @@ -0,0 +1,93 @@ +# 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. + +import unittest + +from mcp_server.client.plain.plain_rest_client_operation import ( + PlainRESTClientOperation, +) +from mcp_server.core.context import GravitinoContext +from mcp_server.core.setting import Setting + + +class TestTokenInjection(unittest.TestCase): + """Verify Bearer token is correctly injected into the httpx client.""" + + def test_token_sets_authorization_header(self): + """When a token is provided, the httpx client carries Authorization: Bearer.""" + client = PlainRESTClientOperation( + "my_metalake", "http://localhost:8090";, token="my-secret-token" + ) + headers = dict(client._catalog_operation.rest_client.headers) + self.assertEqual(headers.get("authorization"), "Bearer my-secret-token") + + def test_empty_token_no_authorization_header(self): + """When token is empty, no Authorization header is added.""" + client = PlainRESTClientOperation( + "my_metalake", "http://localhost:8090";, token="" + ) + headers = dict(client._catalog_operation.rest_client.headers) + self.assertNotIn("authorization", headers) + + def test_no_token_argument_no_authorization_header(self): + """When token argument is omitted entirely, no Authorization header is added.""" + client = PlainRESTClientOperation("my_metalake", "http://localhost:8090";) + headers = dict(client._catalog_operation.rest_client.headers) + self.assertNotIn("authorization", headers) + + +class TestSettingTokenMasking(unittest.TestCase): + """Verify that the token is not exposed in Setting string representation.""" + + def test_token_masked_in_str(self): + """Token value must not appear in Setting.__str__.""" + setting = Setting(metalake="ml", token="super-secret-token-value") + self.assertNotIn("super-secret-token-value", str(setting)) + self.assertIn("***", str(setting)) + + def test_empty_token_shows_empty_in_str(self): + """When no token is set, __str__ shows empty placeholder.""" + setting = Setting(metalake="ml", token="") + self.assertNotIn("***", str(setting)) + + +class TestGravitinoContextTokenPropagation(unittest.TestCase): + """Verify GravitinoContext passes token from Setting to the REST client.""" + + def test_context_propagates_token(self): + """Token from Setting reaches the httpx client Authorization header.""" + setting = Setting( + metalake="ml", + gravitino_uri="http://localhost:8090";, + token="ctx-token-xyz", + ) + ctx = GravitinoContext(setting) + rest_client = ctx.rest_client() + headers = dict(rest_client._catalog_operation.rest_client.headers) + self.assertEqual(headers.get("authorization"), "Bearer ctx-token-xyz") + + def test_context_anonymous_when_no_token(self): + """Empty token in Setting → no Authorization header in REST calls.""" + setting = Setting( + metalake="ml", + gravitino_uri="http://localhost:8090";, + token="", + ) + ctx = GravitinoContext(setting) + rest_client = ctx.rest_client() + headers = dict(rest_client._catalog_operation.rest_client.headers) + self.assertNotIn("authorization", headers) diff --git a/mcp-server/tests/unit/tools/mock_operation.py b/mcp-server/tests/unit/tools/mock_operation.py index d76ba94b60..a985bf7912 100644 --- a/mcp-server/tests/unit/tools/mock_operation.py +++ b/mcp-server/tests/unit/tools/mock_operation.py @@ -31,7 +31,7 @@ from mcp_server.client.statistic_operation import StatisticOperation class MockOperation(GravitinoOperation): - def __init__(self, metalake, uri): + def __init__(self, metalake, uri, token=""): pass def as_table_operation(self) -> TableOperation:
