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")