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]

Reply via email to