Copilot commented on code in PR #11620:
URL: https://github.com/apache/gravitino/pull/11620#discussion_r3401146026
##########
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 = ""
Review Comment:
`Setting` overrides `__str__()` to mask the token, but the
dataclass-generated `__repr__()` will still include the raw `token` value
(e.g., if the object is logged with `%r` or appears in exception/debug output).
This is a real risk of leaking credentials even if `__str__` is safe.
##########
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")
Review Comment:
This assertion converts `httpx.Headers` to a `dict` and then looks up a
lowercase key. Header casing in the resulting `dict` is not guaranteed across
httpx versions; use `Headers.get(...)` (case-insensitive) instead, and close
the underlying `AsyncClient` to avoid ResourceWarnings.
##########
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)
+
Review Comment:
Same as above: avoid relying on header key casing by using
`Headers.get(...)`, and make sure to close the `AsyncClient` created by
`PlainRESTClientOperation` to prevent ResourceWarnings.
##########
mcp-server/mcp_server/main.py:
##########
@@ -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.",
+ )
Review Comment:
The new `--token` argument has important behavior (CLI value vs
`GRAVITINO_TOKEN` fallback) but there’s no unit test covering
parsing/precedence. Given the guideline that every change should have a
corresponding test, please add a small test to validate: (1) env var is used
when `--token` is omitted, and (2) `--token` overrides the env var when both
are set.
##########
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
+
Review Comment:
These unit tests create `httpx.AsyncClient` instances (via
`PlainRESTClientOperation`) but never close them, which can produce
ResourceWarnings and leak sockets across the test run. Import `asyncio` so the
tests can close the client with `aclose()`.
##########
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")
+
Review Comment:
This test also depends on `dict(headers)` key casing and leaves the
underlying `httpx.AsyncClient` unclosed (the context creates a client under the
hood). Use `Headers.get(...)` and close the client at the end to avoid
ResourceWarnings.
##########
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)
Review Comment:
This test leaves the underlying `httpx.AsyncClient` unclosed and checks for
a lowercase key in a `dict` created from `httpx.Headers` (header key casing
isn’t guaranteed). Use `Headers.get(...)` and close the client to keep the unit
test suite clean and non-leaky.
##########
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)
Review Comment:
Same as above: use `Headers.get(...)` rather than converting to `dict` and
checking for a specific key casing, and close the created `AsyncClient` to
avoid leaking resources during the test run.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]