This is an automated email from the ASF dual-hosted git repository.
roryqi 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 c4453641e8 [#12176] feat(mcp-server): support policy-on-tag tools
(#13477)
c4453641e8 is described below
commit c4453641e8ec01c1c78e6c602e0319a054415e05
Author: roryqi <[email protected]>
AuthorDate: Wed Sep 30 09:33:10 2026 +0800
[#12176] feat(mcp-server): support policy-on-tag tools (#13477)
### What changes were proposed in this pull request?
Add MCP tools for managing policy-to-tag associations:
- List policies associated with a tag.
- Associate a policy with a tag and selector.
- Remove a policy association from a tag.
- List tags associated with a policy.
Also add REST client support, URL encoding coverage, unit tests, and MCP
documentation.
### Why are the changes needed?
Policy associations are now managed through tags instead of direct
metadata-object associations. The MCP server needs corresponding tools
so agents can manage and inspect policy-to-tag relationships.
Fix: #12176
### Does this PR introduce _any_ user-facing change?
Yes. The MCP server exposes four new tools:
- `list_policies_for_tag`
- `associate_policy_with_tag`
- `disassociate_policy_from_tag`
- `list_tags_for_policy`
### How was this patch tested?
- `./gradlew :mcp-server:testPython`
- 278 tests passed.
- `./gradlew :mcp-server:formatCheckPython`
- `git diff --check`
---
docs/gravitino-mcp-server.md | 135 +++++++++---------
.../plain/plain_rest_client_policy_operation.py | 42 +++++-
mcp-server/mcp_server/client/policy_operation.py | 54 +++++--
mcp-server/mcp_server/tools/policy.py | 157 +++++++++++++++------
mcp-server/tests/unit/client/test_url_encoding.py | 58 ++++++++
mcp-server/tests/unit/tools/mock_operation.py | 17 ++-
mcp-server/tests/unit/tools/test_policy.py | 83 +++++++++--
mcp-server/tests/unit/tools/test_statistic.py | 1 -
8 files changed, 413 insertions(+), 134 deletions(-)
diff --git a/docs/gravitino-mcp-server.md b/docs/gravitino-mcp-server.md
index fa8b539948..fb25fdafbd 100644
--- a/docs/gravitino-mcp-server.md
+++ b/docs/gravitino-mcp-server.md
@@ -69,72 +69,75 @@ You could start Gravitino MCP server by Docker image,
`docker run -p 8000:8000 -
Gravitino MCP server supports the following tools, and you could export tool
by tag.
-| Tool name | Description
| Tag |
-|-------------------------------------|--------------------------------------------------------------------------------|--------------|
-| `list_metalakes` | Retrieve the metalakes the caller can
access. | `metalake` |
-| `get_list_of_catalogs` | Retrieve a list of all catalogs in the
system. | `catalog` |
-| `create_catalog` | Create a new catalog.
| `catalog` |
-| `alter_catalog` | Alter an existing catalog.
| `catalog` |
-| `drop_catalog` | Drop a catalog.
| `catalog` |
-| `set_catalog_in_use` | Enable or disable a catalog.
| `catalog` |
-| `get_list_of_schemas` | Retrieve a list of schemas belonging
to a specific catalog. | `schema` |
-| `create_schema` | Create a new schema.
| `schema` |
-| `alter_schema` | Alter an existing schema.
| `schema` |
-| `drop_schema` | Drop a schema.
| `schema` |
-| `get_list_of_tables` | Retrieve a list of tables within a
specific catalog and schema. | `table` |
-| `get_table_metadata_details` | Retrieve comprehensive metadata
details for a specific table. | `table` |
-| `create_table` | Create a new table.
| `table` |
-| `alter_table` | Alter an existing table.
| `table` |
-| `drop_table` | Drop a table.
| `table` |
-| `list_of_models` | Retrieve a list of models within a
specific catalog and schema. | `model` |
-| `load_model` | Retrieve comprehensive metadata
details for a specific model. | `model` |
-| `list_model_versions` | Retrieve a list of versions for a
specific model. | `model` |
-| `load_model_version` | Retrieve comprehensive metadata
details for a specific model version. | `model` |
-| `load_model_version_by_alias` | Retrieve comprehensive metadata
details for a specific model version by alias. | `model` |
-| `register_model` | Register a new model.
| `model` |
-| `delete_model` | Delete a model.
| `model` |
-| `link_model_version` | Link a new version to a model.
| `model` |
-| `delete_model_version` | Delete a model version.
| `model` |
-| `delete_model_version_by_alias` | Delete a model version by one of its
aliases. | `model` |
-| `alter_model` | Alter an existing model.
| `model` |
-| `alter_model_version` | Alter a model version.
| `model` |
-| `alter_model_version_by_alias` | Alter a model version by one of its
aliases. | `model` |
-| `metadata_type_to_fullname_formats` | Retrieve the metadata type to fullname
formats mapping. | `metadata` |
-| `list_of_topics` | Retrieve a list of topics within a
specific catalog and schema. | `topic` |
-| `load_topic` | Retrieve comprehensive metadata
details for a specific topic. | `topic` |
-| `create_topic` | Create a new topic.
| `topic` |
-| `alter_topic` | Alter an existing topic.
| `topic` |
-| `delete_topic` | Delete a topic.
| `topic` |
-| `list_of_filesets` | Retrieve a list of filesets within a
specific catalog and schema. | `fileset` |
-| `load_fileset` | Retrieve comprehensive metadata
details for a specific fileset. | `fileset` |
-| `list_files_in_fileset` | Retrieve a list of files within a
specific fileset. | `fileset` |
-| `create_fileset` | Create a new fileset.
| `fileset` |
-| `alter_fileset` | Alter an existing fileset.
| `fileset` |
-| `drop_fileset` | Drop a fileset.
| `fileset` |
-| `list_of_jobs` | Retrieve a list of jobs
| `job` |
-| `get_job_by_id` | Retrieve a job by its ID.
| `job` |
-| `list_of_job_templates` | Retrieve a list of job templates.
| `job` |
-| `get_job_template_by_name` | Retrieve a job template by its name.
| `job` |
-| `run_job` | Run a job with the specified
parameters. | `job` |
-| `cancel_job` | Cancel a running job by its ID.
| `job` |
-| `get_tag_by_name` | Retrieve a tag by its name.
| `tag` |
-| `list_of_tags` | Retrieve a list of tags.
| `tag` |
-| `list_tags_for_metadata` | Retrieve a list of tags associated
with a specific metadata item. | `tag` |
-| `list_metadata_by_tag` | Retrieve a list of metadata items
associated with a specific tag. | `tag` |
-| `associate_tag_with_metadata` | Associate tags with a specific
metadata item. | `tag` |
-| `disassociate_tag_from_metadata` | Disassociate tags from a specific
metadata item. | `tag` |
-| `create_tag` | Create a new tag.
| `tag` |
-| `alter_tag` | Alter an existing tag.
| `tag` |
-| `delete_tag` | Delete a tag.
| `tag` |
-| `list_statistics_for_metadata` | Retrieve a list of statistics
associated with a specific metadata item. | `statistics` |
-| `list_statistics_for_partition` | Retrieve a list of statistics
associated with a specific partition. | `statistics` |
-| `get_list_of_policies` | Retrieve a list of policies in the
system. | `policy` |
-| `get_policy_detail_information` | Retrieve detailed information for a
specific policy by policy name. | `policy` |
-| `list_policies_for_metadata` | List all policies derived for a
specific metadata item. | `policy` |
-| `list_of_partitions` | Retrieve partitions for a table. Only
for catalogs with a partition API. | `partition` |
-| `get_partition` | Retrieve a partition's metadata. Only
for catalogs with a partition API. | `partition` |
-| `list_of_views` | Retrieve a list of views for a schema.
Only for catalogs supporting views. | `view` |
-| `load_view` | Retrieve a view's metadata. Only for
catalogs supporting views. | `view` |
+| Tool name | Description
| Tag |
+|-------------------------------------|--------------------------------------------------------------------------------|-----------------|
+| `list_metalakes` | Retrieve the metalakes the caller can
access. | `metalake` |
+| `get_list_of_catalogs` | Retrieve a list of all catalogs in the
system. | `catalog` |
+| `create_catalog` | Create a new catalog.
| `catalog` |
+| `alter_catalog` | Alter an existing catalog.
| `catalog` |
+| `drop_catalog` | Drop a catalog.
| `catalog` |
+| `set_catalog_in_use` | Enable or disable a catalog.
| `catalog` |
+| `get_list_of_schemas` | Retrieve a list of schemas belonging
to a specific catalog. | `schema` |
+| `create_schema` | Create a new schema.
| `schema` |
+| `alter_schema` | Alter an existing schema.
| `schema` |
+| `drop_schema` | Drop a schema.
| `schema` |
+| `get_list_of_tables` | Retrieve a list of tables within a
specific catalog and schema. | `table` |
+| `get_table_metadata_details` | Retrieve comprehensive metadata
details for a specific table. | `table` |
+| `create_table` | Create a new table.
| `table` |
+| `alter_table` | Alter an existing table.
| `table` |
+| `drop_table` | Drop a table.
| `table` |
+| `list_of_models` | Retrieve a list of models within a
specific catalog and schema. | `model` |
+| `load_model` | Retrieve comprehensive metadata
details for a specific model. | `model` |
+| `list_model_versions` | Retrieve a list of versions for a
specific model. | `model` |
+| `load_model_version` | Retrieve comprehensive metadata
details for a specific model version. | `model` |
+| `load_model_version_by_alias` | Retrieve comprehensive metadata
details for a specific model version by alias. | `model` |
+| `register_model` | Register a new model.
| `model` |
+| `delete_model` | Delete a model.
| `model` |
+| `link_model_version` | Link a new version to a model.
| `model` |
+| `delete_model_version` | Delete a model version.
| `model` |
+| `delete_model_version_by_alias` | Delete a model version by one of its
aliases. | `model` |
+| `alter_model` | Alter an existing model.
| `model` |
+| `alter_model_version` | Alter a model version.
| `model` |
+| `alter_model_version_by_alias` | Alter a model version by one of its
aliases. | `model` |
+| `metadata_type_to_fullname_formats` | Retrieve the metadata type to fullname
formats mapping. | `metadata` |
+| `list_of_topics` | Retrieve a list of topics within a
specific catalog and schema. | `topic` |
+| `load_topic` | Retrieve comprehensive metadata
details for a specific topic. | `topic` |
+| `create_topic` | Create a new topic.
| `topic` |
+| `alter_topic` | Alter an existing topic.
| `topic` |
+| `delete_topic` | Delete a topic.
| `topic` |
+| `list_of_filesets` | Retrieve a list of filesets within a
specific catalog and schema. | `fileset` |
+| `load_fileset` | Retrieve comprehensive metadata
details for a specific fileset. | `fileset` |
+| `list_files_in_fileset` | Retrieve a list of files within a
specific fileset. | `fileset` |
+| `create_fileset` | Create a new fileset.
| `fileset` |
+| `alter_fileset` | Alter an existing fileset.
| `fileset` |
+| `drop_fileset` | Drop a fileset.
| `fileset` |
+| `list_of_jobs` | Retrieve a list of jobs
| `job` |
+| `get_job_by_id` | Retrieve a job by its ID.
| `job` |
+| `list_of_job_templates` | Retrieve a list of job templates.
| `job` |
+| `get_job_template_by_name` | Retrieve a job template by its name.
| `job` |
+| `run_job` | Run a job with the specified
parameters. | `job` |
+| `cancel_job` | Cancel a running job by its ID.
| `job` |
+| `get_tag_by_name` | Retrieve a tag by its name.
| `tag` |
+| `list_of_tags` | Retrieve a list of tags.
| `tag` |
+| `list_tags_for_metadata` | Retrieve a list of tags associated
with a specific metadata item. | `tag` |
+| `list_metadata_by_tag` | Retrieve a list of metadata items
associated with a specific tag. | `tag` |
+| `associate_tag_with_metadata` | Associate tags with a specific
metadata item. | `tag` |
+| `disassociate_tag_from_metadata` | Disassociate tags from a specific
metadata item. | `tag` |
+| `create_tag` | Create a new tag.
| `tag` |
+| `alter_tag` | Alter an existing tag.
| `tag` |
+| `delete_tag` | Delete a tag.
| `tag` |
+| `list_statistics_for_metadata` | Retrieve a list of statistics
associated with a specific metadata item. | `statistics` |
+| `list_statistics_for_partition` | Retrieve a list of statistics
associated with a specific partition. | `statistics` |
+| `get_list_of_policies` | Retrieve a list of policies in the
system. | `policy` |
+| `get_policy_detail_information` | Retrieve detailed information for a
specific policy by policy name. | `policy` |
+| `list_policies_for_tag` | List policies directly associated with
a tag, including selectors. | `policy`, `tag` |
+| `associate_policy_with_tag` | Associate a policy with a tag and
selector. | `policy`, `tag` |
+| `disassociate_policy_from_tag` | Remove a direct policy association
from a tag. | `policy`, `tag` |
+| `list_tags_for_policy` | List tags directly associated with a
policy, including selectors. | `policy`, `tag` |
+| `list_of_partitions` | Retrieve partitions for a table. Only
for catalogs with a partition API. | `partition` |
+| `get_partition` | Retrieve a partition's metadata. Only
for catalogs with a partition API. | `partition` |
+| `list_of_views` | Retrieve a list of views for a schema.
Only for catalogs supporting views. | `view` |
+| `load_view` | Retrieve a view's metadata. Only for
catalogs supporting views. | `view` |
## Configuration
diff --git
a/mcp-server/mcp_server/client/plain/plain_rest_client_policy_operation.py
b/mcp-server/mcp_server/client/plain/plain_rest_client_policy_operation.py
index 8d16e27cca..1c39384555 100644
--- a/mcp-server/mcp_server/client/plain/plain_rest_client_policy_operation.py
+++ b/mcp-server/mcp_server/client/plain/plain_rest_client_policy_operation.py
@@ -15,10 +15,13 @@
# specific language governing permissions and limitations
# under the License.
+import json
+
from mcp_server.client import PolicyOperation
from mcp_server.client.plain.utils import (
encode_path_segment,
extract_content_from_response,
+ extract_response,
)
@@ -46,12 +49,41 @@ class PlainRESTClientPolicyOperation(PolicyOperation):
)
return extract_content_from_response(response, "policy", {})
- async def list_policies_for_metadata(
- self, metadata_full_name: str, metadata_type: str
+ async def list_policies_for_tag(self, tag_name: str) -> str:
+ response = await self.rest_client.get(
+ f"/api/metalakes/{encode_path_segment(self.metalake_name)}"
+ f"/tags/{encode_path_segment(tag_name)}/policies?details=true"
+ )
+ return extract_content_from_response(response, "associations", [])
+
+ async def associate_policy_with_tag(
+ self, tag_name: str, policy_name: str, selector: dict
) -> str:
+ response = await self.rest_client.post(
+ f"/api/metalakes/{encode_path_segment(self.metalake_name)}"
+ f"/tags/{encode_path_segment(tag_name)}"
+ f"/policies/{encode_path_segment(policy_name)}",
+ json={"selector": selector},
+ )
+ return extract_response(response)
+
+ async def disassociate_policy_from_tag(
+ self, tag_name: str, policy_name: str
+ ) -> str:
+ response = await self.rest_client.delete(
+ f"/api/metalakes/{encode_path_segment(self.metalake_name)}"
+ f"/tags/{encode_path_segment(tag_name)}"
+ f"/policies/{encode_path_segment(policy_name)}"
+ )
+ if response.status_code == 204:
+ return json.dumps(
+ {"policy": policy_name, "tag": tag_name, "removed": True}
+ )
+ return extract_response(response)
+
+ async def list_tags_for_policy(self, policy_name: str) -> str:
response = await self.rest_client.get(
f"/api/metalakes/{encode_path_segment(self.metalake_name)}"
- f"/objects/{encode_path_segment(metadata_type)}"
-
f"/{encode_path_segment(metadata_full_name)}/policies?details=true",
+ f"/policies/{encode_path_segment(policy_name)}/tags?details=true"
)
- return extract_content_from_response(response, "policies", [])
+ return extract_content_from_response(response, "associations", [])
diff --git a/mcp-server/mcp_server/client/policy_operation.py
b/mcp-server/mcp_server/client/policy_operation.py
index 00ce3b45c6..0f23d1862f 100644
--- a/mcp-server/mcp_server/client/policy_operation.py
+++ b/mcp-server/mcp_server/client/policy_operation.py
@@ -47,20 +47,56 @@ class PolicyOperation(ABC):
pass
@abstractmethod
- async def list_policies_for_metadata(
- self, metadata_full_name: str, metadata_type: str
+ async def list_policies_for_tag(self, tag_name: str) -> str:
+ """List all policies directly associated with a tag.
+
+ Args:
+ tag_name: Name of the tag
+
+ Returns:
+ str: JSON-formatted list of policy-tag associations, including
selectors
+ """
+ pass
+
+ @abstractmethod
+ async def associate_policy_with_tag(
+ self, tag_name: str, policy_name: str, selector: dict
) -> str:
+ """Associate one policy with a tag.
+
+ Args:
+ tag_name: Name of the tag
+ policy_name: Name of the policy
+ selector: Selector controlling which tag assignments match the
policy
+
+ Returns:
+ str: JSON-formatted policy-tag association
"""
- List all policies associated with a specific metadata item.
+ pass
+
+ @abstractmethod
+ async def disassociate_policy_from_tag(
+ self, tag_name: str, policy_name: str
+ ) -> str:
+ """Remove one policy association from a tag.
Args:
- metadata_full_name: Full name of the metadata object to associate
policies with.
- It's typically in the format "catalog.schema.table" or
"catalog.schema" or "catalog"
- or "catalog.schema.fileset". The "model", "topic" are also
supported and the format
- is the same as for "catalog.schema.table".
- metadata_type: Type of the metadata (e.g., "table", "column")
+ tag_name: Name of the tag
+ policy_name: Name of the policy
+
+ Returns:
+ str: JSON-formatted removal confirmation
+ """
+ pass
+
+ @abstractmethod
+ async def list_tags_for_policy(self, policy_name: str) -> str:
+ """List all tags directly associated with a policy.
+
+ Args:
+ policy_name: Name of the policy
Returns:
- str: JSON formatted string containing list of policy metadata
associated with the metadata
+ str: JSON-formatted list of policy-tag associations, including
selectors
"""
pass
diff --git a/mcp-server/mcp_server/tools/policy.py
b/mcp-server/mcp_server/tools/policy.py
index 0ec1672d49..5508e00d05 100644
--- a/mcp-server/mcp_server/tools/policy.py
+++ b/mcp-server/mcp_server/tools/policy.py
@@ -132,60 +132,133 @@ def load_policy_tools(mcp: FastMCP):
client = ctx.request_context.lifespan_context.rest_client()
return await client.as_policy_operation().load_policy(policy_name)
- @mcp.tool(tags={"policy"})
- async def list_policies_for_metadata(
- ctx: Context, metadata_full_name: str, metadata_type: str
+ @mcp.tool(tags={"policy", "tag"})
+ async def list_policies_for_tag(ctx: Context, tag_name: str) -> str:
+ """List the policies directly associated with a tag.
+
+ The result includes each policy and the selector on its association.
+ ALL_VALUES selectors match any assignment of the tag, while TAG_VALUE
+ selectors match only the specified assignment value.
+
+ Args:
+ ctx (Context): The request context containing the REST client.
+ tag_name (str): Name of the tag.
+
+ Returns:
+ str: JSON-formatted policy-tag associations.
+
+ Example Return Value:
+ [
+ {
+ "policy": {
+ "name": "retention_policy",
+ "policyType": "custom",
+ "enabled": true
+ },
+ "selector": {
+ "type": "TAG_VALUE",
+ "value": "finance"
+ }
+ }
+ ]
+ """
+ client = ctx.request_context.lifespan_context.rest_client()
+ return await client.as_policy_operation().list_policies_for_tag(
+ tag_name
+ )
+
+ @mcp.tool(tags={"policy", "tag"})
+ async def associate_policy_with_tag(
+ ctx: Context,
+ tag_name: str,
+ policy_name: str,
+ selector: dict,
) -> str:
+ """Associate one policy with a tag and a selector.
+
+ Use {"type": "ALL_VALUES"} to select the policy whenever the tag is
+ present. Use {"type": "TAG_VALUE", "value": "finance"} to select it
+ only when the tag is assigned with that value.
+
+ Args:
+ ctx (Context): The request context containing the REST client.
+ tag_name (str): Name of the tag.
+ policy_name (str): Name of the policy.
+ selector (dict): Required association selector.
+
+ Returns:
+ str: JSON-formatted policy-tag association.
+
+ Example Return Value:
+ {
+ "code": 0,
+ "policy": "retention_policy",
+ "tag": "data_domain",
+ "selector": {
+ "type": "TAG_VALUE",
+ "value": "finance"
+ }
+ }
"""
- List all policies associated with a specific metadata item.
+ client = ctx.request_context.lifespan_context.rest_client()
+ return await client.as_policy_operation().associate_policy_with_tag(
+ tag_name, policy_name, selector
+ )
+
+ @mcp.tool(tags={"policy", "tag"})
+ async def disassociate_policy_from_tag(
+ ctx: Context, tag_name: str, policy_name: str
+ ) -> str:
+ """Remove one direct policy association from a tag.
Args:
- ctx (Context): The request context object containing lifespan
context
- and connector information.
- metadata_full_name (str): Full name of the metadata item. For
more, please see tool
- `get_metadata_fullname_formats`.
- metadata_type (str): Type of the metadata (e.g., "table",
"column"). For More information, please see
- tool `list_all_metadata_types`.
+ ctx (Context): The request context containing the REST client.
+ tag_name (str): Name of the tag.
+ policy_name (str): Name of the policy.
- Example input:
- metadata_full_name: "catalog.schema.table"
- metadata_type: "table"
+ Returns:
+ str: JSON-formatted removal confirmation.
+
+ Example Return Value:
+ {
+ "policy": "retention_policy",
+ "tag": "data_domain",
+ "removed": true
+ }
+ """
+ client = ctx.request_context.lifespan_context.rest_client()
+ return await client.as_policy_operation().disassociate_policy_from_tag(
+ tag_name, policy_name
+ )
+
+ @mcp.tool(tags={"policy", "tag"})
+ async def list_tags_for_policy(ctx: Context, policy_name: str) -> str:
+ """List the tags directly associated with a policy.
+
+ The result includes each tag and the selector on its association.
+
+ Args:
+ ctx (Context): The request context containing the REST client.
+ policy_name (str): Name of the policy.
Returns:
- str: JSON-formatted string containing the list of policies
associated with the metadata.
+ str: JSON-formatted policy-tag associations.
Example Return Value:
[
- {
- "name": "my_policy1",
- "comment": "This is a test policy",
- "policyType": "custom",
- "enabled": true,
- "content": {
- "customRules": {
- "rule1": 123
- },
- "properties": {
- "key1": "value1"
- },
- "supportedObjectTypes": [
- "fileset",
- "model",
- "topic",
- "schema",
- "table",
- "catalog"
- ]
- },
- "inherited": false,
- "audit": {
- "creator": "anonymous",
- "createTime": "2025-08-18T08:29:30.016501Z"
- }
+ {
+ "tag": {
+ "name": "data_domain",
+ "comment": "Business data domain"
+ },
+ "selector": {
+ "type": "TAG_VALUE",
+ "value": "finance"
}
+ }
]
"""
client = ctx.request_context.lifespan_context.rest_client()
- return await client.as_policy_operation().list_policies_for_metadata(
- metadata_full_name, metadata_type
+ return await client.as_policy_operation().list_tags_for_policy(
+ policy_name
)
diff --git a/mcp-server/tests/unit/client/test_url_encoding.py
b/mcp-server/tests/unit/client/test_url_encoding.py
index 3b84d32544..e4450b9f01 100644
--- a/mcp-server/tests/unit/client/test_url_encoding.py
+++ b/mcp-server/tests/unit/client/test_url_encoding.py
@@ -496,6 +496,64 @@ class TestPolicyOperationUrlEncoding(unittest.TestCase):
self.assertIn(_ENCODED_PATH_TRAVERSAL, url)
self.assertNotIn("../../", url)
+ def test_list_policies_for_tag_encodes_tag_name(self):
+ client = _make_mock_client({"associations": [{"policy": {}}]})
+ op = PlainRESTClientPolicyOperation(METALAKE, client)
+ result = asyncio.run(op.list_policies_for_tag(_PATH_TRAVERSAL))
+ url = _called_url(client.get)
+ self.assertEqual('[{"policy": {}}]', result)
+ self.assertIn(_ENCODED_PATH_TRAVERSAL, url)
+ self.assertIn("details=true", url)
+ self.assertNotIn("../../", url)
+
+ def test_associate_policy_with_tag_encodes_names_and_sends_selector(self):
+ client = _make_mock_client(
+ {
+ "code": 0,
+ "policy": _QUERY_INJECTION,
+ "tag": _PATH_TRAVERSAL,
+ "selector": {"type": "ALL_VALUES"},
+ }
+ )
+ op = PlainRESTClientPolicyOperation(METALAKE, client)
+ result = asyncio.run(
+ op.associate_policy_with_tag(
+ _PATH_TRAVERSAL,
+ _QUERY_INJECTION,
+ {"type": "ALL_VALUES"},
+ )
+ )
+ url = _called_url(client.post)
+ self.assertIn('"selector": {"type": "ALL_VALUES"}', result)
+ self.assertIn(_ENCODED_PATH_TRAVERSAL, url)
+ self.assertIn(_ENCODED_QUERY_INJECTION, url)
+ self.assertEqual(
+ {"selector": {"type": "ALL_VALUES"}},
+ client.post.call_args.kwargs["json"],
+ )
+
+ def test_disassociate_policy_from_tag_handles_no_content_response(self):
+ client = _make_mock_client({})
+ client.delete.return_value.status_code = 204
+ op = PlainRESTClientPolicyOperation(METALAKE, client)
+ result = asyncio.run(
+ op.disassociate_policy_from_tag(_PATH_TRAVERSAL, _QUERY_INJECTION)
+ )
+ url = _called_url(client.delete)
+ self.assertIn('"removed": true', result)
+ self.assertIn(_ENCODED_PATH_TRAVERSAL, url)
+ self.assertIn(_ENCODED_QUERY_INJECTION, url)
+
+ def test_list_tags_for_policy_encodes_policy_name(self):
+ client = _make_mock_client({"associations": [{"tag": {}}]})
+ op = PlainRESTClientPolicyOperation(METALAKE, client)
+ result = asyncio.run(op.list_tags_for_policy(_QUERY_INJECTION))
+ url = _called_url(client.get)
+ self.assertEqual('[{"tag": {}}]', result)
+ self.assertIn(_ENCODED_QUERY_INJECTION, url)
+ self.assertIn("details=true", url)
+ self.assertNotIn("?admin=true", url)
+
class TestStatisticOperationUrlEncoding(unittest.TestCase):
def test_list_of_statistics_encodes_metadata_fullname(self):
diff --git a/mcp-server/tests/unit/tools/mock_operation.py
b/mcp-server/tests/unit/tools/mock_operation.py
index dfe6125525..371e125509 100644
--- a/mcp-server/tests/unit/tools/mock_operation.py
+++ b/mcp-server/tests/unit/tools/mock_operation.py
@@ -225,13 +225,24 @@ class MockFilesetOperation(FilesetOperation):
class MockPolicyOperation(PolicyOperation):
- async def list_policies_for_metadata(
- self, metadata_full_name: str, metadata_type: str
+ async def list_policies_for_tag(self, tag_name: str) -> str:
+ return f"list_policies_for_tag: {tag_name}"
+
+ async def associate_policy_with_tag(
+ self, tag_name: str, policy_name: str, selector: dict
) -> str:
return (
- f"list_policies_for_metadata: {metadata_full_name},
{metadata_type}"
+ f"associate_policy_with_tag: {tag_name}, {policy_name}, {selector}"
)
+ async def disassociate_policy_from_tag(
+ self, tag_name: str, policy_name: str
+ ) -> str:
+ return f"disassociate_policy_from_tag: {tag_name}, {policy_name}"
+
+ async def list_tags_for_policy(self, policy_name: str) -> str:
+ return f"list_tags_for_policy: {policy_name}"
+
async def get_list_of_policies(self) -> str:
return "mock_policies"
diff --git a/mcp-server/tests/unit/tools/test_policy.py
b/mcp-server/tests/unit/tools/test_policy.py
index fc49b5e312..64c21fd9f5 100644
--- a/mcp-server/tests/unit/tools/test_policy.py
+++ b/mcp-server/tests/unit/tools/test_policy.py
@@ -54,28 +54,95 @@ class TestPolicyTool(unittest.TestCase):
asyncio.run(_test_get_policy_detail_information(self.mcp))
- def test_list_policies_for_metadata(self):
- async def _test_list_policies_for_metadata(mcp_server):
- async with Client(mcp_server) as client:
+ def test_list_policies_for_tag(self):
+ async def _test():
+ async with Client(self.mcp) as client:
result = await client.call_tool(
- "list_policies_for_metadata",
+ "list_policies_for_tag", {"tag_name": "data_domain"}
+ )
+ self.assertEqual(
+ "list_policies_for_tag: data_domain",
+ result.content[0].text,
+ )
+
+ asyncio.run(_test())
+
+ def test_associate_policy_with_tag(self):
+ async def _test():
+ async with Client(self.mcp) as client:
+ result = await client.call_tool(
+ "associate_policy_with_tag",
+ {
+ "tag_name": "data_domain",
+ "policy_name": "retention_policy",
+ "selector": {
+ "type": "TAG_VALUE",
+ "value": "finance",
+ },
+ },
+ )
+ self.assertEqual(
+ "associate_policy_with_tag: data_domain, "
+ "retention_policy, {'type': 'TAG_VALUE', "
+ "'value': 'finance'}",
+ result.content[0].text,
+ )
+
+ asyncio.run(_test())
+
+ def test_disassociate_policy_from_tag(self):
+ async def _test():
+ async with Client(self.mcp) as client:
+ result = await client.call_tool(
+ "disassociate_policy_from_tag",
{
- "metadata_full_name": "catalog.db.table",
- "metadata_type": "table",
+ "tag_name": "data_domain",
+ "policy_name": "retention_policy",
},
)
self.assertEqual(
- "list_policies_for_metadata: catalog.db.table, table",
+ "disassociate_policy_from_tag: data_domain, "
+ "retention_policy",
result.content[0].text,
)
- asyncio.run(_test_list_policies_for_metadata(self.mcp))
+ asyncio.run(_test())
+
+ def test_list_tags_for_policy(self):
+ async def _test():
+ async with Client(self.mcp) as client:
+ result = await client.call_tool(
+ "list_tags_for_policy",
+ {"policy_name": "retention_policy"},
+ )
+ self.assertEqual(
+ "list_tags_for_policy: retention_policy",
+ result.content[0].text,
+ )
+
+ asyncio.run(_test())
+
+ def test_policy_tag_tools_are_exposed(self):
+ async def _test():
+ async with Client(self.mcp) as client:
+ names = {tool.name for tool in await client.list_tools()}
+ self.assertTrue(
+ {
+ "list_policies_for_tag",
+ "associate_policy_with_tag",
+ "disassociate_policy_from_tag",
+ "list_tags_for_policy",
+ }.issubset(names)
+ )
+
+ asyncio.run(_test())
def test_removed_object_policy_tools_are_not_exposed(self):
async def _test():
async with Client(self.mcp) as client:
names = {tool.name for tool in await client.list_tools()}
for name in (
+ "list_policies_for_metadata",
"associate_policy_with_metadata",
"disassociate_policy_from_metadata",
"get_policy_for_metadata",
diff --git a/mcp-server/tests/unit/tools/test_statistic.py
b/mcp-server/tests/unit/tools/test_statistic.py
index 4f9a385baa..7768b9f9e2 100644
--- a/mcp-server/tests/unit/tools/test_statistic.py
+++ b/mcp-server/tests/unit/tools/test_statistic.py
@@ -80,7 +80,6 @@ class TestStatisticTool(unittest.TestCase):
"associate_tag_with_metadata",
"disassociate_tag_from_metadata",
"list_tags_for_metadata",
- "list_policies_for_metadata",
):
with self.subTest(tool=name):
schema = tools[name].inputSchema
