mcp-server: route agent discovery through Gravitino metadata

mcp-server/mcp_server/server.py

@@ -44,6 +44,28 @@
 from mcp_server.core.setting import Setting
 from mcp_server.tools import load_tools
 
+GRAVITINO_MCP_INSTRUCTIONS = """\
+The Gravitino MCP server is the authoritative source for metadata and \
+governance across the metalake. It answers what data exists, where it lives, \
+how it is structured, and how it is governed, spanning catalogs, schemas, \
+tables, columns, models, filesets, topics, statistics, jobs, tags, and \
+policies.
+
+Use these tools first for any question about discovery or governance. That \
+includes locating a catalog, schema, or table; inspecting a table's columns \
+and properties; and finding which tags or policies apply to an object. Do not \
+route these questions to a SQL query engine. A query engine discovers \
+structure only by enumerating it one level at a time and executes against the \
+live source, while the metadata tools answer the same questions directly from \
+Gravitino in a single call.
+
+The Gravitino tools return metadata only and never return table rows. A SQL \
+query engine is the right tool for reading or computing over data, and only \
+after the target object and its governing catalog have been identified here. \
+Before reading data that may hold personal, financial, or otherwise sensitive \
+information, consult the tag and policy tools to determine the object's \
+classification."""
+
 
 def _get_principal_from_request(fallback_authorization: str = "") -> str:
     """Derive a display principal for audit attribution.
@@ -100,12 +122,14 @@ def _create_gravitino_mcp(setting: Setting) -> FastMCP:
     if setting.tags is not None and len(setting.tags) > 0:
         mcp = FastMCP(
             "Gravitino MCP Server",
+            instructions=GRAVITINO_MCP_INSTRUCTIONS,
             lifespan=_create_lifespan_manager(GravitinoContext(setting)),
             include_tags=setting.tags,
         )
     else:
         mcp = FastMCP(
             "Gravitino MCP Server",
+            instructions=GRAVITINO_MCP_INSTRUCTIONS,
             lifespan=_create_lifespan_manager(GravitinoContext(setting)),
         )
 

mcp-server/mcp_server/tools/metadata.py

@@ -15,9 +15,101 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import json
+
 from fastmcp import Context, FastMCP
 
 
+async def _find_table_metadata(
+    client, name: str, case_sensitive: bool
+) -> dict:
+    """Locate tables by name across relational catalogs and schemas.
+
+    Composes the existing catalog, schema, and table listing operations into
+    a single server-side crawl. A catalog or schema that cannot be listed is
+    recorded under "skipped" so one bad node does not abort the search.
+    """
+    if not name or not name.strip():
+        raise ValueError("name must be a non-empty string")
+    needle = name if case_sensitive else name.casefold()
+
+    def _matches(candidate: str) -> bool:
+        hay = candidate if case_sensitive else candidate.casefold()
+        return needle in hay
+
+    matches = []
+    skipped = []
+    searched_catalogs = 0
+    searched_schemas = 0
+
+    catalogs = json.loads(
+        await client.as_catalog_operation().get_list_of_catalogs()
+    )
+    for catalog in catalogs:
+        if not isinstance(catalog, dict):
+            continue
+        if str(catalog.get("type", "")).casefold() != "relational":
+            continue
+        catalog_name = catalog.get("name", "")
+        searched_catalogs += 1
+        try:
+            schemas = json.loads(
+                await client.as_schema_operation().get_list_of_schemas(
+                    catalog_name
+                )
+            )
+        except Exception as exc:  # pylint: disable=broad-except
+            skipped.append({"location": catalog_name, "reason": str(exc)})
+            continue
+        for schema in schemas:
+            schema_name = (
+                schema.get("name", "")
+                if isinstance(schema, dict)
+                else str(schema)
+            )
+            searched_schemas += 1
+            try:
+                tables = json.loads(
+                    await client.as_table_operation().get_list_of_tables(
+                        catalog_name, schema_name
+                    )
+                )
+            except Exception as exc:  # pylint: disable=broad-except
+                skipped.append(
+                    {
+                        "location": f"{catalog_name}.{schema_name}",
+                        "reason": str(exc),
+                    }
+                )
+                continue
+            for table in tables:
+                table_name = (
+                    table.get("name", "")
+                    if isinstance(table, dict)
+                    else str(table)
+                )
+                if _matches(table_name):
+                    matches.append(
+                        {
+                            "catalog": catalog_name,
+                            "schema": schema_name,
+                            "table": table_name,
+                            "fullName": (
+                                f"{catalog_name}.{schema_name}.{table_name}"
+                            ),
+                            "catalogType": catalog.get("type", ""),
+                        }
+                    )
+    return {
+        "matches": matches,
+        "searched": {
+            "catalogs": searched_catalogs,
+            "schemas": searched_schemas,
+        },
+        "skipped": skipped,
+    }
+
+
 def load_metadata_tool(mcp: FastMCP):
     @mcp.tool(tags={"tag", "policy"})
     async def metadata_type_to_fullname_formats(ctx: Context) -> dict:
@@ -42,3 +134,65 @@ async def metadata_type_to_fullname_formats(ctx: Context) -> dict:
             "fileset": "{catalog}.{schema}.{fileset}",
             "column": "{catalog}.{schema}.{table}.{column}",
         }
+
+    @mcp.tool(tags={"table"})
+    async def find_metadata(
+        ctx: Context, name: str, case_sensitive: bool = False
+    ) -> str:
+        """
+        Locate table metadata by name across every relational catalog and
+        schema in the metalake, in a single call.
+
+        Use this to answer "where does table <name> live" or "which catalogs
+        contain a table called <name>" without knowing the catalog or schema
+        in advance. Prefer it over listing catalogs and schemas by hand, and
+        prefer it over discovering table locations through a SQL query engine.
+        The tool returns metadata only and never returns table rows.
+
+        The lookup crawls the metalake server-side: it lists relational
+        catalogs, then schemas within each, then tables within each schema,
+        and returns every table whose name matches. Matching is a
+        case-insensitive substring match by default. A catalog or schema that
+        cannot be listed, for example a schema reported as managed by multiple
+        catalogs, is skipped and recorded under "skipped" rather than aborting
+        the search.
+
+        Args:
+            ctx (Context): The request context object.
+            name (str): The table name or name fragment to search for.
+            case_sensitive (bool): Match the name by case when True. Defaults
+                to False (case-insensitive).
+
+        Returns:
+            str: A JSON string with the following structure:
+                {
+                    "matches": [
+                        {
+                            "catalog": "catalog_postgres",
+                            "schema": "hr",
+                            "table": "employees",
+                            "fullName": "catalog_postgres.hr.employees",
+                            "catalogType": "relational"
+                        }
+                    ],
+                    "searched": {"catalogs": 3, "schemas": 7},
+                    "skipped": [
+                        {
+                            "location": "catalog_hive.sales",
+                            "reason": "Schema managed by multiple catalogs"
+                        }
+                    ]
+                }
+
+        Special Considerations:
+            - Only relational catalogs are searched. Filesets, topics, and
+              models are out of scope for this tool.
+            - On a large metalake the crawl issues one listing call per catalog
+              and per schema, so it is heavier than a direct load when the
+              catalog and schema are already known. In that case use
+              get_list_of_tables or get_table_metadata_details instead.
+        """
+        client = ctx.request_context.lifespan_context.rest_client()
+        return json.dumps(
+            await _find_table_metadata(client, name, case_sensitive)
+        )

mcp-server/tests/unit/tools/test_metadata.py

@@ -0,0 +1,177 @@
+# 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 json
+import unittest
+
+from mcp_server.tools.metadata import _find_table_metadata
+
+
+class _FakeCatalogOperation:
+    def __init__(self, catalogs):
+        self._catalogs = catalogs
+
+    async def get_list_of_catalogs(self):
+        return json.dumps(self._catalogs)
+
+
+class _FakeSchemaOperation:
+    def __init__(self, schemas_by_catalog):
+        self._schemas_by_catalog = schemas_by_catalog
+
+    async def get_list_of_schemas(self, catalog_name):
+        entry = self._schemas_by_catalog[catalog_name]
+        if isinstance(entry, Exception):
+            raise entry
+        return json.dumps(entry)
+
+
+class _FakeTableOperation:
+    def __init__(self, tables_by_schema):
+        self._tables_by_schema = tables_by_schema
+
+    async def get_list_of_tables(self, catalog_name, schema_name):
+        entry = self._tables_by_schema[(catalog_name, schema_name)]
+        if isinstance(entry, Exception):
+            raise entry
+        return json.dumps(entry)
+
+
+class _FakeClient:
+    """Minimal stand-in exposing only the three operations the crawl uses."""
+
+    def __init__(self, catalogs, schemas_by_catalog, tables_by_schema):
+        self._catalog_operation = _FakeCatalogOperation(catalogs)
+        self._schema_operation = _FakeSchemaOperation(schemas_by_catalog)
+        self._table_operation = _FakeTableOperation(tables_by_schema)
+
+    def as_catalog_operation(self):
+        return self._catalog_operation
+
+    def as_schema_operation(self):
+        return self._schema_operation
+
+    def as_table_operation(self):
+        return self._table_operation
+
+
+def _sample_client():
+    """A metalake with two relational catalogs, one fileset catalog to be
+    ignored, and one schema that fails to list (managed by multiple catalogs).
+    """
+    catalogs = [
+        {"name": "catalog_postgres", "type": "relational"},
+        {"name": "catalog_hive", "type": "relational"},
+        {"name": "catalog_fileset", "type": "fileset"},
+    ]
+    schemas_by_catalog = {
+        "catalog_postgres": [{"name": "hr"}, {"name": "public"}],
+        "catalog_hive": [{"name": "product"}, {"name": "sales"}],
+    }
+    tables_by_schema = {
+        ("catalog_postgres", "hr"): [{"name": "employees"}, {"name": "salaries"}],
+        ("catalog_postgres", "public"): [],
+        ("catalog_hive", "product"): [{"name": "Employees"}],
+        ("catalog_hive", "sales"): Exception(
+            "Schema managed by multiple catalogs"
+        ),
+    }
+    return _FakeClient(catalogs, schemas_by_catalog, tables_by_schema)
+
+
+class TestFindTableMetadata(unittest.TestCase):
+    """Unit tests for the find_metadata crawl helper."""
+
+    def test_case_insensitive_match_across_catalogs(self):
+        """A case-insensitive search returns hits from every relational
+        catalog, including a name that differs only by case."""
+        result = asyncio.run(
+            _find_table_metadata(_sample_client(), "employees", False)
+        )
+        full_names = {match["fullName"] for match in result["matches"]}
+        self.assertEqual(
+            full_names,
+            {
+                "catalog_postgres.hr.employees",
+                "catalog_hive.product.Employees",
+            },
+        )
+
+    def test_fileset_catalog_is_ignored(self):
+        """Only relational catalogs are crawled; the fileset catalog is not
+        searched and does not appear in the counts."""
+        result = asyncio.run(
+            _find_table_metadata(_sample_client(), "employees", False)
+        )
+        self.assertEqual(result["searched"]["catalogs"], 2)
+        self.assertEqual(result["searched"]["schemas"], 4)
+
+    def test_unlistable_schema_is_skipped_not_fatal(self):
+        """A schema that cannot be listed is recorded under 'skipped' rather
+        than aborting the crawl."""
+        result = asyncio.run(
+            _find_table_metadata(_sample_client(), "employees", False)
+        )
+        self.assertEqual(
+            result["skipped"],
+            [
+                {
+                    "location": "catalog_hive.sales",
+                    "reason": "Schema managed by multiple catalogs",
+                }
+            ],
+        )
+
+    def test_case_sensitive_excludes_differing_case(self):
+        """A case-sensitive search drops a name that differs only by case."""
+        result = asyncio.run(
+            _find_table_metadata(_sample_client(), "employees", True)
+        )
+        full_names = {match["fullName"] for match in result["matches"]}
+        self.assertEqual(full_names, {"catalog_postgres.hr.employees"})
+
+    def test_substring_match(self):
+        """Matching is a substring match by default."""
+        result = asyncio.run(
+            _find_table_metadata(_sample_client(), "ployee", False)
+        )
+        full_names = {match["fullName"] for match in result["matches"]}
+        self.assertEqual(
+            full_names,
+            {
+                "catalog_postgres.hr.employees",
+                "catalog_hive.product.Employees",
+            },
+        )
+
+    def test_no_match_returns_empty(self):
+        """A name that matches nothing returns no matches but still reports
+        what was searched."""
+        result = asyncio.run(
+            _find_table_metadata(_sample_client(), "no_such_table", False)
+        )
+        self.assertEqual(result["matches"], [])
+        self.assertEqual(result["searched"]["catalogs"], 2)
+
+    def test_empty_name_is_rejected(self):
+        """A blank name raises rather than matching every table."""
+        with self.assertRaises(ValueError):
+            asyncio.run(_find_table_metadata(_sample_client(), "  ", False))
+
+
+if __name__ == "__main__":
+    unittest.main()