This is an automated email from the ASF dual-hosted git repository.
fanng 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 168137633a [#8138] feat(mcp-server): Support statistic operation for
MCP server. (#8214)
168137633a is described below
commit 168137633a7cb9681e828055fd4e2a3c15d27fd0
Author: Mini Yu <[email protected]>
AuthorDate: Mon Aug 25 10:06:06 2025 +0800
[#8138] feat(mcp-server): Support statistic operation for MCP server.
(#8214)
### What changes were proposed in this pull request?
Support statistics related operations for MCP server.
### Why are the changes needed?
It's a big feature for MCP server.
Fix: #8138
### Does this PR introduce _any_ user-facing change?
N/A.
### How was this patch tested?
UTs and test locally.
---
docs/gravitino-mcp-server.md | 48 +++++--
mcp-server/mcp_server/client/fileset_operation.py | 2 +-
.../mcp_server/client/gravitino_operation.py | 11 ++
mcp-server/mcp_server/client/job_operation.py | 4 +-
mcp-server/mcp_server/client/model_operation.py | 4 +-
.../plain/plain_rest_client_fileset_operation.py | 2 +-
.../plain/plain_rest_client_job_operation.py | 4 +-
.../plain/plain_rest_client_model_operation.py | 4 +-
.../client/plain/plain_rest_client_operation.py | 9 ++
.../plain/plain_rest_client_statistic_operation.py | 53 ++++++++
.../plain/plain_rest_client_tag_operation.py | 2 +-
.../plain/plain_rest_client_topic_operation.py | 4 +-
.../mcp_server/client/statistic_operation.py | 72 ++++++++++
mcp-server/mcp_server/client/tag_operation.py | 2 +-
mcp-server/mcp_server/client/topic_operation.py | 4 +-
mcp-server/mcp_server/tools/__init__.py | 2 +
mcp-server/mcp_server/tools/fileset.py | 4 +-
mcp-server/mcp_server/tools/job.py | 10 +-
mcp-server/mcp_server/tools/model.py | 4 +-
mcp-server/mcp_server/tools/statistic.py | 147 +++++++++++++++++++++
mcp-server/mcp_server/tools/tag.py | 10 +-
mcp-server/mcp_server/tools/topic.py | 4 +-
mcp-server/tests/unit/tools/mock_operation.py | 45 +++++--
mcp-server/tests/unit/tools/test_fileset.py | 77 +++++++++++
mcp-server/tests/unit/tools/test_job.py | 10 +-
mcp-server/tests/unit/tools/test_model.py | 108 +++++++++++++++
mcp-server/tests/unit/tools/test_statistic.py | 72 ++++++++++
mcp-server/tests/unit/tools/test_tag.py | 100 ++++++++++++++
mcp-server/tests/unit/tools/test_topic.py | 59 +++++++++
29 files changed, 813 insertions(+), 64 deletions(-)
diff --git a/docs/gravitino-mcp-server.md b/docs/gravitino-mcp-server.md
index 46d3b505cf..ca0f021184 100644
--- a/docs/gravitino-mcp-server.md
+++ b/docs/gravitino-mcp-server.md
@@ -55,17 +55,43 @@ Or start a HTTP MCP server by `uv run mcp_server --metalake
test --uri http://12
Gravitino MCP server supports the following tools, and you could export tool
by tag.
-| Tool name | Description
| Tag | Since version |
-|---------------------------------|---------------------------------------------------------------------|-----------|---------------|
-| `get_list_of_catalogs` | Retrieve a list of all catalogs in the
system. | `catalog` | 1.0.0 |
-| `get_list_of_schemas` | Retrieve a list of schemas belonging to a
specific catalog. | `schema` | 1.0.0 |
-| `get_list_of_tables` | Retrieve a list of tables within a
specific catalog and schema. | `table` | 1.0.0 |
-| `get_table_metadata_details` | Retrieve comprehensive metadata details
for a specific table. | `table` | 1.0.0 |
-| `get_list_of_policies` | Retrieve a list of policies in the system.
| `policy` | 1.0.0 |
-| `get_policy_detail_information` | Retrieve detailed information for a
specific policy by policy name. | `policy` | 1.0.0 |
-| `list_policies_for_metadata` | List all policies associated with a
specific metadata item. | `policy` | 1.0.0 |
-| `list_metadata_by_policy` | List all metadata items associated with a
specific policy. | `policy` | 1.0.0 |
-| `get_policy_for_metadata` | Get a policy associated with a specific
metadata item. | `policy` | 1.0.0 |
+| Tool name | Description
| Tag | Since version |
+|-------------------------------------|--------------------------------------------------------------------------------|--------------|---------------|
+| `get_list_of_catalogs` | Retrieve a list of all catalogs in the
system. | `catalog` | 1.0.0 |
+| `get_list_of_schemas` | Retrieve a list of schemas belonging
to a specific catalog. | `schema` | 1.0.0 |
+| `get_list_of_tables` | Retrieve a list of tables within a
specific catalog and schema. | `table` | 1.0.0 |
+| `get_table_metadata_details` | Retrieve comprehensive metadata
details for a specific table. | `table` | 1.0.0 |
+| `list_of_models` | Retrieve a list of models within a
specific catalog and schema. | `model` | 1.0.0 |
+| `load_model` | Retrieve comprehensive metadata
details for a specific model. | `model` | 1.0.0 |
+| `list_model_versions` | Retrieve a list of versions for a
specific model. | `model` | 1.0.0 |
+| `load_model_version` | Retrieve comprehensive metadata
details for a specific model version. | `model` | 1.0.0 |
+| `load_model_version_by_alias` | Retrieve comprehensive metadata
details for a specific model version by alias. | `model` | 1.0.0 |
+| `metadata_type_to_fullname_formats` | Retrieve the metadata type to fullname
formats mapping. | `metadata` | 1.0.0 |
+| `list_of_topics` | Retrieve a list of topics within a
specific catalog and schema. | `topic` | 1.0.0 |
+| `load_topic` | Retrieve comprehensive metadata
details for a specific topic. | `topic` | 1.0.0 |
+| `list_of_filesets` | Retrieve a list of filesets within a
specific catalog and schema. | `fileset` | 1.0.0 |
+| `load_fileset` | Retrieve comprehensive metadata
details for a specific fileset. | `fileset` | 1.0.0 |
+| `list_files_in_fileset` | Retrieve a list of files within a
specific fileset. | `fileset` | 1.0.0 |
+| `list_of_jobs` | Retrieve a list of jobs
| `job` | 1.0.0 |
+| `get_job_by_id` | Retrieve a job by its ID.
| `job` | 1.0.0 |
+| `list_of_job_templates` | Retrieve a list of job templates.
| `job` | 1.0.0 |
+| `get_job_template_by_name` | Retrieve a job template by its name.
| `job` | 1.0.0 |
+| `run_job` | Run a job with the specified
parameters. | `job` | 1.0.0
|
+| `cancel_job` | Cancel a running job by its ID.
| `job` | 1.0.0 |
+| `get_tag_by_name` | Retrieve a tag by its name.
| `tag` | 1.0.0 |
+| `list_of_tags` | Retrieve a list of tags.
| `tag` | 1.0.0 |
+| `list_tags_for_metadata` | Retrieve a list of tags associated
with a specific metadata item. | `tag` | 1.0.0 |
+| `list_metadata_by_tag` | Retrieve a list of metadata items
associated with a specific tag. | `tag` | 1.0.0 |
+| `associate_tag_with_metadata` | Associate tags with a specific
metadata item. | `tag` | 1.0.0 |
+| `disassociate_tag_from_metadata` | Disassociate tags from a specific
metadata item. | `tag` | 1.0.0 |
+| `list_statistics_for_metadata` | Retrieve a list of statistics
associated with a specific metadata item. | `statistics` | 1.0.0
|
+| `list_statistics_for_partition` | Retrieve a list of statistics
associated with a specific partition. | `statistics` | 1.0.0
|
+| `get_list_of_policies` | Retrieve a list of policies in the
system. | `policy` | 1.0.0 |
+| `get_policy_detail_information` | Retrieve detailed information for a
specific policy by policy name. | `policy` | 1.0.0 |
+| `list_policies_for_metadata` | List all policies associated with a
specific metadata item. | `policy` | 1.0.0 |
+| `list_metadata_by_policy` | List all metadata items associated
with a specific policy. | `policy` | 1.0.0 |
+| `get_policy_for_metadata` | Get a policy associated with a
specific metadata item. | `policy` | 1.0.0 |
+
### Configuration
diff --git a/mcp-server/mcp_server/client/fileset_operation.py
b/mcp-server/mcp_server/client/fileset_operation.py
index e41b3ec74f..a158c11cda 100644
--- a/mcp-server/mcp_server/client/fileset_operation.py
+++ b/mcp-server/mcp_server/client/fileset_operation.py
@@ -24,7 +24,7 @@ class FilesetOperation(ABC):
"""
@abstractmethod
- async def get_list_of_filesets(
+ async def list_of_filesets(
self, catalog_name: str, schema_name: str
) -> str:
"""
diff --git a/mcp-server/mcp_server/client/gravitino_operation.py
b/mcp-server/mcp_server/client/gravitino_operation.py
index 6b0c4fa12b..a2e7fa152e 100644
--- a/mcp-server/mcp_server/client/gravitino_operation.py
+++ b/mcp-server/mcp_server/client/gravitino_operation.py
@@ -23,6 +23,7 @@ from mcp_server.client.job_operation import JobOperation
from mcp_server.client.model_operation import ModelOperation
from mcp_server.client.policy_operation import PolicyOperation
from mcp_server.client.schema_operation import SchemaOperation
+from mcp_server.client.statistic_operation import StatisticOperation
from mcp_server.client.table_operation import TableOperation
from mcp_server.client.tag_operation import TagOperation
from mcp_server.client.topic_operation import TopicOperation
@@ -121,3 +122,13 @@ class GravitinoOperation(ABC):
JobOperation: Interface for performing job-level operations
"""
pass
+
+ @abstractmethod
+ def as_statistic_operation(self) -> StatisticOperation:
+ """
+ Access the statistic operation interface of this Gravitino operation.
+
+ Returns:
+ StatisticOperation: Interface for performing statistic-level
operations
+ """
+ pass
diff --git a/mcp-server/mcp_server/client/job_operation.py
b/mcp-server/mcp_server/client/job_operation.py
index c83f750bfc..1361ce0aa3 100644
--- a/mcp-server/mcp_server/client/job_operation.py
+++ b/mcp-server/mcp_server/client/job_operation.py
@@ -37,7 +37,7 @@ class JobOperation(ABC):
pass
@abstractmethod
- async def get_list_of_jobs(self, job_template_name: str) -> str:
+ async def list_of_jobs(self, job_template_name: str) -> str:
"""
Retrieve the list of jobs within the metalake
@@ -51,7 +51,7 @@ class JobOperation(ABC):
pass
@abstractmethod
- async def get_list_of_job_templates(self) -> str:
+ async def list_of_job_templates(self) -> str:
"""
Retrieve the list of job templates within the metalake
diff --git a/mcp-server/mcp_server/client/model_operation.py
b/mcp-server/mcp_server/client/model_operation.py
index 9f6e4e9278..c63eda326a 100644
--- a/mcp-server/mcp_server/client/model_operation.py
+++ b/mcp-server/mcp_server/client/model_operation.py
@@ -24,9 +24,7 @@ class ModelOperation(ABC):
"""
@abstractmethod
- async def get_list_of_models(
- self, catalog_name: str, schema_name: str
- ) -> str:
+ async def list_of_models(self, catalog_name: str, schema_name: str) -> str:
"""
Retrieve the list of models within a specified catalog.
diff --git
a/mcp-server/mcp_server/client/plain/plain_rest_client_fileset_operation.py
b/mcp-server/mcp_server/client/plain/plain_rest_client_fileset_operation.py
index 9548d7aab2..d0af2ebc9e 100644
--- a/mcp-server/mcp_server/client/plain/plain_rest_client_fileset_operation.py
+++ b/mcp-server/mcp_server/client/plain/plain_rest_client_fileset_operation.py
@@ -23,7 +23,7 @@ class PlainRESTClientFilesetOperation(FilesetOperation):
self.metalake_name = metalake_name
self.rest_client = rest_client
- async def get_list_of_filesets(
+ async def list_of_filesets(
self, catalog_name: str, schema_name: str
) -> str:
response = await self.rest_client.get(
diff --git
a/mcp-server/mcp_server/client/plain/plain_rest_client_job_operation.py
b/mcp-server/mcp_server/client/plain/plain_rest_client_job_operation.py
index 687c974976..964bab841d 100644
--- a/mcp-server/mcp_server/client/plain/plain_rest_client_job_operation.py
+++ b/mcp-server/mcp_server/client/plain/plain_rest_client_job_operation.py
@@ -33,7 +33,7 @@ class PlainRESTClientJobOperation(JobOperation):
)
return extract_content_from_response(response, "job", {})
- async def get_list_of_jobs(self, job_template_name: str) -> str:
+ async def list_of_jobs(self, job_template_name: str) -> str:
url = f"/api/metalakes/{self.metalake_name}/jobs/runs"
if job_template_name:
url += f"?jobTemplateName={job_template_name}"
@@ -47,7 +47,7 @@ class PlainRESTClientJobOperation(JobOperation):
)
return extract_content_from_response(response, "jobTemplate", {})
- async def get_list_of_job_templates(self) -> str:
+ async def list_of_job_templates(self) -> str:
response = await self.rest_client.get(
f"/api/metalakes/{self.metalake_name}/jobs/templates?details=true"
)
diff --git
a/mcp-server/mcp_server/client/plain/plain_rest_client_model_operation.py
b/mcp-server/mcp_server/client/plain/plain_rest_client_model_operation.py
index 266865973b..3419d4fb52 100644
--- a/mcp-server/mcp_server/client/plain/plain_rest_client_model_operation.py
+++ b/mcp-server/mcp_server/client/plain/plain_rest_client_model_operation.py
@@ -27,9 +27,7 @@ class PlainRESTClientModelOperation(ModelOperation):
self.metalake_name = metalake_name
self.rest_client = rest_client
- async def get_list_of_models(
- self, catalog_name: str, schema_name: str
- ) -> str:
+ async def list_of_models(self, catalog_name: str, schema_name: str) -> str:
response = await self.rest_client.get(
f"/api/metalakes/{self.metalake_name}/catalogs/{catalog_name}/schemas/{schema_name}/models"
)
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 28798d699f..205562b6d8 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
@@ -45,6 +45,9 @@ from
mcp_server.client.plain.plain_rest_client_policy_operation import (
from mcp_server.client.plain.plain_rest_client_schema_operation import (
PlainRESTClientSchemaOperation,
)
+from mcp_server.client.plain.plain_rest_client_statistic_operation import (
+ PlainRESTClientStatisticOperation,
+)
from mcp_server.client.plain.plain_rest_client_table_operation import (
PlainRESTClientTableOperation,
)
@@ -88,6 +91,9 @@ class PlainRESTClientOperation(GravitinoOperation):
self._policy_operation = PlainRESTClientPolicyOperation(
metalake_name, _rest_client
)
+ self._statistic_operation = PlainRESTClientStatisticOperation(
+ metalake_name, _rest_client
+ )
def as_catalog_operation(self) -> CatalogOperation:
return self._catalog_operation
@@ -113,5 +119,8 @@ class PlainRESTClientOperation(GravitinoOperation):
def as_job_operation(self) -> JobOperation:
return self._job_operation
+ def as_statistic_operation(self):
+ return self._statistic_operation
+
def as_policy_operation(self) -> PolicyOperation:
return self._policy_operation
diff --git
a/mcp-server/mcp_server/client/plain/plain_rest_client_statistic_operation.py
b/mcp-server/mcp_server/client/plain/plain_rest_client_statistic_operation.py
new file mode 100644
index 0000000000..4417cd8b01
--- /dev/null
+++
b/mcp-server/mcp_server/client/plain/plain_rest_client_statistic_operation.py
@@ -0,0 +1,53 @@
+# 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.
+
+from mcp_server.client.plain.utils import extract_content_from_response
+from mcp_server.client.statistic_operation import StatisticOperation
+
+
+class PlainRESTClientStatisticOperation(StatisticOperation):
+ def __init__(self, metalake_name: str, rest_client):
+ self.metalake_name = metalake_name
+ self.rest_client = rest_client
+
+ async def list_of_statistics(
+ self, metalake_name: str, metadata_type: str, metadata_fullname: str
+ ) -> str:
+ response = await self.rest_client.get(
+
f"/api/metalakes/{metalake_name}/objects/{metadata_type}/{metadata_fullname}/statistics"
+ )
+ return extract_content_from_response(response, "statistics", [])
+
+ # pylint: disable=R0917
+ async def list_statistic_for_partition(
+ self,
+ metalake_name: str,
+ metadata_type: str,
+ metadata_fullname: str,
+ from_partition_name: str,
+ to_partition_name: str,
+ from_inclusive: bool = True,
+ to_inclusive: bool = False,
+ ) -> str:
+ response = await self.rest_client.get(
+
f"/api/metalakes/{metalake_name}/objects/{metadata_type}/{metadata_fullname}/statistics/"
+
f"partitions?from={from_partition_name}&to={to_partition_name}&fromInclusive={from_inclusive}"
+ f"&toInclusive={to_inclusive}"
+ )
+ return extract_content_from_response(
+ response, "partitionStatistics", []
+ )
diff --git
a/mcp-server/mcp_server/client/plain/plain_rest_client_tag_operation.py
b/mcp-server/mcp_server/client/plain/plain_rest_client_tag_operation.py
index cd3a67bcb0..85bbc13c87 100644
--- a/mcp-server/mcp_server/client/plain/plain_rest_client_tag_operation.py
+++ b/mcp-server/mcp_server/client/plain/plain_rest_client_tag_operation.py
@@ -25,7 +25,7 @@ class PlainRESTClientTagOperation(TagOperation):
self.metalake_name = metalake_name
self.rest_client = rest_client
- async def get_list_of_tags(self) -> str:
+ async def list_of_tags(self) -> str:
response = await self.rest_client.get(
f"/api/metalakes/{self.metalake_name}/tags?details=true"
)
diff --git
a/mcp-server/mcp_server/client/plain/plain_rest_client_topic_operation.py
b/mcp-server/mcp_server/client/plain/plain_rest_client_topic_operation.py
index e9a4b26bc9..af8f105ae8 100644
--- a/mcp-server/mcp_server/client/plain/plain_rest_client_topic_operation.py
+++ b/mcp-server/mcp_server/client/plain/plain_rest_client_topic_operation.py
@@ -27,9 +27,7 @@ class PlainRESTClientTopicOperation(TopicOperation):
self.metalake_name = metalake_name
self.rest_client = rest_client
- async def get_list_of_topics(
- self, catalog_name: str, schema_name: str
- ) -> str:
+ async def list_of_topics(self, catalog_name: str, schema_name: str) -> str:
response = await self.rest_client.get(
f"/api/metalakes/{self.metalake_name}/catalogs/{catalog_name}/schemas/{schema_name}/topics"
)
diff --git a/mcp-server/mcp_server/client/statistic_operation.py
b/mcp-server/mcp_server/client/statistic_operation.py
new file mode 100644
index 0000000000..b9c2f1f99d
--- /dev/null
+++ b/mcp-server/mcp_server/client/statistic_operation.py
@@ -0,0 +1,72 @@
+# 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.
+
+from abc import ABC, abstractmethod
+
+
+class StatisticOperation(ABC):
+ """
+ Abstract base class for Gravitino statistic operations.
+ """
+
+ @abstractmethod
+ async def list_of_statistics(
+ self, metalake_name: str, metadata_type: str, metadata_fullname: str
+ ) -> str:
+ """
+ Retrieve the list of statistics for a specific metadata type and
fullname within a metalake.
+ Args:
+ metalake_name: Name of the metalake
+ metadata_type: Type of metadata (e.g., table, column)
+ metadata_fullname: Full name of the metadata item
+
+ Returns:
+ str: JSON-formatted string containing statistics information
+ """
+ pass
+
+ # pylint: disable=R0917
+ @abstractmethod
+ async def list_statistic_for_partition(
+ self,
+ metalake_name: str,
+ metadata_type: str,
+ metadata_fullname: str,
+ from_partition_name: str,
+ to_partition_name: str,
+ from_inclusive: bool = True,
+ to_inclusive: bool = False,
+ ) -> str:
+ """
+ Retrieve statistics for a specific partition range of a metadata item.
+ Note: This method is currently only supported list statistics for
partitions of tables.
+ So `metadata_type` should always be "table".
+
+ Args:
+ metalake_name: Name of the metalake
+ metadata_type: Type of metadata, should be "table" for partition
statistics
+ metadata_fullname: Full name of the metadata item, the format
should be
+ "{catalog}.{schema}.{table}".
+ from_partition_name: Starting partition name
+ to_partition_name: Ending partition name
+ from_inclusive: Whether to include the starting partition
+ to_inclusive: Whether to include the ending partition
+
+ Returns:
+ str: JSON-formatted string containing statistics for the specified
partitions
+ """
+ pass
diff --git a/mcp-server/mcp_server/client/tag_operation.py
b/mcp-server/mcp_server/client/tag_operation.py
index b00a06b8f4..afb4df7b32 100644
--- a/mcp-server/mcp_server/client/tag_operation.py
+++ b/mcp-server/mcp_server/client/tag_operation.py
@@ -54,7 +54,7 @@ class TagOperation(ABC):
pass
@abstractmethod
- async def get_list_of_tags(self) -> str:
+ async def list_of_tags(self) -> str:
"""
Retrieve the list of tags within the metalake
diff --git a/mcp-server/mcp_server/client/topic_operation.py
b/mcp-server/mcp_server/client/topic_operation.py
index 47a2f4f3b6..6bc0bc94ab 100644
--- a/mcp-server/mcp_server/client/topic_operation.py
+++ b/mcp-server/mcp_server/client/topic_operation.py
@@ -24,9 +24,7 @@ class TopicOperation(ABC):
"""
@abstractmethod
- async def get_list_of_topics(
- self, catalog_name: str, schema_name: str
- ) -> str:
+ async def list_of_topics(self, catalog_name: str, schema_name: str) -> str:
"""
Retrieve the list of topics within a specified catalog.
diff --git a/mcp-server/mcp_server/tools/__init__.py
b/mcp-server/mcp_server/tools/__init__.py
index 69fabbf31c..5f754b495b 100644
--- a/mcp-server/mcp_server/tools/__init__.py
+++ b/mcp-server/mcp_server/tools/__init__.py
@@ -24,6 +24,7 @@ from mcp_server.tools.metadata import load_metadata_tool
from mcp_server.tools.model import load_model_tools
from mcp_server.tools.policy import load_policy_tools
from mcp_server.tools.schema import load_schema_tools
+from mcp_server.tools.statistic import load_statistic_tools
from mcp_server.tools.table import load_table_tools
from mcp_server.tools.tag import load_tag_tool
from mcp_server.tools.topic import load_topic_tools
@@ -39,4 +40,5 @@ def load_tools(mcp: FastMCP):
load_fileset_tools(mcp)
load_tag_tool(mcp)
load_metadata_tool(mcp)
+ load_statistic_tools(mcp)
load_policy_tools(mcp)
diff --git a/mcp-server/mcp_server/tools/fileset.py
b/mcp-server/mcp_server/tools/fileset.py
index 27400de207..c3997ac16d 100644
--- a/mcp-server/mcp_server/tools/fileset.py
+++ b/mcp-server/mcp_server/tools/fileset.py
@@ -20,7 +20,7 @@ from fastmcp import Context, FastMCP
def load_fileset_tools(mcp: FastMCP):
@mcp.tool(tags={"fileset"})
- async def get_list_of_filesets(
+ async def list_of_filesets(
ctx: Context,
catalog_name: str,
schema_name: str,
@@ -50,7 +50,7 @@ def load_fileset_tools(mcp: FastMCP):
]
"""
client = ctx.request_context.lifespan_context.rest_client()
- return await client.as_fileset_operation().get_list_of_filesets(
+ return await client.as_fileset_operation().list_of_filesets(
catalog_name, schema_name
)
diff --git a/mcp-server/mcp_server/tools/job.py
b/mcp-server/mcp_server/tools/job.py
index d1e5e24414..c0d141d2ab 100644
--- a/mcp-server/mcp_server/tools/job.py
+++ b/mcp-server/mcp_server/tools/job.py
@@ -20,7 +20,7 @@ from fastmcp import Context, FastMCP
def load_job_tool(mcp: FastMCP):
@mcp.tool(tags={"job"})
- async def get_list_of_jobs(
+ async def list_of_jobs(
ctx: Context,
job_template_name: str = None,
) -> str:
@@ -55,9 +55,7 @@ def load_job_tool(mcp: FastMCP):
audit: An object containing audit information, including
creator and creation time.
"""
client = ctx.request_context.lifespan_context.rest_client()
- return await client.as_job_operation().get_list_of_jobs(
- job_template_name
- )
+ return await client.as_job_operation().list_of_jobs(job_template_name)
@mcp.tool(tags={"job"})
async def get_job_by_id(
@@ -93,7 +91,7 @@ def load_job_tool(mcp: FastMCP):
return await client.as_job_operation().get_job_by_id(job_id)
@mcp.tool(tags={"job"})
- async def get_list_of_job_templates(
+ async def list_of_job_templates(
ctx: Context,
) -> str:
"""
@@ -132,7 +130,7 @@ def load_job_tool(mcp: FastMCP):
scripts: A list of scripts associated with the job template
and can be called by the executable.
"""
client = ctx.request_context.lifespan_context.rest_client()
- return await client.as_job_operation().get_list_of_job_templates()
+ return await client.as_job_operation().list_of_job_templates()
@mcp.tool(tags={"job"})
async def get_job_template_by_name(
diff --git a/mcp-server/mcp_server/tools/model.py
b/mcp-server/mcp_server/tools/model.py
index 08e56343da..4e5f21fdb5 100644
--- a/mcp-server/mcp_server/tools/model.py
+++ b/mcp-server/mcp_server/tools/model.py
@@ -20,7 +20,7 @@ from fastmcp import Context, FastMCP
def load_model_tools(mcp: FastMCP):
@mcp.tool(tags={"model"})
- async def get_list_of_models(
+ async def list_of_models(
ctx: Context,
catalog_name: str,
schema_name: str,
@@ -62,7 +62,7 @@ def load_model_tools(mcp: FastMCP):
]
"""
client = ctx.request_context.lifespan_context.rest_client()
- return await client.as_model_operation().get_list_of_models(
+ return await client.as_model_operation().list_of_models(
catalog_name, schema_name
)
diff --git a/mcp-server/mcp_server/tools/statistic.py
b/mcp-server/mcp_server/tools/statistic.py
new file mode 100644
index 0000000000..76dd1d9a6f
--- /dev/null
+++ b/mcp-server/mcp_server/tools/statistic.py
@@ -0,0 +1,147 @@
+# 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.
+
+from fastmcp import Context, FastMCP
+
+
+def load_statistic_tools(mcp: FastMCP):
+ @mcp.tool(tags={"statistic"})
+ async def list_statistics_for_metadata(
+ ctx: Context,
+ metalake_name: str,
+ metadata_type: str,
+ metadata_fullname: str,
+ ) -> str:
+ """
+ Retrieve a list of statistics for a specific metadata object.
Currently,
+ this tool only supports statistics for tables, so `metadata_type`
+ should always be "table" and metadata_fullname should be in the
format
+ "{catalog}.{schema}.{table}". For more information about the
metadata
+ type and full name formats, please refer to the tool
+ 'metadata_type_to_fullname_formats'.
+
+ Args:
+ ctx (Context): The request context.
+ metalake_name (str): The name of the metalake.
+ metadata_type (str): The type of metadata (e.g., table, column).
For
+ more, please refer to too 'metadata_type_to_fullname_formats'
+ metadata_fullname (str): The full name of the metadata object. For
+ more, please refer to tool 'metadata_type_to_fullname_formats'.
+
+
+ Returns:
+ str: A JSON string containing the list of statistics.
+
+ Example Return Value:
+ [
+ {
+ "name": "custom-key1",
+ "value": "value1",
+ "reserved": false,
+ "modifiable": true,
+ "audit": {
+ "creator": "anonymous",
+ "createTime": "2025-08-20T07:33:41.233089Z",
+ "lastModifier": "anonymous",
+ "lastModifiedTime": "2025-08-20T07:33:41.233104Z"
+ }
+ }
+ ]
+
+ name: The name of the statistic.
+ value: The value of the statistic.
+ reserved: Indicates if the statistic is reserved.
+ modifiable: Indicates if the statistic can be modified.
+ audit: Metadata about the statistic's creation and modification.
+ """
+ client = ctx.request_context.lifespan_context.rest_client()
+ return await client.as_statistic_operation().list_of_statistics(
+ metalake_name, metadata_type, metadata_fullname
+ )
+
+ # pylint: disable=R0917
+ @mcp.tool(tags={"statistic"})
+ async def list_statistics_for_partition(
+ ctx: Context,
+ metalake_name: str,
+ metadata_type: str,
+ metadata_fullname: str,
+ from_partition_name: str,
+ to_partition_name: str,
+ from_inclusive: bool = True,
+ to_inclusive: bool = False,
+ ) -> str:
+ """
+ Retrieve statistics for a specific partition range of a metadata item.
+ Note: This method is currently only supported list statistics for
partitions of tables.
+ So `metadata_type` should always be "table".
+
+ Args:
+ ctx (Context): The request context.
+ metalake_name (str): The name of the metalake.
+ metadata_type (str): The type of metadata, should be "table" for
partition statistics.
+ metadata_fullname (str): The full name of the metadata item, the
format should be
+ "{catalog}.{schema}.{table}".
+ from_partition_name (str): Starting partition name.
+ to_partition_name (str): Ending partition name.
+ from_inclusive (bool): Whether to include the starting partition.
+ to_inclusive (bool): Whether to include the ending partition.
+
+ Returns:
+ str: A JSON string containing statistics for the specified
partitions.
+
+ Example Return Value:
+ [
+ {
+ "partitionName": "partitionName_92f3fa736e47",
+ "statistics": [
+ {
+ "name": "custom-key1",
+ "value": "value1",
+ "reserved": false,
+ "modifiable": true,
+ "audit": {
+ "creator": "anonymous",
+ "createTime": "2025-08-20T07:33:41.233089Z",
+ "lastModifier": "anonymous",
+ "lastModifiedTime": "2025-08-20T07:33:41.233104Z"
+ }
+ }
+ ]
+ }
+ ]
+
+ partitionName: The name of the partition.
+ statistics: A list of statistics for the partition, each statistic
includes:
+ - name: The name of the statistic.
+ - value: The value of the statistic.
+ - reserved: Indicates if the statistic is reserved.
+ - modifiable: Indicates if the statistic can be modified.
+ - audit: Metadata about the statistic's creation and
modification.
+ """
+ client = ctx.request_context.lifespan_context.rest_client()
+ return (
+ await client.as_statistic_operation().list_statistic_for_partition(
+ metalake_name,
+ metadata_type,
+ metadata_fullname,
+ from_partition_name,
+ to_partition_name,
+ from_inclusive,
+ to_inclusive,
+ )
+ )
diff --git a/mcp-server/mcp_server/tools/tag.py
b/mcp-server/mcp_server/tools/tag.py
index d708b27a5b..3a96dfb93a 100644
--- a/mcp-server/mcp_server/tools/tag.py
+++ b/mcp-server/mcp_server/tools/tag.py
@@ -57,14 +57,14 @@ def load_tag_tool(mcp: FastMCP):
)
@mcp.tool(tags={"tag"})
- async def get_tag_by_name(ctx: Context, name: str) -> str:
+ async def get_tag_by_name(ctx: Context, tag_name: str) -> str:
"""
Load a tag by its name.
Args:
ctx (Context): The request context object containing lifespan
context
and connector information.
- name (str): Name of the tag to get
+ tag_name (str): Name of the tag to get
Returns:
str: JSON-formatted string containing the tag information
@@ -85,10 +85,10 @@ def load_tag_tool(mcp: FastMCP):
}
"""
client = ctx.request_context.lifespan_context.rest_client()
- return await client.as_tag_operation().get_tag_by_name(name)
+ return await client.as_tag_operation().get_tag_by_name(tag_name)
@mcp.tool(tags={"tag"})
- async def get_list_of_tags(ctx: Context) -> str:
+ async def list_of_tags(ctx: Context) -> str:
"""
Retrieve the list of tags within the metalake.
@@ -118,7 +118,7 @@ def load_tag_tool(mcp: FastMCP):
]
"""
client = ctx.request_context.lifespan_context.rest_client()
- return await client.as_tag_operation().get_list_of_tags()
+ return await client.as_tag_operation().list_of_tags()
# Disable the alter_tag tool by default as it can be destructive.
@mcp.tool(tags={"tag"}, enabled=False)
diff --git a/mcp-server/mcp_server/tools/topic.py
b/mcp-server/mcp_server/tools/topic.py
index 552b79956a..93d2aa57d5 100644
--- a/mcp-server/mcp_server/tools/topic.py
+++ b/mcp-server/mcp_server/tools/topic.py
@@ -20,7 +20,7 @@ from fastmcp import Context, FastMCP
def load_topic_tools(mcp: FastMCP):
@mcp.tool(tags={"topic"})
- async def get_list_of_topic(
+ async def list_of_topics(
ctx: Context,
catalog_name: str,
schema_name: str,
@@ -80,7 +80,7 @@ def load_topic_tools(mcp: FastMCP):
]
"""
client = ctx.request_context.lifespan_context.rest_client()
- return await client.as_topic_operation().get_list_of_topics(
+ return await client.as_topic_operation().list_of_topics(
catalog_name, schema_name
)
diff --git a/mcp-server/tests/unit/tools/mock_operation.py
b/mcp-server/tests/unit/tools/mock_operation.py
index 0831dfb69b..d76ba94b60 100644
--- a/mcp-server/tests/unit/tools/mock_operation.py
+++ b/mcp-server/tests/unit/tools/mock_operation.py
@@ -27,6 +27,7 @@ from mcp_server.client import (
)
from mcp_server.client.fileset_operation import FilesetOperation
from mcp_server.client.job_operation import JobOperation
+from mcp_server.client.statistic_operation import StatisticOperation
class MockOperation(GravitinoOperation):
@@ -57,6 +58,9 @@ class MockOperation(GravitinoOperation):
def as_job_operation(self) -> JobOperation:
return MockJobOperation()
+ def as_statistic_operation(self) -> StatisticOperation:
+ return MockStatisticOperation()
+
def as_policy_operation(self) -> PolicyOperation:
return MockPolicyOperation()
@@ -84,7 +88,7 @@ class MockTableOperation(TableOperation):
class MockFilesetOperation(FilesetOperation):
- async def get_list_of_filesets(
+ async def list_of_filesets(
self, catalog_name: str, schema_name: str
) -> str:
return "mock_filesets"
@@ -143,9 +147,7 @@ class MockPolicyOperation(PolicyOperation):
class MockModelOperation(ModelOperation):
- async def get_list_of_models(
- self, catalog_name: str, schema_name: str
- ) -> str:
+ async def list_of_models(self, catalog_name: str, schema_name: str) -> str:
return "mock_models"
async def load_model(
@@ -170,9 +172,7 @@ class MockModelOperation(ModelOperation):
class MockTopicOperation(TopicOperation):
- async def get_list_of_topics(
- self, catalog_name: str, schema_name: str
- ) -> str:
+ async def list_of_topics(self, catalog_name: str, schema_name: str) -> str:
return "mock_topics"
async def load_topic(
@@ -182,7 +182,7 @@ class MockTopicOperation(TopicOperation):
class MockTagOperation(TagOperation):
- async def get_list_of_tags(self) -> str:
+ async def list_of_tags(self) -> str:
return "mock_tags"
async def create_tag(
@@ -204,7 +204,7 @@ class MockTagOperation(TagOperation):
metadata_full_name: str,
metadata_type: str,
tags_to_associate: list,
- tags_to_disassociate,
+ tags_to_disassociate: list,
) -> str:
return f"mock_associated_tags: {tags_to_associate} with metadata
{metadata_full_name} of type {metadata_type}"
@@ -218,13 +218,13 @@ class MockTagOperation(TagOperation):
class MockJobOperation(JobOperation):
- async def get_list_of_jobs(self, job_template_name: str = "") -> str:
+ async def list_of_jobs(self, job_template_name: str = "") -> str:
return "mock_jobs"
async def get_job_by_id(self, job_id: str) -> str:
return f"mock_job: {job_id}"
- async def get_list_of_job_templates(self) -> str:
+ async def list_of_job_templates(self) -> str:
return "mock_job_templates"
async def get_job_template_by_name(self, name: str) -> str:
@@ -235,3 +235,26 @@ class MockJobOperation(JobOperation):
async def cancel_job(self, job_id: str) -> str:
return f"mock_job_cancelled: {job_id}"
+
+
+class MockStatisticOperation(StatisticOperation):
+ async def list_of_statistics(
+ self, metalake_name: str, metadata_type: str, metadata_fullname: str
+ ) -> str:
+ return f"mock_statistics: {metalake_name}, {metadata_type},
{metadata_fullname}"
+
+ # pylint: disable=R0917
+ async def list_statistic_for_partition(
+ self,
+ metalake_name: str,
+ metadata_type: str,
+ metadata_fullname: str,
+ from_partition_name: str,
+ to_partition_name: str,
+ from_inclusive: bool = True,
+ to_inclusive: bool = False,
+ ) -> str:
+ return (
+ f"mock_statistics_for_partition: {metalake_name}, {metadata_type},
{metadata_fullname},"
+ f" {from_partition_name}, {to_partition_name}, {from_inclusive},
{to_inclusive}"
+ )
diff --git a/mcp-server/tests/unit/tools/test_fileset.py
b/mcp-server/tests/unit/tools/test_fileset.py
new file mode 100644
index 0000000000..a0c0371491
--- /dev/null
+++ b/mcp-server/tests/unit/tools/test_fileset.py
@@ -0,0 +1,77 @@
+# 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 asyncio
+import unittest
+
+from fastmcp import Client
+
+from mcp_server.client.factory import RESTClientFactory
+from mcp_server.core import Setting
+from mcp_server.server import GravitinoMCPServer
+from tests.unit.tools import MockOperation
+
+
+class TestFilesetTool(unittest.TestCase):
+ def setUp(self):
+ RESTClientFactory.set_rest_client(MockOperation)
+ server = GravitinoMCPServer(Setting("mock_metalake"))
+ self.mcp = server.mcp
+
+ def test_list_filesets(self):
+ async def _test_list_filesets(mcp_server):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "list_of_filesets",
+ {"catalog_name": "mock", "schema_name": "mock"},
+ )
+ self.assertEqual("mock_filesets", result.content[0].text)
+
+ asyncio.run(_test_list_filesets(self.mcp))
+
+ def test_load_fileset(self):
+ async def _test_load_fileset(mcp_server):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "load_fileset",
+ {
+ "catalog_name": "mock",
+ "schema_name": "mock",
+ "fileset_name": "mock",
+ },
+ )
+ self.assertEqual("mock_fileset", result.content[0].text)
+
+ asyncio.run(_test_load_fileset(self.mcp))
+
+ def test_list_files_in_fileset(self):
+ async def _test_list_files_in_fileset(mcp_server):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "list_files_in_fileset",
+ {
+ "catalog_name": "mock",
+ "schema_name": "mock",
+ "fileset_name": "mock",
+ "location_name": "mock_location",
+ },
+ )
+ self.assertEqual(
+ "mock_files_in_fileset", result.content[0].text
+ )
+
+ asyncio.run(_test_list_files_in_fileset(self.mcp))
diff --git a/mcp-server/tests/unit/tools/test_job.py
b/mcp-server/tests/unit/tools/test_job.py
index 22132b7f29..41440e71b9 100644
--- a/mcp-server/tests/unit/tools/test_job.py
+++ b/mcp-server/tests/unit/tools/test_job.py
@@ -35,7 +35,7 @@ class TestJobTool(unittest.TestCase):
def test_list_job_templates(self):
async def _test_list_job_templates(mcp_server):
async with Client(mcp_server) as client:
- result = await client.call_tool("get_list_of_job_templates")
+ result = await client.call_tool("list_of_job_templates")
self.assertEqual("mock_job_templates", result.content[0].text)
asyncio.run(_test_list_job_templates(self.mcp))
@@ -64,13 +64,13 @@ class TestJobTool(unittest.TestCase):
asyncio.run(_test_get_job_by_id(self.mcp))
- def test_get_list_of_jobs(self):
- async def _test_get_list_of_jobs(mcp_server):
+ def test_list_of_jobs(self):
+ async def _test_list_of_jobs(mcp_server):
async with Client(mcp_server) as client:
- result = await client.call_tool("get_list_of_jobs")
+ result = await client.call_tool("list_of_jobs")
self.assertEqual("mock_jobs", result.content[0].text)
- asyncio.run(_test_get_list_of_jobs(self.mcp))
+ asyncio.run(_test_list_of_jobs(self.mcp))
def test_run_job(self):
async def _test_run_job(mcp_server):
diff --git a/mcp-server/tests/unit/tools/test_model.py
b/mcp-server/tests/unit/tools/test_model.py
new file mode 100644
index 0000000000..ff27e64d3b
--- /dev/null
+++ b/mcp-server/tests/unit/tools/test_model.py
@@ -0,0 +1,108 @@
+# 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 asyncio
+import unittest
+
+from fastmcp import Client
+
+from mcp_server.client.factory import RESTClientFactory
+from mcp_server.core import Setting
+from mcp_server.server import GravitinoMCPServer
+from tests.unit.tools import MockOperation
+
+
+class TestModelTool(unittest.TestCase):
+ def setUp(self):
+ RESTClientFactory.set_rest_client(MockOperation)
+ server = GravitinoMCPServer(Setting("mock_metalake"))
+ self.mcp = server.mcp
+
+ def test_list_models(self):
+ async def _test_list_models(mcp_server):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "list_of_models",
+ {"catalog_name": "mock", "schema_name": "mock"},
+ )
+ self.assertEqual("mock_models", result.content[0].text)
+
+ asyncio.run(_test_list_models(self.mcp))
+
+ def test_load_model(self):
+ async def _test_load_model(mcp_server):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "load_model",
+ {
+ "catalog_name": "mock",
+ "schema_name": "mock",
+ "model_name": "mock",
+ },
+ )
+ self.assertEqual("mock_model", result.content[0].text)
+
+ asyncio.run(_test_load_model(self.mcp))
+
+ def test_list_model_versions(self):
+ async def _test_list_model_versions(mcp_server):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "list_model_versions",
+ {
+ "catalog_name": "mock",
+ "schema_name": "mock",
+ "model_name": "mock",
+ },
+ )
+ self.assertEqual("mock_model_versions", result.content[0].text)
+
+ asyncio.run(_test_list_model_versions(self.mcp))
+
+ def test_load_model_version(self):
+ async def _test_load_model_version(mcp_server):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "load_model_version",
+ {
+ "catalog_name": "mock",
+ "schema_name": "mock",
+ "model_name": "mock",
+ "version": 1,
+ },
+ )
+ self.assertEqual("mock_model_version", result.content[0].text)
+
+ asyncio.run(_test_load_model_version(self.mcp))
+
+ def test_load_model_version_by_alias(self):
+ async def _test_load_model_version_by_alias(mcp_server):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "load_model_version_by_alias",
+ {
+ "catalog_name": "mock",
+ "schema_name": "mock",
+ "model_name": "mock",
+ "alias": "latest",
+ },
+ )
+ self.assertEqual(
+ "mock_model_version_by_alias", result.content[0].text
+ )
+
+ asyncio.run(_test_load_model_version_by_alias(self.mcp))
diff --git a/mcp-server/tests/unit/tools/test_statistic.py
b/mcp-server/tests/unit/tools/test_statistic.py
new file mode 100644
index 0000000000..1a8527c6bb
--- /dev/null
+++ b/mcp-server/tests/unit/tools/test_statistic.py
@@ -0,0 +1,72 @@
+# 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 asyncio
+import unittest
+
+from fastmcp import Client
+
+from mcp_server.client.factory import RESTClientFactory
+from mcp_server.core import Setting
+from mcp_server.server import GravitinoMCPServer
+from tests.unit.tools import MockOperation
+
+
+class TestStatisticTool(unittest.TestCase):
+ def setUp(self):
+ RESTClientFactory.set_rest_client(MockOperation)
+ server = GravitinoMCPServer(Setting("mock_job"))
+ self.mcp = server.mcp
+
+ def test_list_of_statistics(self):
+ async def _test_list_of_statistics(mcp_server):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "list_statistics_for_metadata",
+ {
+ "metalake_name": "mock_metalake",
+ "metadata_type": "mock_type",
+ "metadata_fullname": "mock_fullname",
+ },
+ )
+ self.assertEqual(
+ "mock_statistics: mock_metalake, mock_type, mock_fullname",
+ result.content[0].text,
+ )
+
+ asyncio.run(_test_list_of_statistics(self.mcp))
+
+ def test_list_statistics_for_partition(self):
+ async def _test_list_statistics_for_partition(mcp_server):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "list_statistics_for_partition",
+ {
+ "metalake_name": "mock_metalake",
+ "metadata_type": "mock_type",
+ "metadata_fullname": "mock_fullname",
+ "from_partition_name": "from_partition",
+ "to_partition_name": "to_partition",
+ },
+ )
+ self.assertEqual(
+ "mock_statistics_for_partition: mock_metalake, mock_type,
mock_fullname, "
+ "from_partition, to_partition, True, False",
+ result.content[0].text,
+ )
+
+ asyncio.run(_test_list_statistics_for_partition(self.mcp))
diff --git a/mcp-server/tests/unit/tools/test_tag.py
b/mcp-server/tests/unit/tools/test_tag.py
new file mode 100644
index 0000000000..ee5011987e
--- /dev/null
+++ b/mcp-server/tests/unit/tools/test_tag.py
@@ -0,0 +1,100 @@
+# 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 asyncio
+import unittest
+
+from fastmcp import Client
+
+from mcp_server.client.factory import RESTClientFactory
+from mcp_server.core import Setting
+from mcp_server.server import GravitinoMCPServer
+from tests.unit.tools import MockOperation
+
+
+class TestTagTool(unittest.TestCase):
+ def setUp(self):
+ RESTClientFactory.set_rest_client(MockOperation)
+ server = GravitinoMCPServer(Setting("mock_metalake"))
+ self.mcp = server.mcp
+
+ def test_list_tags(self):
+ async def _test_list_tags(mcp_server):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool("list_of_tags")
+ self.assertEqual("mock_tags", result.content[0].text)
+
+ asyncio.run(_test_list_tags(self.mcp))
+
+ def test_get_tag_by_name(self):
+ async def _test_get_tag_by_name(mcp_server):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "get_tag_by_name",
+ {
+ "tag_name": "mock",
+ },
+ )
+ self.assertEqual("mock_tag: mock", result.content[0].text)
+
+ asyncio.run(_test_get_tag_by_name(self.mcp))
+
+ def test_associate_tag_with_metadata(self):
+ async def _test_associate_tag_with_metadata(mcp_server):
+
+ tags_to_associate = ["tag1", "tag2"]
+ metadata_full_name = "catalog.schema.table"
+ metadata_type = "table"
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "associate_tag_with_metadata",
+ {
+ "metadata_full_name": metadata_full_name,
+ "metadata_type": metadata_type,
+ "tags_to_associate": tags_to_associate,
+ },
+ )
+ self.assertEqual(
+ f"mock_associated_tags: {tags_to_associate} with metadata"
+ f" {metadata_full_name} of type {metadata_type}",
+ result.content[0].text,
+ )
+
+ asyncio.run(_test_associate_tag_with_metadata(self.mcp))
+
+ def test_disassociate_tag_from_metadata(self):
+ async def _test_disassociate_tag_from_metadata(mcp_server):
+
+ tags_to_disassociate = ["tag1", "tag2"]
+ metadata_full_name = "catalog.schema.table"
+ metadata_type = "table"
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "disassociate_tag_from_metadata",
+ {
+ "metadata_full_name": metadata_full_name,
+ "metadata_type": metadata_type,
+ "tags_to_disassociate": tags_to_disassociate,
+ },
+ )
+ self.assertEqual(
+ f"mock_associated_tags: [] with metadata"
+ f" {metadata_full_name} of type {metadata_type}",
+ result.content[0].text,
+ )
+
+ asyncio.run(_test_disassociate_tag_from_metadata(self.mcp))
diff --git a/mcp-server/tests/unit/tools/test_topic.py
b/mcp-server/tests/unit/tools/test_topic.py
new file mode 100644
index 0000000000..28422d0634
--- /dev/null
+++ b/mcp-server/tests/unit/tools/test_topic.py
@@ -0,0 +1,59 @@
+# 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 asyncio
+import unittest
+
+from fastmcp import Client
+
+from mcp_server.client.factory import RESTClientFactory
+from mcp_server.core import Setting
+from mcp_server.server import GravitinoMCPServer
+from tests.unit.tools.mock_operation import MockOperation
+
+
+class TestTopicTool(unittest.TestCase):
+ def setUp(self):
+ RESTClientFactory.set_rest_client(MockOperation)
+ server = GravitinoMCPServer(Setting("mock_metalake"))
+ self.mcp = server.mcp
+
+ def test_list_topics(self):
+ async def _test_list_topics(mcp_server):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "list_of_topics",
+ {"catalog_name": "mock", "schema_name": "mock"},
+ )
+ self.assertEqual("mock_topics", result.content[0].text)
+
+ asyncio.run(_test_list_topics(self.mcp))
+
+ def test_load_topic(self):
+ async def _test_load_topic(mcp_server):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "load_topic",
+ {
+ "catalog_name": "mock",
+ "schema_name": "mock",
+ "topic_name": "mock",
+ },
+ )
+ self.assertEqual("mock_topic", result.content[0].text)
+
+ asyncio.run(_test_load_topic(self.mcp))
