From c3ecd9fb33d67fe2f9baaf46e03cdeefd883e62a Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Mon, 28 Sep 2026 14:51:39 +0000 Subject: [PATCH 1/2] Fix postgres filter bug --- python/semantic_kernel/connectors/postgres.py | 164 ++++++++-------- .../memory/postgres/test_postgres_int.py | 80 ++++++++ .../connectors/memory/test_postgres_store.py | 179 +++++++++++++++++- 3 files changed, 335 insertions(+), 88 deletions(-) diff --git a/python/semantic_kernel/connectors/postgres.py b/python/semantic_kernel/connectors/postgres.py index 5a7e66c15880..47fe30a80cd0 100644 --- a/python/semantic_kernel/connectors/postgres.py +++ b/python/semantic_kernel/connectors/postgres.py @@ -308,7 +308,11 @@ class PostgresCollection( VectorSearch[TKey, TModel], Generic[TKey, TModel], ): - """PostgreSQL collection implementation.""" + """PostgreSQL collection implementation. + + Search filters accept Python lambdas or strings containing lambda expressions, not raw SQL. + Filter values are passed separately from the SQL command as query parameters. + """ connection_pool: AsyncConnectionPool | None = None db_schema: str = DEFAULT_SCHEMA @@ -793,12 +797,13 @@ def _construct_vector_query( table=sql.Identifier(self.collection_name), ) - if where_clauses := self._build_filter(options.filter): # type: ignore - query += ( - sql.SQL("WHERE {clause}").format(clause=sql.SQL(" AND ").join(where_clauses)) - if isinstance(where_clauses, list) - else sql.SQL("WHERE {clause}").format(clause=where_clauses) - ) + params: list[Any] = ["[" + ",".join([str(float(v)) for v in vector]) + "]"] + if filters := self._build_filter(options.filter): + clauses: list[sql.Composable] = [] + for clause, values in filters if isinstance(filters, list) else [filters]: + clauses.append(sql.SQL("({})").format(clause)) + params.extend(values) + query += sql.SQL(" WHERE {}").format(sql.SQL(" AND ").join(clauses)) query += sql.SQL(" ORDER BY {dist_col} LIMIT {limit}").format( dist_col=sql.Identifier(self._distance_column_name), @@ -832,9 +837,6 @@ def _construct_vector_query( subquery=query, ) - # Convert the vector to a string for the query - params = ["[" + ",".join([str(float(v)) for v in vector]) + "]"] - return ( query, params, @@ -849,81 +851,73 @@ def _construct_vector_query( ) @override - def _lambda_parser(self, node: ast.AST) -> Any: - # Comparison operations - match node: - case ast.Compare(): - if len(node.ops) > 1: - # Chain comparisons (e.g., 1 < x < 3) become AND of each comparison - values = [] - for idx in range(len(node.ops)): - left = node.left if idx == 0 else node.comparators[idx - 1] - right = node.comparators[idx] - op = node.ops[idx] - values.append(self._lambda_parser(ast.Compare(left=left, ops=[op], comparators=[right]))) - return f"({' AND '.join(values)})" - left = self._lambda_parser(node.left) - right = self._lambda_parser(node.comparators[0]) - op = node.ops[0] - match op: - case ast.In(): - return f"{left} IN {right}" - case ast.NotIn(): - return f"{left} NOT IN {right}" - 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 = node.op # type: ignore - values = [self._lambda_parser(v) for v in node.values] - if isinstance(op, ast.And): - return f"({' AND '.join(values)})" - if isinstance(op, ast.Or): - return f"({' OR '.join(values)})" - raise NotImplementedError(f"Unsupported BoolOp: {type(op)}") - case ast.UnaryOp(): - match node.op: - case ast.Not(): - operand = self._lambda_parser(node.operand) - return f"NOT ({operand})" - case ast.UAdd() | ast.USub() | ast.Invert(): - raise NotImplementedError("Unary +, -, ~ are not supported in PostgreSQL filters.") - case ast.Attribute(): - # Only allow attributes that are in the data model - 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'"{node.attr}"' - case ast.Name(): - # Only allow names that are in the data model - if node.id not in self.definition.storage_names: - raise VectorStoreOperationException( - f"Field '{node.id}' not in data model (storage property names are used)." - ) - return f'"{node.id}"' - case ast.Constant(): - if isinstance(node.value, str): - return "'" + node.value.replace("'", "''") + "'" - if node.value is None: - return "NULL" - if isinstance(node.value, bool): - return "TRUE" if node.value else "FALSE" - return str(node.value) - case ast.List(): - # For IN/NOT IN lists - return "(" + ", ".join(self._lambda_parser(elt) for elt in node.elts) + ")" - raise NotImplementedError(f"Unsupported AST node: {type(node)}") + def _lambda_parser(self, node: ast.AST) -> tuple[sql.Composable, list[Any]]: + """Parse a filter into SQL with placeholders and its ordered parameter values.""" + params: list[Any] = [] + + def parse(node: ast.AST) -> sql.Composable: + match node: + case ast.Compare(): + if len(node.ops) > 1: + comparisons = [] + for idx, op in enumerate(node.ops): + left_node = node.left if idx == 0 else node.comparators[idx - 1] + right_node = node.comparators[idx] + comparisons.append(parse(ast.Compare(left=left_node, ops=[op], comparators=[right_node]))) + return sql.SQL("({})").format(sql.SQL(" AND ").join(comparisons)) + left = parse(node.left) + right = parse(node.comparators[0]) + op = node.ops[0] + match op: + case ast.In(): + return sql.SQL("{} IN {}").format(left, right) + case ast.NotIn(): + return sql.SQL("{} NOT IN {}").format(left, right) + case ast.Eq(): + return sql.SQL("{} = {}").format(left, right) + case ast.NotEq(): + return sql.SQL("{} <> {}").format(left, right) + case ast.Gt(): + return sql.SQL("{} > {}").format(left, right) + case ast.GtE(): + return sql.SQL("{} >= {}").format(left, right) + case ast.Lt(): + return sql.SQL("{} < {}").format(left, right) + case ast.LtE(): + return sql.SQL("{} <= {}").format(left, right) + raise NotImplementedError(f"Unsupported operator: {type(op)}") + case ast.BoolOp(): + values = [parse(value) for value in node.values] + if isinstance(node.op, ast.And): + return sql.SQL("({})").format(sql.SQL(" AND ").join(values)) + if isinstance(node.op, ast.Or): + return sql.SQL("({})").format(sql.SQL(" OR ").join(values)) + raise NotImplementedError(f"Unsupported BoolOp: {type(node.op)}") + case ast.UnaryOp(): + match node.op: + case ast.Not(): + return sql.SQL("NOT ({})").format(parse(node.operand)) + case ast.UAdd() | ast.USub() | ast.Invert(): + raise NotImplementedError("Unary +, -, ~ are not supported in PostgreSQL filters.") + case ast.Attribute() | ast.Name(): + name = node.attr if isinstance(node, ast.Attribute) else node.id + if name not in self.definition.storage_names: + raise VectorStoreOperationException( + f"Field '{name}' not in data model (storage property names are used)." + ) + return sql.Identifier(name) + case ast.Constant(): + if not isinstance(node.value, (str, int, float, bool)) and node.value is not None: + raise VectorStoreOperationException(f"Unsupported constant type: {type(node.value)}") + params.append(node.value) + return sql.Placeholder() + case ast.List(): + if not node.elts: + raise VectorStoreOperationException("Empty lists are not supported in PostgreSQL filters.") + return sql.SQL("({})").format(sql.SQL(", ").join(parse(elt) for elt in node.elts)) + raise NotImplementedError(f"Unsupported AST node: {type(node)}") + + return parse(node), params @override def _get_record_from_result(self, result: dict[str, Any]) -> dict[str, Any]: diff --git a/python/tests/integration/memory/postgres/test_postgres_int.py b/python/tests/integration/memory/postgres/test_postgres_int.py index bbfe0e91db89..ddaa67e69c62 100644 --- a/python/tests/integration/memory/postgres/test_postgres_int.py +++ b/python/tests/integration/memory/postgres/test_postgres_int.py @@ -11,6 +11,7 @@ from pydantic import BaseModel from semantic_kernel.connectors.postgres import PostgresCollection, PostgresSettings, PostgresStore +from semantic_kernel.data._shared import default_dynamic_filter_function from semantic_kernel.data.vector import ( DistanceFunction, IndexKind, @@ -22,6 +23,7 @@ MemoryConnectorConnectionException, MemoryConnectorInitializationError, ) +from semantic_kernel.functions import KernelParameterMetadata try: import psycopg # noqa: F401 @@ -257,3 +259,81 @@ async def test_search(vector_store: PostgresStore): finally: await simple_collection.delete([r.id for r in records]) + + +@vectorstoremodel +class FilteredDataModel(BaseModel): + id: Annotated[int, VectorStoreField("key")] + tenant: Annotated[str, VectorStoreField("data")] + embedding: Annotated[ + list[float] | None, + VectorStoreField( + "vector", dimensions=3, index_kind=IndexKind.HNSW, distance_function=DistanceFunction.COSINE_DISTANCE + ), + ] = None + + +@pytest_asyncio.fixture(params=["on", "off"]) +async def filtered_collection(request) -> AsyncGenerator[PostgresCollection[int, FilteredDataModel], None]: + settings = pg_settings.model_copy(update={"min_pool": 1, "max_pool": 1}) + async with await settings.create_connection_pool(options=f"-c standard_conforming_strings={request.param}") as pool: + async with pool.connection() as conn: + cursor = await conn.execute("SHOW standard_conforming_strings") + assert await cursor.fetchone() == (request.param,) + + collection = PostgresCollection( + record_type=FilteredDataModel, + collection_name=f"test_filters_{uuid.uuid4().hex}", + connection_pool=pool, + ) + await collection.ensure_collection_exists() + try: + yield collection + finally: + await collection.ensure_collection_deleted() + + +@pytest.mark.parametrize("include_total_count", [False, True]) +@pytest.mark.parametrize("filter_kind", ["single", "multiple", "membership", "dynamic"]) +@pytest.mark.parametrize("value", ["tenant_a", "O'Brien", "a\\' OR 1=1 --", "'; SELECT 1; --", "50%_%s\nnext"]) +async def test_search_filters( + filtered_collection: PostgresCollection[int, FilteredDataModel], + include_total_count: bool, + filter_kind: str, + value: str, +): + await filtered_collection.upsert([ + FilteredDataModel(id=1, tenant=value, embedding=[1.0, 0.0, 0.0]), + FilteredDataModel(id=2, tenant="tenant_b", embedding=[1.0, 0.0, 0.0]), + FilteredDataModel(id=3, tenant=value, embedding=[1.0, 0.0, 0.0]), + ]) + filter = f"lambda x: x.tenant == {value!r}" + expected_ids = {1, 3} + if filter_kind == "multiple": + filter = [filter, "lambda x: x.id < 3"] + expected_ids = {1} + elif filter_kind == "membership": + filter = f"lambda x: x.tenant in [{value!r}, 'missing'] and x.id < 3" + expected_ids = {1} + elif filter_kind == "dynamic": + filter = default_dynamic_filter_function( + filter="lambda x: x.id < 3", + parameters=[KernelParameterMetadata(name="tenant")], + tenant=value, + ) + expected_ids = {1} + + results = await filtered_collection.search( + vector=[1.0, 0.0, 0.0], filter=filter, include_total_count=include_total_count, top=10 + ) + + assert {result.record.id async for result in results.results} == expected_ids + assert results.total_count == (len(expected_ids) if include_total_count else None) + + no_matches = await filtered_collection.search( + vector=[1.0, 0.0, 0.0], + filter="lambda x: x.tenant == 'missing'", + include_total_count=include_total_count, + ) + assert [result async for result in no_matches.results] == [] + assert no_matches.total_count == (0 if include_total_count else None) diff --git a/python/tests/unit/connectors/memory/test_postgres_store.py b/python/tests/unit/connectors/memory/test_postgres_store.py index 8f417cf7036b..8e94dc313cee 100644 --- a/python/tests/unit/connectors/memory/test_postgres_store.py +++ b/python/tests/unit/connectors/memory/test_postgres_store.py @@ -1,5 +1,7 @@ # Copyright (c) Microsoft. All rights reserved. +import ast +import asyncio from collections.abc import AsyncGenerator from dataclasses import dataclass from typing import Annotated, Any @@ -7,7 +9,7 @@ import pytest import pytest_asyncio -from psycopg import AsyncConnection, AsyncCursor +from psycopg import AsyncConnection, AsyncCursor, ProgrammingError from psycopg_pool import AsyncConnectionPool from pytest import fixture @@ -17,7 +19,16 @@ PostgresSettings, PostgresStore, ) -from semantic_kernel.data.vector import DistanceFunction, IndexKind, VectorStoreField, vectorstoremodel +from semantic_kernel.data._shared import default_dynamic_filter_function +from semantic_kernel.data.vector import ( + DistanceFunction, + IndexKind, + VectorSearchOptions, + VectorStoreField, + vectorstoremodel, +) +from semantic_kernel.exceptions import VectorSearchExecutionException, VectorStoreOperationException +from semantic_kernel.functions import KernelParameterMetadata @fixture(scope="function") @@ -233,6 +244,18 @@ async def test_get_records(vector_store: PostgresStore, mock_cursor: Mock) -> No # region Vector Search tests +@pytest.mark.parametrize( + "filter, filter_sql, filter_params", + [ + (None, "", []), + ("lambda x: x.id == 1", ' WHERE ("id" = %s)', [1]), + ( + ["lambda x: x.id > 0", "lambda x: x.id == 1 or x.id == 2"], + ' WHERE ("id" > %s) AND (("id" = %s OR "id" = %s))', + [0, 1, 2], + ), + ], +) @pytest.mark.parametrize( "distance_function, operator, subquery_distance, include_vectors, include_total_count", [ @@ -251,6 +274,9 @@ async def test_vector_search( subquery_distance: str | None, include_vectors: bool, include_total_count: bool, + filter: str | list[str] | None, + filter_sql: str, + filter_params: list[Any], ) -> None: @vectorstoremodel @dataclass @@ -284,6 +310,7 @@ def model_post_init(self, context: Any) -> None: skip=5, include_vectors=include_vectors, include_total_count=include_total_count, + filter=filter, ) if include_total_count: # Including total count issues query directly @@ -308,7 +335,7 @@ def model_post_init(self, context: Any) -> None: expected_statement = ( f'SELECT {expected_columns}, "embedding" {operator} %s as "{DISTANCE_COLUMN_NAME}" ' - 'FROM "public"."test_collection" ' + f'FROM "public"."test_collection"{filter_sql} ' f'ORDER BY "{DISTANCE_COLUMN_NAME}" LIMIT 10 OFFSET 5' ) @@ -320,6 +347,152 @@ def model_post_init(self, context: Any) -> None: ) assert statement_str == expected_statement + assert execute_args[1] == ["[1.0,2.0,3.0]", *filter_params] + + +@fixture +def filter_collection(vector_store): + @vectorstoremodel + @dataclass + class FilterRecord: + id: Annotated[int, VectorStoreField("key")] + tenant: Annotated[str, VectorStoreField("data")] + embedding: Annotated[ + list[float], + VectorStoreField("vector", dimensions=3, distance_function=DistanceFunction.COSINE_DISTANCE), + ] + label: Annotated[str, VectorStoreField("data", storage_name='label"name')] = "" + + return vector_store.get_collection(collection_name="filter_records", record_type=FilterRecord) + + +@pytest.mark.parametrize( + "expression, expected_sql, expected_params", + [ + ("x.id == 1", '"id" = %s', [1]), + ("x.id != 1", '"id" <> %s', [1]), + ("x.id > 1", '"id" > %s', [1]), + ("x.id >= 1", '"id" >= %s', [1]), + ("x.id < 1", '"id" < %s', [1]), + ("x.id <= 1", '"id" <= %s', [1]), + ("x.id in [1, 2]", '"id" IN (%s, %s)', [1, 2]), + ("x.id not in [1, 2]", '"id" NOT IN (%s, %s)', [1, 2]), + ("x.id == 1 and x.id < 3", '("id" = %s AND "id" < %s)', [1, 3]), + ("x.id == 1 or x.id == 2", '("id" = %s OR "id" = %s)', [1, 2]), + ("not x.id == 1", 'NOT ("id" = %s)', [1]), + ("0 < x.id < 3", '(%s < "id" AND "id" < %s)', [0, 3]), + ("x.id < 2 < 3", '("id" < %s AND %s < %s)', [2, 2, 3]), + ("x.id == 1.5", '"id" = %s', [1.5]), + ("x.id == True", '"id" = %s', [True]), + ("x.id == False", '"id" = %s', [False]), + ("x.tenant == None", '"tenant" = %s', [None]), + ("x.tenant != None", '"tenant" <> %s', [None]), + ("x.id == x.id", '"id" = "id"', []), + ("id == 1", '"id" = %s', [1]), + ( + "(x.id == 1 or x.id == 2) and not x.tenant == 'other'", + '(("id" = %s OR "id" = %s) AND NOT ("tenant" = %s))', + [1, 2, "other"], + ), + ], +) +def test_filter_parser(filter_collection, expression, expected_sql, expected_params): + clause, params = filter_collection._build_filter(f"lambda x: {expression}") + + assert clause.as_string() == expected_sql + assert params == expected_params + assert [type(value) for value in params] == [type(value) for value in expected_params] + + +@pytest.mark.parametrize( + "value", + ["tenant_a", "", "O'Brien", "\\", "a\\' OR 1=1 --", "'; SELECT 1; --", "50%_%s\nnext", "\u00e9"], +) +def test_filter_values_are_parameters(filter_collection, value): + filter = default_dynamic_filter_function( + parameters=[KernelParameterMetadata(name="tenant")], + tenant=value, + ) + query, params, _ = filter_collection._construct_vector_query( + [1, 0, 0], VectorSearchOptions(filter=filter, include_total_count=True) + ) + + assert query.as_string() == ( + 'SELECT "id", "tenant", "label""name", "embedding" <=> %s as "sk_pg_distance" ' + 'FROM "public"."filter_records" WHERE ("tenant" = %s) ORDER BY "sk_pg_distance" LIMIT 3' + ) + assert params == ["[1.0,0.0,0.0]", value] + + +def test_callable_filter(filter_collection): + clause, params = filter_collection._build_filter(lambda x: x.tenant == "tenant_a") + + assert clause.as_string() == '"tenant" = %s' + assert params == ["tenant_a"] + + +def test_filter_identifier_quoting(filter_collection): + clause, params = filter_collection._lambda_parser(ast.Name(id='label"name')) + + assert clause.as_string() == '"label""name"' + assert params == [] + + +@pytest.mark.parametrize( + "expression, error", + [ + ("x.unknown == 1", VectorStoreOperationException), + ("unknown == 1", VectorStoreOperationException), + ("x.id in []", VectorStoreOperationException), + ("x.id not in []", VectorStoreOperationException), + ("x.id == b'bytes'", VectorStoreOperationException), + ("x.id == 1j", VectorStoreOperationException), + ("x.id is None", NotImplementedError), + ("x.id == -1", NotImplementedError), + ("x.id == +1", NotImplementedError), + ("x.id == ~1", NotImplementedError), + ("x.tenant.startswith('a')", NotImplementedError), + ("x.id in (1, 2)", NotImplementedError), + ], +) +def test_filter_rejects_unsupported_input(filter_collection, expression, error): + with pytest.raises(error): + filter_collection._construct_vector_query([1, 0, 0], VectorSearchOptions(filter=f"lambda x: {expression}")) + + +def test_filter_rejects_raw_sql(filter_collection): + with pytest.raises((SyntaxError, VectorStoreOperationException)): + filter_collection._build_filter("tenant = 'tenant_a'") + + +@pytest.mark.parametrize("include_total_count", [False, True]) +async def test_filtered_search_error_is_not_retried(filter_collection, mock_cursor, include_total_count): + mock_cursor.execute.side_effect = ProgrammingError("test query failure") + + with pytest.raises((ProgrammingError, VectorSearchExecutionException)): + results = await filter_collection.search( + vector=[1, 0, 0], filter="lambda x: x.tenant == 'tenant_a'", include_total_count=include_total_count + ) + async for _ in results.results: + pass + + assert mock_cursor.execute.call_count == 1 + assert ' WHERE ("tenant" = %s)' in mock_cursor.execute.call_args.args[0].as_string() + assert mock_cursor.execute.call_args.args[1] == ["[1.0,0.0,0.0]", "tenant_a"] + + +async def test_concurrent_filtered_searches_keep_separate_parameters(filter_collection, mock_cursor): + async def search(tenant): + results = await filter_collection.search(vector=[1, 0, 0], filter=f"lambda x: x.tenant == {tenant!r}") + await asyncio.sleep(0) + return [result async for result in results.results] + + await asyncio.gather(search("tenant_a"), search("tenant_b")) + + assert [call.args[1] for call in mock_cursor.execute.call_args_list] == [ + ["[1.0,0.0,0.0]", "tenant_a"], + ["[1.0,0.0,0.0]", "tenant_b"], + ] async def test_model_post_init_conflicting_distance_column_name(vector_store: PostgresStore) -> None: From 1f8b78cd698a10d51046d0034d21b9b8aa1d85ce Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Mon, 28 Sep 2026 16:09:39 +0000 Subject: [PATCH 2/2] Address PR comments --- python/semantic_kernel/connectors/postgres.py | 19 +++++- .../memory/postgres/test_postgres_int.py | 45 +++++++++++++- .../connectors/memory/test_postgres_store.py | 58 ++++++++++++++++++- 3 files changed, 117 insertions(+), 5 deletions(-) diff --git a/python/semantic_kernel/connectors/postgres.py b/python/semantic_kernel/connectors/postgres.py index 47fe30a80cd0..42994d75d042 100644 --- a/python/semantic_kernel/connectors/postgres.py +++ b/python/semantic_kernel/connectors/postgres.py @@ -312,6 +312,7 @@ class PostgresCollection( Search filters accept Python lambdas or strings containing lambda expressions, not raw SQL. Filter values are passed separately from the SQL command as query parameters. + Comparisons using ``== None`` and ``!= None`` match null and non-null fields, respectively. """ connection_pool: AsyncConnectionPool | None = None @@ -865,9 +866,23 @@ def parse(node: ast.AST) -> sql.Composable: right_node = node.comparators[idx] comparisons.append(parse(ast.Compare(left=left_node, ops=[op], comparators=[right_node]))) return sql.SQL("({})").format(sql.SQL(" AND ").join(comparisons)) - left = parse(node.left) - right = parse(node.comparators[0]) + left_node = node.left + right_node = node.comparators[0] op = node.ops[0] + if isinstance(op, (ast.Eq, ast.NotEq)): + if isinstance(left_node, ast.Constant) and left_node.value is None: + left_node, right_node = right_node, left_node + if isinstance(right_node, ast.Constant) and right_node.value is None: + if isinstance(left_node, ast.Constant) and ( + left_node.value is None or isinstance(left_node.value, (str, int, float, bool)) + ): + matches = (left_node.value is None) == isinstance(op, ast.Eq) + return sql.SQL("TRUE" if matches else "FALSE") + return sql.SQL("{} IS NULL" if isinstance(op, ast.Eq) else "{} IS NOT NULL").format( + parse(left_node) + ) + left = parse(left_node) + right = parse(right_node) match op: case ast.In(): return sql.SQL("{} IN {}").format(left, right) diff --git a/python/tests/integration/memory/postgres/test_postgres_int.py b/python/tests/integration/memory/postgres/test_postgres_int.py index ddaa67e69c62..cf9df99e967a 100644 --- a/python/tests/integration/memory/postgres/test_postgres_int.py +++ b/python/tests/integration/memory/postgres/test_postgres_int.py @@ -264,7 +264,7 @@ async def test_search(vector_store: PostgresStore): @vectorstoremodel class FilteredDataModel(BaseModel): id: Annotated[int, VectorStoreField("key")] - tenant: Annotated[str, VectorStoreField("data")] + tenant: Annotated[str | None, VectorStoreField("data")] embedding: Annotated[ list[float] | None, VectorStoreField( @@ -337,3 +337,46 @@ async def test_search_filters( ) assert [result async for result in no_matches.results] == [] assert no_matches.total_count == (0 if include_total_count else None) + + +@pytest.mark.parametrize("include_total_count", [False, True]) +@pytest.mark.parametrize( + "filter, expected_ids", + [ + ("lambda x: x.tenant == None", {1, 4}), + ("lambda x: x.tenant != None", {2, 3}), + ("lambda x: None == x.tenant", {1, 4}), + ("lambda x: None != x.tenant", {2, 3}), + ("lambda x: None == None", {1, 2, 3, 4}), + ("lambda x: None != None", set()), + ("lambda x: None != 'value'", {1, 2, 3, 4}), + ("lambda x: False == None", set()), + ("lambda x: not x.tenant == None", {2, 3}), + ("lambda x: not x.tenant != None", {1, 4}), + ("lambda x: x.tenant == None or x.tenant == 'value'", {1, 2, 4}), + ("lambda x: x.id > 1 and x.tenant == None", {4}), + (["lambda x: x.id >= 3", "lambda x: x.tenant == None"], {4}), + (["lambda x: x.id > 1", "lambda x: None != x.tenant"], {2, 3}), + ("lambda x: None == x.tenant == None", {1, 4}), + ("lambda x: x.tenant == ''", {3}), + ], +) +async def test_search_null_filters( + filtered_collection: PostgresCollection[int, FilteredDataModel], + include_total_count: bool, + filter: str | list[str], + expected_ids: set[int], +): + await filtered_collection.upsert([ + FilteredDataModel(id=1, tenant=None, embedding=[1.0, 0.0, 0.0]), + FilteredDataModel(id=2, tenant="value", embedding=[1.0, 0.0, 0.0]), + FilteredDataModel(id=3, tenant="", embedding=[1.0, 0.0, 0.0]), + FilteredDataModel(id=4, tenant=None, embedding=[1.0, 0.0, 0.0]), + ]) + + results = await filtered_collection.search( + vector=[1.0, 0.0, 0.0], filter=filter, include_total_count=include_total_count, top=10 + ) + + assert {result.record.id async for result in results.results} == expected_ids + assert results.total_count == (len(expected_ids) if include_total_count else None) diff --git a/python/tests/unit/connectors/memory/test_postgres_store.py b/python/tests/unit/connectors/memory/test_postgres_store.py index 8e94dc313cee..ad33d864ca9f 100644 --- a/python/tests/unit/connectors/memory/test_postgres_store.py +++ b/python/tests/unit/connectors/memory/test_postgres_store.py @@ -385,8 +385,6 @@ class FilterRecord: ("x.id == 1.5", '"id" = %s', [1.5]), ("x.id == True", '"id" = %s', [True]), ("x.id == False", '"id" = %s', [False]), - ("x.tenant == None", '"tenant" = %s', [None]), - ("x.tenant != None", '"tenant" <> %s', [None]), ("x.id == x.id", '"id" = "id"', []), ("id == 1", '"id" = %s', [1]), ( @@ -404,6 +402,57 @@ def test_filter_parser(filter_collection, expression, expected_sql, expected_par assert [type(value) for value in params] == [type(value) for value in expected_params] +@pytest.mark.parametrize( + "expression, expected_sql, expected_params", + [ + ("x.tenant == None", '"tenant" IS NULL', []), + ("x.tenant != None", '"tenant" IS NOT NULL', []), + ("None == x.tenant", '"tenant" IS NULL', []), + ("None != x.tenant", '"tenant" IS NOT NULL', []), + ("None == None", "TRUE", []), + ("None != None", "FALSE", []), + ("'value' != None", "TRUE", []), + ("None == 'value'", "FALSE", []), + ("0 != None", "TRUE", []), + ("False == None", "FALSE", []), + ("not x.tenant == None", 'NOT ("tenant" IS NULL)', []), + ("not x.tenant != None", 'NOT ("tenant" IS NOT NULL)', []), + ("x.tenant == None or x.id == 1", '("tenant" IS NULL OR "id" = %s)', [1]), + ("None == x.tenant == None", '("tenant" IS NULL AND "tenant" IS NULL)', []), + ("'value' == x.tenant != None", '(%s = "tenant" AND "tenant" IS NOT NULL)', ["value"]), + ], +) +def test_null_filter_parser(filter_collection, expression, expected_sql, expected_params): + clause, params = filter_collection._build_filter(f"lambda x: {expression}") + + assert clause.as_string() == expected_sql + assert params == expected_params + + +def test_null_filter_query_parameter_order(filter_collection): + filters = [ + "lambda x: x.id > 0", + "lambda x: x.tenant == None or x.tenant == 'value'", + "lambda x: x.id < 5", + ] + query, params, _ = filter_collection._construct_vector_query([1, 0, 0], VectorSearchOptions(filter=filters)) + + assert query.as_string() == ( + 'SELECT "id", "tenant", "label""name", "embedding" <=> %s as "sk_pg_distance" ' + 'FROM "public"."filter_records" WHERE ("id" > %s) ' + 'AND (("tenant" IS NULL OR "tenant" = %s)) AND ("id" < %s) ORDER BY "sk_pg_distance" LIMIT 3' + ) + assert params == ["[1.0,0.0,0.0]", 0, "value", 5] + + +def test_dynamic_null_filter(filter_collection): + filter = default_dynamic_filter_function(parameters=[KernelParameterMetadata(name="tenant")], tenant=None) + clause, params = filter_collection._build_filter(filter) + + assert clause.as_string() == '"tenant" IS NULL' + assert params == [] + + @pytest.mark.parametrize( "value", ["tenant_a", "", "O'Brien", "\\", "a\\' OR 1=1 --", "'; SELECT 1; --", "50%_%s\nnext", "\u00e9"], @@ -442,12 +491,17 @@ def test_filter_identifier_quoting(filter_collection): "expression, error", [ ("x.unknown == 1", VectorStoreOperationException), + ("x.unknown == None", VectorStoreOperationException), + ("None != x.unknown", VectorStoreOperationException), ("unknown == 1", VectorStoreOperationException), ("x.id in []", VectorStoreOperationException), ("x.id not in []", VectorStoreOperationException), ("x.id == b'bytes'", VectorStoreOperationException), ("x.id == 1j", VectorStoreOperationException), + ("None == b'bytes'", VectorStoreOperationException), + ("1j != None", VectorStoreOperationException), ("x.id is None", NotImplementedError), + ("x.id is not None", NotImplementedError), ("x.id == -1", NotImplementedError), ("x.id == +1", NotImplementedError), ("x.id == ~1", NotImplementedError),