From 219b352e56a424e12020984ee8f526d281c1044d Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Thu, 23 Jul 2026 16:07:35 +0000 Subject: [PATCH 1/2] Fix cosmosdb vectorstore bug --- .../connectors/azure_cosmos_db.py | 207 +++++++++++------- .../test_azure_cosmos_db_no_sql.py | 49 +++++ .../test_azure_cosmos_db_no_sql_collection.py | 144 +++++++++++- 3 files changed, 306 insertions(+), 94 deletions(-) diff --git a/python/semantic_kernel/connectors/azure_cosmos_db.py b/python/semantic_kernel/connectors/azure_cosmos_db.py index 0da6f52b8346..2d1bb53d9630 100644 --- a/python/semantic_kernel/connectors/azure_cosmos_db.py +++ b/python/semantic_kernel/connectors/azure_cosmos_db.py @@ -2,6 +2,7 @@ import ast import asyncio +import re import sys from collections.abc import Sequence from importlib import metadata @@ -788,14 +789,36 @@ async def _inner_search( if not vector: vector = await self._generate_vector_from_values(values, options) - if where_clauses := self._build_filter(options.filter): # type: ignore + # Translate the filter into a Cosmos SQL clause plus bound parameters. Filter string + # constants are bound as query parameters (@filter_pN) instead of concatenated, preventing + # SQL injection. _build_filter returns None, a single (clause, params) tuple, or a list of + # them; parameter names are renumbered into one global sequence so they stay unique when + # multiple filters are AND-combined. + parsed_filter = self._build_filter(options.filter) + where_clauses = "" + if parsed_filter: + clause_tuples = parsed_filter if isinstance(parsed_filter, list) else [parsed_filter] + rendered_clauses: list[str] = [] + filter_param_count = 0 + for clause, clause_parameters in clause_tuples: + rename: dict[str, str] = {} + for parameter in clause_parameters: + new_name = f"@filter_p{filter_param_count}" + filter_param_count += 1 + rename[parameter["name"]] = new_name + params.append({"name": new_name, "value": parameter["value"]}) + if rename: + + def _substitute(match: "re.Match[str]", mapping: dict[str, str] = rename) -> str: + return mapping[match.group(0)] + + clause = re.sub(r"@filter_p\d+", _substitute, clause) + rendered_clauses.append(clause) where_clauses = ( - f"WHERE {where_clauses} " - if isinstance(where_clauses, str) - else f"WHERE ({' AND '.join(where_clauses)}) " + f"WHERE {rendered_clauses[0]} " + if len(rendered_clauses) == 1 + else f"WHERE ({' AND '.join(rendered_clauses)}) " ) - else: - where_clauses = "" # Empty string instead of None vector_field_name = vector_field.storage_name or vector_field.name select_clause = self._build_select_clause(options.include_vectors) params.append({"name": "@vector", "value": vector}) @@ -849,86 +872,102 @@ def _build_select_clause(self, include_vectors: bool) -> str: return ", ".join(f"c.{field}" for field in included_fields) @override - def _lambda_parser(self, node: ast.AST) -> Any: - match node: - case ast.Compare(): - if len(node.ops) > 1: - # Chain comparisons (e.g., a < b < c) become AND-ed comparisons - values = [] - for idx in range(len(node.ops)): - if idx == 0: - values.append( - ast.Compare( - left=node.left, - ops=[node.ops[idx]], - comparators=[node.comparators[idx]], + def _lambda_parser(self, node: ast.AST) -> tuple[str, list[dict[str, Any]]]: # type: ignore[override] + # The base contract returns just the filter string; Cosmos extends the return type to a + # (clause, parameters) tuple so filter string constants can be bound as query parameters + # (@filter_pN) rather than concatenated into the SQL text. This prevents SQL injection: + # Cosmos DB uses backslash escape sequences (\'), so escaping by doubling apostrophes ('') + # is both unsafe and functionally incorrect. The parameters list is local to this call, so + # no shared state is used and it is safe to call concurrently from multiple tasks/threads. + # This mirrors the composite-return pattern used by the SQL Server and Oracle connectors, + # letting Cosmos reuse the base ``_build_filter`` for AST parsing/visiting. + parameters: list[dict[str, Any]] = [] + + def parse(node: ast.AST) -> str: + match node: + case ast.Compare(): + if len(node.ops) > 1: + # Chain comparisons (e.g., a < b < c) become AND-ed comparisons + values = [] + for idx in range(len(node.ops)): + if idx == 0: + values.append( + ast.Compare( + left=node.left, + ops=[node.ops[idx]], + comparators=[node.comparators[idx]], + ) ) - ) - else: - values.append( - ast.Compare( - left=node.comparators[idx - 1], - ops=[node.ops[idx]], - comparators=[node.comparators[idx]], + else: + values.append( + ast.Compare( + left=node.comparators[idx - 1], + ops=[node.ops[idx]], + comparators=[node.comparators[idx]], + ) ) - ) - return "(" + " AND ".join([self._lambda_parser(v) for v in values]) + ")" - left = self._lambda_parser(node.left) - right = self._lambda_parser(node.comparators[0]) - op = node.ops[0] - match op: - case ast.In(): - # Cosmos DB: ARRAY_CONTAINS(right, left) - return f"ARRAY_CONTAINS({right}, {left})" - case ast.NotIn(): - return f"NOT ARRAY_CONTAINS({right}, {left})" - case ast.Eq(): - return f"{left} = {right}" - case ast.NotEq(): - return f"{left} != {right}" - case ast.Gt(): - return f"{left} > {right}" - case ast.GtE(): - return f"{left} >= {right}" - case ast.Lt(): - return f"{left} < {right}" - case ast.LtE(): - return f"{left} <= {right}" - raise NotImplementedError(f"Unsupported operator: {type(op)}") - case ast.BoolOp(): - op_str = "AND" if isinstance(node.op, ast.And) else "OR" - return "(" + f" {op_str} ".join([self._lambda_parser(v) for v in node.values]) + ")" - case ast.UnaryOp(): - match node.op: - case ast.Not(): - return f"NOT ({self._lambda_parser(node.operand)})" - case ast.UAdd(): - return f"+{self._lambda_parser(node.operand)}" - case ast.USub(): - return f"-{self._lambda_parser(node.operand)}" - case ast.Invert(): - raise NotImplementedError("Invert operation is not supported.") - raise NotImplementedError(f"Unsupported unary operator: {type(node.op)}") - case ast.Attribute(): - # Cosmos DB: c.field_name - if node.attr not in self.definition.storage_names: - raise VectorStoreOperationException( - f"Field '{node.attr}' not in data model (storage property names are used)." - ) - return f"c.{node.attr}" - case ast.Name(): - # Could be a variable or constant; not supported - raise NotImplementedError("Constants or variables are not supported, use a value or attribute.") - case ast.Constant(): - # Quote strings, leave numbers as is - if isinstance(node.value, str): - return "'" + node.value.replace("'", "''") + "'" - if isinstance(node.value, (float, int)): - return str(node.value) - if node.value is None: - return "null" - raise NotImplementedError(f"Unsupported constant type: {type(node.value)}") - raise NotImplementedError(f"Unsupported AST node: {type(node)}") + return "(" + " AND ".join([parse(v) for v in values]) + ")" + left = parse(node.left) + right = parse(node.comparators[0]) + op = node.ops[0] + match op: + case ast.In(): + # Cosmos DB: ARRAY_CONTAINS(right, left) + return f"ARRAY_CONTAINS({right}, {left})" + case ast.NotIn(): + return f"NOT ARRAY_CONTAINS({right}, {left})" + case ast.Eq(): + return f"{left} = {right}" + case ast.NotEq(): + return f"{left} != {right}" + case ast.Gt(): + return f"{left} > {right}" + case ast.GtE(): + return f"{left} >= {right}" + case ast.Lt(): + return f"{left} < {right}" + case ast.LtE(): + return f"{left} <= {right}" + raise NotImplementedError(f"Unsupported operator: {type(op)}") + case ast.BoolOp(): + op_str = "AND" if isinstance(node.op, ast.And) else "OR" + return "(" + f" {op_str} ".join([parse(v) for v in node.values]) + ")" + case ast.UnaryOp(): + match node.op: + case ast.Not(): + return f"NOT ({parse(node.operand)})" + case ast.UAdd(): + return f"+{parse(node.operand)}" + case ast.USub(): + return f"-{parse(node.operand)}" + case ast.Invert(): + raise NotImplementedError("Invert operation is not supported.") + raise NotImplementedError(f"Unsupported unary operator: {type(node.op)}") + case ast.Attribute(): + # Cosmos DB: c.field_name + if node.attr not in self.definition.storage_names: + raise VectorStoreOperationException( + f"Field '{node.attr}' not in data model (storage property names are used)." + ) + return f"c.{node.attr}" + case ast.Name(): + # Could be a variable or constant; not supported + raise NotImplementedError("Constants or variables are not supported, use a value or attribute.") + case ast.Constant(): + # Bind strings as query parameters to avoid SQL injection. Numbers and null + # cannot carry injection, so they are inlined. + if isinstance(node.value, str): + name = f"@filter_p{len(parameters)}" + parameters.append({"name": name, "value": node.value}) + return name + if isinstance(node.value, (float, int)): + return str(node.value) + if node.value is None: + return "null" + raise NotImplementedError(f"Unsupported constant type: {type(node.value)}") + raise NotImplementedError(f"Unsupported AST node: {type(node)}") + + return parse(node), parameters @override def _get_record_from_result(self, result: dict[str, Any]) -> dict[str, Any]: diff --git a/python/tests/integration/memory/azure_cosmos_db/test_azure_cosmos_db_no_sql.py b/python/tests/integration/memory/azure_cosmos_db/test_azure_cosmos_db_no_sql.py index 1b8d1874ee72..1d8082bd9d0a 100644 --- a/python/tests/integration/memory/azure_cosmos_db/test_azure_cosmos_db_no_sql.py +++ b/python/tests/integration/memory/azure_cosmos_db/test_azure_cosmos_db_no_sql.py @@ -229,3 +229,52 @@ async def test_custom_client( assert await collection.collection_exists() is False collection_names = await store.list_collection_names() assert collection_name not in collection_names + + async def test_search_with_filter( + self, + stores: dict[str, Callable[[], VectorStore]], + record_type: type, + ): + """Test vector search with equality filters, including apostrophe and injection payloads. + + This is a regression test for the Cosmos DB NoSQL filter SQL-injection issue: filter + constants must be bound as query parameters, so that: + * an equality filter returns only the matching record(s), + * a value containing an apostrophe (e.g. O'Reilly) works instead of raising HTTP 400, and + * a malicious value (\\' OR true --) is treated as an inert literal and does not bypass + the filter. + """ + vector = [0.1, 0.2, 0.3, 0.4, 0.5] + records = [ + record_type(id="red-1", product_type="red", description="red record", vector=vector), + record_type(id="blue-1", product_type="blue", description="blue record", vector=vector), + record_type(id="oreilly-1", product_type="O'Reilly", description="apostrophe record", vector=vector), + ] + + async with stores["azure_cosmos_db_no_sql"]() as store: + collection_name = "search_with_filter" + collection = store.get_collection(collection_name=collection_name, record_type=record_type) + + try: + await collection.ensure_collection_exists() + await collection.upsert(records) + + # Equality filter returns only the matching record. + results = await collection.search(vector=vector, filter=lambda x: x.product_type == "red", top=10) + matched = [result.record async for result in results.results] + assert {record.id for record in matched} == {"red-1"} + + # A value containing an apostrophe must be handled as a literal (no HTTP 400). + results = await collection.search(vector=vector, filter=lambda x: x.product_type == "O'Reilly", top=10) + matched = [result.record async for result in results.results] + assert {record.id for record in matched} == {"oreilly-1"} + + # A SQL-injection payload must not bypass the filter; it matches no record. + results = await collection.search( + vector=vector, filter=lambda x: x.product_type == "\\' OR true --", top=10 + ) + matched = [result.record async for result in results.results] + assert matched == [] + finally: + await collection.ensure_collection_deleted() + assert await collection.collection_exists() is False diff --git a/python/tests/unit/connectors/memory/azure_cosmos_db/test_azure_cosmos_db_no_sql_collection.py b/python/tests/unit/connectors/memory/azure_cosmos_db/test_azure_cosmos_db_no_sql_collection.py index 617a097053f0..d7dd03f4e01f 100644 --- a/python/tests/unit/connectors/memory/azure_cosmos_db/test_azure_cosmos_db_no_sql_collection.py +++ b/python/tests/unit/connectors/memory/azure_cosmos_db/test_azure_cosmos_db_no_sql_collection.py @@ -68,36 +68,38 @@ def test_azure_cosmos_db_no_sql_collection_init_env( assert vector_collection.create_database is False -def test_azure_cosmos_db_no_sql_build_filter_escapes_apostrophes( +def test_azure_cosmos_db_no_sql_build_filter_parameterizes_apostrophes( azure_cosmos_db_no_sql_unit_test_env, record_type, collection_name: str, ) -> None: - """Test Cosmos DB filter building escapes apostrophes in string literals.""" + """Test Cosmos DB filter building binds string literals as parameters (handles apostrophes).""" vector_collection = CosmosNoSqlCollection( record_type=record_type, collection_name=collection_name, ) - filter_string = vector_collection._build_filter('lambda x: x.content == "O\'Reilly"') + filter_string, parameters = vector_collection._build_filter('lambda x: x.content == "O\'Reilly"') - assert filter_string == "c.content = 'O''Reilly'" + assert filter_string == "c.content = @filter_p0" + assert parameters == [{"name": "@filter_p0", "value": "O'Reilly"}] -def test_azure_cosmos_db_no_sql_build_filter_escapes_injection_payload( +def test_azure_cosmos_db_no_sql_build_filter_parameterizes_injection_payload( azure_cosmos_db_no_sql_unit_test_env, record_type, collection_name: str, ) -> None: - """Test Cosmos DB filter building keeps injection-shaped strings inside the literal.""" + """Test Cosmos DB filter building binds injection-shaped strings as inert parameters.""" vector_collection = CosmosNoSqlCollection( record_type=record_type, collection_name=collection_name, ) - filter_string = vector_collection._build_filter("lambda x: x.content == \"test' OR '1'='1\"") + filter_string, parameters = vector_collection._build_filter('lambda x: x.content == "\\\' OR true --"') - assert filter_string == "c.content = 'test'' OR ''1''=''1'" + assert filter_string == "c.content = @filter_p0" + assert parameters == [{"name": "@filter_p0", "value": "' OR true --"}] def test_azure_cosmos_db_no_sql_dynamic_filter_injection_payload_stays_literal( @@ -121,11 +123,133 @@ def test_azure_cosmos_db_no_sql_dynamic_filter_injection_payload_stays_literal( type_object=str, ) ], - content="test' OR '1'='1", + content="\\' OR true --", ) assert isinstance(generated_filter, str) - assert vector_collection._build_filter(generated_filter) == "c.content = 'test'' OR ''1''=''1'" + filter_string, parameters = vector_collection._build_filter(generated_filter) + assert filter_string == "c.content = @filter_p0" + assert parameters == [{"name": "@filter_p0", "value": "\\' OR true --"}] + + +async def test_azure_cosmos_db_no_sql_translate_filter_is_isolated_across_threads( + azure_cosmos_db_no_sql_unit_test_env, + record_type, + collection_name: str, +) -> None: + """Test _lambda_parser holds no shared state: concurrent calls on one instance don't bleed.""" + import asyncio + + vector_collection = CosmosNoSqlCollection( + record_type=record_type, + collection_name=collection_name, + ) + + def translate(value: str) -> tuple[str, list]: + return vector_collection._build_filter(f'lambda x: x.content == "{value}"') + + values = [f"tenant-{i}" for i in range(50)] + results = await asyncio.gather(*(asyncio.to_thread(translate, value) for value in values)) + + for value, (clause, parameters) in zip(values, results): + # Every call gets its own local accumulator, so each is a single @filter_p0 bound to its + # own value with no cross-contamination between concurrent calls. + assert clause == "c.content = @filter_p0" + assert parameters == [{"name": "@filter_p0", "value": value}] + + +async def test_azure_cosmos_db_no_sql_inner_search_parameterizes_filter( + azure_cosmos_db_no_sql_unit_test_env, + record_type, + collection_name: str, +) -> None: + """Test _inner_search binds filter constants as parameters instead of concatenating them.""" + from semantic_kernel.data.vector import SearchType, VectorSearchOptions + + vector_collection = CosmosNoSqlCollection( + record_type=record_type, + collection_name=collection_name, + ) + + async def _empty_results(): + for item in []: + yield item + + mock_container_proxy = MagicMock() + mock_container_proxy.query_items = MagicMock(return_value=_empty_results()) + + options = VectorSearchOptions( + vector_property_name="vector", + filter=lambda x: x.content == "\\' OR true --", + top=3, + ) + + with patch.object( + CosmosNoSqlCollection, + "_get_container_proxy", + new=AsyncMock(return_value=mock_container_proxy), + ): + await vector_collection._inner_search( + search_type=SearchType.VECTOR, + options=options, + vector=[0.1, 0.2, 0.3, 0.4, 0.5], + ) + + query, kwargs = mock_container_proxy.query_items.call_args + executed_query = query[0] + parameters = kwargs["parameters"] + + # The injection payload must never appear in the query text; it must be a bound parameter. + assert "OR true" not in executed_query + assert "c.content = @filter_p0" in executed_query + assert {"name": "@filter_p0", "value": "\\' OR true --"} in parameters + + +async def test_azure_cosmos_db_no_sql_inner_search_multiple_filters_have_unique_params( + azure_cosmos_db_no_sql_unit_test_env, + record_type, + collection_name: str, +) -> None: + """Test _inner_search renumbers filter params to unique names when multiple filters are given.""" + from semantic_kernel.data.vector import SearchType, VectorSearchOptions + + vector_collection = CosmosNoSqlCollection( + record_type=record_type, + collection_name=collection_name, + ) + + async def _empty_results(): + for item in []: + yield item + + mock_container_proxy = MagicMock() + mock_container_proxy.query_items = MagicMock(return_value=_empty_results()) + + options = VectorSearchOptions( + vector_property_name="vector", + filter=['lambda x: x.content == "alpha"', 'lambda x: x.content == "beta"'], + top=3, + ) + + with patch.object( + CosmosNoSqlCollection, + "_get_container_proxy", + new=AsyncMock(return_value=mock_container_proxy), + ): + await vector_collection._inner_search( + search_type=SearchType.VECTOR, + options=options, + vector=[0.1, 0.2, 0.3, 0.4, 0.5], + ) + + query, kwargs = mock_container_proxy.query_items.call_args + executed_query = query[0] + parameters = kwargs["parameters"] + + # Two filters, each with one string constant, must get distinct parameter names and be AND-ed. + assert "WHERE (c.content = @filter_p0 AND c.content = @filter_p1)" in executed_query + assert {"name": "@filter_p0", "value": "alpha"} in parameters + assert {"name": "@filter_p1", "value": "beta"} in parameters @pytest.mark.parametrize("exclude_list", [["AZURE_COSMOS_DB_NO_SQL_URL"]], indirect=True) From d00d8059ca31d58552c108a724d09676e2e6e555 Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Thu, 23 Jul 2026 17:01:11 +0000 Subject: [PATCH 2/2] Address PR comment --- .../azure_cosmos_db/test_azure_cosmos_db_no_sql_collection.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/tests/unit/connectors/memory/azure_cosmos_db/test_azure_cosmos_db_no_sql_collection.py b/python/tests/unit/connectors/memory/azure_cosmos_db/test_azure_cosmos_db_no_sql_collection.py index d7dd03f4e01f..717c9ee65e7f 100644 --- a/python/tests/unit/connectors/memory/azure_cosmos_db/test_azure_cosmos_db_no_sql_collection.py +++ b/python/tests/unit/connectors/memory/azure_cosmos_db/test_azure_cosmos_db_no_sql_collection.py @@ -96,10 +96,10 @@ def test_azure_cosmos_db_no_sql_build_filter_parameterizes_injection_payload( collection_name=collection_name, ) - filter_string, parameters = vector_collection._build_filter('lambda x: x.content == "\\\' OR true --"') + filter_string, parameters = vector_collection._build_filter('lambda x: x.content == "\\\\\' OR true --"') assert filter_string == "c.content = @filter_p0" - assert parameters == [{"name": "@filter_p0", "value": "' OR true --"}] + assert parameters == [{"name": "@filter_p0", "value": "\\' OR true --"}] def test_azure_cosmos_db_no_sql_dynamic_filter_injection_payload_stays_literal(