This is an automated email from the ASF dual-hosted git repository. FreeOnePlus pushed a commit to branch feat/v1-07-catalog-context in repository https://gitbox.apache.org/repos/asf/doris-mcp-server.git
commit dba7dea670f98a77ea54cbf155bdada1112acbef Author: FreeOnePlus <[email protected]> AuthorDate: Fri Jul 31 12:34:44 2026 +0800 feat: complete Catalog domain context and metadata --- README.md | 18 + doris_mcp_server/state_handles.py | 5 +- doris_mcp_server/tools/capability_detector.py | 17 + doris_mcp_server/tools/capability_registry.py | 9 +- doris_mcp_server/tools/catalog_handlers.py | 217 ++++++++ doris_mcp_server/tools/domain_catalog.py | 34 +- doris_mcp_server/tools/domain_dispatcher.py | 409 +++++++++++++-- doris_mcp_server/tools/doris_feature_matrix.py | 2 +- doris_mcp_server/tools/tools_manager.py | 142 ++---- doris_mcp_server/utils/catalog_metadata.py | 661 +++++++++++++++++++++++++ test/protocol/test_state_handles.py | 14 + test/tools/test_capability_detector.py | 24 + test/tools/test_capability_registry.py | 38 ++ test/tools/test_domain_catalog.py | 11 + test/tools/test_domain_dispatcher.py | 397 ++++++++++++++- test/tools/test_doris_feature_matrix.py | 18 + test/tools/test_tools_manager.py | 120 +++++ test/utils/test_catalog_metadata.py | 369 ++++++++++++++ 18 files changed, 2325 insertions(+), 180 deletions(-) diff --git a/README.md b/README.md index 050e98c..b7e9a79 100644 --- a/README.md +++ b/README.md @@ -772,6 +772,24 @@ The Doris MCP Server supports **catalog federation**, enabling interaction with * **Cross-Catalog SQL Queries**: Use `doris_query.execute_query` with three-part table naming. * **Catalog Discovery**: Use `doris_catalog.list_catalogs`. +* **Typed Object Discovery**: `list_tables` distinguishes tables, views, and + materialized views. Its `page` value is an opaque, expiring continuation + handle bound to the exact request and authorization principal. +* **Deterministic Table Context**: `get_table_context` always reads `schema` + first and can add `comments`, `indexes`, and `basic`. Every returned section + contains `status`, `data`, `warnings`, and `source`; an unsupported optional + section produces a partial result instead of discarding the available + schema. +* **Bounded Size Metadata**: `get_table_size` reads table row and byte + statistics and optionally partition statistics. If partition metadata is + unavailable, the table-level result remains available with + `status: "partial"`. + +Catalog execution failures use the normal child error envelope. The +`error.details.reason_code` value distinguishes +`CATALOG_OBJECT_NOT_FOUND`, `CATALOG_PERMISSION_DENIED`, +`CATALOG_SECTION_UNSUPPORTED`, and `CATALOG_BACKEND_UNAVAILABLE` without +returning raw backend error text. #### Three-Part Naming Requirement: diff --git a/doris_mcp_server/state_handles.py b/doris_mcp_server/state_handles.py index 938daae..4001374 100644 --- a/doris_mcp_server/state_handles.py +++ b/doris_mcp_server/state_handles.py @@ -137,13 +137,16 @@ def _encode_base64url(value: bytes) -> str: def _decode_base64url(value: str) -> bytes: try: padding = "=" * (-len(value) % 4) - return base64.b64decode( + decoded = base64.b64decode( (value + padding).encode("ascii"), altchars=b"-_", validate=True, ) except (UnicodeEncodeError, binascii.Error) as exc: raise StateHandleError("invalid") from exc + if _encode_base64url(decoded) != value: + raise StateHandleError("invalid") + return decoded def _validate_claim(name: str, value: str) -> None: diff --git a/doris_mcp_server/tools/capability_detector.py b/doris_mcp_server/tools/capability_detector.py index 007d1b8..8c3dd25 100644 --- a/doris_mcp_server/tools/capability_detector.py +++ b/doris_mcp_server/tools/capability_detector.py @@ -149,6 +149,23 @@ _DOMAIN_PROBES: Mapping[str, tuple[tuple[str, tuple[str, ...]], ...]] = { ), ("table_metadata_readable",), ), + ( + ( + "SELECT COLUMN_NAME, DATA_TYPE " + "FROM information_schema.columns LIMIT 1" + ), + ( + "information_schema.columns", + "table_context_sections_readable", + ), + ), + ( + ( + "SELECT PARTITION_NAME, TABLE_ROWS, DATA_LENGTH " + "FROM information_schema.partitions LIMIT 1" + ), + ("table_partition_statistics_readable",), + ), ), "doris_query": ( ( diff --git a/doris_mcp_server/tools/capability_registry.py b/doris_mcp_server/tools/capability_registry.py index fc6cb75..80104a9 100644 --- a/doris_mcp_server/tools/capability_registry.py +++ b/doris_mcp_server/tools/capability_registry.py @@ -21,6 +21,7 @@ from __future__ import annotations import asyncio import hashlib import json +import re from collections.abc import Callable, Coroutine, Mapping, Sequence from dataclasses import dataclass from datetime import UTC, datetime, timedelta @@ -418,7 +419,7 @@ class CapabilityEvaluator: limitations=(), ) evidence_sources.extend(probe.evidence_sources) - evidence_sources.append(probe.probe_id) + evidence_sources.append(_public_evidence_source_id(probe.probe_id)) if probe.status is CapabilityProbeStatus.MISCONFIGURED: return _CandidateFailure( status=AvailabilityStatus.MISCONFIGURED, @@ -882,6 +883,12 @@ def _normalized_source_ids(values: tuple[str, ...]) -> tuple[str, ...]: return _ordered_unique([value.lower() for value in values]) +def _public_evidence_source_id(value: str) -> str: + """Normalize private probe keys into public Manifest identifiers.""" + normalized = re.sub(r"[^a-z0-9_]+", "_", value.lower()).strip("_") + return normalized or "runtime_probe" + + def _select_failure( failures: list[_CandidateFailure], ) -> _CandidateFailure: diff --git a/doris_mcp_server/tools/catalog_handlers.py b/doris_mcp_server/tools/catalog_handlers.py new file mode 100644 index 0000000..248426c --- /dev/null +++ b/doris_mcp_server/tools/catalog_handlers.py @@ -0,0 +1,217 @@ +# 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. + +"""Catalog handler routes shared by formal and hidden legacy bindings.""" + +from __future__ import annotations + +from datetime import datetime +from typing import Any + +from ..utils.analysis_tools import SQLAnalyzer +from ..utils.catalog_metadata import ( + CatalogMetadataFailure, + DorisCatalogMetadataReader, +) +from ..utils.data_quality_tools import DataQualityTools +from ..utils.db import DorisConnectionManager +from ..utils.schema_extractor import MetadataExtractor + + +class CatalogToolHandlersMixin: + """Keep Catalog-domain routing outside the manager control plane.""" + + metadata_extractor: MetadataExtractor + catalog_metadata: DorisCatalogMetadataReader + sql_analyzer: SQLAnalyzer + data_quality_tools: DataQualityTools + + def _initialize_catalog_handlers( + self, + connection_manager: DorisConnectionManager, + ) -> None: + self.catalog_metadata = DorisCatalogMetadataReader(connection_manager) + + @staticmethod + def _required_string(arguments: dict[str, Any], name: str) -> str: + raise NotImplementedError + + async def _get_table_schema_tool( + self, + arguments: dict[str, Any], + ) -> dict[str, Any]: + table_name = self._required_string(arguments, "table_name") + db_name = arguments.get("db_name") + catalog_name = arguments.get("catalog_name") + if arguments.get("_formal_catalog"): + return await self.catalog_metadata.get_table_schema( + table_name=table_name, + database_name=self._required_string(arguments, "db_name"), + catalog_name=catalog_name, + ) + return await self.metadata_extractor.get_table_schema_for_mcp( + table_name, + db_name, + catalog_name, + ) + + async def _get_db_table_list_tool( + self, + arguments: dict[str, Any], + ) -> dict[str, Any]: + db_name = arguments.get("db_name") + catalog_name = arguments.get("catalog_name") + if arguments.get("_formal_catalog"): + return await self.catalog_metadata.list_tables( + database_name=self._required_string(arguments, "db_name"), + catalog_name=catalog_name, + ) + return await self.metadata_extractor.get_db_table_list_for_mcp( + db_name, + catalog_name, + ) + + async def _get_db_list_tool( + self, + arguments: dict[str, Any], + ) -> dict[str, Any]: + catalog_name = arguments.get("catalog_name") + if arguments.get("_formal_catalog"): + return await self.catalog_metadata.list_databases( + catalog_name=catalog_name, + ) + return await self.metadata_extractor.get_db_list_for_mcp(catalog_name) + + async def _get_table_comment_tool( + self, + arguments: dict[str, Any], + ) -> dict[str, Any]: + table_name = self._required_string(arguments, "table_name") + db_name = arguments.get("db_name") + catalog_name = arguments.get("catalog_name") + if arguments.get("_formal_catalog"): + return await self.catalog_metadata.get_table_comment( + table_name=table_name, + database_name=self._required_string(arguments, "db_name"), + catalog_name=catalog_name, + ) + return await self.metadata_extractor.get_table_comment_for_mcp( + table_name, + db_name, + catalog_name, + ) + + async def _get_table_column_comments_tool( + self, + arguments: dict[str, Any], + ) -> dict[str, Any]: + table_name = self._required_string(arguments, "table_name") + db_name = arguments.get("db_name") + catalog_name = arguments.get("catalog_name") + if arguments.get("_formal_catalog"): + return await self.catalog_metadata.get_column_comments( + table_name=table_name, + database_name=self._required_string(arguments, "db_name"), + catalog_name=catalog_name, + ) + return await self.metadata_extractor.get_table_column_comments_for_mcp( + table_name, + db_name, + catalog_name, + ) + + async def _get_table_indexes_tool( + self, + arguments: dict[str, Any], + ) -> dict[str, Any]: + table_name = self._required_string(arguments, "table_name") + db_name = arguments.get("db_name") + catalog_name = arguments.get("catalog_name") + if arguments.get("_formal_catalog"): + return await self.catalog_metadata.get_table_indexes( + table_name=table_name, + database_name=self._required_string(arguments, "db_name"), + catalog_name=catalog_name, + ) + return await self.metadata_extractor.get_table_indexes_for_mcp( + table_name, + db_name, + catalog_name, + ) + + async def _get_catalog_list_tool( + self, + arguments: dict[str, Any], + ) -> dict[str, Any]: + if arguments.get("_formal_catalog"): + return await self.catalog_metadata.list_catalogs() + return await self.metadata_extractor.get_catalog_list_for_mcp() + + async def _get_table_data_size_tool( + self, + arguments: dict[str, Any], + ) -> dict[str, Any]: + db_name = arguments.get("db_name") + table_name = arguments.get("table_name") + single_replica = arguments.get("single_replica", False) + if arguments.get("_formal_catalog"): + return await self.catalog_metadata.get_table_size( + table_name=self._required_string(arguments, "table_name"), + database_name=self._required_string(arguments, "db_name"), + catalog_name=arguments.get("catalog_name"), + include_partitions=bool( + arguments.get("include_partitions", False) + ), + ) + return await self.sql_analyzer.get_table_data_size( + db_name, + table_name, + single_replica, + ) + + async def _get_table_basic_info_tool( + self, + arguments: dict[str, Any], + ) -> dict[str, Any]: + try: + table_name = self._required_string(arguments, "table_name") + catalog_name = arguments.get("catalog_name") + db_name = arguments.get("db_name") + if arguments.get("_formal_catalog"): + return await self.catalog_metadata.get_table_basic( + table_name=table_name, + database_name=self._required_string( + arguments, + "db_name", + ), + catalog_name=catalog_name, + ) + return await self.data_quality_tools.get_table_basic_info( + table_name=table_name, + catalog_name=catalog_name, + db_name=db_name, + ) + except CatalogMetadataFailure: + raise + except Exception as exc: + return { + "error": str(exc), + "analysis_type": "table_basic_info", + "timestamp": datetime.now().isoformat(), + } + + +__all__ = ["CatalogToolHandlersMixin"] diff --git a/doris_mcp_server/tools/domain_catalog.py b/doris_mcp_server/tools/domain_catalog.py index 3c9f706..59b4993 100644 --- a/doris_mcp_server/tools/domain_catalog.py +++ b/doris_mcp_server/tools/domain_catalog.py @@ -596,14 +596,31 @@ _QUERY_OUTPUT = _result_schema( "additionalProperties": False, } ) +_TABLE_CONTEXT_SECTION_OUTPUT = { + "type": "object", + "properties": { + "status": { + "type": "string", + "enum": ["success", "partial", "unavailable"], + }, + "data": {"type": "object", "additionalProperties": True}, + "warnings": { + "type": "array", + "items": {"type": "string"}, + }, + "source": {"type": "string"}, + }, + "required": ["status", "data", "warnings", "source"], + "additionalProperties": False, +} _TABLE_CONTEXT_OUTPUT = _result_schema( { "type": "object", "properties": { - "schema": {"type": "object", "additionalProperties": True}, - "comments": {"type": "object", "additionalProperties": True}, - "indexes": {"type": "object", "additionalProperties": True}, - "basic": {"type": "object", "additionalProperties": True}, + "schema": _TABLE_CONTEXT_SECTION_OUTPUT, + "comments": _TABLE_CONTEXT_SECTION_OUTPUT, + "indexes": _TABLE_CONTEXT_SECTION_OUTPUT, + "basic": _TABLE_CONTEXT_SECTION_OUTPUT, }, "required": ["schema"], "additionalProperties": False, @@ -829,7 +846,8 @@ DOMAIN_DEFINITIONS = ( "doris_catalog", "list_tables", "List tables", - "List visible tables, views, and materialized views.", + "List visible tables, views, and materialized views with " + "deterministic filtering and opaque continuation cursors.", _input_schema( { "catalog": _string("Catalog name."), @@ -850,7 +868,8 @@ DOMAIN_DEFINITIONS = ( "get_table_context", "Get table context", "Read schema, comments, indexes, and basic table information " - "through a deterministic four-section composition.", + "through a deterministic four-section composition. Each " + "section reports its own status, warnings, and source.", _input_schema( { "catalog": _string("Catalog name."), @@ -871,7 +890,8 @@ DOMAIN_DEFINITIONS = ( "doris_catalog", "get_table_size", "Get table size", - "Read table and optional partition size and row statistics.", + "Read table and optional partition size and row statistics. " + "Unsupported partition metadata returns a partial result.", _input_schema( { "catalog": _string("Catalog name."), diff --git a/doris_mcp_server/tools/domain_dispatcher.py b/doris_mcp_server/tools/domain_dispatcher.py index a037452..7927efd 100644 --- a/doris_mcp_server/tools/domain_dispatcher.py +++ b/doris_mcp_server/tools/domain_dispatcher.py @@ -19,6 +19,8 @@ from __future__ import annotations import fnmatch +import hashlib +import hmac import json import math import secrets @@ -39,6 +41,8 @@ from ..schema_validation import ( ToolOutputValidationError, ToolSchemaGuard, ) +from ..state_handles import StateHandleCodec, StateHandleError +from ..utils.catalog_metadata import CatalogMetadataFailure from ..utils.logger import get_audit_logger, get_logger from . import domain_catalog as domain_catalog_module from .domain_catalog import ( @@ -99,6 +103,10 @@ class ChildArgumentsAdapterError(ValueError): """Raised when formal arguments cannot use the active handler adapter.""" +_COLLECTION_PAGE_SIZE = 100 +_TABLE_CONTEXT_SECTION_ORDER = ("schema", "comments", "indexes", "basic") + + class ChildHandlerOwner(Protocol): """Structural owner of the existing read-only Doris handler methods.""" @@ -264,6 +272,7 @@ class DomainDispatcher: *, catalog: DorisDomainCatalog | None = None, schema_guard: ToolSchemaGuard | None = None, + state_handle_codec: StateHandleCodec | None = None, ) -> None: catalog = catalog or domain_catalog_module.DORIS_DOMAIN_CATALOG catalog.validate_integrity() @@ -271,6 +280,9 @@ class DomainDispatcher: self._manifest_service = manifest_service self._catalog = catalog self._schema_guard = schema_guard or ToolSchemaGuard() + self._state_handle_codec = state_handle_codec or StateHandleCodec( + secrets.token_urlsafe(32) + ) self._flat_children: dict[str, tuple[str, ChildToolDefinition]] = {} self._schemas: dict[str, CompiledToolSchema] = {} self._bindings = _build_bindings(owner, catalog) @@ -525,6 +537,7 @@ class DomainDispatcher: raw, arguments, adapter_warnings, + auth_context, ) else: data = await self._call_table_context( @@ -552,6 +565,33 @@ class DomainDispatcher: "violations": violations, }, ) + except CatalogMetadataFailure as exc: + self._audit(feature_id, arguments, "error", started) + if exc.reason_code == "CATALOG_ARGUMENT_INVALID": + return self._error( + domain.name, + DomainErrorCode.CHILD_ARGUMENTS_INVALID, + str(exc), + child_tool=child.name, + manifest_version=manifest.manifest_version, + details={ + "rediscover": False, + "reason_code": exc.reason_code, + }, + ) + return self._error( + domain.name, + DomainErrorCode.CHILD_EXECUTION_FAILED, + str(exc), + child_tool=child.name, + manifest_version=manifest.manifest_version, + retryable=exc.retryable, + details={ + "rediscover": False, + "reason_code": exc.reason_code, + "status_code": exc.status_code, + }, + ) except ToolOutputValidationError: logger.exception("Formal child output validation failed for %s", feature_id) self._audit( @@ -631,7 +671,7 @@ class DomainDispatcher: if raw.get("success") is False or ( "error" in raw and raw.get("success") is not True ): - raise RuntimeError("child handler reported failure") + raise _failure_from_legacy_response(raw) return raw, warnings async def _call_table_context( @@ -656,6 +696,7 @@ class DomainDispatcher: "catalog_name": arguments.get("catalog"), "db_name": arguments.get("database"), "table_name": arguments.get("table"), + "_formal_catalog": True, } legacy_arguments = { key: value @@ -663,32 +704,71 @@ class DomainDispatcher: if value is not None } - data: dict[str, Any] = {} + section_states: dict[str, dict[str, Any]] = {} evidence: list[dict[str, Any]] = [] warnings: list[str] = [] - comments: dict[str, Any] = {} for migration in migrations: section = cast(str, migration.target_section) if section not in effective_sections: continue handler = getattr(self._owner, migration.legacy_handler_name) - raw = await handler(dict(legacy_arguments)) - success = isinstance(raw, dict) and not ( - raw.get("success") is False - or ("error" in raw and raw.get("success") is not True) + state = section_states.setdefault( + section, + { + "data": {}, + "sources": [], + "warnings": [], + "successes": 0, + "failures": 0, + }, ) + failure: CatalogMetadataFailure | None = None + try: + raw = await handler(dict(legacy_arguments)) + if not isinstance(raw, dict): + raise TypeError("child handler must return an object") + if raw.get("success") is False or ( + "error" in raw and raw.get("success") is not True + ): + failure = _failure_from_legacy_response(raw) + except CatalogMetadataFailure as exc: + raw = {} + failure = exc + except Exception: + raw = {} + failure = CatalogMetadataFailure( + "Doris catalog metadata execution failed.", + reason_code="CATALOG_EXECUTION_FAILED", + status_code=502, + ) + + success = failure is None evidence.append( { "section": section, "handler": migration.legacy_handler_name, "success": success, + "reason_code": ( + None if failure is None else failure.reason_code + ), } ) - if not success: + state["sources"].append(migration.legacy_handler_name) + if failure is not None: if section == "schema": - raise RuntimeError("required schema section failed") - warnings.append(f"Optional {section} section is unavailable.") + raise failure + state["failures"] += 1 + warning = ( + f"Optional {section} section is unavailable " + f"({failure.reason_code})." + ) + state["warnings"].append(warning) + warnings.append(warning) continue + state["successes"] += 1 + handler_warnings = _legacy_warnings(raw) + state["warnings"].extend(handler_warnings) + warnings.extend(handler_warnings) value = _unwrap_legacy_result(raw) if section == "comments": comment_key = ( @@ -697,18 +777,53 @@ class DomainDispatcher: == "get_table_column_comments" else "table" ) - comments[comment_key] = _as_object(value) + cast(dict[str, Any], state["data"])[comment_key] = _as_object( + value + ) else: - data[section] = _as_object(value) - if comments: - data["comments"] = comments - if "schema" not in data: + state["data"] = _as_object(value) + + schema_state = section_states.get("schema") + if not schema_state or not schema_state["successes"]: raise RuntimeError("required schema section was not executed") + data: dict[str, Any] = {} + for section in _TABLE_CONTEXT_SECTION_ORDER: + if section not in effective_sections: + continue + section_state = section_states.get(section) + if section_state is None: + continue + section_warnings = list( + dict.fromkeys(section_state["warnings"]) + ) + if not section_state["successes"]: + status = "unavailable" + elif section_state["failures"] or section_warnings: + status = "partial" + else: + status = "success" + data[section] = { + "status": status, + "data": section_state["data"], + "warnings": section_warnings, + "source": ",".join(section_state["sources"]), + } return { "status": "partial" if warnings else "success", "data": data, - "warnings": warnings, - "metadata": {"sections": sorted(effective_sections)}, + "warnings": list(dict.fromkeys(warnings)), + "metadata": { + "sections": [ + section + for section in _TABLE_CONTEXT_SECTION_ORDER + if section in effective_sections + ], + "requested_sections": [ + section + for section in _TABLE_CONTEXT_SECTION_ORDER + if section in requested_sections + ], + }, "evidence": evidence, } @@ -718,6 +833,7 @@ class DomainDispatcher: raw: dict[str, Any], arguments: dict[str, Any], adapter_warnings: tuple[str, ...], + auth_context: Any | None = None, ) -> dict[str, Any]: output = cast( dict[str, Any], @@ -732,6 +848,7 @@ class DomainDispatcher: data_schema.get("properties", {}), ) warnings = list(adapter_warnings) + warnings.extend(_legacy_warnings(raw)) metadata = _legacy_metadata(raw) if {"items", "next_cursor", "truncated"} <= set(data_properties): @@ -748,19 +865,69 @@ class DomainDispatcher: for item in items if fnmatch.fnmatchcase(str(item.get("name", "")), pattern) ] - if arguments.get("page") is not None: - warnings.append( - "Pagination is not supported by the active handler." - ) - if arguments.get("types") is not None: - warnings.append( - "Object type filtering is not supported by the active handler." + if child.name == "list_catalogs" and not arguments.get( + "include_external", + True, + ): + items = [ + item + for item in items + if item.get("scope") == "internal" + or item.get("name") == "internal" + ] + requested_types = arguments.get("types") + if isinstance(requested_types, list): + allowed_types = { + str(item) for item in requested_types + } + items = [ + item + for item in items + if str(item.get("type", "table")) in allowed_types + ] + items.sort( + key=lambda item: ( + str(item.get("name", "")).casefold(), + str(item.get("type", "")).casefold(), ) + ) + total_items = len(items) + scope = _collection_cursor_scope(child, arguments) + snapshot = _collection_snapshot(items) + offset = _decode_collection_cursor( + arguments.get("page"), + expected_scope=scope, + expected_snapshot=snapshot, + child=child, + auth_context=auth_context, + handle_codec=self._state_handle_codec, + ) + page_items = items[offset : offset + _COLLECTION_PAGE_SIZE] + next_offset = offset + len(page_items) + truncated = next_offset < total_items child_data: dict[str, Any] = { - "items": items, - "next_cursor": None, - "truncated": False, + "items": page_items, + "next_cursor": ( + _encode_collection_cursor( + next_offset, + scope, + snapshot, + child=child, + auth_context=auth_context, + handle_codec=self._state_handle_codec, + ) + if truncated + else None + ), + "truncated": truncated, } + metadata.update( + { + "page_size": _COLLECTION_PAGE_SIZE, + "returned_items": len(page_items), + "total_visible_items": total_items, + } + ) elif {"columns", "rows", "row_count", "truncated"} <= set( data_properties ): @@ -886,8 +1053,9 @@ def _adapt_arguments( elif adapter_name == "adapt:catalog_object_names": for unsupported in ("pattern", "types", "page"): prepared.pop(unsupported, None) + prepared["_formal_catalog"] = True elif adapter_name == "adapt:remove_random_string": - prepared = {} + prepared = {"_formal_catalog": True} elif adapter_name == "adapt:explain_arguments": level = prepared.pop("level", None) if level == "costs": @@ -906,10 +1074,7 @@ def _adapt_arguments( ) elif adapter_name == "adapt:table_size_arguments": prepared["single_replica"] = False - if prepared.pop("include_partitions", False): - warnings.append( - "Partition detail is not supported by the active handler." - ) + prepared["_formal_catalog"] = True elif adapter_name == "adapt:monitoring_arguments": if prepared: warnings.append( @@ -1038,6 +1203,186 @@ def _legacy_metadata(raw: Mapping[str, Any]) -> dict[str, Any]: return {} +def _legacy_warnings(raw: Mapping[str, Any]) -> list[str]: + warnings = raw.get("warnings") + if not isinstance(warnings, list): + return [] + return [str(item) for item in warnings] + + +def _failure_from_legacy_response( + raw: Mapping[str, Any], +) -> CatalogMetadataFailure: + raw_code = str(raw.get("error_code", "")).upper() + raw_status = raw.get("status_code") + status_code = int(raw_status) if isinstance(raw_status, int) else 500 + message = str(raw.get("error", "")).lower() + if ( + "NOT_VISIBLE" in raw_code + or "NOT_FOUND" in raw_code + or status_code == 404 + or any( + marker in message + for marker in ( + "does not exist", + "doesn't exist", + "not found", + "unknown table", + "unknown database", + ) + ) + ): + return CatalogMetadataFailure( + "The requested Doris metadata object was not found.", + reason_code="CATALOG_OBJECT_NOT_FOUND", + status_code=404, + ) + if ( + "PERMISSION" in raw_code + or "DENIED" in raw_code + or status_code == 403 + or any( + marker in message + for marker in ( + "access denied", + "permission denied", + "not authorized", + "privilege", + ) + ) + ): + return CatalogMetadataFailure( + "Doris denied access to the requested catalog metadata.", + reason_code="CATALOG_PERMISSION_DENIED", + status_code=403, + ) + if ( + "UNSUPPORTED" in raw_code + or status_code == 501 + or "not supported" in message + or "unsupported" in message + ): + return CatalogMetadataFailure( + "The requested Doris metadata section is unsupported.", + reason_code="CATALOG_SECTION_UNSUPPORTED", + status_code=501, + ) + if status_code in {502, 503, 504}: + return CatalogMetadataFailure( + "Doris catalog metadata is temporarily unavailable.", + reason_code="CATALOG_BACKEND_UNAVAILABLE", + status_code=status_code, + retryable=True, + ) + return CatalogMetadataFailure( + "Doris catalog metadata execution failed.", + reason_code="CATALOG_EXECUTION_FAILED", + status_code=status_code, + ) + + +def _collection_cursor_scope( + child: ChildToolDefinition, + arguments: Mapping[str, Any], +) -> str: + payload = { + "child": child.name, + "arguments": { + key: value + for key, value in arguments.items() + if key != "page" + }, + } + canonical = json.dumps( + payload, + ensure_ascii=True, + separators=(",", ":"), + sort_keys=True, + ) + return hashlib.sha256(canonical.encode("utf-8")).hexdigest()[:16] + + +def _collection_snapshot(items: list[dict[str, Any]]) -> str: + identifiers = [ + { + "name": str(item.get("name", "")), + "type": str(item.get("type", "")), + } + for item in items + ] + canonical = json.dumps( + identifiers, + ensure_ascii=True, + separators=(",", ":"), + sort_keys=True, + ) + return hashlib.sha256(canonical.encode("utf-8")).hexdigest() + + +def _encode_collection_cursor( + offset: int, + scope: str, + snapshot: str, + *, + child: ChildToolDefinition, + auth_context: Any | None, + handle_codec: StateHandleCodec, +) -> str: + return handle_codec.issue( + kind="child-list-page", + scope=f"child:page:doris_catalog:{child.name}", + resource=f"mcp://doris_catalog/{child.name}/{scope}", + state={ + "offset": offset, + "snapshot": snapshot, + }, + auth_context=auth_context, + ) + + +def _decode_collection_cursor( + value: Any, + *, + expected_scope: str, + expected_snapshot: str, + child: ChildToolDefinition, + auth_context: Any | None, + handle_codec: StateHandleCodec, +) -> int: + if value is None: + return 0 + if not isinstance(value, str) or not value: + raise ChildArgumentsAdapterError("page cursor must be a non-empty string") + try: + claims = handle_codec.resolve( + value, + expected_kind="child-list-page", + expected_scope=f"child:page:doris_catalog:{child.name}", + expected_resource=( + f"mcp://doris_catalog/{child.name}/{expected_scope}" + ), + auth_context=auth_context, + ) + except StateHandleError as exc: + raise ChildArgumentsAdapterError( + f"page cursor is invalid ({exc.reason})" + ) from exc + if set(claims.state) != {"offset", "snapshot"}: + raise ChildArgumentsAdapterError("page cursor is invalid") + offset = claims.state.get("offset") + snapshot = claims.state.get("snapshot") + if ( + not isinstance(offset, int) + or isinstance(offset, bool) + or offset < 0 + or not isinstance(snapshot, str) + ): + raise ChildArgumentsAdapterError("page cursor is invalid") + if not hmac.compare_digest(snapshot, expected_snapshot): + raise ChildArgumentsAdapterError("page cursor is stale") + return offset + + def _as_object(value: Any) -> dict[str, Any]: if isinstance(value, dict): return dict(value) diff --git a/doris_mcp_server/tools/doris_feature_matrix.py b/doris_mcp_server/tools/doris_feature_matrix.py index bd5f4b9..decbfec 100644 --- a/doris_mcp_server/tools/doris_feature_matrix.py +++ b/doris_mcp_server/tools/doris_feature_matrix.py @@ -948,7 +948,7 @@ FEATURE_DEFINITIONS = ( _feature( "doris_catalog", "get_table_size", - A, + M, _variant( "partition_statistics", probes=("table_partition_statistics_readable",), diff --git a/doris_mcp_server/tools/tools_manager.py b/doris_mcp_server/tools/tools_manager.py index 77d6f86..8d2d072 100644 --- a/doris_mcp_server/tools/tools_manager.py +++ b/doris_mcp_server/tools/tools_manager.py @@ -19,6 +19,7 @@ Apache Doris MCP Tools Manager Responsible for tool registration, management, scheduling and routing, does not contain specific business logic implementation """ +import secrets from collections.abc import Iterable from datetime import datetime from typing import Any @@ -27,6 +28,7 @@ from mcp.types import Tool from ..auth.operation_policy import authorize_operation from ..result_limits import configured_default_result_rows +from ..state_handles import StateHandleCodec from ..utils.adbc_query_tools import DorisADBCQueryTools from ..utils.analysis_tools import MemoryTracker, SQLAnalyzer, TableAnalyzer from ..utils.data_exploration_tools import DataExplorationTools @@ -47,6 +49,7 @@ from .capability_registry import ( CapabilityProviderRegistry, CapabilityRegistry, ) +from .catalog_handlers import CatalogToolHandlersMixin from .domain_catalog import CURRENT_FLAT_TOOL_NAMES from .domain_dispatcher import ( BoundHandlerAvailabilityProvider, @@ -67,7 +70,10 @@ from .tool_registry import ToolRegistryError logger = get_logger(__name__) -class DorisToolsManager(DomainManifestManagerMixin): +class DorisToolsManager( + CatalogToolHandlersMixin, + DomainManifestManagerMixin, +): """Apache Doris Tools Manager""" def __init__( @@ -87,6 +93,7 @@ class DorisToolsManager(DomainManifestManagerMixin): self.metadata_extractor = MetadataExtractor( connection_manager=connection_manager ) + self._initialize_catalog_handlers(connection_manager) self.monitoring_tools = DorisMonitoringTools(connection_manager) self.memory_tracker = MemoryTracker(connection_manager) @@ -151,9 +158,29 @@ class DorisToolsManager(DomainManifestManagerMixin): self._tool_exposure_mode = ToolExposureMode.parse( tool_exposure_mode or configured_mode ) + state_handle_secret = getattr( + config, + "mcp_state_handle_secret", + None, + ) + if not isinstance(state_handle_secret, str | bytes) or len( + state_handle_secret + ) < 32: + state_handle_secret = secrets.token_urlsafe(32) + state_handle_ttl = getattr( + config, + "mcp_state_handle_ttl_seconds", + 300, + ) + if not isinstance(state_handle_ttl, int): + state_handle_ttl = 300 self._domain_dispatcher = DomainDispatcher( self, self._domain_manifest_service, + state_handle_codec=StateHandleCodec( + state_handle_secret, + default_ttl_seconds=state_handle_ttl, + ), ) self._validate_custom_tool_names() @@ -287,75 +314,6 @@ class DorisToolsManager(DomainManifestManagerMixin): max_bytes=max_bytes, ) - async def _get_table_schema_tool(self, arguments: dict[str, Any]) -> dict[str, Any]: - """Get table schema tool routing""" - table_name = self._required_string(arguments, "table_name") - db_name = arguments.get("db_name") - catalog_name = arguments.get("catalog_name") - - # Delegate to metadata extractor for processing - return await self.metadata_extractor.get_table_schema_for_mcp( - table_name, db_name, catalog_name - ) - - async def _get_db_table_list_tool( - self, arguments: dict[str, Any] - ) -> dict[str, Any]: - """Get database table list tool routing""" - db_name = arguments.get("db_name") - catalog_name = arguments.get("catalog_name") - - # Delegate to metadata extractor for processing - return await self.metadata_extractor.get_db_table_list_for_mcp( - db_name, catalog_name - ) - - async def _get_db_list_tool(self, arguments: dict[str, Any]) -> dict[str, Any]: - """Get database list tool routing""" - catalog_name = arguments.get("catalog_name") - - # Delegate to metadata extractor for processing - return await self.metadata_extractor.get_db_list_for_mcp(catalog_name) - - async def _get_table_comment_tool( - self, arguments: dict[str, Any] - ) -> dict[str, Any]: - """Get table comment tool routing""" - table_name = self._required_string(arguments, "table_name") - db_name = arguments.get("db_name") - catalog_name = arguments.get("catalog_name") - - # Delegate to metadata extractor for processing - return await self.metadata_extractor.get_table_comment_for_mcp( - table_name, db_name, catalog_name - ) - - async def _get_table_column_comments_tool( - self, arguments: dict[str, Any] - ) -> dict[str, Any]: - """Get table column comments tool routing""" - table_name = self._required_string(arguments, "table_name") - db_name = arguments.get("db_name") - catalog_name = arguments.get("catalog_name") - - # Delegate to metadata extractor for processing - return await self.metadata_extractor.get_table_column_comments_for_mcp( - table_name, db_name, catalog_name - ) - - async def _get_table_indexes_tool( - self, arguments: dict[str, Any] - ) -> dict[str, Any]: - """Get table indexes tool routing""" - table_name = self._required_string(arguments, "table_name") - db_name = arguments.get("db_name") - catalog_name = arguments.get("catalog_name") - - # Delegate to metadata extractor for processing - return await self.metadata_extractor.get_table_indexes_for_mcp( - table_name, db_name, catalog_name - ) - async def _get_recent_audit_logs_tool( self, arguments: dict[str, Any] ) -> dict[str, Any]: @@ -366,14 +324,6 @@ class DorisToolsManager(DomainManifestManagerMixin): # Delegate to metadata extractor for processing return await self.metadata_extractor.get_recent_audit_logs_for_mcp(days, limit) - async def _get_catalog_list_tool(self, arguments: dict[str, Any]) -> dict[str, Any]: - """Get catalog list tool routing""" - # random_string parameter is required in the source project, but not actually used in business logic - # Here we ignore it and directly call business logic - - # Delegate to metadata extractor for processing - return await self.metadata_extractor.get_catalog_list_for_mcp() - async def _get_sql_explain_tool(self, arguments: dict[str, Any]) -> dict[str, Any]: """SQL Explain tool routing""" sql = self._required_string(arguments, "sql") @@ -398,19 +348,6 @@ class DorisToolsManager(DomainManifestManagerMixin): sql, db_name, catalog_name, timeout ) - async def _get_table_data_size_tool( - self, arguments: dict[str, Any] - ) -> dict[str, Any]: - """Table data size tool routing""" - db_name = arguments.get("db_name") - table_name = arguments.get("table_name") - single_replica = arguments.get("single_replica", False) - - # Delegate to SQL analyzer for processing - return await self.sql_analyzer.get_table_data_size( - db_name, table_name, single_replica - ) - async def _get_monitoring_metrics_tool( self, arguments: dict[str, Any] ) -> dict[str, Any]: @@ -538,29 +475,6 @@ class DorisToolsManager(DomainManifestManagerMixin): # ==================== v0.5.0 Advanced Analytics Tools Private Methods ==================== - async def _get_table_basic_info_tool( - self, arguments: dict[str, Any] - ) -> dict[str, Any]: - """Get table basic information tool routing""" - try: - table_name = self._required_string(arguments, "table_name") - catalog_name = arguments.get("catalog_name") - db_name = arguments.get("db_name") - - # Delegate to atomic data quality tools - result = await self.data_quality_tools.get_table_basic_info( - table_name=table_name, catalog_name=catalog_name, db_name=db_name - ) - - return result - - except Exception as e: - return { - "error": str(e), - "analysis_type": "table_basic_info", - "timestamp": datetime.now().isoformat(), - } - async def _analyze_columns_tool(self, arguments: dict[str, Any]) -> dict[str, Any]: """Analyze columns tool routing""" try: diff --git a/doris_mcp_server/utils/catalog_metadata.py b/doris_mcp_server/utils/catalog_metadata.py new file mode 100644 index 0000000..3e8b4ae --- /dev/null +++ b/doris_mcp_server/utils/catalog_metadata.py @@ -0,0 +1,661 @@ +# 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. + +"""Strict, read-only Doris catalog metadata access for formal domain tools.""" + +from __future__ import annotations + +import uuid +from collections.abc import Mapping +from typing import Any + +from .db import DorisConnectionManager +from .security import get_current_auth_context +from .sql_security_utils import ( + SQLSecurityError, + quote_identifier, + validate_identifier, +) + + +class CatalogMetadataFailure(RuntimeError): + """A sanitized catalog failure that remains machine distinguishable.""" + + def __init__( + self, + message: str, + *, + reason_code: str, + status_code: int, + retryable: bool = False, + ) -> None: + super().__init__(message) + self.reason_code = reason_code + self.status_code = status_code + self.retryable = retryable + + +class DorisCatalogMetadataReader: + """Read formal Catalog-domain metadata without swallowing Doris failures.""" + + def __init__(self, connection_manager: DorisConnectionManager) -> None: + self._connection_manager = connection_manager + self._session_id = f"formal_catalog_{uuid.uuid4().hex[:8]}" + + async def list_catalogs(self) -> dict[str, Any]: + """Return visible catalogs with their reported provider type.""" + rows = await self._execute("SHOW CATALOGS") + items: list[dict[str, Any]] = [] + for row in rows: + name = _row_value(row, "CatalogName") + if name is None: + values = list(row.values()) + name = values[1] if len(values) > 1 else values[0] if values else None + if name is None: + continue + catalog_name = str(name) + reported_type = _row_value(row, "Type") + scope = "internal" if catalog_name == "internal" else "external" + items.append( + { + "name": catalog_name, + "type": ( + str(reported_type).lower() + if reported_type not in (None, "") + else scope + ), + "scope": scope, + "is_current": _as_bool(_row_value(row, "IsCurrent")), + } + ) + return _success(items, "SHOW CATALOGS") + + async def list_databases( + self, + *, + catalog_name: str | None, + ) -> dict[str, Any]: + """Return visible databases in one catalog.""" + effective_catalog = catalog_name or "internal" + if catalog_name and catalog_name != "internal": + safe_catalog = _quote(catalog_name, "catalog name") + sql = f"SHOW DATABASES FROM {safe_catalog}" + else: + sql = "SHOW DATABASES" + rows = await self._execute(sql) + items = [ + { + "name": str(name), + "catalog": effective_catalog, + } + for row in rows + if (name := _first_value(row)) is not None + ] + return _success(items, "SHOW DATABASES") + + async def list_tables( + self, + *, + database_name: str, + catalog_name: str | None, + ) -> dict[str, Any]: + """Return visible tables and views with stable normalized object types.""" + relation = _qualified_name(catalog_name, database_name) + rows = await self._execute(f"SHOW FULL TABLES FROM {relation}") + items: list[dict[str, Any]] = [] + for row in rows: + name = next( + ( + value + for key, value in row.items() + if str(key).lower().startswith("tables_in_") + ), + _first_value(row), + ) + if name is None: + continue + raw_type = _row_value(row, "Table_type", "TABLE_TYPE") + items.append( + { + "name": str(name), + "type": _normalize_table_type(raw_type), + "raw_type": "" if raw_type is None else str(raw_type), + "catalog": catalog_name or "internal", + "database": database_name, + "storage_format": _optional_string( + _row_value(row, "Storage_format") + ), + "inverted_index_storage_format": _optional_string( + _row_value(row, "Inverted_index_storage_format") + ), + } + ) + return _success(items, "SHOW FULL TABLES") + + async def get_table_schema( + self, + *, + table_name: str, + database_name: str, + catalog_name: str | None, + ) -> dict[str, Any]: + """Return a strict schema section and prove that the table is visible.""" + relation = _qualified_name(catalog_name, database_name, table_name) + rows = await self._execute(f"SHOW FULL COLUMNS FROM {relation}") + if not rows: + raise _object_not_found() + columns = [ + { + "column_name": _optional_string(_row_value(row, "Field")), + "data_type": _optional_string(_row_value(row, "Type")), + "is_nullable": str(_row_value(row, "Null") or "NO").upper() + == "YES", + "default_value": _row_value(row, "Default"), + "comment": _optional_string(_row_value(row, "Comment")), + "key": _optional_string(_row_value(row, "Key")), + "extra": _optional_string(_row_value(row, "Extra")), + } + for row in rows + ] + return _success( + { + "catalog": catalog_name or "internal", + "database": database_name, + "table": table_name, + "columns": columns, + "column_count": len(columns), + }, + "SHOW FULL COLUMNS", + ) + + async def get_table_comment( + self, + *, + table_name: str, + database_name: str, + catalog_name: str | None, + ) -> dict[str, Any]: + """Return the table comment from the catalog-local information schema.""" + relation = _information_schema_relation(catalog_name, "tables") + # SQL sink audit: relation is assembled only from quote_identifier + # outputs; filter values stay bound before _execute reaches Doris. + rows = await self._execute( + ( + "SELECT TABLE_COMMENT " + f"FROM {relation} " # nosec B608 + "WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s" + ), + (database_name, table_name), + ) + if not rows: + raise _object_not_found() + return _success( + { + "comment": _optional_string( + _row_value(rows[0], "TABLE_COMMENT") + ) + }, + "information_schema.tables", + ) + + async def get_column_comments( + self, + *, + table_name: str, + database_name: str, + catalog_name: str | None, + ) -> dict[str, Any]: + """Return ordered column comments from the catalog-local metadata.""" + relation = _information_schema_relation(catalog_name, "columns") + # SQL sink audit: relation is assembled only from quote_identifier + # outputs; filter values stay bound before _execute reaches Doris. + rows = await self._execute( + ( + "SELECT COLUMN_NAME, COLUMN_COMMENT " + f"FROM {relation} " # nosec B608 + "WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s " + "ORDER BY ORDINAL_POSITION" + ), + (database_name, table_name), + ) + return _success( + { + str(name): _optional_string( + _row_value(row, "COLUMN_COMMENT") + ) + for row in rows + if (name := _row_value(row, "COLUMN_NAME")) is not None + }, + "information_schema.columns", + ) + + async def get_table_indexes( + self, + *, + table_name: str, + database_name: str, + catalog_name: str | None, + ) -> dict[str, Any]: + """Return table indexes, preserving multi-column index order.""" + relation = _qualified_name(catalog_name, database_name, table_name) + rows = await self._execute(f"SHOW INDEX FROM {relation}") + indexes: dict[str, dict[str, Any]] = {} + for row in rows: + raw_name = _row_value(row, "Key_name", "KEY_NAME") + if raw_name is None: + continue + name = str(raw_name) + index = indexes.setdefault( + name, + { + "name": name, + "columns": [], + "unique": _as_int(_row_value(row, "Non_unique"), 1) == 0, + "type": _optional_string( + _row_value(row, "Index_type", "INDEX_TYPE") + ), + }, + ) + column_name = _row_value(row, "Column_name", "COLUMN_NAME") + if column_name is not None: + index["columns"].append(str(column_name)) + return _success( + {"items": list(indexes.values()), "index_count": len(indexes)}, + "SHOW INDEX", + ) + + async def get_table_basic( + self, + *, + table_name: str, + database_name: str, + catalog_name: str | None, + ) -> dict[str, Any]: + """Return stable basic metadata without scanning table data.""" + row = await self._get_table_metadata_row( + table_name=table_name, + database_name=database_name, + catalog_name=catalog_name, + ) + return _success( + _normalize_basic_row( + row, + catalog_name=catalog_name, + database_name=database_name, + table_name=table_name, + ), + "information_schema.tables", + ) + + async def get_table_size( + self, + *, + table_name: str, + database_name: str, + catalog_name: str | None, + include_partitions: bool, + ) -> dict[str, Any]: + """Return table statistics and optional partition statistics.""" + row = await self._get_table_metadata_row( + table_name=table_name, + database_name=database_name, + catalog_name=catalog_name, + ) + data_size = _optional_int(_row_value(row, "DATA_LENGTH")) + index_size = _optional_int(_row_value(row, "INDEX_LENGTH")) + result: dict[str, Any] = { + "catalog": catalog_name or "internal", + "database": database_name, + "table": table_name, + "row_count": _optional_int(_row_value(row, "TABLE_ROWS")), + "data_size_bytes": data_size, + "index_size_bytes": index_size, + "total_size_bytes": ( + data_size + index_size + if data_size is not None and index_size is not None + else data_size + ), + "partitions": [], + } + warnings: list[str] = [] + metadata: dict[str, Any] = { + "source": "information_schema.tables", + "partition_status": "not_requested", + } + if include_partitions: + try: + partition_rows = await self._get_partition_rows( + table_name=table_name, + database_name=database_name, + catalog_name=catalog_name, + ) + except CatalogMetadataFailure as exc: + warnings.append( + "Partition statistics are unavailable for the requested " + f"table ({exc.reason_code})." + ) + metadata["partition_status"] = "unavailable" + metadata["partition_reason_code"] = exc.reason_code + else: + result["partitions"] = [ + { + "name": _optional_string( + _row_value(partition, "PARTITION_NAME") + ), + "ordinal_position": _optional_int( + _row_value( + partition, + "PARTITION_ORDINAL_POSITION", + ) + ), + "row_count": _optional_int( + _row_value(partition, "TABLE_ROWS") + ), + "data_size_bytes": _optional_int( + _row_value(partition, "DATA_LENGTH") + ), + "index_size_bytes": _optional_int( + _row_value(partition, "INDEX_LENGTH") + ), + } + for partition in partition_rows + ] + metadata["partition_status"] = "available" + metadata["source"] = ( + "information_schema.tables," + "information_schema.partitions" + ) + return _success( + result, + metadata["source"], + warnings=warnings, + metadata=metadata, + ) + + async def _get_table_metadata_row( + self, + *, + table_name: str, + database_name: str, + catalog_name: str | None, + ) -> Mapping[str, Any]: + relation = _information_schema_relation(catalog_name, "tables") + # SQL sink audit: relation is assembled only from quote_identifier + # outputs; filter values stay bound before _execute reaches Doris. + rows = await self._execute( + ( + "SELECT TABLE_CATALOG, TABLE_SCHEMA, TABLE_NAME, TABLE_TYPE, " + "ENGINE, TABLE_ROWS, AVG_ROW_LENGTH, DATA_LENGTH, INDEX_LENGTH, " + "CREATE_TIME, UPDATE_TIME, TABLE_COMMENT " + f"FROM {relation} " # nosec B608 + "WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s" + ), + (database_name, table_name), + ) + if not rows: + raise _object_not_found() + return rows[0] + + async def _get_partition_rows( + self, + *, + table_name: str, + database_name: str, + catalog_name: str | None, + ) -> list[Mapping[str, Any]]: + relation = _information_schema_relation(catalog_name, "partitions") + # SQL sink audit: relation is assembled only from quote_identifier + # outputs; filter values stay bound before _execute reaches Doris. + return await self._execute( + ( + "SELECT PARTITION_NAME, PARTITION_ORDINAL_POSITION, TABLE_ROWS, " + f"DATA_LENGTH, INDEX_LENGTH FROM {relation} " # nosec B608 + "WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s " + "ORDER BY PARTITION_ORDINAL_POSITION" + ), + (database_name, table_name), + ) + + async def _execute( + self, + sql: str, + params: tuple[Any, ...] | None = None, + ) -> list[Mapping[str, Any]]: + try: + result = await self._connection_manager.execute_query( + self._session_id, + sql, + params, + get_current_auth_context(), + ) + except Exception as exc: + raise _classify_failure(exc) from exc + return [ + row + for row in (result.data or []) + if isinstance(row, Mapping) + ] + + +def _success( + result: Any, + source: str, + *, + warnings: list[str] | None = None, + metadata: Mapping[str, Any] | None = None, +) -> dict[str, Any]: + response_metadata = dict(metadata or {}) + response_metadata.setdefault("source", source) + return { + "success": True, + "result": result, + "warnings": list(warnings or ()), + "metadata": response_metadata, + } + + +def _classify_failure(exc: Exception) -> CatalogMetadataFailure: + if isinstance(exc, CatalogMetadataFailure): + return exc + if isinstance(exc, SQLSecurityError): + return CatalogMetadataFailure( + "Catalog arguments contain an invalid Doris identifier.", + reason_code="CATALOG_ARGUMENT_INVALID", + status_code=400, + ) + + numeric_code = next( + ( + value + for value in getattr(exc, "args", ()) + if isinstance(value, int) + ), + None, + ) + message = str(exc).lower() + if numeric_code in {1044, 1045, 1142, 1227} or any( + marker in message + for marker in ( + "access denied", + "permission denied", + "not authorized", + "privilege", + ) + ): + return CatalogMetadataFailure( + "Doris denied access to the requested catalog metadata.", + reason_code="CATALOG_PERMISSION_DENIED", + status_code=403, + ) + if numeric_code in {1049, 1146} or any( + marker in message + for marker in ( + "doesn't exist", + "does not exist", + "unknown database", + "unknown table", + "not found", + ) + ): + return _object_not_found() + if numeric_code in {1064, 1109} or any( + marker in message + for marker in ( + "not supported", + "unsupported", + "unknown system table", + ) + ): + return CatalogMetadataFailure( + "The requested Doris metadata section is unsupported.", + reason_code="CATALOG_SECTION_UNSUPPORTED", + status_code=501, + ) + if isinstance(exc, TimeoutError | ConnectionError | OSError): + return CatalogMetadataFailure( + "Doris catalog metadata is temporarily unavailable.", + reason_code="CATALOG_BACKEND_UNAVAILABLE", + status_code=503, + retryable=True, + ) + return CatalogMetadataFailure( + "Doris catalog metadata execution failed.", + reason_code="CATALOG_EXECUTION_FAILED", + status_code=502, + ) + + +def _object_not_found() -> CatalogMetadataFailure: + return CatalogMetadataFailure( + "The requested Doris metadata object was not found.", + reason_code="CATALOG_OBJECT_NOT_FOUND", + status_code=404, + ) + + +def _quote(value: str, label: str) -> str: + try: + validate_identifier(value, label) + return quote_identifier(value, label) + except SQLSecurityError as exc: + raise _classify_failure(exc) from exc + + +def _qualified_name( + catalog_name: str | None, + database_name: str, + table_name: str | None = None, +) -> str: + parts = [] + if catalog_name and catalog_name != "internal": + parts.append(_quote(catalog_name, "catalog name")) + parts.append(_quote(database_name, "database name")) + if table_name is not None: + parts.append(_quote(table_name, "table name")) + return ".".join(parts) + + +def _information_schema_relation( + catalog_name: str | None, + table_name: str, +) -> str: + parts = [] + if catalog_name and catalog_name != "internal": + parts.append(_quote(catalog_name, "catalog name")) + parts.extend( + ( + quote_identifier("information_schema", "database name"), + _quote(table_name, "metadata table name"), + ) + ) + return ".".join(parts) + + +def _row_value(row: Mapping[str, Any], *names: str) -> Any | None: + lowered = {str(key).lower(): value for key, value in row.items()} + for name in names: + if name.lower() in lowered: + return lowered[name.lower()] + return None + + +def _first_value(row: Mapping[str, Any]) -> Any | None: + return next(iter(row.values()), None) + + +def _optional_string(value: Any) -> str: + return "" if value is None else str(value) + + +def _optional_int(value: Any) -> int | None: + if value in (None, ""): + return None + try: + return int(value) + except (TypeError, ValueError): + return None + + +def _as_int(value: Any, default: int) -> int: + parsed = _optional_int(value) + return default if parsed is None else parsed + + +def _as_bool(value: Any) -> bool: + if isinstance(value, bool): + return value + return str(value or "").strip().lower() in {"1", "true", "yes"} + + +def _normalize_table_type(value: Any) -> str: + normalized = str(value or "").strip().upper().replace("-", " ") + if "MATERIALIZED" in normalized: + return "materialized_view" + if "VIEW" in normalized: + return "view" + return "table" + + +def _normalize_basic_row( + row: Mapping[str, Any], + *, + catalog_name: str | None, + database_name: str, + table_name: str, +) -> dict[str, Any]: + return { + "catalog": catalog_name or "internal", + "database": database_name, + "table": table_name, + "type": _normalize_table_type(_row_value(row, "TABLE_TYPE")), + "raw_type": _optional_string(_row_value(row, "TABLE_TYPE")), + "engine": _optional_string(_row_value(row, "ENGINE")), + "row_count": _optional_int(_row_value(row, "TABLE_ROWS")), + "average_row_length": _optional_int( + _row_value(row, "AVG_ROW_LENGTH") + ), + "data_size_bytes": _optional_int(_row_value(row, "DATA_LENGTH")), + "index_size_bytes": _optional_int(_row_value(row, "INDEX_LENGTH")), + "created_at": _optional_string(_row_value(row, "CREATE_TIME")), + "updated_at": _optional_string(_row_value(row, "UPDATE_TIME")), + "comment": _optional_string(_row_value(row, "TABLE_COMMENT")), + } + + +__all__ = [ + "CatalogMetadataFailure", + "DorisCatalogMetadataReader", +] diff --git a/test/protocol/test_state_handles.py b/test/protocol/test_state_handles.py index f9c1196..1fb004c 100644 --- a/test/protocol/test_state_handles.py +++ b/test/protocol/test_state_handles.py @@ -209,6 +209,20 @@ def test_handle_rejects_expiry_and_tampering_before_exposing_state(): auth_context=_auth_context(), ) + alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_" + parts = handle.split(".") + tail_index = alphabet.index(parts[2][-1]) + noncanonical_tail = alphabet[(tail_index & ~0b11) | ((tail_index + 1) & 0b11)] + parts[2] = parts[2][:-1] + noncanonical_tail + with pytest.raises(StateHandleError, match="invalid"): + codec.resolve( + ".".join(parts), + expected_kind="list-page", + expected_scope="tools:list", + expected_resource="mcp://tools", + auth_context=_auth_context(), + ) + def test_handle_enforces_secret_ttl_claim_and_payload_limits(): with pytest.raises(ValueError, match="at least 32 bytes"): diff --git a/test/tools/test_capability_detector.py b/test/tools/test_capability_detector.py index fab456b..3c0afbe 100644 --- a/test/tools/test_capability_detector.py +++ b/test/tools/test_capability_detector.py @@ -97,6 +97,22 @@ class _ProbeConnection: ): [ {"table_metadata_probe": 1} ], + ( + "SELECT COLUMN_NAME, DATA_TYPE " + "FROM information_schema.columns LIMIT 1" + ): [ + {"COLUMN_NAME": "id", "DATA_TYPE": "bigint"} + ], + ( + "SELECT PARTITION_NAME, TABLE_ROWS, DATA_LENGTH " + "FROM information_schema.partitions LIMIT 1" + ): [ + { + "PARTITION_NAME": "p1", + "TABLE_ROWS": 1, + "DATA_LENGTH": 8, + } + ], } return SimpleNamespace( data=self.row_overrides.get(sql, rows.get(sql, [])) @@ -166,6 +182,14 @@ async def test_detector_builds_version_vector_and_extends_domains_lazily() -> No catalog.probe("database_metadata_readable").status is CapabilityProbeStatus.SUPPORTED ) + assert ( + catalog.probe("table_context_sections_readable").status + is CapabilityProbeStatus.SUPPORTED + ) + assert ( + catalog.probe("table_partition_statistics_readable").status + is CapabilityProbeStatus.SUPPORTED + ) assert connection.statements.count("SELECT @@version_comment;") == 1 diff --git a/test/tools/test_capability_registry.py b/test/tools/test_capability_registry.py index f5fa7ea..0f5fd3d 100644 --- a/test/tools/test_capability_registry.py +++ b/test/tools/test_capability_registry.py @@ -193,6 +193,44 @@ def test_evaluator_requires_version_probes_handler_and_call_permission() -> None assert allowed.callable is True +def test_evaluator_normalizes_system_object_probe_evidence_for_manifest() -> None: + evaluator = CapabilityEvaluator( + matrix=DORIS_FEATURE_MATRIX, + bound_handlers=_BoundHandlers( + "doris_catalog.get_table_context" + ), # type: ignore[arg-type] + ) + domain = DORIS_DOMAIN_CATALOG.resolve_domain("doris_catalog") + child = DORIS_DOMAIN_CATALOG.resolve_child( + "doris_catalog", + "get_table_context", + ) + probes = { + probe_id: CapabilityProbeEvidence( + probe_id=probe_id, + status=CapabilityProbeStatus.SUPPORTED, + reason_code="TEST_PROBE_SUPPORTED", + evidence_sources=("test_probe",), + ) + for probe_id in ( + "information_schema.columns", + "table_context_sections_readable", + ) + } + + availability = evaluator.evaluate( + snapshot=_snapshot(probes=probes), + providers=CapabilityProviderRegistry({}).snapshot(), + domain=domain, + child=child, + auth_context=None, + ) + + assert availability.callable is True + assert "information_schema_columns" in availability.evidence_sources + assert "information_schema.columns" not in availability.evidence_sources + + def test_evaluator_distinguishes_provider_misconfiguration() -> None: bound = _BoundHandlers("doris_query.diagnose_query_performance") evaluator = CapabilityEvaluator( diff --git a/test/tools/test_domain_catalog.py b/test/tools/test_domain_catalog.py index aa6de8e..6a05b78 100644 --- a/test/tools/test_domain_catalog.py +++ b/test/tools/test_domain_catalog.py @@ -226,7 +226,18 @@ def test_table_context_has_four_sections_and_five_deterministic_steps() -> None: assert plan is not None section_enum = _wire_input(child)["properties"]["sections"]["items"]["enum"] + output = child.to_wire()["output_schema"] + data_schema = output["properties"]["data"] assert tuple(section_enum) == ("schema", "comments", "indexes", "basic") + for section in ("schema", "comments", "indexes", "basic"): + section_schema = data_schema["properties"][section] + assert section_schema["required"] == [ + "status", + "data", + "warnings", + "source", + ] + assert section_schema["additionalProperties"] is False assert tuple(step.name for step in plan.steps) == ( "schema", "table_comments", diff --git a/test/tools/test_domain_dispatcher.py b/test/tools/test_domain_dispatcher.py index cad865c..6a51825 100644 --- a/test/tools/test_domain_dispatcher.py +++ b/test/tools/test_domain_dispatcher.py @@ -807,19 +807,33 @@ async def test_table_context_composite_merges_selected_sections() -> None: assert result.mode == "result" child_result = cast_mapping(result.to_wire()["data"]) assert child_result["status"] == "partial" - assert child_result["data"]["schema"] == {"columns": ["id"]} - assert child_result["data"]["comments"] == { + assert child_result["data"]["schema"] == { + "status": "success", + "data": {"columns": ["id"]}, + "warnings": [], + "source": "_get_table_schema_tool", + } + assert child_result["data"]["comments"]["data"] == { "table": {"comment": "orders"}, "columns": {"id": "key"}, } - assert "indexes" not in child_result["data"] - assert child_result["data"]["basic"]["row_count"] == 1 + assert child_result["data"]["indexes"]["status"] == "unavailable" + assert child_result["data"]["indexes"]["data"] == {} + assert ( + child_result["data"]["indexes"]["warnings"] + == [ + "Optional indexes section is unavailable " + "(CATALOG_EXECUTION_FAILED)." + ] + ) + assert child_result["data"]["basic"]["data"]["row_count"] == 1 assert len(child_result["evidence"]) == 5 manager._get_table_schema_tool.assert_awaited_once_with( { "catalog_name": "internal", "db_name": "analytics", "table_name": "orders", + "_formal_catalog": True, } ) @@ -848,6 +862,40 @@ async def test_table_context_required_schema_failure_is_execution_error() -> Non assert result.error.code is DomainErrorCode.CHILD_EXECUTION_FAILED [email protected] +async def test_table_context_required_schema_distinguishes_missing_object() -> None: + manager = _manager("doris_catalog.get_table_context") + manager._get_table_schema_tool = AsyncMock( + return_value={ + "success": False, + "error": "not visible", + "error_code": "DORIS_METADATA_NOT_VISIBLE", + "status_code": 404, + } + ) + + result = await manager.domain_dispatcher.call_domain( + "doris_catalog", + { + "child_tool": "get_table_context", + "arguments": { + "database": "analytics", + "table": "missing", + "sections": ["schema"], + }, + }, + None, + ) + + assert result.mode == "error" + assert result.error.code is DomainErrorCode.CHILD_EXECUTION_FAILED + assert ( + result.error.details["reason_code"] + == "CATALOG_OBJECT_NOT_FOUND" + ) + assert result.error.details["status_code"] == 404 + + @pytest.mark.asyncio async def test_table_context_skips_unrequested_optional_sections() -> None: manager = _manager("doris_catalog.get_table_context") @@ -869,7 +917,14 @@ async def test_table_context_skips_unrequested_optional_sections() -> None: ) assert result.mode == "result" - assert result.to_wire()["data"]["data"] == {"schema": {"columns": ["id"]}} + assert result.to_wire()["data"]["data"] == { + "schema": { + "status": "success", + "data": {"columns": ["id"]}, + "warnings": [], + "source": "_get_table_schema_tool", + } + } @pytest.mark.asyncio @@ -903,13 +958,27 @@ async def test_table_context_includes_schema_for_comments_only_selection() -> No assert result.mode == "result" child_result = cast_mapping(result.to_wire()["data"]) assert child_result["data"] == { - "schema": {"columns": ["id"]}, + "schema": { + "status": "success", + "data": {"columns": ["id"]}, + "warnings": [], + "source": "_get_table_schema_tool", + }, "comments": { - "table": {"comment": "orders"}, - "columns": {"id": "key"}, + "status": "success", + "data": { + "table": {"comment": "orders"}, + "columns": {"id": "key"}, + }, + "warnings": [], + "source": ( + "_get_table_comment_tool," + "_get_table_column_comments_tool" + ), }, } - assert child_result["metadata"]["sections"] == ["comments", "schema"] + assert child_result["metadata"]["sections"] == ["schema", "comments"] + assert child_result["metadata"]["requested_sections"] == ["comments"] manager._get_table_indexes_tool.assert_not_awaited() manager._get_table_basic_info_tool.assert_not_awaited() @@ -946,7 +1015,7 @@ async def test_table_context_internal_guards_reject_invalid_composition() -> Non ) -def test_collection_result_normalization_filters_and_warns() -> None: +def test_collection_result_normalization_filters_and_paginates() -> None: manager = _manager() child = DORIS_DOMAIN_CATALOG.resolve_child( "doris_catalog", @@ -958,26 +1027,300 @@ def test_collection_result_normalization_filters_and_warns() -> None: { "success": True, "result": [ - "internal", - {"name": "iceberg"}, - {"name": "hive"}, + {"name": "internal", "scope": "internal"}, + *[ + {"name": f"external_{index:03d}", "scope": "external"} + for index in range(101) + ], ], }, { - "pattern": "i*", - "page": "cursor", - "types": ["internal", "hms"], + "pattern": "*", + "include_external": True, + }, + (), + ) + + assert result["status"] == "success" + assert len(result["data"]["items"]) == 100 + assert result["data"]["truncated"] is True + assert isinstance(result["data"]["next_cursor"], str) + + next_page = manager.domain_dispatcher._normalize_child_result( + child, + { + "success": True, + "result": [ + {"name": "internal", "scope": "internal"}, + *[ + {"name": f"external_{index:03d}", "scope": "external"} + for index in range(101) + ], + ], + }, + { + "pattern": "*", + "include_external": True, + "page": result["data"]["next_cursor"], + }, + (), + ) + assert len(next_page["data"]["items"]) == 2 + assert next_page["data"]["next_cursor"] is None + assert next_page["data"]["truncated"] is False + + internal_only = manager.domain_dispatcher._normalize_child_result( + child, + { + "success": True, + "result": [ + {"name": "internal", "scope": "internal"}, + {"name": "iceberg", "scope": "external"}, + ], + }, + {"include_external": False}, + (), + ) + assert internal_only["data"]["items"] == [ + {"name": "internal", "scope": "internal"} + ] + + +def test_table_collection_filters_exact_normalized_object_types() -> None: + manager = _manager() + child = DORIS_DOMAIN_CATALOG.resolve_child( + "doris_catalog", + "list_tables", + ) + + result = manager.domain_dispatcher._normalize_child_result( + child, + { + "success": True, + "result": [ + {"name": "orders", "type": "table"}, + {"name": "orders_view", "type": "view"}, + {"name": "orders_mv", "type": "materialized_view"}, + ], + }, + { + "database": "analytics", + "types": ["view", "materialized_view"], }, (), ) - assert result["status"] == "partial" assert result["data"]["items"] == [ - {"name": "internal"}, - {"name": "iceberg"}, + {"name": "orders_mv", "type": "materialized_view"}, + {"name": "orders_view", "type": "view"}, ] - assert result["data"]["next_cursor"] is None - assert len(result["warnings"]) == 2 + + +def test_collection_cursor_is_bound_to_the_exact_request() -> None: + manager = _manager() + child = DORIS_DOMAIN_CATALOG.resolve_child( + "doris_catalog", + "list_tables", + ) + raw = { + "success": True, + "result": [ + {"name": f"table_{index:03d}", "type": "table"} + for index in range(101) + ], + } + first = manager.domain_dispatcher._normalize_child_result( + child, + raw, + {"database": "analytics"}, + (), + ) + + with pytest.raises( + ChildArgumentsAdapterError, + match="invalid", + ): + manager.domain_dispatcher._normalize_child_result( + child, + raw, + { + "database": "other_database", + "page": first["data"]["next_cursor"], + }, + (), + ) + + +def test_collection_cursor_is_bound_to_the_authorization_principal() -> None: + manager = _manager() + child = DORIS_DOMAIN_CATALOG.resolve_child( + "doris_catalog", + "list_tables", + ) + raw = { + "success": True, + "result": [ + {"name": f"table_{index:03d}", "type": "table"} + for index in range(101) + ], + } + alice = AuthContext(auth_method="basic", user_id="alice") + bob = AuthContext(auth_method="basic", user_id="bob") + first = manager.domain_dispatcher._normalize_child_result( + child, + raw, + {"database": "analytics"}, + (), + alice, + ) + + with pytest.raises(ChildArgumentsAdapterError, match="invalid"): + manager.domain_dispatcher._normalize_child_result( + child, + raw, + { + "database": "analytics", + "page": first["data"]["next_cursor"], + }, + (), + bob, + ) + + [email protected] [email protected]( + ("raw", "reason_code"), + [ + ( + { + "success": False, + "error": "hidden backend message", + "error_code": "DORIS_METADATA_PERMISSION_DENIED", + "status_code": 403, + }, + "CATALOG_PERMISSION_DENIED", + ), + ( + { + "success": False, + "error": "table does not exist", + "error_code": "DORIS_METADATA_NOT_VISIBLE", + "status_code": 404, + }, + "CATALOG_OBJECT_NOT_FOUND", + ), + ], +) +async def test_catalog_atomic_failures_have_stable_reason_codes( + raw: dict[str, Any], + reason_code: str, +) -> None: + manager = _manager("doris_catalog.list_tables") + manager._get_db_table_list_tool = AsyncMock(return_value=raw) + + result = await manager.domain_dispatcher.call_domain( + "doris_catalog", + { + "child_tool": "list_tables", + "arguments": {"database": "analytics"}, + }, + None, + ) + + assert result.mode == "error" + assert result.error.code is DomainErrorCode.CHILD_EXECUTION_FAILED + assert result.error.details["reason_code"] == reason_code + assert "hidden backend message" not in result.error.message + + [email protected] +async def test_table_size_preserves_partition_degradation_as_partial_result() -> None: + manager = _manager("doris_catalog.get_table_size") + manager._get_table_data_size_tool = AsyncMock( + return_value={ + "success": True, + "result": { + "row_count": 10, + "data_size_bytes": 100, + "partitions": [], + }, + "warnings": [ + "Partition statistics are unavailable " + "(CATALOG_SECTION_UNSUPPORTED)." + ], + "metadata": { + "partition_status": "unavailable", + "partition_reason_code": "CATALOG_SECTION_UNSUPPORTED", + }, + } + ) + + result = await manager.domain_dispatcher.call_domain( + "doris_catalog", + { + "child_tool": "get_table_size", + "arguments": { + "database": "analytics", + "table": "orders", + "include_partitions": True, + }, + }, + None, + ) + + assert result.mode == "result" + child_result = result.to_wire()["data"] + assert child_result["status"] == "partial" + assert child_result["data"]["row_count"] == 10 + assert child_result["metadata"]["partition_status"] == "unavailable" + manager._get_table_data_size_tool.assert_awaited_once_with( + { + "db_name": "analytics", + "table_name": "orders", + "include_partitions": True, + "single_replica": False, + "_formal_catalog": True, + } + ) + + [email protected] +async def test_table_context_marks_unsupported_optional_section() -> None: + manager = _manager("doris_catalog.get_table_context") + manager._get_table_schema_tool = AsyncMock( + return_value={"success": True, "result": {"columns": ["id"]}} + ) + manager._get_table_indexes_tool = AsyncMock( + return_value={ + "success": False, + "error": "unsupported", + "error_code": "CATALOG_SECTION_UNSUPPORTED", + "status_code": 501, + } + ) + + result = await manager.domain_dispatcher.call_domain( + "doris_catalog", + { + "child_tool": "get_table_context", + "arguments": { + "database": "analytics", + "table": "orders", + "sections": ["indexes"], + }, + }, + None, + ) + + assert result.mode == "result" + indexes = result.to_wire()["data"]["data"]["indexes"] + assert indexes["status"] == "unavailable" + assert indexes["data"] == {} + assert "CATALOG_SECTION_UNSUPPORTED" in indexes["warnings"][0] + assert ( + result.to_wire()["data"]["evidence"][1]["reason_code"] + == "CATALOG_SECTION_UNSUPPORTED" + ) def test_query_result_normalization_handles_malformed_legacy_shapes() -> None: @@ -1058,13 +1401,17 @@ def test_detail_result_normalization_adds_formal_evidence() -> None: "types": ["table"], "page": "cursor", }, - {"catalog_name": "internal", "db_name": "db"}, + { + "catalog_name": "internal", + "db_name": "db", + "_formal_catalog": True, + }, 0, ), ( "adapt:remove_random_string", {"include_external": True}, - {}, + {"_formal_catalog": True}, 0, ), ( @@ -1083,9 +1430,11 @@ def test_detail_result_normalization_adds_formal_evidence() -> None: { "db_name": "db", "table_name": "t", + "include_partitions": True, "single_replica": False, + "_formal_catalog": True, }, - 1, + 0, ), ( "adapt:monitoring_arguments", diff --git a/test/tools/test_doris_feature_matrix.py b/test/tools/test_doris_feature_matrix.py index e9b86a1..33a66fc 100644 --- a/test/tools/test_doris_feature_matrix.py +++ b/test/tools/test_doris_feature_matrix.py @@ -411,6 +411,24 @@ def test_master_fe_scope_does_not_inherit_an_unrelated_old_be_version() -> None: assert result.compatible is True +def test_table_size_uses_master_fe_metadata_with_unknown_backend_version() -> None: + result = DORIS_FEATURE_MATRIX.evaluate( + domain="doris_catalog", + child_name="get_table_size", + versions=DorisClusterVersionVector.from_comments( + master_fe="Doris version doris-4.0.5", + backends=( + "Doris version doris-4.0.5", + "", + ), + ), + ) + + assert result.version_scope is CapabilityVersionScope.MASTER_FE + assert result.effective_version == "4.0.5" + assert result.compatible is True + + @pytest.mark.parametrize( "child_name", [ diff --git a/test/tools/test_tools_manager.py b/test/tools/test_tools_manager.py index 25ed764..dc9cf5a 100644 --- a/test/tools/test_tools_manager.py +++ b/test/tools/test_tools_manager.py @@ -219,6 +219,126 @@ class TestDorisToolsManager: assert result == expected mock_execute.assert_awaited_once_with() + @pytest.mark.asyncio + async def test_formal_catalog_discovery_routes_to_strict_reader( + self, + tools_manager, + ): + expected = {"success": True, "result": [{"name": "orders"}]} + tools_manager.catalog_metadata.list_catalogs = AsyncMock( + return_value=expected + ) + tools_manager.catalog_metadata.list_databases = AsyncMock( + return_value=expected + ) + tools_manager.catalog_metadata.list_tables = AsyncMock( + return_value=expected + ) + + catalogs = await tools_manager._get_catalog_list_tool( + {"_formal_catalog": True} + ) + databases = await tools_manager._get_db_list_tool( + {"catalog_name": "lake", "_formal_catalog": True} + ) + tables = await tools_manager._get_db_table_list_tool( + { + "catalog_name": "lake", + "db_name": "analytics", + "_formal_catalog": True, + } + ) + + assert catalogs == databases == tables == expected + tools_manager.catalog_metadata.list_catalogs.assert_awaited_once_with() + tools_manager.catalog_metadata.list_databases.assert_awaited_once_with( + catalog_name="lake" + ) + tools_manager.catalog_metadata.list_tables.assert_awaited_once_with( + database_name="analytics", + catalog_name="lake", + ) + + @pytest.mark.asyncio + async def test_formal_table_context_routes_all_five_section_handlers( + self, + tools_manager, + ): + expected = {"success": True, "result": {}} + tools_manager.catalog_metadata.get_table_schema = AsyncMock( + return_value=expected + ) + tools_manager.catalog_metadata.get_table_comment = AsyncMock( + return_value=expected + ) + tools_manager.catalog_metadata.get_column_comments = AsyncMock( + return_value=expected + ) + tools_manager.catalog_metadata.get_table_indexes = AsyncMock( + return_value=expected + ) + tools_manager.catalog_metadata.get_table_basic = AsyncMock( + return_value=expected + ) + arguments = { + "catalog_name": "internal", + "db_name": "analytics", + "table_name": "orders", + "_formal_catalog": True, + } + + results = [ + await tools_manager._get_table_schema_tool(dict(arguments)), + await tools_manager._get_table_comment_tool(dict(arguments)), + await tools_manager._get_table_column_comments_tool( + dict(arguments) + ), + await tools_manager._get_table_indexes_tool(dict(arguments)), + await tools_manager._get_table_basic_info_tool(dict(arguments)), + ] + + assert results == [expected] * 5 + for method in ( + tools_manager.catalog_metadata.get_table_schema, + tools_manager.catalog_metadata.get_table_comment, + tools_manager.catalog_metadata.get_column_comments, + tools_manager.catalog_metadata.get_table_indexes, + tools_manager.catalog_metadata.get_table_basic, + ): + method.assert_awaited_once_with( + table_name="orders", + database_name="analytics", + catalog_name="internal", + ) + + @pytest.mark.asyncio + async def test_formal_table_size_preserves_partition_request( + self, + tools_manager, + ): + expected = {"success": True, "result": {"row_count": 1}} + tools_manager.catalog_metadata.get_table_size = AsyncMock( + return_value=expected + ) + + result = await tools_manager._get_table_data_size_tool( + { + "catalog_name": "internal", + "db_name": "analytics", + "table_name": "orders", + "include_partitions": True, + "_formal_catalog": True, + } + ) + + assert result == expected + tools_manager.catalog_metadata.get_table_size.assert_awaited_once_with( + table_name="orders", + database_name="analytics", + catalog_name="internal", + include_partitions=True, + ) + @pytest.mark.asyncio async def test_invalid_tool_name(self, tools_manager): diff --git a/test/utils/test_catalog_metadata.py b/test/utils/test_catalog_metadata.py new file mode 100644 index 0000000..b110045 --- /dev/null +++ b/test/utils/test_catalog_metadata.py @@ -0,0 +1,369 @@ +# 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. + +"""Production contract tests for strict formal Catalog metadata reads.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import pytest + +from doris_mcp_server.utils.catalog_metadata import ( + CatalogMetadataFailure, + DorisCatalogMetadataReader, +) + + +class _ConnectionManager: + def __init__(self, *responses: Any) -> None: + self.responses = list(responses) + self.calls: list[tuple[str, tuple[Any, ...] | None]] = [] + + async def execute_query( + self, + _session_id: str, + sql: str, + params: tuple[Any, ...] | None, + _auth_context: Any, + ) -> SimpleNamespace: + self.calls.append((sql, params)) + response = self.responses.pop(0) + if isinstance(response, Exception): + raise response + return SimpleNamespace(data=response) + + +def _reader(*responses: Any) -> tuple[DorisCatalogMetadataReader, _ConnectionManager]: + manager = _ConnectionManager(*responses) + return ( + DorisCatalogMetadataReader(manager), # type: ignore[arg-type] + manager, + ) + + [email protected] +async def test_discovery_reads_rich_catalog_database_and_table_metadata() -> None: + reader, manager = _reader( + [ + { + "CatalogId": 0, + "CatalogName": "internal", + "Type": "internal", + "IsCurrent": "Yes", + }, + { + "CatalogId": 1, + "CatalogName": "lake", + "Type": "iceberg", + "IsCurrent": "No", + }, + ], + [{"Database": "analytics"}], + [ + { + "Tables_in_analytics": "orders", + "Table_type": "BASE TABLE", + "Storage_format": "V2", + "Inverted_index_storage_format": "V1", + }, + { + "Tables_in_analytics": "orders_view", + "Table_type": "VIEW", + "Storage_format": "NONE", + "Inverted_index_storage_format": "NONE", + }, + { + "Tables_in_analytics": "orders_mv", + "Table_type": "MATERIALIZED VIEW", + }, + ], + ) + + catalogs = await reader.list_catalogs() + databases = await reader.list_databases(catalog_name="lake") + tables = await reader.list_tables( + database_name="analytics", + catalog_name="lake", + ) + + assert catalogs["result"] == [ + { + "name": "internal", + "type": "internal", + "scope": "internal", + "is_current": True, + }, + { + "name": "lake", + "type": "iceberg", + "scope": "external", + "is_current": False, + }, + ] + assert databases["result"] == [ + {"name": "analytics", "catalog": "lake"} + ] + assert [item["type"] for item in tables["result"]] == [ + "table", + "view", + "materialized_view", + ] + assert manager.calls == [ + ("SHOW CATALOGS", None), + ("SHOW DATABASES FROM `lake`", None), + ("SHOW FULL TABLES FROM `lake`.`analytics`", None), + ] + + [email protected] +async def test_table_context_sections_use_strict_read_only_metadata_paths() -> None: + reader, manager = _reader( + [ + { + "Field": "id", + "Type": "bigint", + "Null": "NO", + "Key": "PRI", + "Default": None, + "Comment": "identifier", + "Extra": "", + } + ], + [{"TABLE_COMMENT": "orders"}], + [ + {"COLUMN_NAME": "id", "COLUMN_COMMENT": "identifier"}, + {"COLUMN_NAME": "amount", "COLUMN_COMMENT": ""}, + ], + [ + { + "Key_name": "idx_amount", + "Column_name": "amount", + "Non_unique": 1, + "Index_type": "INVERTED", + }, + { + "Key_name": "idx_compound", + "Column_name": "id", + "Non_unique": 0, + "Index_type": "BTREE", + }, + { + "Key_name": "idx_compound", + "Column_name": "amount", + "Non_unique": 0, + "Index_type": "BTREE", + }, + ], + [ + { + "TABLE_TYPE": "BASE TABLE", + "ENGINE": "OLAP", + "TABLE_ROWS": 7, + "AVG_ROW_LENGTH": 20, + "DATA_LENGTH": 140, + "INDEX_LENGTH": 10, + "CREATE_TIME": "2026-07-01 00:00:00", + "UPDATE_TIME": "2026-07-30 00:00:00", + "TABLE_COMMENT": "orders", + } + ], + ) + + schema = await reader.get_table_schema( + table_name="orders", + database_name="analytics", + catalog_name="internal", + ) + table_comment = await reader.get_table_comment( + table_name="orders", + database_name="analytics", + catalog_name="internal", + ) + column_comments = await reader.get_column_comments( + table_name="orders", + database_name="analytics", + catalog_name="internal", + ) + indexes = await reader.get_table_indexes( + table_name="orders", + database_name="analytics", + catalog_name="internal", + ) + basic = await reader.get_table_basic( + table_name="orders", + database_name="analytics", + catalog_name="internal", + ) + + assert schema["result"]["column_count"] == 1 + assert schema["result"]["columns"][0]["comment"] == "identifier" + assert table_comment["result"] == {"comment": "orders"} + assert column_comments["result"] == { + "id": "identifier", + "amount": "", + } + assert indexes["result"]["items"][1]["columns"] == ["id", "amount"] + assert basic["result"]["row_count"] == 7 + assert basic["result"]["data_size_bytes"] == 140 + assert manager.calls[0] == ( + "SHOW FULL COLUMNS FROM `analytics`.`orders`", + None, + ) + assert all( + statement.lstrip().upper().startswith(("SHOW", "SELECT")) + for statement, _params in manager.calls + ) + + [email protected] +async def test_table_size_returns_optional_partition_statistics() -> None: + reader, manager = _reader( + [ + { + "TABLE_ROWS": "9", + "DATA_LENGTH": "200", + "INDEX_LENGTH": "20", + } + ], + [ + { + "PARTITION_NAME": "p1", + "PARTITION_ORDINAL_POSITION": 1, + "TABLE_ROWS": 4, + "DATA_LENGTH": 80, + "INDEX_LENGTH": 8, + } + ], + ) + + result = await reader.get_table_size( + table_name="orders", + database_name="analytics", + catalog_name=None, + include_partitions=True, + ) + + assert result["result"]["total_size_bytes"] == 220 + assert result["result"]["partitions"] == [ + { + "name": "p1", + "ordinal_position": 1, + "row_count": 4, + "data_size_bytes": 80, + "index_size_bytes": 8, + } + ] + assert result["warnings"] == [] + assert result["metadata"]["partition_status"] == "available" + assert "`information_schema`.`partitions`" in manager.calls[1][0] + + [email protected] +async def test_table_size_preserves_base_result_when_partitions_are_unsupported() -> None: + reader, _manager = _reader( + [{"TABLE_ROWS": 9, "DATA_LENGTH": 200, "INDEX_LENGTH": 20}], + RuntimeError(1109, "Unknown system table"), + ) + + result = await reader.get_table_size( + table_name="orders", + database_name="analytics", + catalog_name="lake", + include_partitions=True, + ) + + assert result["result"]["row_count"] == 9 + assert result["result"]["partitions"] == [] + assert result["metadata"]["partition_status"] == "unavailable" + assert ( + result["metadata"]["partition_reason_code"] + == "CATALOG_SECTION_UNSUPPORTED" + ) + assert result["warnings"] == [ + "Partition statistics are unavailable for the requested table " + "(CATALOG_SECTION_UNSUPPORTED)." + ] + + [email protected] [email protected]( + ("failure", "reason_code", "status_code", "retryable"), + [ + ( + RuntimeError(1142, "SELECT command denied"), + "CATALOG_PERMISSION_DENIED", + 403, + False, + ), + ( + RuntimeError(1146, "Table does not exist"), + "CATALOG_OBJECT_NOT_FOUND", + 404, + False, + ), + ( + RuntimeError(1064, "Syntax is not supported"), + "CATALOG_SECTION_UNSUPPORTED", + 501, + False, + ), + ( + TimeoutError("timed out"), + "CATALOG_BACKEND_UNAVAILABLE", + 503, + True, + ), + ], +) +async def test_catalog_failures_are_sanitized_and_distinguishable( + failure: Exception, + reason_code: str, + status_code: int, + retryable: bool, +) -> None: + reader, _manager = _reader(failure) + + with pytest.raises(CatalogMetadataFailure) as exc_info: + await reader.list_catalogs() + + assert exc_info.value.reason_code == reason_code + assert exc_info.value.status_code == status_code + assert exc_info.value.retryable is retryable + assert "SELECT command denied" not in str(exc_info.value) + + [email protected] +async def test_empty_schema_and_invalid_identifier_fail_before_ambiguous_results() -> None: + empty_reader, _manager = _reader([]) + invalid_reader, invalid_manager = _reader([]) + + with pytest.raises(CatalogMetadataFailure) as missing: + await empty_reader.get_table_schema( + table_name="missing", + database_name="analytics", + catalog_name=None, + ) + with pytest.raises(CatalogMetadataFailure) as invalid: + await invalid_reader.list_tables( + database_name="analytics; DROP TABLE orders", + catalog_name=None, + ) + + assert missing.value.reason_code == "CATALOG_OBJECT_NOT_FOUND" + assert invalid.value.reason_code == "CATALOG_ARGUMENT_INVALID" + assert invalid_manager.calls == [] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
