This is an automated email from the ASF dual-hosted git repository. FreeOnePlus pushed a commit to branch feat/v1-11-search-domain in repository https://gitbox.apache.org/repos/asf/doris-mcp-server.git
commit 953ea5679f9983ef1151bc7353097e8a8acbe903 Author: FreeOnePlus <[email protected]> AuthorDate: Fri Jul 31 18:22:53 2026 +0800 feat: add Search domain runtime --- CHANGELOG.md | 11 + doris_mcp_server/tools/capability_detector.py | 266 +++- doris_mcp_server/tools/domain_catalog.py | 189 ++- doris_mcp_server/tools/domain_dispatcher.py | 26 + doris_mcp_server/tools/doris_feature_matrix.py | 23 +- doris_mcp_server/tools/search_handlers.py | 110 ++ doris_mcp_server/tools/tools_manager.py | 6 + doris_mcp_server/utils/search_runtime.py | 1658 ++++++++++++++++++++++++ test/integration/test_real_doris_transports.py | 241 ++++ test/tools/test_capability_detector.py | 113 ++ test/tools/test_domain_dispatcher.py | 69 +- test/tools/test_doris_feature_matrix.py | 25 + test/utils/test_search_runtime.py | 543 ++++++++ 13 files changed, 3211 insertions(+), 69 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7ec61bb..2a96620 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -54,6 +54,9 @@ under **Unreleased** until a new version is selected and published. - A read-only Pipeline domain with five capability-gated children for ingestion status and diagnosis, materialized-view refresh state, recorded table freshness, and bounded upstream or downstream dependency evidence. +- A read-only Search domain with four capability-gated children for + target-index-validated text, vector, and hybrid retrieval, Doris-native + tokenizer previews, search-index inspection, and evidence-based diagnosis. - Real Doris process tests covering Streamable HTTP and stdio. ### Changed @@ -98,6 +101,14 @@ under **Unreleased** until a new version is selected and published. - Isolated Doris domain probe statements in independent route-aware connection contexts so one unsupported version-specific statement cannot poison later capability evidence. +- Kept Search filters, identifiers, vectors, and result fields structured and + bounded, with caller values remaining driver-bound instead of accepting raw + search SQL or Doris `SEARCH` DSL through the structured retrieval child. +- Classified missing Doris Search functions as unsupported capability evidence + and preserved analyzer terms without colliding with credential-token + redaction. +- Bound hybrid Search vector, text, and structured-filter parameters in exact + SQL placeholder order. - Excluded explicitly dead Doris components from active version gating while preserving them in node inventory, and kept live runtime manifests within the 16 KiB domain budget. diff --git a/doris_mcp_server/tools/capability_detector.py b/doris_mcp_server/tools/capability_detector.py index 467c4e5..df67609 100644 --- a/doris_mcp_server/tools/capability_detector.py +++ b/doris_mcp_server/tools/capability_detector.py @@ -38,6 +38,11 @@ from ..utils.doris_http_client import ( configured_fe_http_hosts, database_config_for_request, ) +from ..utils.sql_security_utils import ( + SQLSecurityError, + quote_identifier, + validate_identifier, +) from .doris_feature_matrix import DorisClusterVersionVector from .doris_version import ( DorisVersion, @@ -255,6 +260,34 @@ _DOMAIN_PROBES: Mapping[str, tuple[tuple[str, tuple[str, ...]], ...]] = { ), ), ), + "doris_search": ( + ( + "SHOW INDEX FROM information_schema.tables", + ( + "inverted_index_metadata_readable", + "ann_index_metadata_readable", + "search_index_metadata_readable", + ), + ), + ( + ( + "SELECT TOKENIZE('Doris search', " + "'\"parser\"=\"english\"') AS tokens" + ), + ("tokenize_function_or_analyzer_ready",), + ), + ( + ( + "SELECT l2_distance_approximate([0.0], [0.0]) " + "AS distance" + ), + ("ann_distance_function_readable",), + ), + ( + "EXPLAIN SELECT 1", + ("search_explain_readable",), + ), + ), } @@ -381,6 +414,12 @@ class DorisCapabilityDetector: probes.update(await self._safe_probe_cluster_services(auth_context)) elif domain_name == "doris_pipeline": probes.update(_combine_pipeline_evidence_probes(probes)) + elif domain_name == "doris_search": + match_probe = await self._probe_search_match_syntax( + auth_context + ) + probes[match_probe.probe_id] = match_probe + probes.update(_combine_search_evidence_probes(probes)) completed_route = self.route_identity(auth_context) if completed_route.fingerprint != base.route.fingerprint: raise CapabilityRouteChangedError( @@ -591,6 +630,116 @@ class DorisCapabilityDetector: ) return probes + async def _probe_search_match_syntax( + self, + auth_context: Any | None, + ) -> CapabilityProbeEvidence: + """Probe MATCH against one visible OLAP string column without writes.""" + discovery_sql = ( + "SELECT TABLE_SCHEMA, TABLE_NAME, COLUMN_NAME " + "FROM information_schema.columns " + "WHERE DATA_TYPE IN ('char', 'varchar', 'string', 'text') " + "AND TABLE_SCHEMA NOT IN ('information_schema', 'mysql') " + "ORDER BY TABLE_SCHEMA, TABLE_NAME, ORDINAL_POSITION LIMIT 8" + ) + route = self.route_identity(auth_context) + try: + session_id = ( + "capability-search:match-target:" + f"{route.fingerprint[:12]}" + ) + async with ( + self._connection_manager.get_connection_context_for_auth_context( + session_id, + auth_context, + ) as connection + ): + result = await connection.execute( + discovery_sql, + mask_result=False, + max_rows=8, + max_bytes=32 * 1024, + ) + rows = tuple( + row + for row in (result.data or ()) + if isinstance(row, Mapping) + ) + except Exception as exc: + status, reason = _classify_probe_error(exc) + return CapabilityProbeEvidence( + probe_id="text_match_syntax_readable", + status=status, + reason_code=reason, + ) + + if not rows: + return CapabilityProbeEvidence( + probe_id="text_match_syntax_readable", + status=CapabilityProbeStatus.DEGRADED, + reason_code="SEARCH_MATCH_PROBE_TARGET_NOT_VISIBLE", + ) + + last_evidence: CapabilityProbeEvidence | None = None + for offset, row in enumerate(rows): + try: + database = validate_identifier( + str(_row_value(row, "TABLE_SCHEMA")), + "database name", + ) + table = validate_identifier( + str(_row_value(row, "TABLE_NAME")), + "table name", + ) + column = validate_identifier( + str(_row_value(row, "COLUMN_NAME")), + "column name", + ) + except SQLSecurityError: + continue + # SQL sink audit: all three metadata identifiers pass the strict + # identifier validator and are quoted before _probe_statement sends + # the read-only EXPLAIN to connection.execute. + statement = ( + f"EXPLAIN SELECT {quote_identifier(column, 'column name')} " # nosec B608 + f"FROM {quote_identifier(database, 'database name')}." + f"{quote_identifier(table, 'table name')} " + f"WHERE {quote_identifier(column, 'column name')} " + "MATCH_ANY 'doris' LIMIT 1" + ) + session_id = ( + f"capability-search:match:{offset}:" + f"{route.fingerprint[:12]}" + ) + try: + async with ( + self._connection_manager + .get_connection_context_for_auth_context( + session_id, + auth_context, + ) as connection + ): + evidence = await self._probe_statement( + connection, + statement, + ("text_match_syntax_readable",), + ) + last_evidence = evidence["text_match_syntax_readable"] + except Exception as exc: + status, reason = _classify_probe_error(exc) + last_evidence = CapabilityProbeEvidence( + probe_id="text_match_syntax_readable", + status=status, + reason_code=reason, + ) + if last_evidence.status is CapabilityProbeStatus.SUPPORTED: + return last_evidence + return last_evidence or CapabilityProbeEvidence( + probe_id="text_match_syntax_readable", + status=CapabilityProbeStatus.DEGRADED, + reason_code="SEARCH_MATCH_PROBE_TARGET_INVALID", + ) + async def _safe_probe_profile_api( self, auth_context: Any | None, @@ -923,7 +1072,7 @@ def _classify_probe_error( CapabilityProbeStatus.UNKNOWN, "PROBE_PERMISSION_DENIED", ) - if error_code in {1064, 1109, 1146}: + if error_code in {1064, 1109, 1146, 1305}: return ( CapabilityProbeStatus.UNSUPPORTED, "PROBE_OBJECT_OR_SYNTAX_UNSUPPORTED", @@ -1387,6 +1536,121 @@ def _combine_pipeline_evidence_probes( } +def _combine_search_evidence_probes( + probes: Mapping[str, CapabilityProbeEvidence], +) -> dict[str, CapabilityProbeEvidence]: + text = _combine_all_runtime_probes( + "inverted_index_and_search_syntax_ready", + probes, + ( + "inverted_index_metadata_readable", + "tokenize_function_or_analyzer_ready", + "text_match_syntax_readable", + ), + supported_reason="INVERTED_INDEX_AND_MATCH_READY", + ) + ann = _combine_all_runtime_probes( + "ann_index_and_metric_compatible", + probes, + ( + "ann_index_metadata_readable", + "ann_distance_function_readable", + ), + supported_reason="ANN_INDEX_AND_DISTANCE_READY", + ) + hybrid = _combine_all_evidence( + "hybrid_search", + (text, ann), + supported_reason="TEXT_AND_ANN_SEARCH_READY", + ) + diagnosis = _combine_all_runtime_probes( + "search_index_and_explain_readable", + probes, + ( + "search_index_metadata_readable", + "search_explain_readable", + ), + supported_reason="SEARCH_INDEX_AND_EXPLAIN_READABLE", + ) + plan = CapabilityProbeEvidence( + probe_id="search_plan_facets_readable", + status=diagnosis.status, + reason_code=diagnosis.reason_code, + evidence_sources=diagnosis.evidence_sources, + ) + ann_feature = CapabilityProbeEvidence( + probe_id="ann_index", + status=ann.status, + reason_code=ann.reason_code, + evidence_sources=ann.evidence_sources, + ) + return { + text.probe_id: text, + ann.probe_id: ann, + ann_feature.probe_id: ann_feature, + hybrid.probe_id: hybrid, + diagnosis.probe_id: diagnosis, + plan.probe_id: plan, + } + + +def _combine_all_runtime_probes( + probe_id: str, + probes: Mapping[str, CapabilityProbeEvidence], + candidate_ids: Sequence[str], + *, + supported_reason: str, +) -> CapabilityProbeEvidence: + candidates = tuple( + probe + for candidate_id in candidate_ids + if (probe := probes.get(candidate_id)) is not None + ) + return _combine_all_evidence( + probe_id, + candidates, + supported_reason=supported_reason, + ) + + +def _combine_all_evidence( + probe_id: str, + candidates: Sequence[CapabilityProbeEvidence], + *, + supported_reason: str, +) -> CapabilityProbeEvidence: + if candidates and all( + probe.status is CapabilityProbeStatus.SUPPORTED + for probe in candidates + ): + status = CapabilityProbeStatus.SUPPORTED + reason = supported_reason + else: + status = next( + ( + candidate_status + for candidate_status in ( + CapabilityProbeStatus.MISCONFIGURED, + CapabilityProbeStatus.UNSUPPORTED, + CapabilityProbeStatus.UNKNOWN, + CapabilityProbeStatus.DEGRADED, + ) + if any( + probe.status is candidate_status + for probe in candidates + ) + ), + CapabilityProbeStatus.UNKNOWN, + ) + reason = f"{probe_id.upper()}_{status.value.upper()}" + return CapabilityProbeEvidence( + probe_id=probe_id, + status=status, + reason_code=reason, + evidence_sources=("runtime_probe",), + ) + + def _combine_any_runtime_probe( probe_id: str, probes: Mapping[str, CapabilityProbeEvidence], diff --git a/doris_mcp_server/tools/domain_catalog.py b/doris_mcp_server/tools/domain_catalog.py index 8c55845..6bd4d58 100644 --- a/doris_mcp_server/tools/domain_catalog.py +++ b/doris_mcp_server/tools/domain_catalog.py @@ -492,6 +492,108 @@ def _array(description: str) -> dict[str, Any]: } +def _search_filters() -> dict[str, Any]: + scalar = {"type": ["string", "number", "integer", "boolean", "null"]} + return { + "type": "object", + "description": ( + "Structured filters keyed by column. A scalar means equality; an " + "object uses operator plus value or values. Supported operators: " + "eq, ne, gt, gte, lt, lte, in, not_in, is_null, is_not_null." + ), + "maxProperties": 32, + "propertyNames": { + "pattern": r"^[A-Za-z_\u4e00-\u9fff]" + r"[A-Za-z0-9_\u4e00-\u9fff]{0,63}$" + }, + "additionalProperties": { + "oneOf": [ + scalar, + { + "type": "object", + "properties": { + "operator": _string( + "Allowlisted filter operator.", + enum=( + "eq", + "ne", + "gt", + "gte", + "lt", + "lte", + "in", + "not_in", + "is_null", + "is_not_null", + ), + ), + "value": scalar, + "values": { + "type": "array", + "items": scalar, + "minItems": 1, + "maxItems": 100, + }, + }, + "required": ["operator"], + "additionalProperties": False, + }, + ] + }, + } + + +def _search_input_schema() -> dict[str, Any]: + return _input_schema( + { + "database": _string("Database name."), + "table": _string("Table name."), + "query": _string( + "Text query for text or hybrid mode.", + max_length=16_384, + ), + "mode": _string( + "Search mode.", + enum=("text", "vector", "hybrid"), + ), + "fields": { + **_string_array( + "Inverted-indexed text fields.", + ), + "maxItems": 32, + }, + "vector": { + "type": "array", + "description": ( + "Finite query-vector values. The length must match the " + "selected ANN index dimension." + ), + "items": {"type": "number"}, + "minItems": 1, + "maxItems": 4_096, + }, + "vector_field": _string( + "ANN-indexed vector field; optional only when one is visible." + ), + "text_operator": _string( + "Token matching behavior.", + enum=("any", "all", "phrase", "phrase_prefix"), + ), + "top_k": _integer( + "Maximum matches.", + minimum=1, + maximum=1_000, + ), + "filters": _search_filters(), + "return_fields": { + **_string_array("Visible return fields."), + "maxItems": 64, + }, + }, + required=("database", "table", "mode"), + ) + + def _input_schema( properties: dict[str, Any], *, @@ -591,31 +693,31 @@ _COLLECTION_OUTPUT = _result_schema( "additionalProperties": False, } ) -_QUERY_OUTPUT = _result_schema( - { - "type": "object", - "properties": { - "columns": { - "type": "array", - "items": { - "type": "object", - "additionalProperties": True, - }, +_QUERY_DATA_SCHEMA = { + "type": "object", + "properties": { + "columns": { + "type": "array", + "items": { + "type": "object", + "additionalProperties": True, }, - "rows": { - "type": "array", - "items": { - "type": "object", - "additionalProperties": True, - }, + }, + "rows": { + "type": "array", + "items": { + "type": "object", + "additionalProperties": True, }, - "row_count": {"type": "integer", "minimum": 0}, - "truncated": {"type": "boolean"}, }, - "required": ["columns", "rows", "row_count", "truncated"], - "additionalProperties": False, - } -) + "row_count": {"type": "integer", "minimum": 0}, + "truncated": {"type": "boolean"}, + }, + "required": ["columns", "rows", "row_count", "truncated"], + "additionalProperties": False, +} +_QUERY_OUTPUT = _result_schema(_QUERY_DATA_SCHEMA) +_SEARCH_QUERY_OUTPUT = _result_schema(_QUERY_DATA_SCHEMA, evidence=True) _TABLE_CONTEXT_SECTION_OUTPUT = { "type": "object", "properties": { @@ -1456,28 +1558,8 @@ DOMAIN_DEFINITIONS = ( "search_data", "Search data", "Run bounded structured text, vector, or hybrid retrieval.", - _input_schema( - { - "database": _string("Database name."), - "table": _string("Table name."), - "query": _string("Text query."), - "mode": _string( - "Search mode.", - enum=("text", "vector", "hybrid"), - ), - "fields": _string_array("Searchable fields."), - "vector": _array("Query vector."), - "top_k": _integer( - "Maximum matches.", - minimum=1, - maximum=1000, - ), - "filters": _object("Structured filters."), - "return_fields": _string_array("Visible return fields."), - }, - required=("database", "table", "mode"), - ), - _QUERY_OUTPUT, + _search_input_schema(), + _SEARCH_QUERY_OUTPUT, ), _child( "doris_search", @@ -1487,12 +1569,22 @@ DOMAIN_DEFINITIONS = ( _input_schema( { "text": _string("Input text."), - "analyzer": _string("Analyzer name."), - "tokenizer": _string("Tokenizer name."), - "token_filters": _string_array("Token filters."), + "analyzer": _string( + "Existing built-in or custom analyzer name." + ), + "tokenizer": _string( + "Built-in tokenizer or expected custom component." + ), + "token_filters": { + **_string_array( + "Expected filters on an existing custom analyzer." + ), + "maxItems": 16, + }, }, required=("text",), ), + _DIAGNOSTIC_OUTPUT, ), _child( "doris_search", @@ -1507,6 +1599,7 @@ DOMAIN_DEFINITIONS = ( }, required=("database", "table"), ), + _DIAGNOSTIC_OUTPUT, ), _child( "doris_search", @@ -1516,7 +1609,7 @@ DOMAIN_DEFINITIONS = ( _input_schema( { "sql": _string("Read-only search SQL."), - "search_request": _object("Structured search request."), + "search_request": _search_input_schema(), "include_profile": _boolean("Include query-profile evidence."), }, any_of=( diff --git a/doris_mcp_server/tools/domain_dispatcher.py b/doris_mcp_server/tools/domain_dispatcher.py index 41e1ae3..898fcce 100644 --- a/doris_mcp_server/tools/domain_dispatcher.py +++ b/doris_mcp_server/tools/domain_dispatcher.py @@ -47,6 +47,7 @@ from ..utils.cluster_runtime import ClusterRuntimeFailure from ..utils.logger import get_audit_logger, get_logger from ..utils.pipeline_runtime import PipelineRuntimeFailure from ..utils.query_runtime import QueryRuntimeFailure +from ..utils.search_runtime import SearchRuntimeFailure from . import domain_catalog as domain_catalog_module from .domain_catalog import ( DorisDomainCatalog, @@ -685,6 +686,31 @@ class DomainDispatcher: "status_code": exc.status_code, }, ) + except SearchRuntimeFailure as exc: + self._audit(feature_id, arguments, "error", started) + return self._error( + domain.name, + ( + DomainErrorCode.CHILD_ARGUMENTS_INVALID + if exc.reason_code + in { + "SEARCH_ARGUMENT_INVALID", + "SEARCH_ANALYZER_NOT_FOUND", + "SEARCH_INDEX_NOT_FOUND", + "SEARCH_TABLE_NOT_FOUND", + } + else 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( diff --git a/doris_mcp_server/tools/doris_feature_matrix.py b/doris_mcp_server/tools/doris_feature_matrix.py index 1067e80..1b09ac4 100644 --- a/doris_mcp_server/tools/doris_feature_matrix.py +++ b/doris_mcp_server/tools/doris_feature_matrix.py @@ -1254,17 +1254,22 @@ FEATURE_DEFINITIONS = ( "doris_search", "search_data", A, - _variant( - "inverted_text_search", - probes=("inverted_index_and_search_syntax_ready",), - ), _variant( "vector_hybrid_search", ranges=(">=4.0.0",), features=("ann_index", "hybrid_search"), - probes=("ann_index_and_metric_compatible",), + probes=( + "ann_index_and_metric_compatible", + "inverted_index_and_search_syntax_ready", + ), + callable_when_degraded=True, sources=("DORIS_RELEASE_4_0_0",), ), + _variant( + "inverted_text_search", + probes=("inverted_index_and_search_syntax_ready",), + callable_when_degraded=True, + ), ), _feature( "doris_search", @@ -1279,10 +1284,6 @@ FEATURE_DEFINITIONS = ( "doris_search", "inspect_search_indexes", M, - _variant( - "inverted_index_metadata", - probes=("inverted_index_metadata_readable",), - ), _variant( "ann_index_metadata", ranges=(">=4.0.0",), @@ -1290,6 +1291,10 @@ FEATURE_DEFINITIONS = ( probes=("ann_index_metadata_readable",), sources=("DORIS_RELEASE_4_0_0",), ), + _variant( + "inverted_index_metadata", + probes=("inverted_index_metadata_readable",), + ), ), _feature( "doris_search", diff --git a/doris_mcp_server/tools/search_handlers.py b/doris_mcp_server/tools/search_handlers.py new file mode 100644 index 0000000..9b22800 --- /dev/null +++ b/doris_mcp_server/tools/search_handlers.py @@ -0,0 +1,110 @@ +# 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. + +"""Formal Search-domain handlers backed by one strict runtime.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any, Protocol, cast + +from ..utils.db import DorisConnectionManager +from ..utils.query_runtime import DorisQueryRuntime +from ..utils.search_runtime import DorisSearchRuntime + + +class _SearchHandlerOwner(Protocol): + """State supplied by ``DorisToolsManager`` to the mixin.""" + + search_runtime: DorisSearchRuntime + + +class SearchToolHandlersMixin: + """Route every Search child through the Doris-native runtime.""" + + def _initialize_search_handlers( + self: _SearchHandlerOwner, + connection_manager: DorisConnectionManager, + query_runtime: DorisQueryRuntime, + ) -> None: + self.search_runtime = DorisSearchRuntime( + connection_manager, + query_runtime, + ) + + async def _formal_doris_search_search_data_tool( + self: _SearchHandlerOwner, + arguments: dict[str, Any], + ) -> dict[str, Any]: + return await self.search_runtime.search_data( + database=cast(str, arguments.get("database")), + table=cast(str, arguments.get("table")), + query=cast(str | None, arguments.get("query")), + mode=cast(str, arguments.get("mode")), + fields=cast(list[str] | None, arguments.get("fields")), + vector=cast(list[Any] | None, arguments.get("vector")), + vector_field=cast(str | None, arguments.get("vector_field")), + text_operator=cast(str | None, arguments.get("text_operator")), + top_k=cast(int | None, arguments.get("top_k")), + filters=cast( + Mapping[str, Any] | None, + arguments.get("filters"), + ), + return_fields=cast( + list[str] | None, + arguments.get("return_fields"), + ), + ) + + async def _formal_doris_search_preview_text_analysis_tool( + self: _SearchHandlerOwner, + arguments: dict[str, Any], + ) -> dict[str, Any]: + return await self.search_runtime.preview_text_analysis( + text=cast(str, arguments.get("text")), + analyzer=cast(str | None, arguments.get("analyzer")), + tokenizer=cast(str | None, arguments.get("tokenizer")), + token_filters=cast( + list[str] | None, + arguments.get("token_filters"), + ), + ) + + async def _formal_doris_search_inspect_search_indexes_tool( + self: _SearchHandlerOwner, + arguments: dict[str, Any], + ) -> dict[str, Any]: + return await self.search_runtime.inspect_search_indexes( + database=cast(str, arguments.get("database")), + table=cast(str, arguments.get("table")), + index=cast(str | None, arguments.get("index")), + ) + + async def _formal_doris_search_diagnose_search_query_tool( + self: _SearchHandlerOwner, + arguments: dict[str, Any], + ) -> dict[str, Any]: + return await self.search_runtime.diagnose_search_query( + sql=cast(str | None, arguments.get("sql")), + search_request=cast( + Mapping[str, Any] | None, + arguments.get("search_request"), + ), + include_profile=bool(arguments.get("include_profile", False)), + ) + + +__all__ = ["SearchToolHandlersMixin"] diff --git a/doris_mcp_server/tools/tools_manager.py b/doris_mcp_server/tools/tools_manager.py index 78dfa73..e92b109 100644 --- a/doris_mcp_server/tools/tools_manager.py +++ b/doris_mcp_server/tools/tools_manager.py @@ -65,6 +65,7 @@ from .domain_manifest import ( from .doris_feature_matrix import DORIS_FEATURE_MATRIX from .pipeline_handlers import PipelineToolHandlersMixin from .query_handlers import QueryToolHandlersMixin +from .search_handlers import SearchToolHandlersMixin from .tool_provider import CustomToolProvider, ToolProviderRuntime from .tool_registry import ToolRegistryError @@ -76,6 +77,7 @@ class DorisToolsManager( CatalogToolHandlersMixin, ClusterToolHandlersMixin, PipelineToolHandlersMixin, + SearchToolHandlersMixin, DomainManifestManagerMixin, ): """Apache Doris Tools Manager""" @@ -120,6 +122,10 @@ class DorisToolsManager( connection_manager, self.adbc_query_tools, ) + self._initialize_search_handlers( + connection_manager, + self.query_runtime, + ) self._capability_registry: CapabilityRegistry | None = None if domain_availability_provider is None: bound_handlers = BoundHandlerAvailabilityProvider(self) diff --git a/doris_mcp_server/utils/search_runtime.py b/doris_mcp_server/utils/search_runtime.py new file mode 100644 index 0000000..921c510 --- /dev/null +++ b/doris_mcp_server/utils/search_runtime.py @@ -0,0 +1,1658 @@ +# 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, evidence-bearing runtime for the read-only Search domain.""" + +from __future__ import annotations + +import json +import math +import re +import uuid +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Any + +from .db import DorisConnectionManager, QueryResult +from .query_runtime import ( + DorisQueryRuntime, + QueryRuntimeFailure, + ReadOnlySQLGuard, +) +from .redaction import redact_sensitive_data +from .security import get_current_auth_context +from .sql_security_utils import ( + SQLSecurityError, + build_table_reference, + quote_identifier, + validate_identifier, +) + +_MAX_BYTES = 2 * 1024 * 1024 +_MAX_QUERY_TEXT_BYTES = 16 * 1024 +_MAX_PREVIEW_TEXT_BYTES = 64 * 1024 +_MAX_TOP_K = 1_000 +_DEFAULT_TOP_K = 10 +_MAX_VECTOR_DIMENSION = 4_096 +_MAX_FIELDS = 32 +_MAX_RETURN_FIELDS = 64 +_MAX_FILTERS = 32 +_MAX_FILTER_VALUES = 100 +_MAX_EXPLAIN_ROWS = 2_000 +_BUILT_IN_ANALYZERS = frozenset( + { + "none", + "standard", + "english", + "chinese", + "unicode", + "icu", + "basic", + "ik", + } +) +_BACKWARD_COMPATIBLE_PARSERS = frozenset( + { + "english", + "chinese", + "unicode", + } +) +_MATCH_OPERATORS = { + "any": "MATCH_ANY", + "all": "MATCH_ALL", + "phrase": "MATCH_PHRASE", + "phrase_prefix": "MATCH_PHRASE_PREFIX", +} +_FILTER_OPERATORS = { + "eq": "=", + "ne": "!=", + "gt": ">", + "gte": ">=", + "lt": "<", + "lte": "<=", +} +_NULL_FILTER_OPERATORS = { + "is_null": "IS NULL", + "is_not_null": "IS NOT NULL", +} +_SET_FILTER_OPERATORS = { + "in": "IN", + "not_in": "NOT IN", +} +_VECTOR_METRICS = { + "l2_distance": ("l2_distance_approximate", "ASC"), + "inner_product": ("inner_product_approximate", "DESC"), +} +_PROPERTY_PAIR = re.compile( + r'"(?P<key>[A-Za-z_][A-Za-z0-9_]*)"\s*=\s*"(?P<value>(?:[^"\\]|\\.)*)"' +) +_SIMPLE_SOURCE = re.compile( + r"\bFROM\s+" + r"(?:(?P<database>`?[A-Za-z_][A-Za-z0-9_]*`?)\s*\.\s*)?" + r"(?P<table>`?[A-Za-z_][A-Za-z0-9_]*`?)", + re.IGNORECASE, +) + + +class SearchRuntimeFailure(RuntimeError): + """Sanitized Search-domain failure with a stable reason code.""" + + 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 + + +@dataclass(frozen=True, slots=True) +class _SearchIndex: + name: str + index_type: str + columns: tuple[str, ...] + properties: Mapping[str, str] + comment: str | None + + @property + def is_inverted(self) -> bool: + return self.index_type == "INVERTED" + + @property + def is_ann(self) -> bool: + return self.index_type == "ANN" + + def to_wire(self) -> dict[str, Any]: + return { + "name": self.name, + "index_type": self.index_type, + "columns": list(self.columns), + "properties": dict(self.properties), + "comment": self.comment, + "parser": self.properties.get("parser"), + "analyzer": ( + self.properties.get("analyzer") + or self.properties.get("built_in_analyzer") + ), + "metric_type": self.properties.get("metric_type"), + "dimension": _as_int(self.properties.get("dim")), + } + + +@dataclass(frozen=True, slots=True) +class _CompiledSearch: + sql: str + params: tuple[Any, ...] + database: str + table: str + mode: str + top_k: int + text_fields: tuple[str, ...] + vector_field: str | None + vector_metric: str | None + vector_dimension: int | None + return_fields: tuple[str, ...] + indexes: tuple[_SearchIndex, ...] + + +class DorisSearchRuntime: + """Run bounded Doris-native text, vector, hybrid, and diagnostic reads.""" + + def __init__( + self, + connection_manager: DorisConnectionManager, + query_runtime: DorisQueryRuntime, + ) -> None: + self._connection_manager = connection_manager + self._query_runtime = query_runtime + self._session_prefix = f"formal_search_{uuid.uuid4().hex[:8]}" + + async def search_data( + self, + *, + database: str, + table: str, + query: str | None, + mode: str, + fields: Sequence[str] | None, + vector: Sequence[Any] | None, + vector_field: str | None, + text_operator: str | None, + top_k: int | None, + filters: Mapping[str, Any] | None, + return_fields: Sequence[str] | None, + ) -> dict[str, Any]: + """Execute one target-index-validated, bounded search request.""" + compiled = await self._compile_search( + { + "database": database, + "table": table, + "query": query, + "mode": mode, + "fields": fields, + "vector": vector, + "vector_field": vector_field, + "text_operator": text_operator, + "top_k": top_k, + "filters": filters, + "return_fields": return_fields, + } + ) + result = await self._execute( + compiled.sql, + params=compiled.params, + max_rows=compiled.top_k + 1, + mask_result=True, + ) + rows = [ + dict(row) + for row in (result.data or ()) + if isinstance(row, Mapping) + ] + driver_truncated = bool(result.metadata.get("truncated")) + overfetched = len(rows) > compiled.top_k + visible_rows = rows[: compiled.top_k] + columns = [ + {"name": str(name)} + for name in result.metadata.get( + "columns", + list(visible_rows[0]) if visible_rows else [], + ) + ] + warnings: list[str] = [] + if driver_truncated: + warnings.append( + "Doris stopped reading at the configured response boundary." + ) + evidence: list[Mapping[str, Any]] = [ + { + "source": "SHOW INDEX", + "kind": "target_index_metadata", + "indexes": [ + index.name + for index in compiled.indexes + if ( + index.is_inverted + and set(index.columns) & set(compiled.text_fields) + ) + or ( + index.is_ann + and compiled.vector_field in index.columns + ) + ], + }, + { + "source": "doris_sql", + "kind": "bounded_search", + "mode": compiled.mode, + }, + ] + return _result( + { + "columns": columns, + "rows": visible_rows, + "row_count": len(visible_rows), + "truncated": driver_truncated or overfetched, + }, + source="doris_native_search", + warnings=warnings, + metadata={ + "database": compiled.database, + "table": compiled.table, + "mode": compiled.mode, + "top_k": compiled.top_k, + "text_fields": list(compiled.text_fields), + "vector_field": compiled.vector_field, + "vector_metric": compiled.vector_metric, + "vector_dimension": compiled.vector_dimension, + "return_fields": list(compiled.return_fields), + "invented_scores": False, + "execution_time_seconds": result.execution_time, + }, + evidence=evidence, + ) + + async def preview_text_analysis( + self, + *, + text: str, + analyzer: str | None, + tokenizer: str | None, + token_filters: Sequence[str] | None, + ) -> dict[str, Any]: + """Run Doris ``TOKENIZE`` without creating analyzer components.""" + normalized_text = _required_text( + text, + "text", + maximum_bytes=_MAX_PREVIEW_TEXT_BYTES, + ) + requested_filters = _identifier_sequence( + token_filters, + "token filter", + maximum=16, + ) + properties: dict[str, str] + component_evidence: list[dict[str, Any]] = [] + + if analyzer is not None: + analyzer_name = _required_identifier(analyzer, "analyzer name") + analyzer_key = analyzer_name.casefold() + if analyzer_key in _BUILT_IN_ANALYZERS: + if tokenizer is not None or requested_filters: + raise _argument_failure( + "Built-in analyzers cannot be combined with custom " + "tokenizer or token-filter names." + ) + properties = _built_in_analyzer_properties(analyzer_key) + else: + definition = await self._custom_analyzer_definition( + analyzer_name + ) + if definition is None: + raise SearchRuntimeFailure( + "The requested custom Doris analyzer does not exist.", + reason_code="SEARCH_ANALYZER_NOT_FOUND", + status_code=404, + ) + observed_tokenizer = _optional_identifier_value( + _value(definition, "tokenizer") + ) + observed_filters = _split_component_names( + _value( + definition, + "token_filter", + "token_filters", + "tokenfilters", + ) + ) + if tokenizer is not None: + tokenizer_name = _required_identifier( + tokenizer, + "tokenizer name", + ) + if tokenizer_name.casefold() != ( + observed_tokenizer or "" + ).casefold(): + raise _argument_failure( + "The requested tokenizer does not match the " + "recorded custom analyzer definition." + ) + if requested_filters and tuple( + item.casefold() for item in requested_filters + ) != tuple(item.casefold() for item in observed_filters): + raise _argument_failure( + "The requested token filters do not match the recorded " + "custom analyzer definition." + ) + properties = {"analyzer": analyzer_name} + component_evidence.append( + { + "source": "SHOW INVERTED INDEX ANALYZER", + "analyzer": analyzer_name, + "tokenizer": observed_tokenizer, + "token_filters": list(observed_filters), + } + ) + else: + tokenizer_name = ( + "unicode" + if tokenizer is None + else _required_identifier(tokenizer, "tokenizer name").casefold() + ) + if tokenizer_name not in _BUILT_IN_ANALYZERS: + raise _argument_failure( + "A custom tokenizer must be referenced through an existing " + "custom analyzer." + ) + if requested_filters: + raise _argument_failure( + "Token filters require an existing custom analyzer; this " + "read-only tool does not create analyzer components." + ) + properties = _built_in_analyzer_properties(tokenizer_name) + + property_string = ",".join( + f'"{key}"="{_escape_property_value(value)}"' + for key, value in properties.items() + ) + result = await self._execute( + "SELECT TOKENIZE(%s, %s) AS tokens", + params=(normalized_text, property_string), + max_rows=1, + mask_result=False, + ) + raw_tokens = ( + _value(result.data[0], "tokens") + if result.data and isinstance(result.data[0], Mapping) + else None + ) + tokens = _normalize_tokens(raw_tokens) + return _result( + { + "tokens": tokens, + "token_count": len(tokens), + "analysis": { + "properties": properties, + "custom_analyzer": properties.get("analyzer"), + }, + }, + source="doris_tokenize", + metadata={ + "input_bytes": len(normalized_text.encode("utf-8")), + "invented_tokens": False, + }, + evidence=[ + { + "source": "doris_sql", + "kind": "TOKENIZE", + "properties": properties, + }, + *component_evidence, + ], + ) + + async def inspect_search_indexes( + self, + *, + database: str, + table: str, + index: str | None, + ) -> dict[str, Any]: + """Return normalized inverted/ANN metadata and recorded build tasks.""" + database_name = _required_identifier(database, "database name") + table_name = _required_identifier(table, "table name") + index_name = ( + None + if index is None + else _required_identifier(index, "index name") + ) + indexes = await self._read_indexes(database_name, table_name) + selected = tuple( + candidate + for candidate in indexes + if index_name is None + or candidate.name.casefold() == index_name.casefold() + ) + if index_name is not None and not selected: + raise SearchRuntimeFailure( + "The requested Doris search index does not exist.", + reason_code="SEARCH_INDEX_NOT_FOUND", + status_code=404, + ) + + build_tasks: list[dict[str, Any]] = [] + warnings: list[str] = [] + try: + build_result = await self._execute( + "SHOW BUILD INDEX WHERE TableName = %s", + params=(table_name,), + max_rows=200, + mask_result=False, + database_context=database_name, + ) + except SearchRuntimeFailure as exc: + warnings.append( + "Index build-task evidence is unavailable " + f"({exc.reason_code})." + ) + else: + build_tasks = [ + _normalized_build_task(row) + for row in (build_result.data or ()) + if isinstance(row, Mapping) + ] + + inverted = [candidate for candidate in selected if candidate.is_inverted] + ann = [candidate for candidate in selected if candidate.is_ann] + capabilities = { + "text": bool(inverted), + "vector": bool(ann), + "hybrid": bool(inverted and ann), + "metrics": sorted( + { + metric + for candidate in ann + if ( + metric := candidate.properties.get("metric_type") + ) + } + ), + } + return _result( + { + "items": [candidate.to_wire() for candidate in selected], + "build_tasks": build_tasks, + "capabilities": capabilities, + "truncated": False, + }, + source="doris_search_index_metadata", + warnings=warnings, + metadata={ + "database": database_name, + "table": table_name, + "index_filter": index_name, + "index_count": len(selected), + "build_task_count": len(build_tasks), + }, + evidence=[ + { + "source": "SHOW INDEX", + "rows_observed": len(indexes), + }, + { + "source": "SHOW BUILD INDEX", + "rows_observed": len(build_tasks), + "available": not warnings, + }, + ], + ) + + async def diagnose_search_query( + self, + *, + sql: str | None, + search_request: Mapping[str, Any] | None, + include_profile: bool, + ) -> dict[str, Any]: + """Combine target index metadata, EXPLAIN, and optional profile facts.""" + if bool(sql) == bool(search_request): + raise _argument_failure( + "Provide exactly one of sql or search_request." + ) + + compiled: _CompiledSearch | None = None + database: str | None = None + table: str | None = None + params: tuple[Any, ...] | None = None + warnings: list[str] = [] + evidence: list[dict[str, Any]] = [] + indexes: tuple[_SearchIndex, ...] = () + + if search_request is not None: + compiled = await self._compile_search(search_request) + query_sql = compiled.sql + params = compiled.params + database = compiled.database + table = compiled.table + indexes = compiled.indexes + else: + try: + statement = ReadOnlySQLGuard.validate( + str(sql), + query_target=True, + ) + except QueryRuntimeFailure as exc: + raise SearchRuntimeFailure( + str(exc), + reason_code="SEARCH_ARGUMENT_INVALID", + status_code=400, + ) from exc + if statement.operation != "SELECT": + raise _argument_failure( + "Search diagnosis accepts a SELECT query only." + ) + query_sql = statement.sql + source = _simple_source(statement.sql) + if source is not None: + database, table = source + if database is not None: + try: + indexes = await self._read_indexes(database, table) + except SearchRuntimeFailure as exc: + warnings.append( + "Target index metadata is unavailable " + f"({exc.reason_code})." + ) + if not indexes: + warnings.append( + "Raw SQL diagnosis could not bind authoritative target " + "index metadata; provide search_request for exact coverage." + ) + + explain = await self._execute( + f"EXPLAIN {query_sql}", + params=params, + max_rows=_MAX_EXPLAIN_ROWS, + mask_result=False, + database_context=database, + ) + plan_rows = [ + dict(row) + for row in (explain.data or ()) + if isinstance(row, Mapping) + ] + plan_text = "\n".join( + " | ".join(str(value) for value in row.values()) + for row in plan_rows + ) + facets = _search_plan_facets(plan_text) + evidence.append( + { + "source": "doris_sql", + "kind": "EXPLAIN", + "plan_rows": len(plan_rows), + } + ) + if indexes: + evidence.append( + { + "source": "SHOW INDEX", + "kind": "target_index_metadata", + "indexes": [candidate.name for candidate in indexes], + } + ) + + findings = _diagnostic_findings( + compiled=compiled, + indexes=indexes, + facets=facets, + query_sql=query_sql, + ) + profile: Mapping[str, Any] | None = None + if include_profile: + if search_request is not None: + warnings.append( + "Profile execution is not performed for parameter-bound " + "structured requests; EXPLAIN and index evidence are " + "returned without weakening parameter safety." + ) + else: + try: + profile_result = await self._query_runtime.get_query_profile( + sql=query_sql, + database=database, + include_operator_tree=False, + ) + except QueryRuntimeFailure as exc: + warnings.append( + "Query-profile evidence is unavailable " + f"({exc.reason_code})." + ) + else: + profile = profile_result.get("data") + evidence.extend(profile_result.get("evidence", [])) + warnings.extend(profile_result.get("warnings", [])) + + if not facets["ann_pushdown_observed"] and compiled is not None and ( + compiled.mode in {"vector", "hybrid"} + ): + warnings.append( + "The plan did not expose Doris ANN SORT INFO for the " + "requested vector path." + ) + return _result( + { + "scope": { + "database": database, + "table": table, + "mode": compiled.mode if compiled is not None else "raw_sql", + }, + "indexes": [candidate.to_wire() for candidate in indexes], + "explain": { + "facets": facets, + "plan_rows": plan_rows, + "truncated": bool(explain.metadata.get("truncated")), + }, + "profile": profile, + "findings": findings, + }, + source="deterministic_search_diagnosis", + warnings=warnings, + metadata={ + "include_profile": include_profile, + "profile_observed": profile is not None, + "invented_index_hits": False, + "execution_time_seconds": explain.execution_time, + }, + evidence=evidence, + ) + + async def _compile_search( + self, + request: Mapping[str, Any], + ) -> _CompiledSearch: + if not isinstance(request, Mapping): + raise _argument_failure("search_request must be an object.") + database = _required_identifier( + request.get("database"), + "database name", + ) + table = _required_identifier(request.get("table"), "table name") + mode = str(request.get("mode", "")).casefold() + if mode not in {"text", "vector", "hybrid"}: + raise _argument_failure( + "mode must be text, vector, or hybrid." + ) + top_k = _bounded_integer( + request.get("top_k"), + default=_DEFAULT_TOP_K, + minimum=1, + maximum=_MAX_TOP_K, + label="top_k", + ) + columns = await self._read_columns(database, table) + column_names = set(columns) + indexes = await self._read_indexes(database, table) + + text_fields: tuple[str, ...] = () + query: str | None = None + match_operator = _MATCH_OPERATORS.get( + str(request.get("text_operator") or "any").casefold() + ) + if match_operator is None: + raise _argument_failure("text_operator is invalid.") + if mode in {"text", "hybrid"}: + query = _required_text( + request.get("query"), + "query", + maximum_bytes=_MAX_QUERY_TEXT_BYTES, + ) + text_fields = _identifier_sequence( + request.get("fields"), + "search field", + maximum=_MAX_FIELDS, + required=True, + ) + _require_known_columns(text_fields, column_names) + indexed_text_fields = { + column + for candidate in indexes + if candidate.is_inverted + for column in candidate.columns + } + missing_indexes = [ + field + for field in text_fields + if field not in indexed_text_fields + ] + if missing_indexes: + raise SearchRuntimeFailure( + "Every requested text field must have a visible inverted " + "index.", + reason_code="SEARCH_TEXT_INDEX_REQUIRED", + status_code=409, + ) + elif request.get("query") not in (None, ""): + raise _argument_failure( + "query is only valid for text or hybrid mode." + ) + + vector_field: str | None = None + vector_values: tuple[float, ...] = () + metric: str | None = None + vector_dimension: int | None = None + if mode in {"vector", "hybrid"}: + vector_values = _vector(request.get("vector")) + vector_field = _resolve_vector_field( + request.get("vector_field"), + indexes, + ) + if vector_field not in column_names: + raise _argument_failure( + "The requested vector field is not visible on the table." + ) + ann_index = _ann_index_for_field(indexes, vector_field) + if ann_index is None: + raise SearchRuntimeFailure( + "The requested vector field must have a visible ANN index.", + reason_code="SEARCH_ANN_INDEX_REQUIRED", + status_code=409, + ) + metric = str( + ann_index.properties.get("metric_type", "") + ).casefold() + if metric not in _VECTOR_METRICS: + raise SearchRuntimeFailure( + "The ANN index uses an unsupported or missing metric.", + reason_code="SEARCH_ANN_METRIC_UNSUPPORTED", + status_code=501, + ) + vector_dimension = _as_int(ann_index.properties.get("dim")) + if vector_dimension is None: + raise SearchRuntimeFailure( + "The ANN index does not expose a valid vector dimension.", + reason_code="SEARCH_ANN_DIMENSION_UNKNOWN", + status_code=409, + ) + if len(vector_values) != vector_dimension: + raise _argument_failure( + "The query vector dimension does not match the ANN index." + ) + elif request.get("vector") not in (None, []): + raise _argument_failure( + "vector is only valid for vector or hybrid mode." + ) + + return_fields = _return_fields( + request.get("return_fields"), + columns, + vector_field=vector_field, + ) + select_parts = [ + quote_identifier(field, "return field") + for field in return_fields + ] + where_parts: list[str] = [] + select_params: list[Any] = [] + where_params: list[Any] = [] + if text_fields and query is not None: + field_predicates = [ + f"{quote_identifier(field, 'search field')} " + f"{match_operator} %s" + for field in text_fields + ] + where_parts.append("(" + " OR ".join(field_predicates) + ")") + where_params.extend(query for _ in field_predicates) + + filter_sql, filter_params = _compile_filters( + request.get("filters"), + column_names, + ) + where_parts.extend(filter_sql) + where_params.extend(filter_params) + + order_sql = "" + if vector_field is not None and metric is not None: + function_name, direction = _VECTOR_METRICS[metric] + distance_alias = "__mcp_vector_distance" + if distance_alias in column_names: + raise SearchRuntimeFailure( + "The table uses a reserved Search result alias.", + reason_code="SEARCH_RESULT_ALIAS_CONFLICT", + status_code=409, + ) + vector_expression = "CAST(%s AS ARRAY<FLOAT>)" + # SQL sink audit: field is validated and quoted; function, alias, + # and ordering are selected from fixed local maps; the vector is a + # bound JSON value before DorisConnection.execute. + select_parts.append( # nosec B608 + f"{function_name}(" + f"{quote_identifier(vector_field, 'vector field')}, " + f"{vector_expression}) AS `{distance_alias}`" + ) + select_params.append( + json.dumps( + list(vector_values), + ensure_ascii=True, + separators=(",", ":"), + ) + ) + order_sql = f" ORDER BY `{distance_alias}` {direction}" + + table_reference = build_table_reference( + table, + db_name=database, + ) + where_sql = ( + " WHERE " + " AND ".join(where_parts) + if where_parts + else "" + ) + # SQL sink audit: all identifiers pass quote_identifier or + # build_table_reference; predicates/operators are local allowlists; + # every caller value remains bound at connection.execute; limit is a + # bounded integer. + sql = ( + f"SELECT {', '.join(select_parts)} FROM {table_reference}" # nosec B608 + f"{where_sql}{order_sql} LIMIT {top_k + 1}" + ) + return _CompiledSearch( + sql=sql, + params=(*select_params, *where_params), + database=database, + table=table, + mode=mode, + top_k=top_k, + text_fields=text_fields, + vector_field=vector_field, + vector_metric=metric, + vector_dimension=vector_dimension, + return_fields=return_fields, + indexes=indexes, + ) + + async def _read_columns( + self, + database: str, + table: str, + ) -> dict[str, str]: + result = await self._execute( + ( + "SELECT COLUMN_NAME, DATA_TYPE " + "FROM information_schema.columns " + "WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s " + "ORDER BY ORDINAL_POSITION LIMIT 2048" + ), + params=(database, table), + max_rows=2_048, + mask_result=False, + ) + columns = { + str(_value(row, "column_name", "COLUMN_NAME")): str( + _value(row, "data_type", "DATA_TYPE") or "" + ) + for row in (result.data or ()) + if isinstance(row, Mapping) + and _value(row, "column_name", "COLUMN_NAME") not in (None, "") + } + if not columns: + raise SearchRuntimeFailure( + "The requested Doris table does not exist or is not visible.", + reason_code="SEARCH_TABLE_NOT_FOUND", + status_code=404, + ) + return columns + + async def _read_indexes( + self, + database: str, + table: str, + ) -> tuple[_SearchIndex, ...]: + table_reference = build_table_reference( + table, + db_name=database, + ) + # SQL sink audit: build_table_reference validates and quotes both + # identifiers before _execute sends this metadata read to + # connection.execute; no caller values are interpolated. + result = await self._execute( + f"SHOW INDEX FROM {table_reference}", # nosec B608 + max_rows=512, + mask_result=False, + ) + grouped: dict[tuple[str, str], dict[str, Any]] = {} + for row in result.data or (): + index_type = str( + _value(row, "index_type", "Index_type") or "" + ).upper() + if index_type not in {"INVERTED", "ANN"}: + continue + name = str(_value(row, "key_name", "Key_name") or "") + column = str( + _value(row, "column_name", "Column_name") or "" + ) + if not name or not column: + continue + key = (name, index_type) + state = grouped.setdefault( + key, + { + "columns": [], + "properties": {}, + "comment": None, + }, + ) + if column not in state["columns"]: + state["columns"].append(column) + state["properties"].update( + _parse_properties( + _value(row, "properties", "Properties") + ) + ) + comment = _value(row, "comment", "Comment") + if comment not in (None, ""): + state["comment"] = str(comment) + return tuple( + _SearchIndex( + name=name, + index_type=index_type, + columns=tuple(state["columns"]), + properties=dict(state["properties"]), + comment=state["comment"], + ) + for (name, index_type), state in sorted(grouped.items()) + ) + + async def _custom_analyzer_definition( + self, + analyzer_name: str, + ) -> Mapping[str, Any] | None: + result = await self._execute( + "SHOW INVERTED INDEX ANALYZER", + max_rows=512, + mask_result=False, + ) + for row in result.data or (): + normalized = _normalized_row(row) + name = _value( + normalized, + "name", + "analyzer_name", + "analyzer", + ) + if ( + isinstance(name, str) + and name.casefold() == analyzer_name.casefold() + ): + properties = _parse_properties( + _value(normalized, "properties") + ) + return {**normalized, **properties} + return None + + async def _execute( + self, + sql: str, + *, + params: Mapping[str, Any] | tuple[Any, ...] | None = None, + max_rows: int, + mask_result: bool, + database_context: str | None = None, + ) -> QueryResult: + auth_context = get_current_auth_context() + session_id = f"{self._session_prefix}:{uuid.uuid4().hex[:8]}" + try: + async with ( + self._connection_manager.get_connection_context_for_auth_context( + session_id, + auth_context, + ) as connection + ): + if database_context is not None: + safe_database = quote_identifier( + database_context, + "database name", + ) + await connection.execute( + f"USE {safe_database}", + auth_context=auth_context, + mask_result=False, + max_rows=1, + max_bytes=1_024, + ) + return await connection.execute( + sql, + params=params, + auth_context=auth_context, + mask_result=mask_result, + max_rows=max_rows, + max_bytes=_MAX_BYTES, + ) + except Exception as exc: + raise _classify_failure(exc) from exc + + +def _result( + data: Mapping[str, Any], + *, + source: str, + warnings: Sequence[str] = (), + metadata: Mapping[str, Any] | None = None, + evidence: Sequence[Mapping[str, Any]] | None = None, +) -> dict[str, Any]: + unique_warnings = list(dict.fromkeys(str(item) for item in warnings)) + response: dict[str, Any] = { + "status": "partial" if unique_warnings else "success", + "data": redact_sensitive_data(dict(data)), + "warnings": unique_warnings, + "metadata": { + "source": source, + **dict(metadata or {}), + }, + } + if evidence is not None: + response["evidence"] = [ + redact_sensitive_data(dict(item)) for item in evidence + ] + return response + + +def _required_identifier(value: Any, label: str) -> str: + if not isinstance(value, str): + raise _argument_failure(f"{label} must be a non-empty string.") + try: + return validate_identifier(value, label) + except SQLSecurityError as exc: + raise _argument_failure(f"{label} is invalid.") from exc + + +def _required_text( + value: Any, + label: str, + *, + maximum_bytes: int, +) -> str: + if not isinstance(value, str) or not value.strip(): + raise _argument_failure(f"{label} must be a non-empty string.") + normalized = value.strip() + if len(normalized.encode("utf-8")) > maximum_bytes: + raise _argument_failure(f"{label} exceeds the maximum accepted size.") + return normalized + + +def _bounded_integer( + value: Any, + *, + default: int, + minimum: int, + maximum: int, + label: str, +) -> int: + if value is None: + return default + if ( + isinstance(value, bool) + or not isinstance(value, int) + or value < minimum + or value > maximum + ): + raise _argument_failure( + f"{label} must be between {minimum} and {maximum}." + ) + return int(value) + + +def _identifier_sequence( + values: Any, + label: str, + *, + maximum: int, + required: bool = False, +) -> tuple[str, ...]: + if values is None: + if required: + raise _argument_failure(f"At least one {label} is required.") + return () + if ( + not isinstance(values, Sequence) + or isinstance(values, str | bytes) + or not values + or len(values) > maximum + ): + raise _argument_failure( + f"{label} values must be a non-empty bounded array." + ) + normalized = tuple( + _required_identifier(value, label) + for value in values + ) + if len(set(normalized)) != len(normalized): + raise _argument_failure(f"{label} values must be unique.") + return normalized + + +def _vector(value: Any) -> tuple[float, ...]: + if ( + not isinstance(value, Sequence) + or isinstance(value, str | bytes) + or not value + or len(value) > _MAX_VECTOR_DIMENSION + ): + raise _argument_failure( + "vector must be a non-empty bounded numeric array." + ) + normalized: list[float] = [] + for item in value: + if isinstance(item, bool) or not isinstance(item, int | float): + raise _argument_failure("vector values must be numbers.") + number = float(item) + if not math.isfinite(number): + raise _argument_failure("vector values must be finite.") + normalized.append(number) + return tuple(normalized) + + +def _return_fields( + values: Any, + columns: Mapping[str, str], + *, + vector_field: str | None, +) -> tuple[str, ...]: + if values is None: + selected = tuple( + name + for name, data_type in columns.items() + if name != vector_field + and not data_type.casefold().startswith("array<float") + )[:_MAX_RETURN_FIELDS] + if not selected: + raise _argument_failure( + "return_fields is required when no scalar columns are visible." + ) + return selected + selected = _identifier_sequence( + values, + "return field", + maximum=_MAX_RETURN_FIELDS, + required=True, + ) + _require_known_columns(selected, set(columns)) + return selected + + +def _require_known_columns( + requested: Sequence[str], + available: set[str], +) -> None: + if any(column not in available for column in requested): + raise _argument_failure( + "One or more requested fields are not visible on the table." + ) + + +def _resolve_vector_field( + requested: Any, + indexes: Sequence[_SearchIndex], +) -> str: + if requested is not None: + return _required_identifier(requested, "vector field") + candidates = sorted( + { + column + for index in indexes + if index.is_ann + for column in index.columns + } + ) + if len(candidates) != 1: + raise _argument_failure( + "vector_field is required unless exactly one ANN-indexed field " + "is visible." + ) + return candidates[0] + + +def _ann_index_for_field( + indexes: Sequence[_SearchIndex], + field: str, +) -> _SearchIndex | None: + candidates = [ + index + for index in indexes + if index.is_ann and field in index.columns + ] + if len(candidates) != 1: + return None + return candidates[0] + + +def _compile_filters( + value: Any, + columns: set[str], +) -> tuple[list[str], list[Any]]: + if value is None: + return [], [] + if not isinstance(value, Mapping) or len(value) > _MAX_FILTERS: + raise _argument_failure("filters must be a bounded object.") + predicates: list[str] = [] + params: list[Any] = [] + for raw_field, raw_rule in value.items(): + field = _required_identifier(raw_field, "filter field") + if field not in columns: + raise _argument_failure( + "One or more filter fields are not visible on the table." + ) + quoted = quote_identifier(field, "filter field") + if not isinstance(raw_rule, Mapping): + _validate_scalar(raw_rule) + predicates.append(f"{quoted} = %s") + params.append(raw_rule) + continue + + unknown = set(raw_rule) - {"operator", "value", "values"} + if unknown: + raise _argument_failure( + "Filter rules contain unsupported properties." + ) + operator = str(raw_rule.get("operator", "")).casefold() + if operator in _FILTER_OPERATORS: + if "value" not in raw_rule or "values" in raw_rule: + raise _argument_failure( + "Scalar filter operators require exactly one value." + ) + scalar = raw_rule["value"] + _validate_scalar(scalar) + predicates.append( + f"{quoted} {_FILTER_OPERATORS[operator]} %s" + ) + params.append(scalar) + continue + if operator in _SET_FILTER_OPERATORS: + values = raw_rule.get("values") + if ( + not isinstance(values, Sequence) + or isinstance(values, str | bytes) + or not values + or len(values) > _MAX_FILTER_VALUES + or "value" in raw_rule + ): + raise _argument_failure( + "Set filter operators require a bounded values array." + ) + for scalar in values: + _validate_scalar(scalar) + placeholders = ", ".join("%s" for _ in values) + predicates.append( + f"{quoted} {_SET_FILTER_OPERATORS[operator]} " + f"({placeholders})" + ) + params.extend(values) + continue + if operator in _NULL_FILTER_OPERATORS: + if "value" in raw_rule or "values" in raw_rule: + raise _argument_failure( + "Null filter operators do not accept values." + ) + predicates.append( + f"{quoted} {_NULL_FILTER_OPERATORS[operator]}" + ) + continue + raise _argument_failure("Filter operator is invalid.") + return predicates, params + + +def _validate_scalar(value: Any) -> None: + if value is not None and not isinstance( + value, + str | int | float | bool, + ): + raise _argument_failure("Filter values must be JSON scalars.") + if isinstance(value, float) and not math.isfinite(value): + raise _argument_failure("Filter numbers must be finite.") + if isinstance(value, str) and len(value.encode("utf-8")) > 8_192: + raise _argument_failure("Filter text exceeds the maximum accepted size.") + + +def _parse_properties(value: Any) -> dict[str, str]: + if isinstance(value, Mapping): + return { + str(key).casefold(): str(item) + for key, item in value.items() + } + if not isinstance(value, str): + return {} + return { + match.group("key").casefold(): match.group("value") + .replace('\\"', '"') + .replace("\\\\", "\\") + for match in _PROPERTY_PAIR.finditer(value) + } + + +def _normalized_build_task(row: Mapping[str, Any]) -> dict[str, Any]: + normalized = _normalized_row(row) + return { + "job_id": normalized.get("job_id"), + "table": normalized.get("table_name"), + "partition": normalized.get("partition_name"), + "state": normalized.get("state"), + "progress": normalized.get("progress"), + "message": normalized.get("msg"), + "create_time": _json_value(normalized.get("create_time")), + "finish_time": _json_value(normalized.get("finish_time")), + "alter_indexes": normalized.get("alter_inverted_indexes"), + } + + +def _normalize_tokens(value: Any) -> list[dict[str, Any]]: + parsed = value + if isinstance(value, str): + try: + parsed = json.loads(value) + except json.JSONDecodeError as exc: + raise SearchRuntimeFailure( + "Doris TOKENIZE returned an invalid payload.", + reason_code="SEARCH_TOKENIZE_INVALID_RESPONSE", + status_code=502, + ) from exc + if not isinstance(parsed, Sequence) or isinstance(parsed, str | bytes): + raise SearchRuntimeFailure( + "Doris TOKENIZE returned an invalid payload.", + reason_code="SEARCH_TOKENIZE_INVALID_RESPONSE", + status_code=502, + ) + tokens: list[dict[str, Any]] = [] + for offset, item in enumerate(parsed): + if isinstance(item, Mapping): + token = item.get("token") + if token is None: + continue + normalized: dict[str, Any] = {"term": str(token)} + if item.get("position") is not None: + normalized["position"] = item["position"] + if item.get("type") is not None: + normalized["type"] = item["type"] + else: + normalized = {"term": str(item), "position": offset} + tokens.append(normalized) + if len(tokens) >= 10_000: + break + return tokens + + +def _built_in_analyzer_properties(name: str) -> dict[str, str]: + key = "parser" if name in _BACKWARD_COMPATIBLE_PARSERS else ( + "built_in_analyzer" + ) + return {key: name} + + +def _split_component_names(value: Any) -> tuple[str, ...]: + parsed = value + if isinstance(value, str): + stripped = value.strip() + if stripped.startswith("["): + try: + parsed = json.loads(stripped) + except json.JSONDecodeError: + parsed = value + if isinstance(parsed, str): + return tuple( + item.strip() + for item in parsed.split(",") + if item.strip() + ) + if isinstance(parsed, Sequence) and not isinstance(parsed, str | bytes): + return tuple( + str(item).strip() + for item in parsed + if str(item).strip() + ) + return () + + +def _optional_identifier_value(value: Any) -> str | None: + if not isinstance(value, str) or not value: + return None + try: + return validate_identifier(value, "component name") + except SQLSecurityError: + return None + + +def _escape_property_value(value: str) -> str: + return value.replace("\\", "\\\\").replace('"', '\\"') + + +def _search_plan_facets(plan_text: str) -> dict[str, Any]: + upper = plan_text.upper() + match_operators = sorted( + { + operator + for operator in ( + "MATCH_ANY", + "MATCH_ALL", + "MATCH_PHRASE", + "MATCH_PHRASE_PREFIX", + "SEARCH", + ) + if operator in upper + } + ) + return { + "ann_pushdown_observed": "ANN SORT INFO" in upper, + "ann_sort_limit_observed": "ANN SORT LIMIT" in upper, + "text_match_predicate_observed": bool(match_operators), + "match_operators": match_operators, + "olap_scan_observed": "OLAPSCANNODE" in upper, + "profile_required_for_inverted_hit_confirmation": True, + } + + +def _diagnostic_findings( + *, + compiled: _CompiledSearch | None, + indexes: Sequence[_SearchIndex], + facets: Mapping[str, Any], + query_sql: str, +) -> list[dict[str, Any]]: + findings: list[dict[str, Any]] = [] + if compiled is not None and compiled.text_fields: + findings.append( + { + "severity": "info", + "code": "TEXT_INDEX_CONFIGURED_AND_MATCH_PLANNED", + "message": ( + "Visible inverted indexes cover the requested text fields " + "and EXPLAIN preserves a MATCH predicate. Query Profile is " + "required to prove runtime inverted-index filtering." + ), + } + ) + if compiled is not None and compiled.vector_field is not None: + if facets.get("ann_pushdown_observed"): + findings.append( + { + "severity": "info", + "code": "ANN_PUSHDOWN_OBSERVED", + "message": "Doris EXPLAIN reports ANN SORT INFO.", + } + ) + else: + findings.append( + { + "severity": "high", + "code": "ANN_PUSHDOWN_NOT_OBSERVED", + "message": ( + "The vector query has ANN metadata, but EXPLAIN did " + "not expose ANN SORT INFO." + ), + } + ) + if compiled is None: + upper = query_sql.upper() + if "MATCH" not in upper and "SEARCH(" not in upper and ( + "_DISTANCE" not in upper + ): + findings.append( + { + "severity": "medium", + "code": "SEARCH_OPERATOR_NOT_OBSERVED", + "message": ( + "The submitted SQL does not expose a recognized Doris " + "text or vector search operator." + ), + } + ) + if not indexes: + findings.append( + { + "severity": "medium", + "code": "INDEX_METADATA_NOT_BOUND", + "message": ( + "No authoritative target index metadata was bound to " + "the raw SQL diagnosis." + ), + } + ) + return findings + + +def _simple_source(sql: str) -> tuple[str | None, str] | None: + match = _SIMPLE_SOURCE.search(sql) + if match is None: + return None + database = _unquote_identifier(match.group("database")) + table = _unquote_identifier(match.group("table")) + if table is None: + return None + try: + table = validate_identifier(table, "table name") + if database is not None: + database = validate_identifier(database, "database name") + except SQLSecurityError: + return None + return database, table + + +def _unquote_identifier(value: str | None) -> str | None: + if value is None: + return None + return value[1:-1] if value.startswith("`") and value.endswith("`") else value + + +def _normalized_row(row: Mapping[str, Any]) -> dict[str, Any]: + return { + str(key).strip().replace(" ", "_").casefold(): value + for key, value in row.items() + } + + +def _value(row: Mapping[str, Any], *names: str) -> Any: + normalized = _normalized_row(row) + for name in names: + key = name.strip().replace(" ", "_").casefold() + if key in normalized: + return normalized[key] + return None + + +def _json_value(value: Any) -> Any: + if value is None or isinstance(value, str | int | float | bool): + return value + if hasattr(value, "isoformat"): + return value.isoformat() + return str(value) + + +def _as_int(value: Any) -> int | None: + if value in (None, ""): + return None + try: + return int(value) + except (TypeError, ValueError): + return None + + +def _argument_failure(message: str) -> SearchRuntimeFailure: + return SearchRuntimeFailure( + message, + reason_code="SEARCH_ARGUMENT_INVALID", + status_code=400, + ) + + +def _classify_failure(exc: Exception) -> SearchRuntimeFailure: + if isinstance(exc, SearchRuntimeFailure): + return exc + if isinstance(exc, QueryRuntimeFailure | SQLSecurityError): + return _argument_failure("Search arguments are invalid.") + numeric_code = next( + ( + value + for value in getattr(exc, "args", ()) + if isinstance(value, int) + ), + None, + ) + message = str(exc).casefold() + if numeric_code in {1044, 1045, 1142, 1227} or any( + marker in message + for marker in ("access denied", "permission denied", "privilege") + ): + return SearchRuntimeFailure( + "Doris denied access to Search data or metadata.", + reason_code="SEARCH_PERMISSION_DENIED", + status_code=403, + ) + if numeric_code in {1064, 1109, 1146} or any( + marker in message + for marker in ( + "doesn't exist", + "does not exist", + "not supported", + "unsupported", + "unknown table", + "unknown function", + "no viable alternative", + ) + ): + return SearchRuntimeFailure( + "The requested Doris Search capability is unsupported.", + reason_code="SEARCH_CAPABILITY_UNSUPPORTED", + status_code=501, + ) + if isinstance(exc, TimeoutError | ConnectionError | OSError): + return SearchRuntimeFailure( + "Doris Search is temporarily unavailable.", + reason_code="SEARCH_BACKEND_UNAVAILABLE", + status_code=503, + retryable=True, + ) + return SearchRuntimeFailure( + "Doris Search execution failed.", + reason_code="SEARCH_EXECUTION_FAILED", + status_code=502, + ) + + +__all__ = ["DorisSearchRuntime", "SearchRuntimeFailure"] diff --git a/test/integration/test_real_doris_transports.py b/test/integration/test_real_doris_transports.py index cfbfa50..43300ef 100644 --- a/test/integration/test_real_doris_transports.py +++ b/test/integration/test_real_doris_transports.py @@ -91,6 +91,12 @@ PIPELINE_CHILD_NAMES = ( "monitor_data_freshness", "analyze_data_dependencies", ) +SEARCH_CHILD_NAMES = ( + "search_data", + "preview_text_analysis", + "inspect_search_indexes", + "diagnose_search_query", +) @dataclass(frozen=True) @@ -120,6 +126,17 @@ class DorisSandbox: cursor.execute(f"KILL CONNECTION {int(connection_id)}") +@dataclass +class DorisSearchSandbox: + settings: RealDorisSettings + admin_connection: pymysql.Connection + table: str + + @property + def qualified_table(self) -> str: + return f"`{self.settings.database}`.`{self.table}`" + + def _real_doris_settings() -> RealDorisSettings: required = { name: os.getenv(name, "").strip() @@ -199,6 +216,66 @@ def doris_sandbox() -> DorisSandbox: admin_connection.close() [email protected] +def doris_search_sandbox() -> DorisSearchSandbox: + settings = _real_doris_settings() + table = f"mcp_search_it_{secrets.token_hex(6)}" + qualified_table = f"`{settings.database}`.`{table}`" + admin_connection = pymysql.connect( + host=settings.host, + port=settings.port, + user=settings.user, + password=settings.password, + database=settings.database, + autocommit=True, + ) + + try: + with admin_connection.cursor() as cursor: + cursor.execute( + f""" + CREATE TABLE {qualified_table} ( + id BIGINT NOT NULL, + title STRING NOT NULL, + category VARCHAR(64) NOT NULL, + embedding ARRAY<FLOAT> NOT NULL, + INDEX idx_title (title) USING INVERTED + PROPERTIES ( + "parser" = "english", + "support_phrase" = "true" + ), + INDEX idx_embedding (embedding) USING ANN + PROPERTIES ( + "index_type" = "hnsw", + "metric_type" = "l2_distance", + "dim" = "3" + ) + ) + DUPLICATE KEY(id) + DISTRIBUTED BY HASH(id) BUCKETS 1 + PROPERTIES ("replication_num" = "1") + """ + ) + cursor.execute( + f""" + INSERT INTO {qualified_table} VALUES + (1, 'Apache Doris search', 'database', [0.1, 0.2, 0.3]), + (2, 'Hybrid retrieval', 'search', [0.2, 0.1, 0.4]), + (3, 'Warehouse analytics', 'analytics', [0.8, 0.7, 0.6]) + """ + ) + + yield DorisSearchSandbox( + settings=settings, + admin_connection=admin_connection, + table=table, + ) + finally: + with suppress(Exception), admin_connection.cursor() as cursor: + cursor.execute(f"DROP TABLE IF EXISTS {qualified_table}") + admin_connection.close() + + def _server_environment( settings: RealDorisSettings, *, @@ -1310,6 +1387,170 @@ async def test_real_doris_hierarchical_pipeline_domain_is_read_only_and_live( assert row == (1,) [email protected]("transport", ["http", "stdio"]) +async def test_real_doris_hierarchical_search_domain_is_read_only_and_live( + transport: str, + doris_search_sandbox: DorisSearchSandbox, +) -> None: + environment = _server_environment( + doris_search_sandbox.settings, + user=doris_search_sandbox.settings.user, + password=doris_search_sandbox.settings.password, + ) + environment["MCP_TOOL_EXPOSURE_MODE"] = "hierarchical" + with doris_search_sandbox.admin_connection.cursor() as cursor: + cursor.execute( + f"SELECT COUNT(*) FROM {doris_search_sandbox.qualified_table}" + ) + row_count_before = int(cursor.fetchone()[0]) + + async with _transport_client( + transport, + environment, + read_timeout_seconds=60, + ) as client: + search_result = await client.call_tool("doris_search", {}) + assert search_result.is_error is False + assert isinstance(search_result.structured_content, dict) + manifest = search_result.structured_content + assert manifest["mode"] == "manifest" + assert manifest["domain"] == "doris_search" + children = {child["name"]: child for child in manifest["children"]} + assert tuple(children) == SEARCH_CHILD_NAMES + assert all(child["availability"]["callable"] for child in children.values()) + manifest_version = manifest["manifest_version"] + + indexes = await _call_domain_child( + client, + domain="doris_search", + child_tool="inspect_search_indexes", + arguments={ + "database": doris_search_sandbox.settings.database, + "table": doris_search_sandbox.table, + }, + manifest_version=manifest_version, + ) + assert { + item["index_type"] for item in indexes["data"]["items"] + } == {"INVERTED", "ANN"} + assert indexes["data"]["capabilities"]["hybrid"] is True + + analysis = await _call_domain_child( + client, + domain="doris_search", + child_tool="preview_text_analysis", + arguments={ + "text": "Apache Doris search", + "analyzer": "english", + }, + manifest_version=manifest_version, + ) + terms = [item["term"] for item in analysis["data"]["tokens"]] + assert {"apache", "doris", "search"} <= set(terms) + assert "[REDACTED]" not in terms + + text = await _call_domain_child( + client, + domain="doris_search", + child_tool="search_data", + arguments={ + "database": doris_search_sandbox.settings.database, + "table": doris_search_sandbox.table, + "query": "Doris", + "mode": "text", + "fields": ["title"], + "top_k": 5, + "return_fields": ["id", "title", "category"], + }, + manifest_version=manifest_version, + ) + assert text["data"]["rows"][0]["id"] == 1 + assert text["metadata"]["invented_scores"] is False + + vector = await _call_domain_child( + client, + domain="doris_search", + child_tool="search_data", + arguments={ + "database": doris_search_sandbox.settings.database, + "table": doris_search_sandbox.table, + "mode": "vector", + "vector": [0.1, 0.2, 0.3], + "vector_field": "embedding", + "top_k": 2, + "return_fields": ["id", "title", "category"], + }, + manifest_version=manifest_version, + ) + assert vector["data"]["rows"][0]["id"] == 1 + assert vector["metadata"]["vector_metric"] == "l2_distance" + + hybrid_request = { + "database": doris_search_sandbox.settings.database, + "table": doris_search_sandbox.table, + "query": "Doris", + "mode": "hybrid", + "fields": ["title"], + "vector": [0.1, 0.2, 0.3], + "vector_field": "embedding", + "top_k": 2, + "filters": {"category": "database"}, + "return_fields": ["id", "title", "category"], + } + hybrid = await _call_domain_child( + client, + domain="doris_search", + child_tool="search_data", + arguments=hybrid_request, + manifest_version=manifest_version, + ) + assert hybrid["data"]["rows"][0]["id"] == 1 + + diagnosis = await _call_domain_child( + client, + domain="doris_search", + child_tool="diagnose_search_query", + arguments={ + "search_request": hybrid_request, + "include_profile": False, + }, + manifest_version=manifest_version, + ) + assert ( + diagnosis["data"]["explain"]["facets"][ + "ann_pushdown_observed" + ] + is True + ) + assert diagnosis["metadata"]["invented_index_hits"] is False + + injected_identifier = await client.call_tool( + "doris_search", + { + "child_tool": "search_data", + "arguments": { + "database": doris_search_sandbox.settings.database, + "table": doris_search_sandbox.table, + "query": "Doris", + "mode": "text", + "fields": ["title"], + "return_fields": [ + f"title; DROP TABLE {doris_search_sandbox.table}" + ], + }, + "manifest_version": manifest_version, + }, + ) + assert injected_identifier.is_error is True + + with doris_search_sandbox.admin_connection.cursor() as cursor: + cursor.execute( + f"SELECT COUNT(*) FROM {doris_search_sandbox.qualified_table}" + ) + row_count_after = int(cursor.fetchone()[0]) + assert row_count_after == row_count_before + + @pytest.mark.skipif( os.getenv("DORIS_REAL_HTTP_INTEGRATION") != "1", reason="set DORIS_REAL_HTTP_INTEGRATION=1 with independent FE/BE HTTP endpoints", diff --git a/test/tools/test_capability_detector.py b/test/tools/test_capability_detector.py index 025aa54..ecdce73 100644 --- a/test/tools/test_capability_detector.py +++ b/test/tools/test_capability_detector.py @@ -37,6 +37,14 @@ from doris_mcp_server.utils.db import DorisRouteIdentity from doris_mcp_server.utils.doris_http_client import DorisHTTPResponse from doris_mcp_server.utils.security import AuthContext +_SEARCH_TARGET_DISCOVERY_SQL = ( + "SELECT TABLE_SCHEMA, TABLE_NAME, COLUMN_NAME " + "FROM information_schema.columns " + "WHERE DATA_TYPE IN ('char', 'varchar', 'string', 'text') " + "AND TABLE_SCHEMA NOT IN ('information_schema', 'mysql') " + "ORDER BY TABLE_SCHEMA, TABLE_NAME, ORDINAL_POSITION LIMIT 8" +) + class _ProbeConnection: def __init__(self) -> None: @@ -206,6 +214,111 @@ async def test_detector_builds_version_vector_and_extends_domains_lazily() -> No assert connection.statements.count("SELECT @@version_comment;") == 1 [email protected] +async def test_search_probes_use_visible_target_and_isolated_connections() -> None: + connection = _ProbeConnection() + connection.row_overrides[_SEARCH_TARGET_DISCOVERY_SQL] = [ + { + "TABLE_SCHEMA": "analytics", + "TABLE_NAME": "documents", + "COLUMN_NAME": "title", + } + ] + manager = _ProbeConnectionManager(connection) + detector = DorisCapabilityDetector(manager) # type: ignore[arg-type] + base = await detector.detect_base( + None, + capability_generation=1, + provider_generation="provider.search", + ) + contexts_before = len(manager.context_sessions) + + search = await detector.detect_domain(base, "doris_search", None) + + assert ( + search.probe("text_match_syntax_readable").status + is CapabilityProbeStatus.SUPPORTED + ) + assert ( + search.probe("inverted_index_and_search_syntax_ready").status + is CapabilityProbeStatus.SUPPORTED + ) + assert ( + search.probe("ann_index_and_metric_compatible").status + is CapabilityProbeStatus.SUPPORTED + ) + assert search.probe("hybrid_search").status is CapabilityProbeStatus.SUPPORTED + assert ( + "EXPLAIN SELECT `title` FROM `analytics`.`documents` " + "WHERE `title` MATCH_ANY 'doris' LIMIT 1" + ) in connection.statements + assert len(manager.context_sessions) - contexts_before == 6 + assert len(set(manager.context_sessions[contexts_before:])) == 6 + + [email protected] +async def test_search_probe_without_visible_target_is_degraded_not_unsupported() -> None: + connection = _ProbeConnection() + manager = _ProbeConnectionManager(connection) + detector = DorisCapabilityDetector(manager) # type: ignore[arg-type] + base = await detector.detect_base( + None, + capability_generation=1, + provider_generation="provider.search", + ) + + search = await detector.detect_domain(base, "doris_search", None) + + assert ( + search.probe("text_match_syntax_readable").status + is CapabilityProbeStatus.DEGRADED + ) + assert ( + search.probe("inverted_index_and_search_syntax_ready").status + is CapabilityProbeStatus.DEGRADED + ) + assert search.probe("hybrid_search").status is CapabilityProbeStatus.DEGRADED + + [email protected] +async def test_search_probe_keeps_text_ready_when_ann_function_is_unsupported() -> None: + connection = _ProbeConnection() + connection.row_overrides[_SEARCH_TARGET_DISCOVERY_SQL] = [ + { + "TABLE_SCHEMA": "analytics", + "TABLE_NAME": "documents", + "COLUMN_NAME": "title", + } + ] + ann_sql = ( + "SELECT l2_distance_approximate([0.0], [0.0]) " + "AS distance" + ) + connection.failures[ann_sql] = RuntimeError( + 1305, + "Unknown function l2_distance_approximate", + ) + manager = _ProbeConnectionManager(connection) + detector = DorisCapabilityDetector(manager) # type: ignore[arg-type] + base = await detector.detect_base( + None, + capability_generation=1, + provider_generation="provider.search", + ) + + search = await detector.detect_domain(base, "doris_search", None) + + assert ( + search.probe("inverted_index_and_search_syntax_ready").status + is CapabilityProbeStatus.SUPPORTED + ) + assert ( + search.probe("ann_index_and_metric_compatible").status + is CapabilityProbeStatus.UNSUPPORTED + ) + assert search.probe("hybrid_search").status is CapabilityProbeStatus.UNSUPPORTED + + @pytest.mark.asyncio async def test_pipeline_probes_isolate_an_unsupported_source_connection() -> None: base_connection = _ProbeConnection() diff --git a/test/tools/test_domain_dispatcher.py b/test/tools/test_domain_dispatcher.py index ac094e2..c628289 100644 --- a/test/tools/test_domain_dispatcher.py +++ b/test/tools/test_domain_dispatcher.py @@ -501,7 +501,7 @@ async def test_child_argument_violation_list_is_bounded_and_marked() -> None: @pytest.mark.asyncio -async def test_manifest_version_unavailable_and_unbound_fail_closed() -> None: +async def test_manifest_version_and_unavailable_child_fail_closed() -> None: callable_manager = _manager("doris_query.execute_query") stale = json.loads( await callable_manager.call_tool( @@ -523,28 +523,75 @@ async def test_manifest_version_unavailable_and_unbound_fail_closed() -> None: }, ) ) - unbound_manager = _manager("doris_search.search_data") - unbound = json.loads( - await unbound_manager.call_tool( + assert stale["error"]["code"] == DomainErrorCode.CHILD_MANIFEST_STALE + assert stale["error"]["details"]["rediscover"] is True + assert unavailable["error"]["code"] == ( + DomainErrorCode.CHILD_CAPABILITY_UNAVAILABLE + ) + assert unavailable["error"]["details"]["reason_code"] == ("TEST_HANDLER_PENDING") + + +def test_search_domain_binds_all_four_children() -> None: + manager = _manager() + bound = BoundHandlerAvailabilityProvider(manager) + search = DORIS_DOMAIN_CATALOG.resolve_domain("doris_search") + + assert len(search.children) == 4 + assert all( + bound.is_bound(search.name, child.name) + for child in search.children + ) + + [email protected] +async def test_formal_search_handler_uses_strict_runtime() -> None: + manager = _manager("doris_search.search_data") + formal_result = { + "status": "success", + "data": { + "columns": [{"name": "id"}], + "rows": [{"id": 1}], + "row_count": 1, + "truncated": False, + }, + "warnings": [], + "metadata": {"source": "doris_native_search"}, + "evidence": [], + } + manager.search_runtime.search_data = AsyncMock(return_value=formal_result) + + response = json.loads( + await manager.call_tool( "doris_search", { "child_tool": "search_data", "arguments": { "database": "analytics", - "table": "events", + "table": "documents", + "query": "Doris", "mode": "text", + "fields": ["title"], + "filters": {"category": "database"}, }, }, ) ) - assert stale["error"]["code"] == DomainErrorCode.CHILD_MANIFEST_STALE - assert stale["error"]["details"]["rediscover"] is True - assert unavailable["error"]["code"] == ( - DomainErrorCode.CHILD_CAPABILITY_UNAVAILABLE + assert response["mode"] == "result" + assert response["data"]["metadata"]["source"] == "doris_native_search" + manager.search_runtime.search_data.assert_awaited_once_with( + database="analytics", + table="documents", + query="Doris", + mode="text", + fields=["title"], + vector=None, + vector_field=None, + text_operator=None, + top_k=None, + filters={"category": "database"}, + return_fields=None, ) - assert unavailable["error"]["details"]["reason_code"] == ("TEST_HANDLER_PENDING") - assert unbound["error"]["details"]["reason_code"] == "HANDLER_NOT_BOUND" @pytest.mark.asyncio diff --git a/test/tools/test_doris_feature_matrix.py b/test/tools/test_doris_feature_matrix.py index 84eecde..edf4ae6 100644 --- a/test/tools/test_doris_feature_matrix.py +++ b/test/tools/test_doris_feature_matrix.py @@ -113,6 +113,31 @@ def test_adbc_is_inside_query_and_variant_is_inside_lakehouse() -> None: ) +def test_search_prefers_vector_hybrid_and_keeps_text_fallback() -> None: + search = DORIS_FEATURE_MATRIX.get_feature( + "doris_search", + "search_data", + ) + + assert tuple( + variant.name for variant in search.support_contract.variants + ) == ( + "vector_hybrid_search", + "inverted_text_search", + ) + vector, text = search.support_contract.variants + assert vector.supported_ranges == (">=4.0.0",) + assert vector.required_probes == ( + "ann_index_and_metric_compatible", + "inverted_index_and_search_syntax_ready", + ) + assert vector.callable_when_degraded is True + assert text.required_probes == ( + "inverted_index_and_search_syntax_ready", + ) + assert text.callable_when_degraded is True + + def test_every_contract_is_fail_closed_and_has_resolvable_sources() -> None: known_sources = {source.source_id for source in DORIS_FEATURE_MATRIX.sources} diff --git a/test/utils/test_search_runtime.py b/test/utils/test_search_runtime.py new file mode 100644 index 0000000..70b73b1 --- /dev/null +++ b/test/utils/test_search_runtime.py @@ -0,0 +1,543 @@ +# 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 contracts for the formal read-only Search runtime.""" + +from __future__ import annotations + +import json +from collections.abc import AsyncIterator, Mapping +from contextlib import asynccontextmanager +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock + +import pytest + +from doris_mcp_server.utils.db import QueryResult +from doris_mcp_server.utils.search_runtime import ( + DorisSearchRuntime, + SearchRuntimeFailure, +) + +_COLUMN_SQL = ( + "SELECT COLUMN_NAME, DATA_TYPE " + "FROM information_schema.columns " + "WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s " + "ORDER BY ORDINAL_POSITION LIMIT 2048" +) +_INDEX_SQL = "SHOW INDEX FROM `analytics`.`documents`" + + +def _columns() -> list[dict[str, Any]]: + return [ + {"COLUMN_NAME": "id", "DATA_TYPE": "BIGINT"}, + {"COLUMN_NAME": "title", "DATA_TYPE": "TEXT"}, + {"COLUMN_NAME": "category", "DATA_TYPE": "VARCHAR"}, + {"COLUMN_NAME": "embedding", "DATA_TYPE": "ARRAY<FLOAT>"}, + ] + + +def _indexes( + *, + inverted: bool = True, + ann: bool = True, + metric: str = "l2_distance", + dimension: int = 3, +) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + if inverted: + rows.append( + { + "Key_name": "idx_title", + "Column_name": "title", + "Index_type": "INVERTED", + "Properties": ( + '"parser" = "english", ' + '"support_phrase" = "true", ' + '"comment_text" = "café"' + ), + "Comment": "title search", + } + ) + if ann: + rows.append( + { + "Key_name": "idx_embedding", + "Column_name": "embedding", + "Index_type": "ANN", + "Properties": ( + '"index_type" = "hnsw", ' + f'"metric_type" = "{metric}", ' + f'"dim" = "{dimension}"' + ), + "Comment": "vector search", + } + ) + return rows + + +class _ConnectionManager: + def __init__( + self, + *, + columns: list[dict[str, Any]] | None = None, + indexes: list[dict[str, Any]] | None = None, + search_rows: list[dict[str, Any]] | None = None, + tokenize_payload: Any = None, + analyzers: list[dict[str, Any]] | None = None, + build_tasks: list[dict[str, Any]] | None = None, + plan_rows: list[dict[str, Any]] | None = None, + failures: dict[str, Exception] | None = None, + ) -> None: + self.columns = _columns() if columns is None else columns + self.indexes = _indexes() if indexes is None else indexes + self.search_rows = ( + [{"id": 1, "title": "Apache Doris", "category": "database"}] + if search_rows is None + else search_rows + ) + self.tokenize_payload = ( + json.dumps( + [ + {"token": "apache", "position": 0, "type": "word"}, + {"token": "doris", "position": 1, "type": "word"}, + ] + ) + if tokenize_payload is None + else tokenize_payload + ) + self.analyzers = analyzers or [] + self.build_tasks = build_tasks or [] + self.plan_rows = plan_rows or [ + { + "Explain String": ( + "0:VTOP-N\nANN SORT INFO: " + "l2_distance_approximate(embedding, [0.1,0.2,0.3])\n" + "ANN SORT LIMIT: 5\n" + "1:VOlapScanNode\nPREDICATES: title MATCH_ANY 'doris'" + ) + } + ] + self.failures = failures or {} + self.calls: list[dict[str, Any]] = [] + self.context_count = 0 + + @asynccontextmanager + async def get_connection_context_for_auth_context( + self, + session_id: str, + _auth_context: Any, + ) -> AsyncIterator[Any]: + self.context_count += 1 + manager = self + + class _Connection: + async def execute( + self, + sql: str, + params: Mapping[str, Any] | tuple[Any, ...] | None = None, + **kwargs: Any, + ) -> QueryResult: + manager.calls.append( + { + "session_id": session_id, + "sql": sql, + "params": params, + "kwargs": kwargs, + } + ) + for prefix, failure in manager.failures.items(): + if sql.startswith(prefix): + raise failure + if sql == _COLUMN_SQL: + rows = manager.columns + elif sql == _INDEX_SQL: + rows = manager.indexes + elif sql == "SHOW INVERTED INDEX ANALYZER": + rows = manager.analyzers + elif sql.startswith("SHOW BUILD INDEX"): + rows = manager.build_tasks + elif sql.startswith("SELECT TOKENIZE"): + rows = [{"tokens": manager.tokenize_payload}] + elif sql.startswith("EXPLAIN "): + rows = manager.plan_rows + elif sql.startswith("USE "): + rows = [] + else: + rows = manager.search_rows + columns = list(rows[0]) if rows else [] + return QueryResult( + data=rows, + metadata={"columns": columns, "truncated": False}, + execution_time=0.01, + row_count=len(rows), + sql=sql, + ) + + yield _Connection() + + +def _runtime( + manager: _ConnectionManager | None = None, +) -> tuple[DorisSearchRuntime, _ConnectionManager, Any]: + connection_manager = manager or _ConnectionManager() + query_runtime = SimpleNamespace( + get_query_profile=AsyncMock( + return_value={ + "status": "success", + "data": {"query_id": "query-1"}, + "warnings": [], + "evidence": [{"source": "runtime_profile"}], + } + ) + ) + return ( + DorisSearchRuntime( # type: ignore[arg-type] + connection_manager, + query_runtime, + ), + connection_manager, + query_runtime, + ) + + +async def _search( + runtime: DorisSearchRuntime, + **overrides: Any, +) -> dict[str, Any]: + request = { + "database": "analytics", + "table": "documents", + "query": "Doris", + "mode": "text", + "fields": ["title"], + "vector": None, + "vector_field": None, + "text_operator": "any", + "top_k": 5, + "filters": {"category": "database"}, + "return_fields": ["id", "title", "category"], + } + request.update(overrides) + return await runtime.search_data(**request) + + [email protected] +async def test_text_search_binds_values_and_uses_target_index_metadata() -> None: + runtime, manager, _ = _runtime() + + result = await _search( + runtime, + query="Doris' OR 1=1 --", + filters={ + "category": { + "operator": "in", + "values": ["database", "analytics"], + } + }, + ) + + assert result["status"] == "success" + assert result["data"]["rows"][0]["title"] == "Apache Doris" + assert result["metadata"]["invented_scores"] is False + query_call = manager.calls[-1] + assert "`title` MATCH_ANY %s" in query_call["sql"] + assert "`category` IN (%s, %s)" in query_call["sql"] + assert "Doris' OR 1=1 --" not in query_call["sql"] + assert query_call["params"] == ( + "Doris' OR 1=1 --", + "database", + "analytics", + ) + assert query_call["kwargs"]["mask_result"] is True + assert manager.context_count == 3 + assert len({call["session_id"] for call in manager.calls}) == 3 + + [email protected] +async def test_search_rejects_identifier_injection_before_execution() -> None: + runtime, manager, _ = _runtime() + + with pytest.raises(SearchRuntimeFailure) as failure: + await _search(runtime, table="documents; DROP TABLE accounts") + + assert failure.value.reason_code == "SEARCH_ARGUMENT_INVALID" + assert manager.calls == [] + + [email protected] +async def test_text_search_requires_visible_inverted_index() -> None: + runtime, manager, _ = _runtime( + _ConnectionManager(indexes=_indexes(inverted=False)) + ) + + with pytest.raises(SearchRuntimeFailure) as failure: + await _search(runtime) + + assert failure.value.reason_code == "SEARCH_TEXT_INDEX_REQUIRED" + assert len(manager.calls) == 2 + + [email protected] +async def test_vector_search_binds_json_vector_and_uses_l2_order() -> None: + runtime, manager, _ = _runtime() + + result = await _search( + runtime, + query=None, + mode="vector", + fields=None, + vector=[0.1, 0.2, 0.3], + vector_field="embedding", + filters=None, + ) + + assert result["status"] == "success" + assert result["metadata"]["vector_metric"] == "l2_distance" + query_call = manager.calls[-1] + assert "CAST(%s AS ARRAY<FLOAT>)" in query_call["sql"] + assert "ORDER BY `__mcp_vector_distance` ASC" in query_call["sql"] + assert "[0.1,0.2,0.3]" not in query_call["sql"] + assert query_call["params"] == ("[0.1,0.2,0.3]",) + + [email protected] +async def test_inner_product_vector_search_orders_descending() -> None: + runtime, manager, _ = _runtime( + _ConnectionManager(indexes=_indexes(metric="inner_product")) + ) + + await _search( + runtime, + query=None, + mode="vector", + fields=None, + vector=[0.1, 0.2, 0.3], + vector_field=None, + filters=None, + ) + + assert "inner_product_approximate" in manager.calls[-1]["sql"] + assert "ORDER BY `__mcp_vector_distance` DESC" in manager.calls[-1]["sql"] + + [email protected] +async def test_hybrid_search_binds_select_parameter_before_predicates() -> None: + runtime, manager, _ = _runtime() + + await _search( + runtime, + mode="hybrid", + vector=[0.1, 0.2, 0.3], + vector_field="embedding", + ) + + assert manager.calls[-1]["params"] == ( + "[0.1,0.2,0.3]", + "Doris", + "database", + ) + + [email protected] +async def test_vector_dimension_mismatch_fails_before_search_execution() -> None: + runtime, manager, _ = _runtime() + + with pytest.raises(SearchRuntimeFailure) as failure: + await _search( + runtime, + query=None, + mode="vector", + fields=None, + vector=[0.1, 0.2], + vector_field="embedding", + filters=None, + ) + + assert failure.value.reason_code == "SEARCH_ARGUMENT_INVALID" + assert len(manager.calls) == 2 + + [email protected] +async def test_filter_operator_is_allowlisted() -> None: + runtime, manager, _ = _runtime() + + with pytest.raises(SearchRuntimeFailure) as failure: + await _search( + runtime, + filters={ + "category": { + "operator": "eq) OR 1=1 --", + "value": "database", + } + }, + ) + + assert failure.value.reason_code == "SEARCH_ARGUMENT_INVALID" + assert len(manager.calls) == 2 + + [email protected] [email protected]( + ("payload", "expected"), + [ + ('["apache", "doris"]', ["apache", "doris"]), + ( + '[{"token":"apache","position":3},{"token":"doris","type":"word"}]', + ["apache", "doris"], + ), + ], +) +async def test_tokenize_normalizes_legacy_and_structured_payloads( + payload: str, + expected: list[str], +) -> None: + runtime, manager, _ = _runtime( + _ConnectionManager(tokenize_payload=payload) + ) + + result = await runtime.preview_text_analysis( + text="Apache Doris", + analyzer="english", + tokenizer=None, + token_filters=None, + ) + + assert [item["term"] for item in result["data"]["tokens"]] == expected + call = manager.calls[-1] + assert call["sql"] == "SELECT TOKENIZE(%s, %s) AS tokens" + assert call["params"] == ( + "Apache Doris", + '"parser"="english"', + ) + + [email protected] +async def test_custom_analyzer_components_must_match_recorded_definition() -> None: + runtime, manager, _ = _runtime( + _ConnectionManager( + analyzers=[ + { + "Name": "customer_text", + "Tokenizer": "standard", + "TokenFilters": '["lowercase", "asciifolding"]', + } + ] + ) + ) + + with pytest.raises(SearchRuntimeFailure) as failure: + await runtime.preview_text_analysis( + text="Apache Doris", + analyzer="customer_text", + tokenizer="standard", + token_filters=["lowercase"], + ) + + assert failure.value.reason_code == "SEARCH_ARGUMENT_INVALID" + assert manager.calls[-1]["sql"] == "SHOW INVERTED INDEX ANALYZER" + + [email protected] +async def test_index_inspection_normalizes_properties_and_build_tasks() -> None: + runtime, _, _ = _runtime( + _ConnectionManager( + build_tasks=[ + { + "JobId": 11, + "TableName": "documents", + "State": "FINISHED", + "Progress": "100%", + } + ] + ) + ) + + result = await runtime.inspect_search_indexes( + database="analytics", + table="documents", + index=None, + ) + + assert result["status"] == "success" + assert result["data"]["capabilities"] == { + "text": True, + "vector": True, + "hybrid": True, + "metrics": ["l2_distance"], + } + items = {item["name"]: item for item in result["data"]["items"]} + assert items["idx_embedding"]["dimension"] == 3 + assert items["idx_title"]["parser"] == "english" + assert items["idx_title"]["properties"]["comment_text"] == "café" + assert result["data"]["build_tasks"][0]["state"] == "FINISHED" + + [email protected] +async def test_diagnosis_reports_ann_plan_without_inventing_text_hits() -> None: + runtime, _, query_runtime = _runtime() + + result = await runtime.diagnose_search_query( + sql=None, + search_request={ + "database": "analytics", + "table": "documents", + "query": "Doris", + "mode": "hybrid", + "fields": ["title"], + "vector": [0.1, 0.2, 0.3], + "vector_field": "embedding", + "top_k": 5, + "return_fields": ["id", "title"], + }, + include_profile=True, + ) + + facets = result["data"]["explain"]["facets"] + assert facets["ann_pushdown_observed"] is True + assert facets["text_match_predicate_observed"] is True + assert facets["profile_required_for_inverted_hit_confirmation"] is True + assert result["metadata"]["invented_index_hits"] is False + assert result["metadata"]["profile_observed"] is False + assert result["status"] == "partial" + query_runtime.get_query_profile.assert_not_awaited() + + [email protected] +async def test_raw_diagnosis_rejects_write_and_profiles_read_only_sql() -> None: + runtime, manager, query_runtime = _runtime() + + with pytest.raises(SearchRuntimeFailure) as failure: + await runtime.diagnose_search_query( + sql="DELETE FROM analytics.documents", + search_request=None, + include_profile=False, + ) + assert failure.value.reason_code == "SEARCH_ARGUMENT_INVALID" + assert manager.calls == [] + + result = await runtime.diagnose_search_query( + sql=( + "SELECT id FROM analytics.documents " + "WHERE title MATCH_ANY 'Doris' LIMIT 5" + ), + search_request=None, + include_profile=True, + ) + + assert result["metadata"]["profile_observed"] is True + query_runtime.get_query_profile.assert_awaited_once() --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
