mirror of
https://github.com/microsoft/semantic-kernel.git
synced 2026-10-12 02:19:06 +00:00
Fix cosmosdb vectorstore bug (#14182)
### Motivation and Context <!-- Thank you for your contribution to the semantic-kernel repo! Please help reviewers and future users, providing the following information: 1. Why is this change required? 2. What problem does it solve? 3. What scenario does it contribute to? 4. If it fixes an open issue, please link to the issue here. --> ### Description <!-- Describe your changes, the overall approach, the underlying design. These notes will help understanding how your code works. Thanks! --> ### Contribution Checklist <!-- Before submitting this PR, please make sure: --> - [ ] The code builds clean without any errors or warnings - [ ] The PR follows the [SK Contribution Guidelines](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md) and the [pre-submission formatting script](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md#development-scripts) raises no violations - [ ] All unit tests pass, and I have added new tests where possible - [ ] I didn't break anyone 😄
This commit is contained in:
@@ -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 @@ class CosmosNoSqlCollection(
|
||||
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 @@ class CosmosNoSqlCollection(
|
||||
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]:
|
||||
|
||||
@@ -229,3 +229,52 @@ class TestCosmosDBNoSQL(VectorStoreTestBase):
|
||||
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
|
||||
|
||||
+134
-10
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user