Python: Fix postgres filter bug

python/semantic_kernel/connectors/postgres.py

@@ -308,7 +308,12 @@ 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.
+    Comparisons using ``== None`` and ``!= None`` match null and non-null fields, respectively.
+    """
 
     connection_pool: AsyncConnectionPool | None = None
     db_schema: str = DEFAULT_SCHEMA
@@ -793,12 +798,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 +838,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 +852,87 @@ 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_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)
+                        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]:

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,124 @@ 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 | None, 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)
+
+
+@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)