Skip to content

Fix cosmosdb vectorstore bug #14182

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
westey-m merged 2 commits into from
Jul 27, 2026
Merged

Fix cosmosdb vectorstore bug #14182

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
219b352
Fix cosmosdb vectorstore bug
westey-m Jul 23, 2026
d00d805
Address PR comment
westey-m Jul 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
207 changes: 123 additions & 84 deletions python/semantic_kernel/connectors/azure_cosmos_db.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import ast
import asyncio
import re
import sys
from collections.abc import Sequence
from importlib import metadata
Expand Down Expand Up @@ -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})
Expand Down Expand Up @@ -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]:
Expand Down
49 changes: 49 additions & 0 deletions python/tests/integration/memory/azure_cosmos_db/test_azure_cosmos_db_no_sql.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Loading
Loading