This is an automated email from the ASF dual-hosted git repository.
bharos 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 9e790c649d [#12099] feat(mcp-server): add read tools for table
partitions (#12128)
9e790c649d is described below
commit 9e790c649ddc421bf389f2105641c99f27fdc8e9
Author: Bharath Krishna <[email protected]>
AuthorDate: Thu Jul 23 10:26:05 2026 -0700
[#12099] feat(mcp-server): add read tools for table partitions (#12128)
### What changes were proposed in this pull request?
Add two read-only MCP tools for table partitions:
- `list_of_partitions(catalog, schema, table, details=false)` — returns
partition
names by default, or full partition metadata when `details=true`.
- `get_partition(catalog, schema, table, partition)` — returns the
metadata of a
single partition.
They follow the existing metadata-type pattern (interface → plain REST
client →
FastMCP tool → factory wiring → mock + unit tests). All user-supplied
path
segments are URL-encoded, and authorization is enforced by Gravitino
(`LOAD_TABLE`), so no filtering logic is added in the MCP layer.
### Why are the changes needed?
The server already exposes `list_statistics_for_partition`, which
requires
partition names as input, but there was no tool to discover those names.
These
tools close that gap and round out read coverage for partitioned tables.
Fix: #12099
### Does this PR introduce _any_ user-facing change?
Yes — two new MCP tools (`list_of_partitions`, `get_partition`), tagged
`partition`.
### How was this patch tested?
Added unit tests for the tools (`test_partition.py`) and for
URL-encoding /
query-param handling (`test_url_encoding.py`). Full unit suite passes
(163 tests); `black` and `pylint` (10/10) are clean.
---
docs/gravitino-mcp-server.md | 2 +
.../mcp_server/client/gravitino_operation.py | 11 ++
.../mcp_server/client/partition_operation.py | 70 ++++++++++++
.../client/plain/plain_rest_client_operation.py | 9 ++
.../plain/plain_rest_client_partition_operation.py | 65 +++++++++++
mcp-server/mcp_server/tools/__init__.py | 2 +
mcp-server/mcp_server/tools/partition.py | 119 +++++++++++++++++++++
mcp-server/tests/unit/client/test_url_encoding.py | 31 ++++++
mcp-server/tests/unit/tools/mock_operation.py | 30 ++++++
mcp-server/tests/unit/tools/test_partition.py | 92 ++++++++++++++++
10 files changed, 431 insertions(+)
diff --git a/docs/gravitino-mcp-server.md b/docs/gravitino-mcp-server.md
index 7d09a53730..769ffdc91f 100644
--- a/docs/gravitino-mcp-server.md
+++ b/docs/gravitino-mcp-server.md
@@ -124,6 +124,8 @@ Gravitino MCP server supports the following tools, and you
could export tool by
| `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 |
+| `list_of_partitions` | Retrieve partitions for a table. Only
for catalogs with a partition API. | `partition` | 2.0.0 |
+| `get_partition` | Retrieve a partition's metadata. Only
for catalogs with a partition API. | `partition` | 2.0.0 |
## Configuration
diff --git a/mcp-server/mcp_server/client/gravitino_operation.py
b/mcp-server/mcp_server/client/gravitino_operation.py
index a2e7fa152e..c8ff7fa719 100644
--- a/mcp-server/mcp_server/client/gravitino_operation.py
+++ b/mcp-server/mcp_server/client/gravitino_operation.py
@@ -21,6 +21,7 @@ from mcp_server.client.catalog_operation import
CatalogOperation
from mcp_server.client.fileset_operation import FilesetOperation
from mcp_server.client.job_operation import JobOperation
from mcp_server.client.model_operation import ModelOperation
+from mcp_server.client.partition_operation import PartitionOperation
from mcp_server.client.policy_operation import PolicyOperation
from mcp_server.client.schema_operation import SchemaOperation
from mcp_server.client.statistic_operation import StatisticOperation
@@ -132,3 +133,13 @@ class GravitinoOperation(ABC):
StatisticOperation: Interface for performing statistic-level
operations
"""
pass
+
+ @abstractmethod
+ def as_partition_operation(self) -> PartitionOperation:
+ """
+ Access the partition operation interface of this Gravitino operation.
+
+ Returns:
+ PartitionOperation: Interface for performing partition-level
operations
+ """
+ pass
diff --git a/mcp-server/mcp_server/client/partition_operation.py
b/mcp-server/mcp_server/client/partition_operation.py
new file mode 100644
index 0000000000..6831cb655d
--- /dev/null
+++ b/mcp-server/mcp_server/client/partition_operation.py
@@ -0,0 +1,70 @@
+# 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 PartitionOperation(ABC):
+ """
+ Abstract base class for Gravitino partition operations.
+ """
+
+ @abstractmethod
+ async def list_of_partitions(
+ self,
+ catalog_name: str,
+ schema_name: str,
+ table_name: str,
+ details: bool = False,
+ ) -> str:
+ """
+ Retrieve the partitions of a specific table.
+
+ Args:
+ catalog_name: Name of the catalog
+ schema_name: Name of the schema
+ table_name: Name of the table
+ details: When False, return partition names only. When True, return
+ full partition metadata.
+
+ Returns:
+ str: JSON-formatted string containing partition names or partition
+ metadata
+ """
+ pass
+
+ @abstractmethod
+ async def get_partition(
+ self,
+ catalog_name: str,
+ schema_name: str,
+ table_name: str,
+ partition_name: str,
+ ) -> str:
+ """
+ Load detailed information of a specific partition.
+
+ Args:
+ catalog_name: Name of the catalog
+ schema_name: Name of the schema
+ table_name: Name of the table
+ partition_name: Name of the partition
+
+ Returns:
+ str: JSON-formatted string containing full partition metadata
+ """
+ pass
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 a6adc0e099..a9e09bf2de 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
@@ -39,6 +39,9 @@ from mcp_server.client.plain.plain_rest_client_job_operation
import (
from mcp_server.client.plain.plain_rest_client_model_operation import (
PlainRESTClientModelOperation,
)
+from mcp_server.client.plain.plain_rest_client_partition_operation import (
+ PlainRESTClientPartitionOperation,
+)
from mcp_server.client.plain.plain_rest_client_policy_operation import (
PlainRESTClientPolicyOperation,
)
@@ -109,6 +112,9 @@ class PlainRESTClientOperation(GravitinoOperation):
self._statistic_operation = PlainRESTClientStatisticOperation(
metalake_name, _rest_client
)
+ self._partition_operation = PlainRESTClientPartitionOperation(
+ metalake_name, _rest_client
+ )
async def close(self) -> None:
"""Close the shared httpx client and release its connection pool."""
@@ -143,3 +149,6 @@ class PlainRESTClientOperation(GravitinoOperation):
def as_policy_operation(self) -> PolicyOperation:
return self._policy_operation
+
+ def as_partition_operation(self):
+ return self._partition_operation
diff --git
a/mcp-server/mcp_server/client/plain/plain_rest_client_partition_operation.py
b/mcp-server/mcp_server/client/plain/plain_rest_client_partition_operation.py
new file mode 100644
index 0000000000..33898413d4
--- /dev/null
+++
b/mcp-server/mcp_server/client/plain/plain_rest_client_partition_operation.py
@@ -0,0 +1,65 @@
+# 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.partition_operation import PartitionOperation
+from mcp_server.client.plain.utils import (
+ encode_path_segment,
+ extract_content_from_response,
+)
+
+
+class PlainRESTClientPartitionOperation(PartitionOperation):
+ """
+ Implementation of PartitionOperation using a plain REST client.
+ """
+
+ def __init__(self, metalake_name: str, rest_client):
+ self.metalake_name = metalake_name
+ self.rest_client = rest_client
+
+ async def list_of_partitions(
+ self,
+ catalog_name: str,
+ schema_name: str,
+ table_name: str,
+ details: bool = False,
+ ) -> str:
+ response = await self.rest_client.get(
+ f"/api/metalakes/{encode_path_segment(self.metalake_name)}"
+ f"/catalogs/{encode_path_segment(catalog_name)}"
+ f"/schemas/{encode_path_segment(schema_name)}"
+ f"/tables/{encode_path_segment(table_name)}/partitions",
+ params={"details": details},
+ )
+ field = "partitions" if details else "names"
+ return extract_content_from_response(response, field, [])
+
+ async def get_partition(
+ self,
+ catalog_name: str,
+ schema_name: str,
+ table_name: str,
+ partition_name: str,
+ ) -> str:
+ response = await self.rest_client.get(
+ f"/api/metalakes/{encode_path_segment(self.metalake_name)}"
+ f"/catalogs/{encode_path_segment(catalog_name)}"
+ f"/schemas/{encode_path_segment(schema_name)}"
+ f"/tables/{encode_path_segment(table_name)}"
+ f"/partitions/{encode_path_segment(partition_name)}"
+ )
+ return extract_content_from_response(response, "partition", {})
diff --git a/mcp-server/mcp_server/tools/__init__.py
b/mcp-server/mcp_server/tools/__init__.py
index 5f754b495b..c44726b038 100644
--- a/mcp-server/mcp_server/tools/__init__.py
+++ b/mcp-server/mcp_server/tools/__init__.py
@@ -22,6 +22,7 @@ from mcp_server.tools.fileset import load_fileset_tools
from mcp_server.tools.job import load_job_tool
from mcp_server.tools.metadata import load_metadata_tool
from mcp_server.tools.model import load_model_tools
+from mcp_server.tools.partition import load_partition_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
@@ -42,3 +43,4 @@ def load_tools(mcp: FastMCP):
load_metadata_tool(mcp)
load_statistic_tools(mcp)
load_policy_tools(mcp)
+ load_partition_tools(mcp)
diff --git a/mcp-server/mcp_server/tools/partition.py
b/mcp-server/mcp_server/tools/partition.py
new file mode 100644
index 0000000000..dbe906e294
--- /dev/null
+++ b/mcp-server/mcp_server/tools/partition.py
@@ -0,0 +1,119 @@
+# 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_partition_tools(mcp: FastMCP):
+ @mcp.tool(tags={"partition"})
+ async def list_of_partitions(
+ ctx: Context,
+ catalog_name: str,
+ schema_name: str,
+ table_name: str,
+ details: bool = False,
+ ):
+ """
+ Retrieve the partitions of a specific table.
+
+ Only catalogs that support the partition API (for example Hive) expose
+ partitions. Catalogs that use hidden partitioning, such as
+ lakehouse-iceberg, do not support this operation and return an
+ "unsupported operation" error; for those catalogs the partitioning is
+ described by the table's partition spec, available via
+ 'get_table_metadata_details'. Use the returned partition names as input
+ to partition-scoped tools such as 'list_statistics_for_partition' and
+ 'get_partition'.
+
+ Parameters:
+ ctx (Context): The request context object containing lifespan
context
+ and connector information.
+ catalog_name (str): The name of the catalog containing the table.
+ schema_name (str): The name of the schema containing the table.
+ table_name (str): The name of the table to list partitions for.
+ details (bool): When False (default), return partition names only.
+ When True, return the full metadata of every partition.
+
+ Returns:
+ str: When details is False, a JSON string containing an array of
+ partition names:
+ ["dt=2025-01-01", "dt=2025-01-02"]
+ When details is True, a JSON string containing an array of
+ partition objects, each with the following structure:
+ [
+ {
+ "type": "identity",
+ "name": "dt=2025-01-01",
+ "fieldNames": [["dt"]],
+ "values": [
+ {"type": "literal", "dataType": "date", "value":
"2025-01-01"}
+ ],
+ "properties": {}
+ }
+ ]
+
+ type: The partition type (identity, range, or list).
+ name: The unique name of the partition.
+ properties: A dictionary of partition-specific properties.
+ """
+ client = ctx.request_context.lifespan_context.rest_client()
+ return await client.as_partition_operation().list_of_partitions(
+ catalog_name, schema_name, table_name, details
+ )
+
+ @mcp.tool(tags={"partition"})
+ async def get_partition(
+ ctx: Context,
+ catalog_name: str,
+ schema_name: str,
+ table_name: str,
+ partition_name: str,
+ ):
+ """
+ Load detailed information of a specific partition.
+
+ Only catalogs that support the partition API (for example Hive) support
+ this operation; catalogs with hidden partitioning such as
+ lakehouse-iceberg return an "unsupported operation" error. Partition
+ names can be discovered with the 'list_of_partitions' tool.
+
+ Parameters:
+ ctx (Context): The request context object containing lifespan
context
+ and connector information.
+ catalog_name (str): The name of the catalog containing the table.
+ schema_name (str): The name of the schema containing the table.
+ table_name (str): The name of the table containing the partition.
+ partition_name (str): The name of the partition to load.
+
+ Returns:
+ str: A JSON string containing full partition metadata.
+
+ Example Return Value:
+ {
+ "type": "identity",
+ "name": "dt=2025-01-01",
+ "fieldNames": [["dt"]],
+ "values": [
+ {"type": "literal", "dataType": "date", "value": "2025-01-01"}
+ ],
+ "properties": {}
+ }
+ """
+ client = ctx.request_context.lifespan_context.rest_client()
+ return await client.as_partition_operation().get_partition(
+ catalog_name, schema_name, table_name, partition_name
+ )
diff --git a/mcp-server/tests/unit/client/test_url_encoding.py
b/mcp-server/tests/unit/client/test_url_encoding.py
index 70de5dd831..75fc38c008 100644
--- a/mcp-server/tests/unit/client/test_url_encoding.py
+++ b/mcp-server/tests/unit/client/test_url_encoding.py
@@ -34,6 +34,9 @@ from mcp_server.client.plain.plain_rest_client_job_operation
import (
from mcp_server.client.plain.plain_rest_client_model_operation import (
PlainRESTClientModelOperation,
)
+from mcp_server.client.plain.plain_rest_client_partition_operation import (
+ PlainRESTClientPartitionOperation,
+)
from mcp_server.client.plain.plain_rest_client_policy_operation import (
PlainRESTClientPolicyOperation,
)
@@ -517,3 +520,31 @@ class TestStatisticOperationUrlEncoding(unittest.TestCase):
self.assertNotIn("../../", url)
self.assertIn("from", params)
self.assertIn("to", params)
+
+
+class TestPartitionOperationUrlEncoding(unittest.TestCase):
+ def test_list_of_partitions_encodes_table_name(self):
+ client = _make_mock_client({"names": []})
+ op = PlainRESTClientPartitionOperation(METALAKE, client)
+ asyncio.run(op.list_of_partitions("catalog", "schema",
_PATH_TRAVERSAL))
+ url = _called_url(client.get)
+ self.assertIn(_ENCODED_PATH_TRAVERSAL, url)
+ self.assertNotIn("../../", url)
+
+ def test_list_of_partitions_passes_details_as_query_param(self):
+ client = _make_mock_client({"partitions": []})
+ op = PlainRESTClientPartitionOperation(METALAKE, client)
+ asyncio.run(
+ op.list_of_partitions("catalog", "schema", "table", details=True)
+ )
+ self.assertTrue(_called_params(client.get)["details"])
+
+ def test_get_partition_encodes_partition_name(self):
+ client = _make_mock_client({"partition": {}})
+ op = PlainRESTClientPartitionOperation(METALAKE, client)
+ asyncio.run(
+ op.get_partition("catalog", "schema", "table", _QUERY_INJECTION)
+ )
+ url = _called_url(client.get)
+ self.assertIn(_ENCODED_QUERY_INJECTION, url)
+ self.assertNotIn("?admin=true", url)
diff --git a/mcp-server/tests/unit/tools/mock_operation.py
b/mcp-server/tests/unit/tools/mock_operation.py
index e1953bc474..2abcd6fe57 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.partition_operation import PartitionOperation
from mcp_server.client.statistic_operation import StatisticOperation
@@ -64,6 +65,9 @@ class MockOperation(GravitinoOperation):
def as_policy_operation(self) -> PolicyOperation:
return MockPolicyOperation()
+ def as_partition_operation(self) -> PartitionOperation:
+ return MockPartitionOperation()
+
class MockCatalogOperation(CatalogOperation):
async def get_list_of_catalogs(self) -> str:
@@ -440,3 +444,29 @@ class MockStatisticOperation(StatisticOperation):
f"mock_statistics_for_partition: {metalake_name}, {metadata_type},
{metadata_fullname},"
f" {from_partition_name}, {to_partition_name}, {from_inclusive},
{to_inclusive}"
)
+
+
+class MockPartitionOperation(PartitionOperation):
+ async def list_of_partitions(
+ self,
+ catalog_name: str,
+ schema_name: str,
+ table_name: str,
+ details: bool = False,
+ ) -> str:
+ return (
+ f"mock_partitions: {catalog_name}, {schema_name}, {table_name}, "
+ f"{details}"
+ )
+
+ async def get_partition(
+ self,
+ catalog_name: str,
+ schema_name: str,
+ table_name: str,
+ partition_name: str,
+ ) -> str:
+ return (
+ f"mock_partition: {catalog_name}, {schema_name}, {table_name}, "
+ f"{partition_name}"
+ )
diff --git a/mcp-server/tests/unit/tools/test_partition.py
b/mcp-server/tests/unit/tools/test_partition.py
new file mode 100644
index 0000000000..080cbf9221
--- /dev/null
+++ b/mcp-server/tests/unit/tools/test_partition.py
@@ -0,0 +1,92 @@
+# 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 TestPartitionTool(unittest.TestCase):
+ def setUp(self):
+ RESTClientFactory.set_rest_client(MockOperation)
+ server = GravitinoMCPServer(Setting("mock_partition"))
+ self.mcp = server.mcp
+
+ def test_list_of_partitions(self):
+ async def _test_list_of_partitions(mcp_server):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "list_of_partitions",
+ {
+ "catalog_name": "mock_catalog",
+ "schema_name": "mock_schema",
+ "table_name": "mock_table",
+ },
+ )
+ self.assertEqual(
+ "mock_partitions: mock_catalog, mock_schema, mock_table, "
+ "False",
+ result.content[0].text,
+ )
+
+ asyncio.run(_test_list_of_partitions(self.mcp))
+
+ def test_list_of_partitions_with_details(self):
+ async def _test_list_of_partitions_with_details(mcp_server):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "list_of_partitions",
+ {
+ "catalog_name": "mock_catalog",
+ "schema_name": "mock_schema",
+ "table_name": "mock_table",
+ "details": True,
+ },
+ )
+ self.assertEqual(
+ "mock_partitions: mock_catalog, mock_schema, mock_table, "
+ "True",
+ result.content[0].text,
+ )
+
+ asyncio.run(_test_list_of_partitions_with_details(self.mcp))
+
+ def test_get_partition(self):
+ async def _test_get_partition(mcp_server):
+ async with Client(mcp_server) as client:
+ result = await client.call_tool(
+ "get_partition",
+ {
+ "catalog_name": "mock_catalog",
+ "schema_name": "mock_schema",
+ "table_name": "mock_table",
+ "partition_name": "mock_partition",
+ },
+ )
+ self.assertEqual(
+ "mock_partition: mock_catalog, mock_schema, mock_table, "
+ "mock_partition",
+ result.content[0].text,
+ )
+
+ asyncio.run(_test_get_partition(self.mcp))
