From 37192950b40349826f52673de8712b7f3955b13b Mon Sep 17 00:00:00 2001 From: MisterRaindrop <278811821@qq.com> Date: Sun, 13 Sep 2026 11:19:18 +0800 Subject: [PATCH] mcp-server: add vector, full-text and hybrid search tools The MCP server could describe a database and run a query, but it could not search one. An agent had no way to find rows by meaning or by keyword, which is the capability an analytical backend has to expose before it is useful for retrieval workloads. Add four tools: - list_searchable_columns reports embedding columns with their declared dimension, tsvector columns, and text columns carrying a full-text index, each with the indexes defined on it. Without it an agent has to guess a target from column names. - vector_search ranks rows by embedding distance over vector, halfvec or sparsevec columns. Filters are applied inside the same scan that walks the vector index, so the search is pre-filtered rather than filtered after truncation. - fulltext_search ranks rows by ts_rank over a text or tsvector column. - hybrid_search runs both arms and fuses them by reciprocal rank. Ranks are fused rather than raw scores because a distance and a ts_rank share no common scale, but their orderings do. Identifiers are resolved against the catalog before they reach any statement and quoted afterwards; every caller supplied value is bound as a parameter. A filter therefore cannot inject SQL, and an unknown column produces an error that lists the ones that exist. A few details are worth calling out, because each one is a silent wrong answer rather than an error if it is got wrong. Recall is exposed through probes and ef_search. At the default ivfflat.probes of 1 an index scan reads a single list and can miss the true nearest neighbour outright, and a selective pre-filter can then return nothing at all while matching rows exist. The settings are session level and reset afterwards: Cloudberry does not dispatch SET LOCAL to the segments where the scan runs, and a session level SET only reaches them once a query executor gang exists, so a transaction scoped setting would quietly do nothing. Their bounds are read from pg_settings rather than assumed, since pgvector stops hnsw.ef_search at 1000 while ivfflat.probes runs to 32768. The document expression is spelled exactly as a to_tsvector() index is, with no coalesce() wrapper and with the text search configuration written as its catalog OID. Either would read as harmless and both stop the planner matching a GIN expression index, turning every full-text search into a sequential scan. to_tsvector is strict, so a NULL document is already excluded without the wrapper. PostgreSQL's query constructors require every term to be present, so a natural sentence often matches nothing. match_mode defaults to requiring all terms and widening to any of them only when that matched nothing, and the result reports which reading produced the rows. Widening is refused when the query carries a negation or a phrase, because ORing the lexemes of `latency -dashboard` would return precisely the rows the caller excluded. Hybrid fusion identifies a row by gp_segment_id and ctid, since a ctid repeats across segments and fusing on it alone would merge unrelated rows. Plain PostgreSQL has no gp_segment_id and uses ctid by itself. Computed columns are renamed when the table already has a column of the same name, and the result reports the name each one landed on, because rows come back as positional lists. Also fix three defects in the existing query tools that this work depends on: - execute_query and explain_query passed parameters to asyncpg as keyword arguments, which asyncpg does not accept, so any call with parameters failed. They now take a list bound to $1, $2, ... in order. - explain_query ran EXPLAIN ANALYZE, which executes the statement, with no read-only validation at all. - get_table_info used a bind parameter as a FROM target, which is not allowed, so the function always failed. Tested against a three-segment Apache Cloudberry cluster with pgvector 0.8.0: 82 new tests covering the tools, their rejection paths, recall and widening behaviour, and the distributed plan shape. --- mcp-server/README.md | 158 ++- mcp-server/pyproject.toml | 4 +- mcp-server/src/cbmcp/__init__.py | 12 + mcp-server/src/cbmcp/database.py | 963 ++++++++++++++++- mcp-server/src/cbmcp/search.py | 546 ++++++++++ mcp-server/src/cbmcp/security.py | 20 + mcp-server/src/cbmcp/server.py | 157 ++- mcp-server/tests/test_search_tools.py | 1423 +++++++++++++++++++++++++ 8 files changed, 3259 insertions(+), 24 deletions(-) create mode 100644 mcp-server/src/cbmcp/search.py create mode 100644 mcp-server/tests/test_search_tools.py diff --git a/mcp-server/README.md b/mcp-server/README.md index 76c8776e722..f94c58bd506 100644 --- a/mcp-server/README.md +++ b/mcp-server/README.md @@ -24,6 +24,7 @@ A Model Communication Protocol (MCP) server for Apache Cloudberry database inter ## Features - **Database Metadata Resources**: Access schemas, tables, views, indexes, and column information +- **Retrieval Tools**: Vector, full-text, and hybrid search for agent and RAG workloads - **Safe Query Tools**: Execute parameterized SQL queries with security validation - **Administrative Tools**: Table statistics, large table analysis, and query optimization - **Context-Aware Prompts**: Predefined prompts for common database tasks @@ -114,9 +115,15 @@ python -m cbmcp.client ### Tools +#### Retrieval Tools +- `list_searchable_columns(schema, table, include_unindexed_text)` - Discover embedding columns, tsvector columns, and text columns with a full-text index +- `vector_search(schema, table, vector_column, query_vector, limit, select_columns, filters, metric, probes, ef_search)` - Nearest neighbour search, optionally pre-filtered +- `fulltext_search(schema, table, text_column, query, limit, select_columns, filters, language, query_mode, match_mode)` - Keyword search ranked by `ts_rank` +- `hybrid_search(schema, table, vector_column, text_column, query_vector, query, limit, select_columns, filters, metric, language, query_mode, match_mode, rrf_k, candidates, probes, ef_search)` - Both of the above, fused by reciprocal rank + #### Query Tools -- `execute_query(query, params, readonly)` - Execute a SQL query -- `explain_query(query, params)` - Get query execution plan +- `execute_query(query, params, readonly)` - Execute a SQL query. `params` is a list bound to `$1, $2, ...` in order +- `explain_query(query, params)` - Get query execution plan. Read-only statements only, because `EXPLAIN ANALYZE` runs what it is given - `get_table_stats(schema, table)` - Get table statistics - `list_large_tables(limit)` - List largest tables @@ -151,6 +158,152 @@ python -m cbmcp.client - `suggest_indexes` - Index recommendation guidance - `database_health_check` - Database health assessment +## Retrieval + +The retrieval tools let an agent search rather than only query. They run over +Apache Cloudberry's own storage: embeddings in `pgvector` columns, documents in +`text` or `tsvector` columns. Nothing is cached or indexed outside the database. + +### Prerequisites + +Vector search covers `vector`, `halfvec` and `sparsevec` columns. A query +vector is always supplied as a plain list of numbers; for a `sparsevec` column +it is converted to that type's own text form, which lists only the non-zero +entries. + +Vector search needs the `pgvector` extension in the target database: + +```sql +CREATE EXTENSION IF NOT EXISTS vector; +``` + +Full-text search needs no extension. A GIN index over `to_tsvector(...)` is +what makes it fast, and is also what `list_searchable_columns` looks for when +deciding whether a text column is worth reporting. + +### Discovering what to search + +An agent that cannot see which columns hold embeddings has to guess from +names. `list_searchable_columns` reports the column kind, the declared +dimension, and the indexes defined on it: + +```json +[ + {"schema": "public", "table": "docs", "column": "embedding", "kind": "vector", + "type": "vector(8)", "dimension": 8, "full_text_indexed": false, + "indexes": [{"name": "docs_embedding_ivf", "method": "ivfflat", "on_expression": false}]}, + {"schema": "public", "table": "docs", "column": "content", "kind": "text", + "type": "text", "dimension": null, "full_text_indexed": true, + "indexes": [{"name": "docs_content_gin", "method": "gin", "on_expression": true}]} +] +``` + +### Filters + +`filters` is a list of conditions applied *before* ranking, so an approximate +nearest neighbour search ranks only rows that already satisfy them: + +```json +{"filters": [{"column": "customer_id", "operator": "in", "value": [1, 2, 3]}]} +``` + +Operators are `eq`, `ne`, `lt`, `lte`, `gt`, `gte`, `in`, `not_in`, `like`, +`ilike`, `is_null`, `is_not_null`. Column names are checked against the catalog +and values are bound as query parameters, so a filter cannot inject SQL. For +anything these do not express, use `execute_query`. + +### Matching every term, or any of them + +PostgreSQL's query constructors require *every* term to be present, so +`query latency slow` becomes `'queri' & 'latenc' & 'slow'`. A document holding +the first two but not the third does not match. The longer and more natural the +query, the likelier it is that nothing matches at all, and the caller gets an +empty result rather than an error. To an agent that reads as "this database +holds nothing on the subject", which is the wrong conclusion when most of the +terms did match something. + +`match_mode` decides how the terms are combined: + +| Mode | Behaviour | +| --- | --- | +| `all` | Every term must appear. Precise, and empty when one term is missing | +| `any` | Any term may appear. `ts_rank` still puts the fuller matches first | +| `all_then_any` | Requires every term, and widens to any of them only if that matched nothing. The default | + +A query that carries an explicit operator is never widened. ORing the lexemes +of `latency -dashboard` would return exactly the rows the caller excluded, and +flattening a phrase is the opposite of asking for one, so a negation or a +phrase is read as an instruction to leave the query alone. The result then +reports `widening_refused: true` alongside an empty row set. + +The result always says which reading produced the rows, so widening is never +silent: + +```json +{"search": {"match_mode": "all_then_any", "matched_with": "any", "widened": true}} +``` + +Under `all_then_any` the existence check is how the decision is made, so it +runs on every call, not only when the strict reading fails. With a full-text +index in place it is an index probe. Widening cannot rescue a query made +entirely of stop words, because the parser produces no lexemes to widen to. + +This matters most in `hybrid_search`. Without it the keyword arm can go empty +while the vector arm still returns rows, so the fused result is pure vector +search presenting itself as hybrid, and nothing in the response says so. + +### Recall + +An IVFFlat index reads `ivfflat.probes` lists per scan, and the default of 1 +can miss the true nearest neighbour outright. `probes` (IVFFlat) and +`ef_search` (HNSW) trade latency for recall, and the settings that were applied +come back in the result so a run can be reproduced: + +```json +{"search": {"mode": "vector", "settings": {"ivfflat.probes": 20}}} +``` + +The settings are session level and reset once the statement finishes, because +Apache Cloudberry does not dispatch `SET LOCAL` to the segments where the index +scan runs. Their upper bounds come from the installed pgvector rather than from +this server, so a value it would refuse is rejected before anything is applied: +`hnsw.ef_search` stops at 1000 while `ivfflat.probes` runs to 32768. + +Recall also decides whether a pre-filtered search finds anything at all. The +filter is applied inside the index scan, so a selective filter combined with a +low probe count can return no rows even though matching rows exist. Raise +`probes` before concluding that a filter matched nothing. + +### Results + +Every retrieval tool returns the rows, the generated SQL, and what it searched: + +```json +{ + "columns": ["id", "content", "distance"], + "rows": [[4138, "...", 0.27]], + "row_count": 1, + "sql": "SELECT ... ORDER BY \"embedding\" <-> $1::text::\"vector\" LIMIT $2", + "search": {"mode": "vector", "metric": "l2", "limit": 1} +} +``` + +Embedding columns are left out of the default projection because they are large +and of no use to a reader; name one in `select_columns` and it comes back as +text. + +A table is free to have a column of its own called `rank` or `score`. Because +rows come back as positional lists, the computed column is renamed rather than +duplicated, and `search` names it: `distance_column`, `rank_column`, +`score_column`, `vector_rank_column`, `text_rank_column`. Read the name from +there rather than assuming it. + +`hybrid_search` adds `score`, `vector_rank` and `text_rank`. A `null` rank means +that arm did not return the row. A row scores +`1 / (rrf_k + rank)` from each arm it appears in, so a row that only one arm +found still places. Ranks are fused rather than raw scores because a distance +and a `ts_rank` share no scale, but their orderings do. + ## Security Features - **SQL Injection Prevention**: Comprehensive query validation @@ -158,6 +311,7 @@ python -m cbmcp.client - **Parameterized Queries**: Safe parameter handling - **Connection Pooling**: Secure connection management - **Sensitive Table Protection**: Blocks access to system tables +- **Catalog-Checked Identifiers**: Retrieval tools resolve every schema, table and column name against the catalog before it reaches a statement, and quote it afterwards ## Quick Start with Cloudberry Demo Cluster diff --git a/mcp-server/pyproject.toml b/mcp-server/pyproject.toml index 984cb5e12a1..e1f17ae466d 100644 --- a/mcp-server/pyproject.toml +++ b/mcp-server/pyproject.toml @@ -57,7 +57,9 @@ dependencies = [ [project.optional-dependencies] dev = [ "pytest>=7.0.0", - "pytest-asyncio>=0.21.0", + # 0.24 introduced loop_scope=, which the retrieval tests rely on to share + # one event loop across a module-scoped fixture. + "pytest-asyncio>=0.24.0", "pytest-cov>=4.0.0", ] diff --git a/mcp-server/src/cbmcp/__init__.py b/mcp-server/src/cbmcp/__init__.py index a584c415ee6..173f44129f3 100644 --- a/mcp-server/src/cbmcp/__init__.py +++ b/mcp-server/src/cbmcp/__init__.py @@ -19,6 +19,13 @@ from .client import CloudberryMCPClient from .config import DatabaseConfig, ServerConfig from .database import DatabaseManager +from .search import ( + FilterOperator, + SearchFilter, + SearchMetric, + TextMatchMode, + TextQueryMode, +) from .security import SQLValidator __version__ = "0.1.0" @@ -28,5 +35,10 @@ "DatabaseConfig", "ServerConfig", "DatabaseManager", + "FilterOperator", + "SearchFilter", + "SearchMetric", + "TextMatchMode", + "TextQueryMode", "SQLValidator", ] \ No newline at end of file diff --git a/mcp-server/src/cbmcp/database.py b/mcp-server/src/cbmcp/database.py index 77cdd2b308b..6eea1e1bbe4 100644 --- a/mcp-server/src/cbmcp/database.py +++ b/mcp-server/src/cbmcp/database.py @@ -22,11 +22,39 @@ """ import logging -from typing import Any, Dict, Optional +import re +from typing import Any, Dict, List, Optional, Sequence from contextlib import asynccontextmanager import asyncpg from .config import DatabaseConfig +from .search import ( + MAX_CANDIDATES, + SearchMetric, + SqlParams, + TEXT_TYPES, + TSVECTOR_TYPE, + TextMatchMode, + TextQueryMode, + VECTOR_TYPES, + build_filter_sql, + build_select_list, + build_tsquery, + coerce_filters, + format_vector_literal, + quote_column, + resolve_column, + resolve_match_mode, + resolve_metric, + resolve_query_mode, + tsquery_has_blocking_operator, + unique_alias, + validate_limit, + validate_positive_int, + validate_query_vector, + validate_vector_column, + warmup_literal, +) from .security import SQLValidator @@ -53,19 +81,52 @@ async def get_connection(self): min_size=1, max_size=10, command_timeout=60.0, + init=self._register_vector_codecs, ) try: async with self._connection_pool.acquire() as conn: yield conn except Exception as e: - logger.error(f"Database connection error: {e}") + logger.error(f"Error while using database connection: {e}") raise + @staticmethod + async def _register_vector_codecs(conn) -> None: + """Decode pgvector columns as their text literal. + + asyncpg ships no codec for the types pgvector adds, so selecting one + of those columns otherwise fails with an unknown-type error. The types + exist only where the extension is installed, so a database without it + simply registers nothing. + """ + try: + records = await conn.fetch( + "SELECT t.typname, n.nspname FROM pg_type t " + "JOIN pg_namespace n ON n.oid = t.typnamespace " + "WHERE t.typname = ANY($1::text[])", + sorted(VECTOR_TYPES), + ) + except Exception as e: + logger.debug(f"Could not look up pgvector types: {e}") + return + + for record in records: + try: + await conn.set_type_codec( + record["typname"], + schema=record["nspname"], + encoder=str, + decoder=str, + format="text", + ) + except Exception as e: + logger.debug(f"Could not register codec for {record['typname']}: {e}") + async def execute_query( self, query: str, - params: Optional[Dict[str, Any]] = None, + params: Optional[Sequence[Any]] = None, readonly: bool = True ) -> Dict[str, Any]: """Execute a SQL query with safety validation""" @@ -81,12 +142,16 @@ async def execute_query( try: async with self.get_connection() as conn: if params: - # Sanitize parameter names - sanitized_params = { - SQLValidator.sanitize_parameter_name(k): v - for k, v in params.items() - } - result = await conn.fetch(query, **sanitized_params) + # asyncpg binds by position, so the query uses $1, $2, ... + # and the values arrive in the same order. + if not isinstance(params, (list, tuple)): + return { + "error": ( + f"params must be a list of values bound to $1, $2, ... " + f"in order, got {type(params).__name__}" + ) + } + result = await conn.fetch(query, *params) else: result = await conn.fetch(query) @@ -129,13 +194,20 @@ async def get_table_info(self, schema: str, table: str) -> Dict[str, Any]: ) # Get table statistics + # The size functions take a regclass, which can be bound as a + # parameter. COUNT(*) needs the relation in the FROM clause, + # where a parameter is not allowed, so the name is quoted + # instead. Quoting is what makes that safe: an embedded quote + # is doubled, so the name can only ever read as one identifier. + qualified_name = SQLValidator.quote_qualified_name(schema, table) stats = await conn.fetchrow( "SELECT " - "pg_size_pretty(pg_total_relation_size($1)) as total_size, " - "pg_size_pretty(pg_relation_size($1)) as table_size, " - "pg_size_pretty(pg_total_relation_size($1) - pg_relation_size($1)) as indexes_size, " - "(SELECT COUNT(*) FROM $1) as row_count", - f"{schema}.{table}" + "pg_size_pretty(pg_total_relation_size($1::regclass)) as total_size, " + "pg_size_pretty(pg_relation_size($1::regclass)) as table_size, " + "pg_size_pretty(pg_total_relation_size($1::regclass) " + "- pg_relation_size($1::regclass)) as indexes_size, " + f"(SELECT COUNT(*) FROM {qualified_name}) as row_count", + qualified_name ) return { @@ -477,12 +549,31 @@ async def list_referenced_tables(self, schema: str, table: str) -> list[dict]: for r in records ] - async def explain_query(self, query: str, params: Optional[dict] = None) -> str: - """Get the execution plan for a query""" + async def explain_query(self, query: str, params: Optional[Sequence[Any]] = None) -> str: + """Get the execution plan for a query + + EXPLAIN ANALYZE runs the statement it is given, so the query is held to + the same read-only rule as execute_query rather than passed straight + through. + """ + is_valid, error_msg = SQLValidator.validate_query(query) + if not is_valid: + return f"Query validation failed: {error_msg}" + if not SQLValidator.is_readonly_query(query): + return ( + "Only read-only queries can be explained, because EXPLAIN ANALYZE " + "executes the statement" + ) + try: async with self.get_connection() as conn: if params: - result = await conn.fetch(f"EXPLAIN (ANALYZE, BUFFERS) {query}", **params) + if not isinstance(params, (list, tuple)): + return ( + f"params must be a list of values bound to $1, $2, ... in " + f"order, got {type(params).__name__}" + ) + result = await conn.fetch(f"EXPLAIN (ANALYZE, BUFFERS) {query}", *params) else: result = await conn.fetch(f"EXPLAIN (ANALYZE, BUFFERS) {query}") @@ -770,4 +861,840 @@ async def list_active_connections(self) -> list[dict]: "query_start": str(r["query_start"]) if r["query_start"] else None } for r in records - ] \ No newline at end of file + ] + + # ------------------------------------------------------------------ + # Retrieval: vector, full-text, and hybrid search + # ------------------------------------------------------------------ + + async def _describe_relation(self, conn, schema: str, table: str) -> dict: + """Read a relation's columns from the catalog. + + Every retrieval method starts here. Resolving caller supplied names + against the catalog is what lets the builders put identifiers into + statement text at all: a name that no column matches never gets there. + """ + records = await conn.fetch( + "SELECT c.oid AS relid, a.attname AS column_name, t.typname AS type_name, " + "format_type(a.atttypid, a.atttypmod) AS type_display, " + "a.atttypmod AS type_modifier, a.attnum AS attribute_number " + "FROM pg_class c " + "JOIN pg_namespace n ON n.oid = c.relnamespace " + "JOIN pg_attribute a ON a.attrelid = c.oid " + "JOIN pg_type t ON t.oid = a.atttypid " + "WHERE n.nspname = $1 AND c.relname = $2 " + "AND a.attnum > 0 AND NOT a.attisdropped " + "ORDER BY a.attnum", + schema, table + ) + if not records: + raise ValueError( + f'Relation "{schema}"."{table}" was not found, or it has no readable columns' + ) + + return { + "relid": records[0]["relid"], + "schema": schema, + "table": table, + "qualified_name": SQLValidator.quote_qualified_name(schema, table), + "columns": [ + { + "column_name": record["column_name"], + "type_name": record["type_name"], + "type_display": record["type_display"], + "type_modifier": record["type_modifier"], + "attribute_number": record["attribute_number"], + } + for record in records + ], + } + + async def _row_key_columns(self, conn, relid: int) -> List[str]: + """Return the system columns that identify a row while fusing results. + + Hybrid search ranks the same table twice and has to recognise a row + that both arms returned. On Cloudberry a ctid is unique only within one + segment, so the segment id belongs in the key; plain PostgreSQL has no + gp_segment_id and ctid alone suffices. + """ + has_segment_id = await conn.fetchval( + "SELECT EXISTS (SELECT 1 FROM pg_attribute " + "WHERE attrelid = $1 AND attname = 'gp_segment_id' AND attnum < 0)", + relid, + ) + return ["gp_segment_id", "ctid"] if has_segment_id else ["ctid"] + + async def _resolve_language(self, conn, language: str) -> tuple: + """Resolve a text search configuration to its name and OID. + + The OID matters for more than validation. A full-text index is built + over to_tsvector('english', col), which the catalog stores with the + configuration already resolved to an OID constant. Writing that same + OID into the query reproduces the constant exactly, so the expression + index still matches; binding the name as a parameter does not, and the + index is quietly skipped once the plan is cached and generic. + """ + if not isinstance(language, str) or not language.strip(): + raise ValueError( + f"language must be a text search configuration name, got {language!r}" + ) + + oid = await conn.fetchval( + "SELECT oid FROM pg_ts_config WHERE cfgname = $1", language + ) + if oid is None: + records = await conn.fetch("SELECT cfgname FROM pg_ts_config ORDER BY cfgname") + available = ", ".join(record["cfgname"] for record in records) + raise ValueError( + f"language '{language}' is not a known text search configuration; " + f"available configurations are: {available}" + ) + # An OID is an integer read out of the catalog, so interpolating it + # cannot inject anything. + return language, f"{int(oid)}::regconfig" + + # Cloudberry dispatches a SET to the segments only when a query executor + # gang already exists. Issued on an otherwise idle connection the value + # reaches the coordinator and is echoed back by current_setting(), yet the + # segment scans that actually read it keep the built-in default. A trivial + # distributed statement creates the gang first. + GANG_WARMUP_SQL = "SELECT count(*) FROM gp_dist_random('gp_id')" + + async def _prepare_vector_gucs( + self, + conn, + type_name: Optional[str], + probes: Optional[int], + ef_search: Optional[int], + ) -> Dict[str, int]: + """Set the pgvector recall knobs for the next statement on this connection. + + Recall is not a property of the query alone: at the default + ivfflat.probes of 1 an index scan reads a single list and can miss the + true nearest neighbour outright. Exposing these lets a caller trade + latency for recall, and measure the difference. + + The settings are session level rather than transaction local because + SET LOCAL and set_config(..., true) are not dispatched to the segments, + so a transaction-scoped setting would silently do nothing. The caller + resets them once the statement is done. + """ + requested: Dict[str, int] = {} + if probes is not None: + requested["ivfflat.probes"] = probes + if ef_search is not None: + requested["hnsw.ef_search"] = ef_search + if not requested: + return requested + + if type_name: + # Load pgvector on the coordinator, which matters when the scan + # does not go out to the segments at all. + literal = warmup_literal(type_name) + quoted_type = SQLValidator.quote_identifier(type_name) + await conn.fetchval(f"SELECT '{literal}'::{quoted_type} IS NOT NULL") + + try: + await conn.fetchval(self.GANG_WARMUP_SQL) + except Exception as e: + # Plain PostgreSQL has no segments and no gp_dist_random(). + logger.debug(f"Skipping segment warm-up: {e}") + + await self._check_guc_bounds(conn, requested) + + applied: Dict[str, int] = {} + try: + for name, value in requested.items(): + # Both the name and the value are server-controlled: the names + # are module constants and the values were checked to be + # integers. + await conn.execute(f"SET {name} = {int(value)}") + applied[name] = value + except Exception: + # Undo whatever already took effect rather than leaving it to the + # driver's reset when the connection goes back to the pool. + await self._reset_vector_gucs(conn, applied) + raise + return applied + + @staticmethod + async def _check_guc_bounds(conn, requested: Dict[str, int]) -> None: + """Reject a recall setting the server would not accept. + + The real limits belong to the installed pgvector, not to this wrapper: + hnsw.ef_search stops at 1000 while ivfflat.probes runs to 32768. Asking + the server for its own declared range gives an accurate error at the + API boundary instead of a raw failure part way through applying the + settings. + """ + records = await conn.fetch( + "SELECT name, min_val, max_val FROM pg_settings WHERE name = ANY($1::text[])", + sorted(requested), + ) + bounds = { + record["name"]: (int(record["min_val"]), int(record["max_val"])) + for record in records + } + + for name, value in requested.items(): + if name not in bounds: + raise ValueError( + f"{name} is not available on this server; it is provided by the " + f"pgvector extension, which may not be installed" + ) + low, high = bounds[name] + if not low <= value <= high: + raise ValueError( + f"{name} must be between {low} and {high} on this server, got {value}" + ) + + async def _reset_vector_gucs(self, conn, applied: Dict[str, int]) -> None: + """Undo _prepare_vector_gucs so a pooled connection is handed back clean.""" + for name in applied: + try: + await conn.execute(f"RESET {name}") + except Exception as e: + logger.warning(f"Could not reset {name}: {e}") + + async def _run_search( + self, + conn, + sql: str, + params: SqlParams, + vector_type: Optional[str] = None, + probes: Optional[int] = None, + ef_search: Optional[int] = None, + meta: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Execute a built statement and shape the result for an agent.""" + applied = await self._prepare_vector_gucs(conn, vector_type, probes, ef_search) + try: + records = await conn.fetch(sql, *params.values) + finally: + if applied: + await self._reset_vector_gucs(conn, applied) + + search_meta = dict(meta or {}) + if applied: + search_meta["settings"] = applied + + return { + "columns": list(records[0].keys()) if records else [], + "rows": [list(record.values()) for record in records], + "row_count": len(records), + "sql": sql, + "search": search_meta, + } + + @staticmethod + def _expression_names_column(expression: Optional[str], column_name: str) -> bool: + """Report whether an index expression references a given column.""" + if not expression: + return False + pattern = rf'(? List[dict]: + """Report the columns an agent can actually run a search against. + + An agent that cannot see which columns hold embeddings, and which text + columns carry a full-text index, has no way to choose a target other + than guessing from names. + """ + searchable_types = sorted(VECTOR_TYPES | TEXT_TYPES | {TSVECTOR_TYPE}) + + conditions = [ + "c.relkind IN ('r', 'p', 'f', 'm', 'v')", + "a.attnum > 0", + "NOT a.attisdropped", + "n.nspname NOT IN ('pg_catalog', 'information_schema', 'pg_toast')", + "t.typname = ANY($1::text[])", + ] + values: List[Any] = [searchable_types] + if schema: + values.append(schema) + conditions.append(f"n.nspname = ${len(values)}") + if table: + values.append(table) + conditions.append(f"c.relname = ${len(values)}") + + async with self.get_connection() as conn: + records = await conn.fetch( + "SELECT c.oid AS relid, n.nspname AS schema_name, c.relname AS table_name, " + "a.attname AS column_name, a.attnum AS attribute_number, " + "t.typname AS type_name, " + "format_type(a.atttypid, a.atttypmod) AS type_display, " + "a.atttypmod AS type_modifier " + "FROM pg_class c " + "JOIN pg_namespace n ON n.oid = c.relnamespace " + "JOIN pg_attribute a ON a.attrelid = c.oid " + "JOIN pg_type t ON t.oid = a.atttypid " + f"WHERE {' AND '.join(conditions)} " + "ORDER BY n.nspname, c.relname, a.attnum", + *values, + ) + if not records: + return [] + + relids = sorted({record["relid"] for record in records}) + index_records = await conn.fetch( + "SELECT i.indrelid AS relid, ic.relname AS index_name, " + "am.amname AS method, pg_get_indexdef(i.indexrelid) AS definition, " + "pg_get_expr(i.indexprs, i.indrelid) AS expression, " + "string_to_array(i.indkey::text, ' ')::int[] AS index_keys " + "FROM pg_index i " + "JOIN pg_class ic ON ic.oid = i.indexrelid " + "JOIN pg_am am ON am.oid = ic.relam " + "WHERE i.indrelid = ANY($1::oid[]) " + "ORDER BY ic.relname", + relids, + ) + + indexes_by_relid: Dict[int, List[dict]] = {} + for record in index_records: + indexes_by_relid.setdefault(record["relid"], []).append(dict(record)) + + results: List[dict] = [] + for record in records: + column_name = record["column_name"] + attribute_number = record["attribute_number"] + type_name = record["type_name"] + + matching: List[dict] = [] + for index in indexes_by_relid.get(record["relid"], []): + keys = index["index_keys"] or [] + on_column = attribute_number in keys + # An index key of 0 marks an expression. pg_get_indexdef only + # quotes an identifier that needs it, so a to_tsvector(content) + # index names the column bare; match it either way. + on_expression = 0 in keys and self._expression_names_column( + index.get("expression"), column_name + ) + if not on_column and not on_expression: + continue + matching.append( + { + "name": index["index_name"], + "method": index["method"], + "definition": index["definition"], + "on_expression": on_expression and not on_column, + } + ) + + if type_name in VECTOR_TYPES: + kind = "vector" + elif type_name == TSVECTOR_TYPE: + kind = "tsvector" + else: + kind = "text" + + full_text_indexed = any( + index["method"] in ("gin", "gist") + and ("to_tsvector" in index["definition"] or kind == "tsvector") + for index in matching + ) + + # Plain text columns without a full-text index are the bulk of any + # schema and drown out the columns worth searching. + if kind == "text" and not full_text_indexed and not include_unindexed_text: + continue + + dimension = None + if kind == "vector" and record["type_modifier"] and record["type_modifier"] > 0: + dimension = record["type_modifier"] + + results.append( + { + "schema": record["schema_name"], + "table": record["table_name"], + "column": column_name, + "kind": kind, + "type": record["type_display"], + "dimension": dimension, + "full_text_indexed": full_text_indexed, + "indexes": matching, + } + ) + + return results + + async def vector_search( + self, + schema: str, + table: str, + vector_column: str, + query_vector: Sequence[float], + limit: int = 10, + select_columns: Optional[Sequence[str]] = None, + filters: Optional[Sequence[Any]] = None, + metric: str = SearchMetric.L2.value, + probes: Optional[int] = None, + ef_search: Optional[int] = None, + ) -> Dict[str, Any]: + """Rank rows by embedding distance, optionally pre-filtered. + + Filters are applied in the same scan that walks the vector index, so + this is a pre-filtered nearest neighbour search rather than a filter + over an already truncated result. + """ + limit = validate_limit(limit) + if probes is not None: + validate_positive_int(probes, "probes") + if ef_search is not None: + validate_positive_int(ef_search, "ef_search") + conditions = coerce_filters(filters) + operator = resolve_metric(metric) + # Checked before the column lookup so a malformed vector is reported + # as such, rather than as a dimension mismatch. + validate_query_vector(query_vector) + + async with self.get_connection() as conn: + relation = await self._describe_relation(conn, schema, table) + columns = relation["columns"] + column = resolve_column(vector_column, columns, "vector_column") + validate_vector_column(column, query_vector) + vector_literal = format_vector_literal(query_vector, column["type_name"]) + + params = SqlParams() + quoted_type = SQLValidator.quote_identifier(column["type_name"]) + # The literal is bound as text and cast, because asyncpg has no + # wire format for an embedding type. + distance = ( + f"{quote_column(column)} {operator} " + f"{params.add(vector_literal)}::text::{quoted_type}" + ) + + projection, output_names = build_select_list(select_columns, columns) + distance_column = unique_alias("distance", output_names) + projection.append( + f"({distance})::float8 AS " + f"{SQLValidator.quote_identifier(distance_column)}" + ) + + where = build_filter_sql(conditions, columns, params) + sql = ( + f"SELECT {', '.join(projection)}\n" + f"FROM {relation['qualified_name']}\n" + + (f"WHERE {where}\n" if where else "") + + f"ORDER BY {distance}\n" + f"LIMIT {params.add(limit)}" + ) + + return await self._run_search( + conn, + sql, + params, + vector_type=column["type_name"], + probes=probes, + ef_search=ef_search, + meta={ + "mode": "vector", + "schema": schema, + "table": table, + "vector_column": column["column_name"], + "metric": metric.value if isinstance(metric, SearchMetric) else metric, + "distance_column": distance_column, + "dimensions": len(query_vector), + "limit": limit, + "filters": len(conditions), + }, + ) + + async def fulltext_search( + self, + schema: str, + table: str, + text_column: str, + query: str, + limit: int = 10, + select_columns: Optional[Sequence[str]] = None, + filters: Optional[Sequence[Any]] = None, + language: str = "english", + query_mode: str = TextQueryMode.WEBSEARCH.value, + match_mode: str = TextMatchMode.ALL_THEN_ANY.value, + ) -> Dict[str, Any]: + """Rank rows by full-text relevance over one text or tsvector column.""" + limit = validate_limit(limit) + conditions = coerce_filters(filters) + query_function = resolve_query_mode(query_mode) + match_mode = resolve_match_mode(match_mode) + if not isinstance(query, str) or not query.strip(): + raise ValueError(f"query must be a non-empty search string, got {query!r}") + + async with self.get_connection() as conn: + relation = await self._describe_relation(conn, schema, table) + columns = relation["columns"] + column = resolve_column(text_column, columns, "text_column") + if column["type_name"] != TSVECTOR_TYPE and column["type_name"] not in TEXT_TYPES: + supported = ", ".join(sorted(TEXT_TYPES | {TSVECTOR_TYPE})) + raise ValueError( + f"text_column '{column['column_name']}' has type " + f"{column['type_display']}, which cannot be searched as text; " + f"expected one of: {supported}" + ) + language, language_sql = await self._resolve_language(conn, language) + matched_with, widening_refused = await self._choose_match_mode( + conn, relation, column, language_sql, query, query_function, + match_mode, conditions, + ) + + params = SqlParams() + query_placeholder = params.add(query) + + document, tsquery, rank = self._build_text_expressions( + column, language_sql, query_placeholder, query_function, + matched_with, + ) + + projection, output_names = build_select_list(select_columns, columns) + rank_column = unique_alias("rank", output_names) + projection.append( + f"({rank})::float8 AS {SQLValidator.quote_identifier(rank_column)}" + ) + + where_parts = [f"{document} @@ {tsquery}"] + filter_sql = build_filter_sql(conditions, columns, params) + if filter_sql: + where_parts.append(filter_sql) + + sql = ( + f"SELECT {', '.join(projection)}\n" + f"FROM {relation['qualified_name']}\n" + f"WHERE {' AND '.join(where_parts)}\n" + f"ORDER BY {rank} DESC\n" + f"LIMIT {params.add(limit)}" + ) + + return await self._run_search( + conn, + sql, + params, + meta={ + "mode": "fulltext", + "schema": schema, + "table": table, + "text_column": column["column_name"], + "language": language, + "query_mode": ( + query_mode.value + if isinstance(query_mode, TextQueryMode) + else query_mode + ), + "match_mode": match_mode, + "matched_with": matched_with, + # True only when a strict reading was wanted and gave up, + # which is the signal a caller has to notice. + "widened": ( + matched_with == TextMatchMode.ANY.value + and match_mode != TextMatchMode.ANY.value + ), + "widening_refused": widening_refused, + "rank_column": rank_column, + "limit": limit, + "filters": len(conditions), + }, + ) + + @staticmethod + def _build_text_expressions( + column: Dict[str, Any], + language_sql: str, + query_placeholder: str, + query_function: str, + match_mode: str = TextMatchMode.ALL.value, + alias: Optional[str] = None, + ) -> tuple: + """Build the document, tsquery, and rank expressions for a column. + + A column already stored as tsvector is used as-is; anything else is + parsed on the fly, spelled exactly as a to_tsvector() expression index + would be so that such an index still matches. + + The column is deliberately not wrapped in coalesce(). It would read as + harmless NULL handling, but to_tsvector is strict, so a NULL document + already yields NULL and NULL @@ query already drops the row. The + wrapper buys nothing and costs the index: matching compares expression + trees, and coalesce(col, '') is not col, so every search would fall + back to a sequential scan. + """ + reference = quote_column(column, alias) + if column["type_name"] == TSVECTOR_TYPE: + document = reference + else: + document = f"to_tsvector({language_sql}, {reference})" + tsquery = build_tsquery( + query_function, language_sql, query_placeholder, match_mode + ) + return document, tsquery, f"ts_rank({document}, {tsquery})" + + async def _choose_match_mode( + self, + conn, + relation: Dict[str, Any], + column: Dict[str, Any], + language_sql: str, + query: str, + query_function: str, + match_mode: str, + conditions: Sequence[Any], + ) -> tuple: + """Decide which term combination the search will actually use. + + Requiring every term is the precise reading of a query, and it is the + right one whenever it matches something. When it matches nothing the + caller gets an empty result rather than an error, which reads as "this + database holds nothing on the subject" and is usually wrong. So check + first, and widen only when the strict reading found nothing. + + Widening is refused outright when the query carries a negation or a + phrase. ORing the lexemes of `latency -dashboard` returns precisely the + rows the caller excluded, and flattening a phrase is the opposite of + asking for one; an empty result is the better answer there. + + Returns the mode to use and whether widening was refused. + """ + if match_mode != TextMatchMode.ALL_THEN_ANY.value: + return match_mode, False + + probe = SqlParams() + document, tsquery, _ = self._build_text_expressions( + column, + language_sql, + probe.add(query), + query_function, + TextMatchMode.ALL.value, + ) + where = [f"{document} @@ {tsquery}"] + filter_sql = build_filter_sql(conditions, relation["columns"], probe) + if filter_sql: + where.append(filter_sql) + + record = await conn.fetchrow( + f"SELECT EXISTS (SELECT 1 FROM {relation['qualified_name']} " + f"WHERE {' AND '.join(where)}) AS matched, " + f"({tsquery})::text AS strict_query", + *probe.values, + ) + if record["matched"]: + return TextMatchMode.ALL.value, False + if tsquery_has_blocking_operator(record["strict_query"]): + return TextMatchMode.ALL.value, True + return TextMatchMode.ANY.value, False + + async def hybrid_search( + self, + schema: str, + table: str, + vector_column: str, + text_column: str, + query_vector: Sequence[float], + query: str, + limit: int = 10, + select_columns: Optional[Sequence[str]] = None, + filters: Optional[Sequence[Any]] = None, + metric: str = SearchMetric.L2.value, + language: str = "english", + query_mode: str = TextQueryMode.WEBSEARCH.value, + match_mode: str = TextMatchMode.ALL_THEN_ANY.value, + rrf_k: int = 60, + candidates: Optional[int] = None, + probes: Optional[int] = None, + ef_search: Optional[int] = None, + ) -> Dict[str, Any]: + """Combine vector and full-text ranking with reciprocal rank fusion. + + Each arm contributes 1 / (rrf_k + rank). Fusing ranks rather than + scores is what makes the two comparable: a distance and a ts_rank have + no common scale, but their orderings do. + """ + limit = validate_limit(limit) + if candidates is None: + candidates = min(max(limit * 5, 50), MAX_CANDIDATES) + candidates = validate_limit(candidates, "candidates", maximum=MAX_CANDIDATES) + if candidates < limit: + raise ValueError( + f"candidates must be at least limit; got candidates={candidates} " + f"and limit={limit}" + ) + rrf_k = validate_limit(rrf_k, "rrf_k", maximum=MAX_CANDIDATES) + if probes is not None: + validate_positive_int(probes, "probes") + if ef_search is not None: + validate_positive_int(ef_search, "ef_search") + + conditions = coerce_filters(filters) + operator = resolve_metric(metric) + query_function = resolve_query_mode(query_mode) + match_mode = resolve_match_mode(match_mode) + validate_query_vector(query_vector) + if not isinstance(query, str) or not query.strip(): + raise ValueError(f"query must be a non-empty search string, got {query!r}") + + async with self.get_connection() as conn: + relation = await self._describe_relation(conn, schema, table) + columns = relation["columns"] + qualified_name = relation["qualified_name"] + + embedding = resolve_column(vector_column, columns, "vector_column") + validate_vector_column(embedding, query_vector) + vector_literal = format_vector_literal(query_vector, embedding["type_name"]) + document_column = resolve_column(text_column, columns, "text_column") + if ( + document_column["type_name"] != TSVECTOR_TYPE + and document_column["type_name"] not in TEXT_TYPES + ): + supported = ", ".join(sorted(TEXT_TYPES | {TSVECTOR_TYPE})) + raise ValueError( + f"text_column '{document_column['column_name']}' has type " + f"{document_column['type_display']}, which cannot be searched as " + f"text; expected one of: {supported}" + ) + language, language_sql = await self._resolve_language(conn, language) + # Decided before the fused statement is built: re-running it after + # the fact would redo the vector arm's work for nothing. + matched_with, widening_refused = await self._choose_match_mode( + conn, relation, document_column, language_sql, query, + query_function, match_mode, conditions, + ) + + params = SqlParams() + quoted_type = SQLValidator.quote_identifier(embedding["type_name"]) + distance = ( + f"{quote_column(embedding)} {operator} " + f"{params.add(vector_literal)}::text::{quoted_type}" + ) + document, tsquery, rank = self._build_text_expressions( + document_column, + language_sql, + params.add(query), + query_function, + matched_with, + ) + + # Built once and referenced by both arms: the placeholders are + # already assigned, so reusing the text reuses the same values. + filter_sql = build_filter_sql(conditions, columns, params) + candidates_placeholder = params.add(candidates) + rrf_placeholder = params.add(rrf_k) + + key_columns = await self._row_key_columns(conn, relation["relid"]) + key_aliases = [f"rowkey_{index}" for index in range(len(key_columns))] + key_select = ", ".join( + f"{SQLValidator.quote_identifier(name)} AS {alias}" + for name, alias in zip(key_columns, key_aliases) + ) + + vector_arm = ( + f" SELECT {key_select},\n" + f" row_number() OVER (ORDER BY {distance}) AS rnk\n" + f" FROM {qualified_name}\n" + + (f" WHERE {filter_sql}\n" if filter_sql else "") + + f" ORDER BY {distance}\n" + f" LIMIT {candidates_placeholder}" + ) + + text_where = [f"{document} @@ {tsquery}"] + if filter_sql: + text_where.append(filter_sql) + text_arm = ( + f" SELECT {key_select},\n" + f" row_number() OVER (ORDER BY {rank} DESC) AS rnk\n" + f" FROM {qualified_name}\n" + f" WHERE {' AND '.join(text_where)}\n" + f" ORDER BY {rank} DESC\n" + f" LIMIT {candidates_placeholder}" + ) + + fused_keys = ", ".join( + f"COALESCE(v.{alias}, f.{alias}) AS {alias}" for alias in key_aliases + ) + join_condition = " AND ".join(f"v.{alias} = f.{alias}" for alias in key_aliases) + rejoin_condition = " AND ".join( + f"d.{SQLValidator.quote_identifier(name)} = fused.{alias}" + for name, alias in zip(key_columns, key_aliases) + ) + fusion = ( + f" SELECT {fused_keys},\n" + f" COALESCE(1.0 / ({rrf_placeholder} + v.rnk), 0)\n" + f" + COALESCE(1.0 / ({rrf_placeholder} + f.rnk), 0) AS score,\n" + f" v.rnk AS vector_rank, f.rnk AS text_rank\n" + f" FROM vector_hits v FULL JOIN text_hits f ON {join_condition}" + ) + + projection, output_names = build_select_list( + select_columns, columns, alias="d" + ) + score_column = unique_alias("score", output_names) + vector_rank_column = unique_alias( + "vector_rank", output_names + [score_column] + ) + text_rank_column = unique_alias( + "text_rank", output_names + [score_column, vector_rank_column] + ) + projection.extend( + [ + f"fused.score::float8 AS " + f"{SQLValidator.quote_identifier(score_column)}", + f"fused.vector_rank AS " + f"{SQLValidator.quote_identifier(vector_rank_column)}", + f"fused.text_rank AS " + f"{SQLValidator.quote_identifier(text_rank_column)}", + ] + ) + + sql = ( + f"WITH vector_hits AS (\n{vector_arm}\n),\n" + f"text_hits AS (\n{text_arm}\n),\n" + f"fused AS (\n{fusion}\n)\n" + f"SELECT {', '.join(projection)}\n" + f"FROM fused JOIN {qualified_name} d ON {rejoin_condition}\n" + f"ORDER BY fused.score DESC\n" + f"LIMIT {params.add(limit)}" + ) + + return await self._run_search( + conn, + sql, + params, + vector_type=embedding["type_name"], + probes=probes, + ef_search=ef_search, + meta={ + "mode": "hybrid", + "schema": schema, + "table": table, + "vector_column": embedding["column_name"], + "text_column": document_column["column_name"], + "metric": metric.value if isinstance(metric, SearchMetric) else metric, + "language": language, + "query_mode": ( + query_mode.value + if isinstance(query_mode, TextQueryMode) + else query_mode + ), + "match_mode": match_mode, + "matched_with": matched_with, + # True only when a strict reading was wanted and gave up, + # which is the signal a caller has to notice. + "widened": ( + matched_with == TextMatchMode.ANY.value + and match_mode != TextMatchMode.ANY.value + ), + "widening_refused": widening_refused, + "score_column": score_column, + "vector_rank_column": vector_rank_column, + "text_rank_column": text_rank_column, + "rrf_k": rrf_k, + "candidates": candidates, + "row_key": key_columns, + "limit": limit, + "filters": len(conditions), + }, + ) diff --git a/mcp-server/src/cbmcp/search.py b/mcp-server/src/cbmcp/search.py new file mode 100644 index 00000000000..2679db3a7df --- /dev/null +++ b/mcp-server/src/cbmcp/search.py @@ -0,0 +1,546 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# 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. + +""" +Retrieval SQL builders for the Apache Cloudberry MCP server. + +These helpers turn catalog-verified identifiers plus caller supplied values +into parameterized SQL for vector, full-text, and hybrid search. + +Two rules hold throughout: + +* Identifiers (schema, table, column) are never interpolated before they have + been matched against the catalog, and they are always quoted afterwards. +* Every caller supplied *value* is bound as a ``$n`` query parameter, never + formatted into the statement text. +""" + +import math +from enum import Enum +from typing import Any, Dict, Iterable, List, Optional, Sequence + +from pydantic import BaseModel, Field + +from .security import SQLValidator + +# Column types that hold embeddings. All are provided by pgvector. +VECTOR_TYPES = frozenset({"vector", "halfvec", "sparsevec"}) + +# Column types that full-text search can build a tsvector from. +TEXT_TYPES = frozenset({"text", "varchar", "bpchar", "name", "citext"}) + +# A column already stored as a parsed document needs no to_tsvector() call. +TSVECTOR_TYPE = "tsvector" + +# The one embedding type whose text form is not a dense bracket list. +SPARSEVEC_TYPE = "sparsevec" + +# Literal used to force the pgvector shared library to load so that its GUCs +# (ivfflat.probes, hnsw.ef_search) become settable in the current session. +# sparsevec has its own text representation and cannot reuse the '[...]' form. +WARMUP_LITERALS: Dict[str, str] = { + "vector": "[1]", + "halfvec": "[1]", + "sparsevec": "{1:1}/1", +} + +# Upper bounds that keep a single tool call from returning or scanning an +# unbounded amount of data. +MAX_LIMIT = 1000 +MAX_CANDIDATES = 10000 + + +class SearchMetric(str, Enum): + """Distance metric used to rank vector search results.""" + + L2 = "l2" + COSINE = "cosine" + INNER_PRODUCT = "inner_product" + L1 = "l1" + + +# pgvector distance operators. Every one of them orders ascending, including +# inner_product: '<#>' returns the *negative* inner product precisely so that +# "smaller is better" holds for all four metrics. +METRIC_OPERATORS: Dict[str, str] = { + SearchMetric.L2.value: "<->", + SearchMetric.COSINE.value: "<=>", + SearchMetric.INNER_PRODUCT.value: "<#>", + SearchMetric.L1.value: "<+>", +} + + +class TextQueryMode(str, Enum): + """How a full-text query string is parsed into a tsquery.""" + + PLAIN = "plain" + PHRASE = "phrase" + WEBSEARCH = "websearch" + + +TEXT_QUERY_FUNCTIONS: Dict[str, str] = { + TextQueryMode.PLAIN.value: "plainto_tsquery", + TextQueryMode.PHRASE.value: "phraseto_tsquery", + TextQueryMode.WEBSEARCH.value: "websearch_to_tsquery", +} + + +class TextMatchMode(str, Enum): + """How the terms of a full-text query are combined.""" + + ALL = "all" + ANY = "any" + ALL_THEN_ANY = "all_then_any" + + +# Escapes one lexeme into the quoted form tsquery's own input function reads. +# The result is cast straight to tsquery rather than passed through +# to_tsquery(): to_tsquery re-parses what it is given, so a lexeme carrying +# punctuation would be split into a phrase again, and these lexemes came out +# of to_tsvector already normalised. +_ESCAPE_LEXEME = "'''' || replace(replace({lexeme}, '\\', '\\\\'), '''', '''''') || ''''" + + +def build_tsquery( + query_function: str, + language_sql: str, + query_placeholder: str, + match_mode: str, +) -> str: + """Build the tsquery expression for one match mode. + + PostgreSQL's query constructors require every term to be present, so a + caller who writes a natural sentence usually matches nothing. The 'any' + form re-parses the same string into its lexemes and ORs them, which keeps + the configuration's own stemming and stop-word handling rather than + splitting words here. + + Ranking still favours the fuller match: ts_rank scores a document that + matched more of the terms above one that matched fewer. + """ + if match_mode == TextMatchMode.ANY.value: + escaped = _ESCAPE_LEXEME.format(lexeme="lexemes.lexeme") + return ( + f"(SELECT string_agg({escaped}, ' | ') " + f"FROM unnest(to_tsvector({language_sql}, {query_placeholder})) " + f"AS lexemes)::tsquery" + ) + return f"{query_function}({language_sql}, {query_placeholder})" + + +def tsquery_has_blocking_operator(rendered: Optional[str]) -> bool: + """Report whether a tsquery carries an operator that widening would undo. + + Widening ORs the query's lexemes together. That is harmless for a query + that only ANDs its terms, but it inverts a negation and flattens a phrase, + handing the caller exactly the rows they asked to exclude. Negation and + phrase distance are therefore read as an instruction not to widen. + + Only operators outside a quoted lexeme count, because a lexeme may itself + contain any character. + """ + if not rendered: + return False + + index = 0 + length = len(rendered) + while index < length: + char = rendered[index] + if char == "'": + # Walk to the closing quote, where '' is an escaped quote. + index += 1 + while index < length: + if rendered[index] == "'": + if index + 1 < length and rendered[index + 1] == "'": + index += 2 + continue + break + index += 1 + index += 1 + continue + if char == "!" or char == "<": + return True + index += 1 + return False + + +def unique_alias(base: str, taken: Sequence[str]) -> str: + """Pick an output name that does not collide with a selected column. + + A table is free to have a column of its own called rank or score, and + results come back as positional rows, so two columns sharing a name leave + the caller reading the wrong one with nothing to warn them. + """ + existing = set(taken) + if base not in existing: + return base + suffix = 2 + while f"{base}_{suffix}" in existing: + suffix += 1 + return f"{base}_{suffix}" + + +def resolve_match_mode(match_mode: Any) -> str: + """Validate a term combination mode.""" + value = match_mode.value if isinstance(match_mode, TextMatchMode) else str(match_mode) + supported = {item.value for item in TextMatchMode} + if value not in supported: + raise ValueError( + f"match_mode '{value}' is not supported; supported modes are: " + f"{', '.join(sorted(supported))}" + ) + return value + + +class FilterOperator(str, Enum): + """Comparison used by a single pre-filter condition.""" + + EQ = "eq" + NE = "ne" + LT = "lt" + LTE = "lte" + GT = "gt" + GTE = "gte" + IN = "in" + NOT_IN = "not_in" + LIKE = "like" + ILIKE = "ilike" + IS_NULL = "is_null" + IS_NOT_NULL = "is_not_null" + + +BINARY_OPERATORS: Dict[str, str] = { + FilterOperator.EQ.value: "=", + FilterOperator.NE.value: "<>", + FilterOperator.LT.value: "<", + FilterOperator.LTE.value: "<=", + FilterOperator.GT.value: ">", + FilterOperator.GTE.value: ">=", + FilterOperator.LIKE.value: "LIKE", + FilterOperator.ILIKE.value: "ILIKE", +} + +LIST_OPERATORS: Dict[str, str] = { + FilterOperator.IN.value: "IN", + FilterOperator.NOT_IN.value: "NOT IN", +} + +NULLARY_OPERATORS: Dict[str, str] = { + FilterOperator.IS_NULL.value: "IS NULL", + FilterOperator.IS_NOT_NULL.value: "IS NOT NULL", +} + + +class SearchFilter(BaseModel): + """One pre-filter condition, applied before ranking. + + Pre-filtering matters for approximate nearest neighbour search: the + condition is pushed into the same scan that walks the vector index, so the + engine ranks only rows that already satisfy it. + """ + + column: str = Field( + description="Column to filter on. Must exist in the table being searched." + ) + operator: FilterOperator = Field( + default=FilterOperator.EQ, + description=( + "Comparison to apply. 'in'/'not_in' take a list value; " + "'is_null'/'is_not_null' take no value." + ), + ) + value: Any = Field( + default=None, + description=( + "Value to compare against. A list for 'in'/'not_in', omitted for " + "'is_null'/'is_not_null'." + ), + ) + + +class SqlParams: + """Collects bind values and hands out the matching ``$n`` placeholders.""" + + def __init__(self) -> None: + self.values: List[Any] = [] + + def add(self, value: Any) -> str: + """Bind one value and return the placeholder that refers to it.""" + self.values.append(value) + return f"${len(self.values)}" + + +def coerce_filters(filters: Optional[Iterable[Any]]) -> List[SearchFilter]: + """Accept filters as models or plain dicts and return validated models.""" + if not filters: + return [] + + coerced: List[SearchFilter] = [] + for index, item in enumerate(filters): + if isinstance(item, SearchFilter): + coerced.append(item) + elif isinstance(item, dict): + try: + coerced.append(SearchFilter(**item)) + except Exception as exc: + raise ValueError(f"filters[{index}] is not a valid filter: {exc}") from exc + else: + raise ValueError( + f"filters[{index}] must be an object with 'column', 'operator' and " + f"'value', got {type(item).__name__}" + ) + return coerced + + +def resolve_column(name: str, columns: Sequence[Dict[str, Any]], label: str) -> Dict[str, Any]: + """Look a column up in catalog metadata, or raise a descriptive error. + + Resolving through the catalog is what makes it safe to interpolate the + name into SQL afterwards: only names the server itself read back from + ``pg_attribute`` ever reach the statement text. + """ + if not isinstance(name, str) or not name.strip(): + raise ValueError(f"{label} must be a non-empty column name, got {name!r}") + + for column in columns: + if column["column_name"] == name: + return column + + available = ", ".join(column["column_name"] for column in columns) + raise ValueError(f"{label} '{name}' does not exist; available columns are: {available}") + + +def quote_column(column: Dict[str, Any], alias: Optional[str] = None) -> str: + """Quote a catalog-resolved column, optionally qualified by a table alias.""" + quoted = SQLValidator.quote_identifier(column["column_name"]) + return f"{alias}.{quoted}" if alias else quoted + + +def build_select_list( + select_columns: Optional[Sequence[str]], + columns: Sequence[Dict[str, Any]], + alias: Optional[str] = None, +) -> tuple: + """Build the projection for a search result, and the names it produces. + + Embedding columns are cast to text so that a client with no pgvector type + codec still receives a readable value. When the caller names no columns, + embeddings are dropped entirely: they are large and of no use to an agent + that is reading the rows. + """ + if select_columns is None: + chosen = [column for column in columns if column["type_name"] not in VECTOR_TYPES] + if not chosen: + chosen = list(columns) + else: + if not isinstance(select_columns, (list, tuple)) or not select_columns: + raise ValueError( + "select_columns must be a non-empty list of column names, or omitted " + "to select every non-embedding column" + ) + chosen = [ + resolve_column(name, columns, f"select_columns[{index}]") + for index, name in enumerate(select_columns) + ] + + projection: List[str] = [] + names: List[str] = [] + for column in chosen: + reference = quote_column(column, alias) + output_name = SQLValidator.quote_identifier(column["column_name"]) + names.append(column["column_name"]) + if column["type_name"] in VECTOR_TYPES: + projection.append(f"{reference}::text AS {output_name}") + else: + projection.append(f"{reference} AS {output_name}") + return projection, names + + +def build_filter_sql( + filters: Sequence[SearchFilter], + columns: Sequence[Dict[str, Any]], + params: SqlParams, + alias: Optional[str] = None, +) -> str: + """Build a parameterized boolean expression, without the WHERE keyword.""" + clauses: List[str] = [] + + for index, condition in enumerate(filters): + label = f"filters[{index}].column" + column = resolve_column(condition.column, columns, label) + reference = quote_column(column, alias) + operator = ( + condition.operator.value + if isinstance(condition.operator, FilterOperator) + else str(condition.operator) + ) + + if operator in NULLARY_OPERATORS: + clauses.append(f"{reference} {NULLARY_OPERATORS[operator]}") + elif operator in LIST_OPERATORS: + if not isinstance(condition.value, (list, tuple)) or not condition.value: + raise ValueError( + f"filters[{index}] operator '{operator}' on column " + f"'{condition.column}' requires a non-empty list value, got " + f"{condition.value!r}" + ) + placeholders = ", ".join(params.add(item) for item in condition.value) + clauses.append(f"{reference} {LIST_OPERATORS[operator]} ({placeholders})") + elif operator in BINARY_OPERATORS: + if condition.value is None: + raise ValueError( + f"filters[{index}] operator '{operator}' on column " + f"'{condition.column}' requires a value; use 'is_null' to test " + f"for NULL" + ) + clauses.append( + f"{reference} {BINARY_OPERATORS[operator]} {params.add(condition.value)}" + ) + else: + supported = ", ".join(sorted(item.value for item in FilterOperator)) + raise ValueError( + f"filters[{index}] operator '{operator}' is not supported; " + f"supported operators are: {supported}" + ) + + return " AND ".join(clauses) + + +def validate_query_vector( + query_vector: Sequence[float], label: str = "query_vector" +) -> List[float]: + """Check an embedding's shape and values, and return it as floats. + + Kept separate from rendering so that a malformed vector is rejected on its + own terms, before the column lookup turns the complaint into a dimension + mismatch. + """ + if not isinstance(query_vector, (list, tuple)) or not query_vector: + raise ValueError(f"{label} must be a non-empty list of numbers, got {query_vector!r}") + + values: List[float] = [] + for index, item in enumerate(query_vector): + if isinstance(item, bool) or not isinstance(item, (int, float)): + raise ValueError( + f"{label}[{index}] must be a number, got {type(item).__name__} ({item!r})" + ) + value = float(item) + if not math.isfinite(value): + raise ValueError(f"{label}[{index}] must be finite, got {value}") + values.append(value) + return values + + +def format_vector_literal( + query_vector: Sequence[float], + type_name: str = "vector", + label: str = "query_vector", +) -> str: + """Render an embedding as the text literal its pgvector type accepts. + + The caller always supplies a dense list, because that is what an embedding + model produces. A sparsevec column stores the same vector as its non-zero + entries, so the literal has to be written in that form: dense brackets are + rejected outright by its input function. + """ + values = validate_query_vector(query_vector, label) + + if type_name == SPARSEVEC_TYPE: + # pgvector's sparse form lists only the non-zero entries, indexed from + # one, with the full width after the slash. An all-zero vector is + # written as an empty entry list. + entries = [ + f"{index}:{value!r}" + for index, value in enumerate(values, start=1) + if value != 0.0 + ] + return "{" + ",".join(entries) + "}/" + str(len(values)) + + return "[" + ",".join(repr(value) for value in values) + "]" + + +def resolve_metric(metric: Any) -> str: + """Map a metric name onto its pgvector operator.""" + value = metric.value if isinstance(metric, SearchMetric) else str(metric) + try: + return METRIC_OPERATORS[value] + except KeyError: + supported = ", ".join(sorted(METRIC_OPERATORS)) + raise ValueError( + f"metric '{value}' is not supported; supported metrics are: {supported}" + ) from None + + +def resolve_query_mode(query_mode: Any) -> str: + """Map a query mode onto its tsquery constructor.""" + value = query_mode.value if isinstance(query_mode, TextQueryMode) else str(query_mode) + try: + return TEXT_QUERY_FUNCTIONS[value] + except KeyError: + supported = ", ".join(sorted(TEXT_QUERY_FUNCTIONS)) + raise ValueError( + f"query_mode '{value}' is not supported; supported modes are: {supported}" + ) from None + + +def validate_positive_int(value: Any, label: str) -> int: + """Reject anything that is not a positive integer. + + Used for the pgvector recall knobs, whose real upper bounds belong to the + server and are read from pg_settings rather than guessed here. + """ + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError(f"{label} must be an integer, got {type(value).__name__} ({value!r})") + if value < 1: + raise ValueError(f"{label} must be at least 1, got {value}") + return value + + +def validate_limit(limit: Any, label: str = "limit", maximum: int = MAX_LIMIT) -> int: + """Reject a row limit that is not a positive integer within bounds.""" + if isinstance(limit, bool) or not isinstance(limit, int): + raise ValueError(f"{label} must be an integer, got {type(limit).__name__} ({limit!r})") + if limit < 1 or limit > maximum: + raise ValueError(f"{label} must be between 1 and {maximum}, got {limit}") + return limit + + +def validate_vector_column(column: Dict[str, Any], query_vector: Sequence[float]) -> None: + """Check that a column holds embeddings and that the dimensions agree.""" + if column["type_name"] not in VECTOR_TYPES: + supported = ", ".join(sorted(VECTOR_TYPES)) + raise ValueError( + f"vector_column '{column['column_name']}' has type " + f"{column['type_display']}, which is not an embedding type; expected one " + f"of: {supported}" + ) + + # pgvector stores the declared dimension directly in atttypmod; -1 means the + # column was declared without one and accepts any width. + declared = column["type_modifier"] + if declared and declared > 0 and len(query_vector) != declared: + raise ValueError( + f"query_vector has {len(query_vector)} dimensions but column " + f"'{column['column_name']}' is {column['type_display']}; they must match" + ) + + +def warmup_literal(type_name: str) -> str: + """Return a literal of the given embedding type, used to load pgvector.""" + return WARMUP_LITERALS.get(type_name, "[1]") diff --git a/mcp-server/src/cbmcp/security.py b/mcp-server/src/cbmcp/security.py index d9cd44a7e17..cf2e5476eda 100644 --- a/mcp-server/src/cbmcp/security.py +++ b/mcp-server/src/cbmcp/security.py @@ -83,6 +83,26 @@ def validate_query(cls, query: str) -> tuple[bool, str]: return True, "Query is valid" + @classmethod + def quote_identifier(cls, name: str) -> str: + """Quote a SQL identifier so it can be embedded in a statement. + + This is the last step of identifier handling, not the only one. Callers + resolve a name against the catalog first; quoting then guarantees that + whatever came back is read as a single identifier and nothing else. + """ + if not isinstance(name, str) or not name: + raise ValueError(f"Identifier must be a non-empty string, got {name!r}") + if "\x00" in name: + raise ValueError("Identifier must not contain a NUL character") + escaped = name.replace('"', '""') + return f'"{escaped}"' + + @classmethod + def quote_qualified_name(cls, schema: str, table: str) -> str: + """Quote a schema-qualified relation name.""" + return f"{cls.quote_identifier(schema)}.{cls.quote_identifier(table)}" + @classmethod def sanitize_parameter_name(cls, param_name: str) -> str: """Sanitize parameter names to prevent injection""" diff --git a/mcp-server/src/cbmcp/server.py b/mcp-server/src/cbmcp/server.py index edf0b539545..c28ad423a89 100644 --- a/mcp-server/src/cbmcp/server.py +++ b/mcp-server/src/cbmcp/server.py @@ -31,6 +31,7 @@ from .config import DatabaseConfig, ServerConfig from .database import DatabaseManager +from .search import SearchFilter, SearchMetric, TextMatchMode, TextQueryMode from .prompt import ( ANALYZE_QUERY_PERFORMANCE_PROMPT, SUGGEST_INDEXES_PROMPT, @@ -128,7 +129,7 @@ async def list_columns( @self.mcp.tool() async def execute_query( query: Annotated[str, Field(description="The SQL query to execute")], - params: Annotated[Optional[Dict[str, Any]], Field(description="The parameters for the query")] = None, + params: Annotated[Optional[List[Any]], Field(description="Values bound to $1, $2, ... in the query, in order")] = None, readonly: Annotated[bool, Field(description="Whether the query is read-only")] = True ) -> Dict[str, Any]: """ @@ -143,7 +144,7 @@ async def execute_query( @self.mcp.tool() async def explain_query( query: Annotated[str, Field(description="The SQL query to explain")], - params: Annotated[Optional[Dict[str, Any]], Field(description="The parameters for the query")] = None + params: Annotated[Optional[List[Any]], Field(description="Values bound to $1, $2, ... in the query, in order")] = None ) -> str: """ Get the execution plan for a query @@ -393,7 +394,157 @@ async def list_active_connections() -> List[Dict[str, Any]]: return await self.db_manager.list_active_connections() except Exception as e: return f"Error listing active connections: {str(e)}" - + + # -------------------------------------------------------------- + # Retrieval tools: what an agent needs to search, not just query + # -------------------------------------------------------------- + + @self.mcp.tool() + async def list_searchable_columns( + schema: Annotated[Optional[str], Field(description="Restrict the result to one schema")] = None, + table: Annotated[Optional[str], Field(description="Restrict the result to one table")] = None, + include_unindexed_text: Annotated[bool, Field(description="Also report text columns that have no full-text index")] = False + ) -> List[Dict[str, Any]]: + """Discover which columns can be searched, and how. + + Reports embedding columns with their dimension, tsvector columns, + and text columns carrying a full-text index, each with the indexes + defined on it. Call this before vector_search, fulltext_search or + hybrid_search to choose a target instead of guessing from names. + """ + logger.info(f"Listing searchable columns in schema: {schema}, table: {table}") + try: + return await self.db_manager.list_searchable_columns( + schema, table, include_unindexed_text + ) + except Exception as e: + return f"Error listing searchable columns: {str(e)}" + + @self.mcp.tool() + async def vector_search( + schema: Annotated[str, Field(description="The schema name")], + table: Annotated[str, Field(description="The table name")], + vector_column: Annotated[str, Field(description="The embedding column to rank by")], + query_vector: Annotated[List[float], Field(description="The query embedding; its length must match the column's declared dimension")], + limit: Annotated[int, Field(description="How many rows to return")] = 10, + select_columns: Annotated[Optional[List[str]], Field(description="Columns to return; defaults to every non-embedding column")] = None, + filters: Annotated[Optional[List[SearchFilter]], Field(description="Conditions applied before ranking, pushed into the index scan")] = None, + metric: Annotated[SearchMetric, Field(description="Distance metric; all four order ascending, so smaller is always better")] = SearchMetric.L2, + probes: Annotated[Optional[int], Field(description="ivfflat lists to probe. The default of 1 can miss the true nearest neighbour; raise it to trade latency for recall")] = None, + ef_search: Annotated[Optional[int], Field(description="HNSW candidate list size; the same recall-for-latency trade on an HNSW index")] = None + ) -> Dict[str, Any]: + """Find the rows whose embedding is nearest to a query vector. + + Filters are applied inside the same scan that walks the vector + index, so this is a pre-filtered nearest neighbour search, not a + filter over an already truncated result. The generated SQL is + returned alongside the rows. + """ + logger.info(f"Vector search on {schema}.{table}.{vector_column}, limit: {limit}") + try: + return await self.db_manager.vector_search( + schema=schema, + table=table, + vector_column=vector_column, + query_vector=query_vector, + limit=limit, + select_columns=select_columns, + filters=filters, + metric=metric, + probes=probes, + ef_search=ef_search, + ) + except Exception as e: + return {"error": f"Error running vector search: {str(e)}"} + + @self.mcp.tool() + async def fulltext_search( + schema: Annotated[str, Field(description="The schema name")], + table: Annotated[str, Field(description="The table name")], + text_column: Annotated[str, Field(description="The text or tsvector column to search")], + query: Annotated[str, Field(description="The search string")], + limit: Annotated[int, Field(description="How many rows to return")] = 10, + select_columns: Annotated[Optional[List[str]], Field(description="Columns to return; defaults to every non-embedding column")] = None, + filters: Annotated[Optional[List[SearchFilter]], Field(description="Conditions applied before ranking")] = None, + language: Annotated[str, Field(description="Text search configuration used to parse both document and query")] = "english", + query_mode: Annotated[TextQueryMode, Field(description="How the search string is parsed: websearch accepts quotes and OR, plain treats it as words, phrase requires the exact sequence")] = TextQueryMode.WEBSEARCH, + match_mode: Annotated[TextMatchMode, Field(description="Whether every term must appear. all_then_any requires all terms and widens to any of them only when that matched nothing, so a natural sentence does not come back empty; the result reports which was used")] = TextMatchMode.ALL_THEN_ANY + ) -> Dict[str, Any]: + """Rank rows by full-text relevance over one text column. + + A tsvector column is searched as stored; a text column is parsed on + the fly, which matches an expression index built over to_tsvector(). + Results are ordered by ts_rank. + """ + logger.info(f"Full-text search on {schema}.{table}.{text_column}, limit: {limit}") + try: + return await self.db_manager.fulltext_search( + schema=schema, + table=table, + text_column=text_column, + query=query, + limit=limit, + select_columns=select_columns, + filters=filters, + language=language, + query_mode=query_mode, + match_mode=match_mode, + ) + except Exception as e: + return {"error": f"Error running full-text search: {str(e)}"} + + @self.mcp.tool() + async def hybrid_search( + schema: Annotated[str, Field(description="The schema name")], + table: Annotated[str, Field(description="The table name")], + vector_column: Annotated[str, Field(description="The embedding column to rank by")], + text_column: Annotated[str, Field(description="The text or tsvector column to search")], + query_vector: Annotated[List[float], Field(description="The query embedding; its length must match the column's declared dimension")], + query: Annotated[str, Field(description="The search string")], + limit: Annotated[int, Field(description="How many rows to return")] = 10, + select_columns: Annotated[Optional[List[str]], Field(description="Columns to return; defaults to every non-embedding column")] = None, + filters: Annotated[Optional[List[SearchFilter]], Field(description="Conditions applied before ranking, to both arms")] = None, + metric: Annotated[SearchMetric, Field(description="Distance metric for the vector arm")] = SearchMetric.L2, + language: Annotated[str, Field(description="Text search configuration for the full-text arm")] = "english", + query_mode: Annotated[TextQueryMode, Field(description="How the search string is parsed")] = TextQueryMode.WEBSEARCH, + match_mode: Annotated[TextMatchMode, Field(description="Whether every term must appear. all_then_any requires all terms and widens to any of them only when that matched nothing, so a natural sentence does not come back empty; the result reports which was used")] = TextMatchMode.ALL_THEN_ANY, + rrf_k: Annotated[int, Field(description="Reciprocal rank fusion constant; larger flattens the weight given to top ranks")] = 60, + candidates: Annotated[Optional[int], Field(description="Rows each arm contributes before fusion; defaults to five times limit, at least 50")] = None, + probes: Annotated[Optional[int], Field(description="ivfflat lists to probe in the vector arm")] = None, + ef_search: Annotated[Optional[int], Field(description="HNSW candidate list size for the vector arm")] = None + ) -> Dict[str, Any]: + """Search by meaning and by keyword at once, then fuse the rankings. + + Each arm returns its own ranked candidates and a row scores + 1 / (rrf_k + rank) from each arm it appears in. Fusing ranks rather + than raw scores is what makes the two comparable: a distance and a + ts_rank share no scale, but their orderings do. Rows found by only + one arm still place, which is the point of running both. + """ + logger.info(f"Hybrid search on {schema}.{table}, limit: {limit}") + try: + return await self.db_manager.hybrid_search( + schema=schema, + table=table, + vector_column=vector_column, + text_column=text_column, + query_vector=query_vector, + query=query, + limit=limit, + select_columns=select_columns, + filters=filters, + metric=metric, + language=language, + query_mode=query_mode, + match_mode=match_mode, + rrf_k=rrf_k, + candidates=candidates, + probes=probes, + ef_search=ef_search, + ) + except Exception as e: + return {"error": f"Error running hybrid search: {str(e)}"} + def _setup_prompts(self): """Setup MCP prompts for common database tasks""" diff --git a/mcp-server/tests/test_search_tools.py b/mcp-server/tests/test_search_tools.py new file mode 100644 index 00000000000..4a0592ac5bd --- /dev/null +++ b/mcp-server/tests/test_search_tools.py @@ -0,0 +1,1423 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# 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. + +""" +Retrieval tool test module + +Covers vector_search, fulltext_search, hybrid_search and +list_searchable_columns against a live Apache Cloudberry cluster. + +These are integration tests. They build their own fixture tables, so they need +a reachable cluster with the pgvector extension available; without one the +whole module skips. +""" + +import re + +import pytest +import pytest_asyncio + +from cbmcp.client import CloudberryMCPClient +from cbmcp.config import DatabaseConfig +from cbmcp.search import ( + format_vector_literal, + tsquery_has_blocking_operator, + unique_alias, +) + +SCHEMA = "cbmcp_search_test" +DOCS = "docs" +EMPTY = "docs_empty" +COLLIDING = "docs_colliding" +SPARSE = "docs_sparse" +DIMENSIONS = 8 +ROWS = 2000 +LISTS = 20 + +# One embedding per row, spread over the unit cube by a handful of coprime +# moduli so that neighbours are not simply adjacent ids. +EMBEDDING_SQL = ( + "ARRAY[(i%7)/7.0, (i%11)/11.0, (i%13)/13.0, (i%3)/3.0, " + "(i%5)/5.0, (i%17)/17.0, (i%19)/19.0, (i%23)/23.0]::vector(8)" +) + +QUERY_VECTOR = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8] +QUERY_VECTOR_LITERAL = "[0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8]" + + +def unwrap(result): + """Return the payload of a tool call across fastmcp result shapes.""" + return result.data if hasattr(result, "data") else result + + +async def open_connection(): + import asyncpg + + config = DatabaseConfig.from_env() + return await asyncpg.connect( + host=config.host, + port=config.port, + database=config.database, + user=config.username, + password=config.password or None, + ) + + +@pytest_asyncio.fixture(scope="module", loop_scope="module") +async def fixture_tables(): + """Build the tables the retrieval tools are exercised against. + + The docs table is distributed by id so that the rows land on several + segments, which is the case the fusion row key and the distributed top-K + merge actually have to handle. + """ + try: + conn = await open_connection() + except Exception as e: + pytest.skip(f"Skipping retrieval tests - no reachable cluster: {e}") + + try: + try: + await conn.execute("CREATE EXTENSION IF NOT EXISTS vector") + except Exception as e: + await conn.close() + pytest.skip(f"Skipping retrieval tests - pgvector unavailable: {e}") + + await conn.execute(f"DROP SCHEMA IF EXISTS {SCHEMA} CASCADE") + await conn.execute(f"CREATE SCHEMA {SCHEMA}") + await conn.execute( + f""" + CREATE TABLE {SCHEMA}.{DOCS} ( + id bigint, + customer_id bigint, + content text, + embedding vector({DIMENSIONS}) + ) DISTRIBUTED BY (id) + """ + ) + await conn.execute( + f""" + INSERT INTO {SCHEMA}.{DOCS} + SELECT i, + (i % 10) + 1, + CASE WHEN i % 3 = 0 THEN 'query latency is slow and timeouts happen' + WHEN i % 3 = 1 THEN 'billing invoice payment problem' + ELSE 'connection refused during upgrade' END, + {EMBEDDING_SQL} + FROM generate_series(1, {ROWS}) i + """ + ) + await conn.execute( + f"CREATE INDEX docs_embedding_ivf ON {SCHEMA}.{DOCS} " + f"USING ivfflat (embedding vector_l2_ops) WITH (lists = {LISTS})" + ) + await conn.execute( + f"CREATE INDEX docs_content_gin ON {SCHEMA}.{DOCS} " + f"USING gin (to_tsvector('english', content))" + ) + await conn.execute( + f""" + CREATE TABLE {SCHEMA}.{EMPTY} ( + id bigint, + content text, + embedding vector({DIMENSIONS}) + ) DISTRIBUTED BY (id) + """ + ) + await conn.execute( + f""" + CREATE TABLE {SCHEMA}.{COLLIDING} ( + id bigint, + rank int, + score int, + distance int, + content text, + embedding vector({DIMENSIONS}) + ) DISTRIBUTED BY (id) + """ + ) + await conn.execute( + f"INSERT INTO {SCHEMA}.{COLLIDING} VALUES " + f"(1, 99, 98, 97, 'query latency is slow', " + f"'[0,0,0,0,0,0,0,0]')" + ) + # sparsevec stores only the non-zero entries, so its text form differs + # from every other embedding type. + await conn.execute( + f""" + CREATE TABLE {SCHEMA}.{SPARSE} ( + id bigint, + embedding sparsevec({DIMENSIONS}) + ) DISTRIBUTED BY (id) + """ + ) + await conn.execute( + f"INSERT INTO {SCHEMA}.{SPARSE} VALUES " + f"(1, '{{1:0.9,2:0.1}}/{DIMENSIONS}'), " + f"(2, '{{3:0.8,7:0.2}}/{DIMENSIONS}'), " + f"(3, '{{1:0.2,5:0.7}}/{DIMENSIONS}')" + ) + await conn.execute(f"ANALYZE {SCHEMA}.{DOCS}") + + yield conn + finally: + try: + await conn.execute(f"DROP SCHEMA IF EXISTS {SCHEMA} CASCADE") + finally: + await conn.close() + + +@pytest_asyncio.fixture(loop_scope="module") +async def client(fixture_tables): + """An in-process MCP client. + + Only stdio is exercised: it needs no separately started server, so the + retrieval behaviour under test is what fails, not the transport. + """ + instance = await CloudberryMCPClient.create(mode="stdio") + yield instance + await instance.close() + + +async def exact_nearest(conn, limit=5, predicate=None): + """Compute the true nearest neighbours with the index scan disabled.""" + where = f"WHERE {predicate}" if predicate else "" + async with conn.transaction(): + await conn.execute("SET LOCAL enable_indexscan = off") + await conn.execute("SET LOCAL enable_indexonlyscan = off") + records = await conn.fetch( + f"SELECT id, (embedding <-> '{QUERY_VECTOR_LITERAL}'::vector)::float8 AS distance " + f"FROM {SCHEMA}.{DOCS} {where} ORDER BY 2, 1 LIMIT {limit}" + ) + return [(r["id"], r["distance"]) for r in records] + + +@pytest.mark.asyncio(loop_scope="module") +class TestListSearchableColumns: + """Discovery of what can be searched.""" + + async def test_reports_vector_column_with_dimension(self, client): + rows = unwrap( + await client.call_tool( + "list_searchable_columns", {"schema": SCHEMA, "table": DOCS} + ) + ) + by_column = {row["column"]: row for row in rows} + + assert "embedding" in by_column + embedding = by_column["embedding"] + assert embedding["kind"] == "vector" + assert embedding["dimension"] == DIMENSIONS + assert any(index["name"] == "docs_embedding_ivf" for index in embedding["indexes"]) + + async def test_reports_text_column_behind_an_expression_index(self, client): + rows = unwrap( + await client.call_tool( + "list_searchable_columns", {"schema": SCHEMA, "table": DOCS} + ) + ) + by_column = {row["column"]: row for row in rows} + + # The GIN index is on to_tsvector(content), and pg_get_indexdef leaves + # such a column name unquoted, so this is the case that a naive + # quoted-name match misses. + assert "content" in by_column + content = by_column["content"] + assert content["kind"] == "text" + assert content["full_text_indexed"] is True + assert any(index["on_expression"] for index in content["indexes"]) + + async def test_unindexed_text_is_hidden_by_default(self, client): + default = unwrap( + await client.call_tool( + "list_searchable_columns", {"schema": SCHEMA, "table": EMPTY} + ) + ) + assert all(row["column"] != "content" for row in default) + + included = unwrap( + await client.call_tool( + "list_searchable_columns", + {"schema": SCHEMA, "table": EMPTY, "include_unindexed_text": True}, + ) + ) + assert any(row["column"] == "content" for row in included) + + async def test_unknown_relation_yields_no_rows(self, client): + rows = unwrap( + await client.call_tool( + "list_searchable_columns", {"schema": SCHEMA, "table": "no_such_table"} + ) + ) + assert rows == [] + + +@pytest.mark.asyncio(loop_scope="module") +class TestVectorSearch: + """Nearest neighbour search, with and without a pre-filter.""" + + async def test_returns_rows_ordered_by_distance(self, client): + result = unwrap( + await client.call_tool( + "vector_search", + { + "schema": SCHEMA, + "table": DOCS, + "vector_column": "embedding", + "query_vector": QUERY_VECTOR, + "limit": 5, + "select_columns": ["id"], + }, + ) + ) + assert result["row_count"] == 5 + assert result["columns"] == ["id", "distance"] + + distances = [row[1] for row in result["rows"]] + assert distances == sorted(distances) + + async def test_embedding_column_is_dropped_from_the_default_projection(self, client): + result = unwrap( + await client.call_tool( + "vector_search", + { + "schema": SCHEMA, + "table": DOCS, + "vector_column": "embedding", + "query_vector": QUERY_VECTOR, + "limit": 1, + }, + ) + ) + assert "embedding" not in result["columns"] + assert {"id", "customer_id", "content", "distance"} <= set(result["columns"]) + + async def test_embedding_column_comes_back_as_text_when_asked_for(self, client): + result = unwrap( + await client.call_tool( + "vector_search", + { + "schema": SCHEMA, + "table": DOCS, + "vector_column": "embedding", + "query_vector": QUERY_VECTOR, + "limit": 1, + "select_columns": ["id", "embedding"], + }, + ) + ) + embedding = result["rows"][0][1] + assert isinstance(embedding, str) + assert embedding.startswith("[") + + async def test_pre_filter_restricts_the_candidates(self, client): + """The pre-filtered ANN shape: the filter rides along with the scan.""" + wanted = [1, 2, 3] + result = unwrap( + await client.call_tool( + "vector_search", + { + "schema": SCHEMA, + "table": DOCS, + "vector_column": "embedding", + "query_vector": QUERY_VECTOR, + "limit": 10, + "select_columns": ["id", "customer_id"], + "probes": LISTS, + "filters": [ + {"column": "customer_id", "operator": "in", "value": wanted} + ], + }, + ) + ) + assert result["row_count"] > 0 + assert all(row[1] in wanted for row in result["rows"]) + + async def test_pre_filtered_result_matches_an_exact_filtered_search( + self, client, fixture_tables + ): + result = unwrap( + await client.call_tool( + "vector_search", + { + "schema": SCHEMA, + "table": DOCS, + "vector_column": "embedding", + "query_vector": QUERY_VECTOR, + "limit": 5, + "select_columns": ["id"], + "probes": LISTS, + "filters": [ + {"column": "customer_id", "operator": "eq", "value": 4} + ], + }, + ) + ) + exact = await exact_nearest(fixture_tables, limit=5, predicate="customer_id = 4") + assert [row[0] for row in result["rows"]] == [row[0] for row in exact] + + async def test_probes_reaches_the_segments(self, client, fixture_tables): + """Regression test for a recall setting that never left the coordinator. + + Probing every list must reproduce the exact answer. When the setting + fails to reach the segment scans the search silently falls back to one + probe, which is a wrong answer rather than an error, so only comparing + against the exact result catches it. + """ + exact = await exact_nearest(fixture_tables, limit=5) + result = unwrap( + await client.call_tool( + "vector_search", + { + "schema": SCHEMA, + "table": DOCS, + "vector_column": "embedding", + "query_vector": QUERY_VECTOR, + "limit": 5, + "select_columns": ["id"], + "probes": LISTS, + }, + ) + ) + assert result["search"]["settings"] == {"ivfflat.probes": LISTS} + assert [row[0] for row in result["rows"]] == [row[0] for row in exact] + + async def test_recall_setting_does_not_leak_to_later_searches(self, client): + """A pooled connection must be handed back without the session setting.""" + await client.call_tool( + "vector_search", + { + "schema": SCHEMA, + "table": DOCS, + "vector_column": "embedding", + "query_vector": QUERY_VECTOR, + "limit": 1, + "select_columns": ["id"], + "probes": LISTS, + }, + ) + result = unwrap( + await client.call_tool( + "vector_search", + { + "schema": SCHEMA, + "table": DOCS, + "vector_column": "embedding", + "query_vector": QUERY_VECTOR, + "limit": 1, + "select_columns": ["id"], + }, + ) + ) + assert "settings" not in result["search"] + + async def test_empty_table_returns_no_rows(self, client): + result = unwrap( + await client.call_tool( + "vector_search", + { + "schema": SCHEMA, + "table": EMPTY, + "vector_column": "embedding", + "query_vector": QUERY_VECTOR, + "limit": 5, + }, + ) + ) + assert result["row_count"] == 0 + assert result["rows"] == [] + + async def test_dimension_mismatch_is_rejected(self, client): + result = unwrap( + await client.call_tool( + "vector_search", + { + "schema": SCHEMA, + "table": DOCS, + "vector_column": "embedding", + "query_vector": [1.0, 2.0], + }, + ) + ) + assert "2 dimensions" in result["error"] + assert "vector(8)" in result["error"] + + async def test_non_embedding_column_is_rejected(self, client): + result = unwrap( + await client.call_tool( + "vector_search", + { + "schema": SCHEMA, + "table": DOCS, + "vector_column": "content", + "query_vector": QUERY_VECTOR, + }, + ) + ) + assert "not an embedding type" in result["error"] + + async def test_unknown_relation_is_rejected(self, client): + result = unwrap( + await client.call_tool( + "vector_search", + { + "schema": SCHEMA, + "table": "no_such_table", + "vector_column": "embedding", + "query_vector": QUERY_VECTOR, + }, + ) + ) + assert "was not found" in result["error"] + + async def test_empty_query_vector_is_rejected(self, client): + result = unwrap( + await client.call_tool( + "vector_search", + { + "schema": SCHEMA, + "table": DOCS, + "vector_column": "embedding", + "query_vector": [], + }, + ) + ) + assert "non-empty list" in result["error"] + + async def test_limit_outside_the_allowed_range_is_rejected(self, client): + result = unwrap( + await client.call_tool( + "vector_search", + { + "schema": SCHEMA, + "table": DOCS, + "vector_column": "embedding", + "query_vector": QUERY_VECTOR, + "limit": 0, + }, + ) + ) + assert "must be between 1 and" in result["error"] + + async def test_empty_filter_list_value_is_rejected(self, client): + result = unwrap( + await client.call_tool( + "vector_search", + { + "schema": SCHEMA, + "table": DOCS, + "vector_column": "embedding", + "query_vector": QUERY_VECTOR, + "filters": [ + {"column": "customer_id", "operator": "in", "value": []} + ], + }, + ) + ) + assert "non-empty list value" in result["error"] + + async def test_unknown_filter_column_is_rejected(self, client): + result = unwrap( + await client.call_tool( + "vector_search", + { + "schema": SCHEMA, + "table": DOCS, + "vector_column": "embedding", + "query_vector": QUERY_VECTOR, + "filters": [ + {"column": "no_such_column", "operator": "eq", "value": 1} + ], + }, + ) + ) + assert "does not exist" in result["error"] + + async def test_identifier_injection_is_rejected_and_changes_nothing( + self, client, fixture_tables + ): + before = await fixture_tables.fetchval(f"SELECT count(*) FROM {SCHEMA}.{DOCS}") + result = unwrap( + await client.call_tool( + "vector_search", + { + "schema": SCHEMA, + "table": DOCS, + "vector_column": f'embedding"; DROP TABLE {SCHEMA}.{DOCS}; --', + "query_vector": QUERY_VECTOR, + }, + ) + ) + assert "does not exist" in result["error"] + assert await fixture_tables.fetchval(f"SELECT count(*) FROM {SCHEMA}.{DOCS}") == before + + +@pytest.mark.asyncio(loop_scope="module") +class TestFullTextSearch: + """Keyword search over a text column.""" + + async def test_ranks_matching_rows(self, client): + result = unwrap( + await client.call_tool( + "fulltext_search", + { + "schema": SCHEMA, + "table": DOCS, + "text_column": "content", + "query": "query latency", + "limit": 5, + "select_columns": ["id", "content"], + }, + ) + ) + assert result["row_count"] == 5 + assert result["columns"] == ["id", "content", "rank"] + assert all("latency" in row[1] for row in result["rows"]) + + ranks = [row[2] for row in result["rows"]] + assert ranks == sorted(ranks, reverse=True) + + async def test_filter_applies_to_the_text_arm(self, client): + result = unwrap( + await client.call_tool( + "fulltext_search", + { + "schema": SCHEMA, + "table": DOCS, + "text_column": "content", + "query": "billing invoice", + "limit": 5, + "select_columns": ["id", "customer_id"], + "filters": [ + {"column": "customer_id", "operator": "eq", "value": 5} + ], + }, + ) + ) + assert result["row_count"] > 0 + assert all(row[1] == 5 for row in result["rows"]) + + async def test_no_match_returns_no_rows(self, client): + result = unwrap( + await client.call_tool( + "fulltext_search", + { + "schema": SCHEMA, + "table": DOCS, + "text_column": "content", + "query": "zzzzzunmatchablezzzzz", + }, + ) + ) + assert result["row_count"] == 0 + + async def test_blank_query_is_rejected(self, client): + result = unwrap( + await client.call_tool( + "fulltext_search", + { + "schema": SCHEMA, + "table": DOCS, + "text_column": "content", + "query": " ", + }, + ) + ) + assert "non-empty search string" in result["error"] + + async def test_unknown_language_is_rejected(self, client): + result = unwrap( + await client.call_tool( + "fulltext_search", + { + "schema": SCHEMA, + "table": DOCS, + "text_column": "content", + "query": "latency", + "language": "no_such_config", + }, + ) + ) + assert "not a known text search configuration" in result["error"] + + async def test_strict_match_stays_strict_when_it_finds_rows(self, client): + """Both terms appear in the fixture, so nothing should be widened.""" + result = unwrap( + await client.call_tool( + "fulltext_search", + { + "schema": SCHEMA, + "table": DOCS, + "text_column": "content", + "query": "query latency", + "limit": 3, + }, + ) + ) + assert result["row_count"] == 3 + assert result["search"]["matched_with"] == "all" + assert result["search"]["widened"] is False + + async def test_widens_when_every_term_together_matches_nothing(self, client): + """A term that appears nowhere must not empty out the whole result. + + Requiring every term is the precise reading, but an empty result reads + to a caller as "this database holds nothing on the subject", which is + the wrong conclusion when most of the terms do match something. + """ + result = unwrap( + await client.call_tool( + "fulltext_search", + { + "schema": SCHEMA, + "table": DOCS, + "text_column": "content", + "query": "query latency zzzznotpresentzzzz", + "limit": 3, + "select_columns": ["id", "content"], + }, + ) + ) + assert result["row_count"] == 3 + assert result["search"]["match_mode"] == "all_then_any" + assert result["search"]["matched_with"] == "any" + assert result["search"]["widened"] is True + assert all("latency" in row[1] for row in result["rows"]) + + async def test_strict_mode_can_be_forced(self, client): + result = unwrap( + await client.call_tool( + "fulltext_search", + { + "schema": SCHEMA, + "table": DOCS, + "text_column": "content", + "query": "query latency zzzznotpresentzzzz", + "match_mode": "all", + }, + ) + ) + assert result["row_count"] == 0 + assert result["search"]["matched_with"] == "all" + assert result["search"]["widened"] is False + + async def test_any_mode_does_not_report_widening(self, client): + """Asking for 'any' up front is not a fallback, so nothing widened.""" + result = unwrap( + await client.call_tool( + "fulltext_search", + { + "schema": SCHEMA, + "table": DOCS, + "text_column": "content", + "query": "query latency zzzznotpresentzzzz", + "match_mode": "any", + "limit": 3, + }, + ) + ) + assert result["row_count"] == 3 + assert result["search"]["matched_with"] == "any" + assert result["search"]["widened"] is False + + async def test_widening_ranks_fuller_matches_first(self, client): + """ts_rank scores a row matching more of the terms above one matching fewer.""" + result = unwrap( + await client.call_tool( + "fulltext_search", + { + "schema": SCHEMA, + "table": DOCS, + "text_column": "content", + "query": "query latency billing zzzznotpresentzzzz", + "limit": 50, + "select_columns": ["id", "content"], + }, + ) + ) + assert result["search"]["widened"] is True + # 'query latency ...' carries two of the terms, 'billing invoice ...' + # only one, so the two-term rows have to come first. + ranks = [row[2] for row in result["rows"]] + assert ranks == sorted(ranks, reverse=True) + assert "latency" in result["rows"][0][1] + + async def test_negation_is_not_undone_by_widening(self, client): + """Widening must never return the rows the caller asked to exclude. + + ORing the lexemes of a negated query inverts it, so the strict reading + is kept and an empty result is the honest answer. + """ + result = unwrap( + await client.call_tool( + "fulltext_search", + { + "schema": SCHEMA, + "table": DOCS, + "text_column": "content", + "query": "latency -timeouts", + "limit": 5, + "select_columns": ["id", "content"], + }, + ) + ) + assert result["search"]["widening_refused"] is True + assert result["search"]["widened"] is False + assert result["search"]["matched_with"] == "all" + assert all("timeouts" not in row[1] for row in result["rows"]) + + async def test_a_phrase_is_not_flattened_by_widening(self, client): + result = unwrap( + await client.call_tool( + "fulltext_search", + { + "schema": SCHEMA, + "table": DOCS, + "text_column": "content", + "query": "timeouts latency", + "query_mode": "phrase", + "limit": 5, + }, + ) + ) + # The two words exist but not adjacent in that order, so a phrase + # query correctly matches nothing and must stay that way. + assert result["row_count"] == 0 + assert result["search"]["widening_refused"] is True + + async def test_document_expression_keeps_an_index_match_possible(self, client): + """The generated SQL must be shaped the way a full-text index is. + + A coalesce() wrapper or a parameterised configuration both read as + harmless, and both stop the planner matching a to_tsvector() index, so + the search silently degrades to a sequential scan. + """ + result = unwrap( + await client.call_tool( + "fulltext_search", + { + "schema": SCHEMA, + "table": DOCS, + "text_column": "content", + "query": "latency", + "limit": 1, + }, + ) + ) + sql = result["sql"] + assert "coalesce" not in sql.lower() + assert "::regconfig" in sql + # The configuration is written as the catalog OID, which reproduces + # the constant an expression index was built with. + assert re.search(r"to_tsvector\(\d+::regconfig", sql) + + async def test_a_query_of_only_stop_words_still_returns_nothing(self, client): + """Widening cannot invent lexemes where the parser found none.""" + result = unwrap( + await client.call_tool( + "fulltext_search", + { + "schema": SCHEMA, + "table": DOCS, + "text_column": "content", + "query": "the and of", + }, + ) + ) + assert result["row_count"] == 0 + + async def test_unknown_match_mode_is_rejected(self, client): + with pytest.raises(Exception) as excinfo: + await client.call_tool( + "fulltext_search", + { + "schema": SCHEMA, + "table": DOCS, + "text_column": "content", + "query": "latency", + "match_mode": "sometimes", + }, + ) + assert "all_then_any" in str(excinfo.value) + + async def test_non_text_column_is_rejected(self, client): + result = unwrap( + await client.call_tool( + "fulltext_search", + { + "schema": SCHEMA, + "table": DOCS, + "text_column": "embedding", + "query": "latency", + }, + ) + ) + assert "cannot be searched as text" in result["error"] + + +@pytest.mark.asyncio(loop_scope="module") +class TestHybridSearch: + """Reciprocal rank fusion across both arms.""" + + async def test_fuses_rows_from_both_arms(self, client): + result = unwrap( + await client.call_tool( + "hybrid_search", + { + "schema": SCHEMA, + "table": DOCS, + "vector_column": "embedding", + "text_column": "content", + "query_vector": QUERY_VECTOR, + "query": "query latency", + "limit": 10, + "select_columns": ["id"], + "probes": LISTS, + }, + ) + ) + assert result["row_count"] == 10 + assert result["columns"] == ["id", "score", "vector_rank", "text_rank"] + + scores = [row[1] for row in result["rows"]] + assert scores == sorted(scores, reverse=True) + + # Both arms have to contribute, otherwise fusion is doing nothing. + assert any(row[2] is not None for row in result["rows"]) + assert any(row[3] is not None for row in result["rows"]) + + async def test_row_key_includes_the_segment_on_cloudberry(self, client): + result = unwrap( + await client.call_tool( + "hybrid_search", + { + "schema": SCHEMA, + "table": DOCS, + "vector_column": "embedding", + "text_column": "content", + "query_vector": QUERY_VECTOR, + "query": "latency", + "limit": 3, + "select_columns": ["id"], + }, + ) + ) + # A ctid repeats across segments, so fusing on it alone would merge + # unrelated rows. + assert result["search"]["row_key"] == ["gp_segment_id", "ctid"] + + async def test_rows_are_distinct(self, client): + result = unwrap( + await client.call_tool( + "hybrid_search", + { + "schema": SCHEMA, + "table": DOCS, + "vector_column": "embedding", + "text_column": "content", + "query_vector": QUERY_VECTOR, + "query": "query latency", + "limit": 20, + "select_columns": ["id"], + }, + ) + ) + ids = [row[0] for row in result["rows"]] + assert len(ids) == len(set(ids)) + + async def test_filter_applies_to_both_arms(self, client): + result = unwrap( + await client.call_tool( + "hybrid_search", + { + "schema": SCHEMA, + "table": DOCS, + "vector_column": "embedding", + "text_column": "content", + "query_vector": QUERY_VECTOR, + "query": "query latency", + "limit": 10, + "select_columns": ["id", "customer_id"], + "filters": [ + {"column": "customer_id", "operator": "in", "value": [2, 7]} + ], + }, + ) + ) + assert result["row_count"] > 0 + assert all(row[1] in (2, 7) for row in result["rows"]) + + async def test_empty_table_returns_no_rows(self, client): + result = unwrap( + await client.call_tool( + "hybrid_search", + { + "schema": SCHEMA, + "table": EMPTY, + "vector_column": "embedding", + "text_column": "content", + "query_vector": QUERY_VECTOR, + "query": "anything", + "limit": 5, + }, + ) + ) + assert result["row_count"] == 0 + + async def test_text_arm_still_contributes_after_widening(self, client): + """Without widening the text arm can go silently empty. + + The fused result would then be pure vector search while still + presenting itself as hybrid, which the caller cannot see. + """ + result = unwrap( + await client.call_tool( + "hybrid_search", + { + "schema": SCHEMA, + "table": DOCS, + "vector_column": "embedding", + "text_column": "content", + "query_vector": QUERY_VECTOR, + "query": "query latency zzzznotpresentzzzz", + "limit": 10, + "select_columns": ["id"], + "probes": LISTS, + }, + ) + ) + assert result["search"]["widened"] is True + assert result["search"]["matched_with"] == "any" + assert any(row[3] is not None for row in result["rows"]) + + async def test_strict_mode_can_leave_the_text_arm_empty(self, client): + result = unwrap( + await client.call_tool( + "hybrid_search", + { + "schema": SCHEMA, + "table": DOCS, + "vector_column": "embedding", + "text_column": "content", + "query_vector": QUERY_VECTOR, + "query": "query latency zzzznotpresentzzzz", + "limit": 5, + "select_columns": ["id"], + "match_mode": "all", + }, + ) + ) + assert result["search"]["widened"] is False + assert all(row[3] is None for row in result["rows"]) + + async def test_negation_is_not_undone_in_the_text_arm(self, client): + result = unwrap( + await client.call_tool( + "hybrid_search", + { + "schema": SCHEMA, + "table": DOCS, + "vector_column": "embedding", + "text_column": "content", + "query_vector": QUERY_VECTOR, + "query": "latency -timeouts", + "limit": 5, + "select_columns": ["id", "content"], + }, + ) + ) + assert result["search"]["widening_refused"] is True + + # The vector arm has no opinion about words and legitimately returns + # rows containing the excluded term, so only rows the text arm ranked + # are checked. Columns are read by the names the result reports rather + # than by position. + text_rank_column = result["search"]["text_rank_column"] + rows = [dict(zip(result["columns"], row)) for row in result["rows"]] + from_text = [row for row in rows if row[text_rank_column] is not None] + assert all("timeouts" not in row["content"] for row in from_text) + + async def test_candidates_below_limit_is_rejected(self, client): + result = unwrap( + await client.call_tool( + "hybrid_search", + { + "schema": SCHEMA, + "table": DOCS, + "vector_column": "embedding", + "text_column": "content", + "query_vector": QUERY_VECTOR, + "query": "latency", + "limit": 10, + "candidates": 5, + }, + ) + ) + assert "candidates must be at least limit" in result["error"] + + +@pytest.mark.asyncio(loop_scope="module") +class TestQueryTools: + """The query tools the retrieval work depends on.""" + + async def test_positional_parameters_are_bound(self, client): + """Regression test: parameters used to be passed as keyword arguments.""" + result = unwrap( + await client.call_tool( + "execute_query", + { + "query": f"SELECT id FROM {SCHEMA}.{DOCS} WHERE id = $1 OR id = $2 ORDER BY id", + "params": [42, 77], + }, + ) + ) + assert result["rows"] == [[42], [77]] + + async def test_parameters_must_be_a_list(self, client): + """A mapping is refused at the tool boundary, before any SQL is built.""" + with pytest.raises(Exception) as excinfo: + await client.call_tool( + "execute_query", + { + "query": f"SELECT id FROM {SCHEMA}.{DOCS} WHERE id = $1", + "params": {"id": 42}, + }, + ) + assert "list" in str(excinfo.value) + + async def test_explain_refuses_a_write_statement(self, client): + """EXPLAIN ANALYZE executes what it is given, so writes must not pass.""" + result = unwrap( + await client.call_tool( + "explain_query", + {"query": f"DELETE FROM {SCHEMA}.{DOCS} WHERE id = 1"}, + ) + ) + assert "Blocked SQL operation" in result or "read-only" in result + + async def test_explain_returns_a_plan_for_a_select(self, client): + result = unwrap( + await client.call_tool( + "explain_query", + {"query": f"SELECT count(*) FROM {SCHEMA}.{DOCS}"}, + ) + ) + assert "Aggregate" in result or "Gather" in result + + +@pytest.mark.asyncio(loop_scope="module") +class TestResultAliases: + """A computed column must not shadow one the table already has.""" + + async def test_rank_alias_avoids_a_real_rank_column(self, client): + result = unwrap( + await client.call_tool( + "fulltext_search", + { + "schema": SCHEMA, + "table": COLLIDING, + "text_column": "content", + "query": "latency", + "limit": 1, + }, + ) + ) + assert len(result["columns"]) == len(set(result["columns"])) + rank_column = result["search"]["rank_column"] + assert rank_column != "rank" + # The stored column keeps its value; the score lands in the new name. + row = dict(zip(result["columns"], result["rows"][0])) + assert row["rank"] == 99 + assert isinstance(row[rank_column], float) + + async def test_distance_alias_avoids_a_real_distance_column(self, client): + result = unwrap( + await client.call_tool( + "vector_search", + { + "schema": SCHEMA, + "table": COLLIDING, + "vector_column": "embedding", + "query_vector": QUERY_VECTOR, + "limit": 1, + }, + ) + ) + assert len(result["columns"]) == len(set(result["columns"])) + distance_column = result["search"]["distance_column"] + assert distance_column != "distance" + row = dict(zip(result["columns"], result["rows"][0])) + assert row["distance"] == 97 + + async def test_hybrid_aliases_avoid_real_columns(self, client): + result = unwrap( + await client.call_tool( + "hybrid_search", + { + "schema": SCHEMA, + "table": COLLIDING, + "vector_column": "embedding", + "text_column": "content", + "query_vector": QUERY_VECTOR, + "query": "latency", + "limit": 1, + }, + ) + ) + assert len(result["columns"]) == len(set(result["columns"])) + assert result["search"]["score_column"] != "score" + row = dict(zip(result["columns"], result["rows"][0])) + assert row["score"] == 98 + + +class TestQueryShapeHelpers: + """Pure helpers, exercised without a cluster.""" + + @pytest.mark.parametrize( + "rendered,blocking", + [ + ("'queri' & 'latenc' & 'slow'", False), + ("'latenc' & !'dashboard'", True), + ("'queri' <-> 'latenc'", True), + ("'a' | 'b'", False), + ("'latenc'", False), + ("", False), + (None, False), + # An operator inside a lexeme is data, not an operator. + ("'ex!clam' & 'other'", False), + ("'a<->b' & 'other'", False), + ("'has''quote' & 'other'", False), + ], + ) + def test_blocking_operator_detection(self, rendered, blocking): + assert tsquery_has_blocking_operator(rendered) is blocking + + @pytest.mark.parametrize( + "base,taken,expected", + [ + ("rank", ["id", "body"], "rank"), + ("rank", ["id", "rank"], "rank_2"), + ("rank", ["rank", "rank_2"], "rank_3"), + ("score", [], "score"), + ], + ) + def test_unique_alias(self, base, taken, expected): + assert unique_alias(base, taken) == expected + + +@pytest.mark.asyncio(loop_scope="module") +class TestSparseEmbeddings: + """A sparsevec column is reported as searchable, so it has to work. + + Its text form is not the dense bracket list every other embedding type + uses, and pgvector's input function rejects the dense form outright. + """ + + SPARSE_QUERY = [0.9, 0.1, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] + + async def test_vector_search_works_on_a_sparsevec_column(self, client): + result = unwrap( + await client.call_tool( + "vector_search", + { + "schema": SCHEMA, + "table": SPARSE, + "vector_column": "embedding", + "query_vector": self.SPARSE_QUERY, + "limit": 3, + "select_columns": ["id"], + }, + ) + ) + assert "error" not in result + assert result["row_count"] == 3 + # Row 1 is the query itself, so it has to come back first. + assert result["rows"][0][0] == 1 + distances = [row[1] for row in result["rows"]] + assert distances == sorted(distances) + + async def test_result_matches_a_hand_written_sparse_query(self, client, fixture_tables): + result = unwrap( + await client.call_tool( + "vector_search", + { + "schema": SCHEMA, + "table": SPARSE, + "vector_column": "embedding", + "query_vector": self.SPARSE_QUERY, + "limit": 3, + "select_columns": ["id"], + }, + ) + ) + records = await fixture_tables.fetch( + f"SELECT id FROM {SCHEMA}.{SPARSE} " + f"ORDER BY embedding <-> '{{1:0.9,2:0.1}}/{DIMENSIONS}'::sparsevec, id" + ) + assert [row[0] for row in result["rows"]] == [r["id"] for r in records] + + async def test_dimension_mismatch_is_caught_on_sparsevec(self, client): + result = unwrap( + await client.call_tool( + "vector_search", + { + "schema": SCHEMA, + "table": SPARSE, + "vector_column": "embedding", + "query_vector": [0.1, 0.2], + }, + ) + ) + assert "2 dimensions" in result["error"] + assert "sparsevec(8)" in result["error"] + + async def test_an_all_zero_query_vector_is_accepted(self, client): + """The sparse form of an all-zero vector is an empty entry list.""" + result = unwrap( + await client.call_tool( + "vector_search", + { + "schema": SCHEMA, + "table": SPARSE, + "vector_column": "embedding", + "query_vector": [0.0] * DIMENSIONS, + "limit": 1, + "select_columns": ["id"], + }, + ) + ) + assert "error" not in result + assert result["row_count"] == 1 + + +@pytest.mark.asyncio(loop_scope="module") +class TestRecallSettingBounds: + """The real limits belong to the installed pgvector, not to this wrapper.""" + + async def test_ef_search_above_the_server_maximum_is_rejected(self, client): + result = unwrap( + await client.call_tool( + "vector_search", + { + "schema": SCHEMA, + "table": DOCS, + "vector_column": "embedding", + "query_vector": QUERY_VECTOR, + "limit": 1, + "ef_search": 5000, + }, + ) + ) + # pgvector caps hnsw.ef_search at 1000; the error has to say so rather + # than let the SET fail half way through applying the settings. + assert "hnsw.ef_search" in result["error"] + assert "1000" in result["error"] + + async def test_probes_above_the_server_maximum_is_rejected(self, client): + result = unwrap( + await client.call_tool( + "vector_search", + { + "schema": SCHEMA, + "table": DOCS, + "vector_column": "embedding", + "query_vector": QUERY_VECTOR, + "limit": 1, + "probes": 99999, + }, + ) + ) + assert "ivfflat.probes" in result["error"] + assert "32768" in result["error"] + + async def test_a_rejected_setting_leaves_the_connection_clean(self, client): + """A refused pair must not leave the first setting applied.""" + await client.call_tool( + "vector_search", + { + "schema": SCHEMA, + "table": DOCS, + "vector_column": "embedding", + "query_vector": QUERY_VECTOR, + "limit": 1, + "probes": LISTS, + "ef_search": 5000, + }, + ) + result = unwrap( + await client.call_tool("execute_query", {"query": "SHOW ivfflat.probes"}) + ) + assert result["rows"][0][0] == "1" + + async def test_both_settings_apply_together_when_valid(self, client): + result = unwrap( + await client.call_tool( + "vector_search", + { + "schema": SCHEMA, + "table": DOCS, + "vector_column": "embedding", + "query_vector": QUERY_VECTOR, + "limit": 1, + "probes": LISTS, + "ef_search": 100, + }, + ) + ) + assert result["search"]["settings"] == { + "ivfflat.probes": LISTS, + "hnsw.ef_search": 100, + } + + @pytest.mark.parametrize("value", [0, -1]) + async def test_a_non_positive_setting_is_rejected(self, client, value): + result = unwrap( + await client.call_tool( + "vector_search", + { + "schema": SCHEMA, + "table": DOCS, + "vector_column": "embedding", + "query_vector": QUERY_VECTOR, + "probes": value, + }, + ) + ) + assert "at least 1" in result["error"] + + +class TestVectorLiteralForms: + """Pure helper, exercised without a cluster.""" + + @pytest.mark.parametrize( + "values,type_name,expected", + [ + ([0.5, 0.25], "vector", "[0.5,0.25]"), + ([0.5, 0.25], "halfvec", "[0.5,0.25]"), + ([0.9, 0.1, 0.0], "sparsevec", "{1:0.9,2:0.1}/3"), + ([0.0, 0.0], "sparsevec", "{}/2"), + ([0.0, 4.0], "sparsevec", "{2:4.0}/2"), + ], + ) + def test_literal_matches_the_type(self, values, type_name, expected): + assert format_vector_literal(values, type_name) == expected + + def test_a_non_finite_value_is_rejected(self): + with pytest.raises(ValueError, match="finite"): + format_vector_literal([1.0, float("inf")], "sparsevec")