Skip to content

Commit 9974625

Browse files
westey-mwestey-m
authored
Python: Fix postgres filter bug (#14497)
### Motivation and Context ### Description - Improve PostgreSQL query construction for search filters. - Expand unit and integration coverage for filtered searches. - Clarify supported filter input in the connector documentation. ### 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 😄
1 parent ee23e07 commit 9974625

3 files changed

Lines changed: 447 additions & 88 deletions

File tree

‎python/semantic_kernel/connectors/postgres.py‎

Lines changed: 94 additions & 85 deletions
Original file line numberDiff line numberDiff line change
@@ -308,7 +308,12 @@ class PostgresCollection(
308308
VectorSearch[TKey, TModel],
309309
Generic[TKey, TModel],
310310
):
311-
"""PostgreSQL collection implementation."""
311+
"""PostgreSQL collection implementation.
312+
313+
Search filters accept Python lambdas or strings containing lambda expressions, not raw SQL.
314+
Filter values are passed separately from the SQL command as query parameters.
315+
Comparisons using ``== None`` and ``!= None`` match null and non-null fields, respectively.
316+
"""
312317

313318
connection_pool: AsyncConnectionPool | None = None
314319
db_schema: str = DEFAULT_SCHEMA
@@ -793,12 +798,13 @@ def _construct_vector_query(
793798
table=sql.Identifier(self.collection_name),
794799
)
795800

796-
if where_clauses := self._build_filter(options.filter): # type: ignore
797-
query += (
798-
sql.SQL("WHERE {clause}").format(clause=sql.SQL(" AND ").join(where_clauses))
799-
if isinstance(where_clauses, list)
800-
else sql.SQL("WHERE {clause}").format(clause=where_clauses)
801-
)
801+
params: list[Any] = ["[" + ",".join([str(float(v)) for v in vector]) + "]"]
802+
if filters := self._build_filter(options.filter):
803+
clauses: list[sql.Composable] = []
804+
for clause, values in filters if isinstance(filters, list) else [filters]:
805+
clauses.append(sql.SQL("({})").format(clause))
806+
params.extend(values)
807+
query += sql.SQL(" WHERE {}").format(sql.SQL(" AND ").join(clauses))
802808

803809
query += sql.SQL(" ORDER BY {dist_col} LIMIT {limit}").format(
804810
dist_col=sql.Identifier(self._distance_column_name),
@@ -832,9 +838,6 @@ def _construct_vector_query(
832838
subquery=query,
833839
)
834840

835-
# Convert the vector to a string for the query
836-
params = ["[" + ",".join([str(float(v)) for v in vector]) + "]"]
837-
838841
return (
839842
query,
840843
params,
@@ -849,81 +852,87 @@ def _construct_vector_query(
849852
)
850853

851854
@override
852-
def _lambda_parser(self, node: ast.AST) -> Any:
853-
# Comparison operations
854-
match node:
855-
case ast.Compare():
856-
if len(node.ops) > 1:
857-
# Chain comparisons (e.g., 1 < x < 3) become AND of each comparison
858-
values = []
859-
for idx in range(len(node.ops)):
860-
left = node.left if idx == 0 else node.comparators[idx - 1]
861-
right = node.comparators[idx]
862-
op = node.ops[idx]
863-
values.append(self._lambda_parser(ast.Compare(left=left, ops=[op], comparators=[right])))
864-
return f"({' AND '.join(values)})"
865-
left = self._lambda_parser(node.left)
866-
right = self._lambda_parser(node.comparators[0])
867-
op = node.ops[0]
868-
match op:
869-
case ast.In():
870-
return f"{left} IN {right}"
871-
case ast.NotIn():
872-
return f"{left} NOT IN {right}"
873-
case ast.Eq():
874-
return f"{left} = {right}"
875-
case ast.NotEq():
876-
return f"{left} <> {right}"
877-
case ast.Gt():
878-
return f"{left} > {right}"
879-
case ast.GtE():
880-
return f"{left} >= {right}"
881-
case ast.Lt():
882-
return f"{left} < {right}"
883-
case ast.LtE():
884-
return f"{left} <= {right}"
885-
raise NotImplementedError(f"Unsupported operator: {type(op)}")
886-
case ast.BoolOp():
887-
op = node.op # type: ignore
888-
values = [self._lambda_parser(v) for v in node.values]
889-
if isinstance(op, ast.And):
890-
return f"({' AND '.join(values)})"
891-
if isinstance(op, ast.Or):
892-
return f"({' OR '.join(values)})"
893-
raise NotImplementedError(f"Unsupported BoolOp: {type(op)}")
894-
case ast.UnaryOp():
895-
match node.op:
896-
case ast.Not():
897-
operand = self._lambda_parser(node.operand)
898-
return f"NOT ({operand})"
899-
case ast.UAdd() | ast.USub() | ast.Invert():
900-
raise NotImplementedError("Unary +, -, ~ are not supported in PostgreSQL filters.")
901-
case ast.Attribute():
902-
# Only allow attributes that are in the data model
903-
if node.attr not in self.definition.storage_names:
904-
raise VectorStoreOperationException(
905-
f"Field '{node.attr}' not in data model (storage property names are used)."
906-
)
907-
return f'"{node.attr}"'
908-
case ast.Name():
909-
# Only allow names that are in the data model
910-
if node.id not in self.definition.storage_names:
911-
raise VectorStoreOperationException(
912-
f"Field '{node.id}' not in data model (storage property names are used)."
913-
)
914-
return f'"{node.id}"'
915-
case ast.Constant():
916-
if isinstance(node.value, str):
917-
return "'" + node.value.replace("'", "''") + "'"
918-
if node.value is None:
919-
return "NULL"
920-
if isinstance(node.value, bool):
921-
return "TRUE" if node.value else "FALSE"
922-
return str(node.value)
923-
case ast.List():
924-
# For IN/NOT IN lists
925-
return "(" + ", ".join(self._lambda_parser(elt) for elt in node.elts) + ")"
926-
raise NotImplementedError(f"Unsupported AST node: {type(node)}")
855+
def _lambda_parser(self, node: ast.AST) -> tuple[sql.Composable, list[Any]]:
856+
"""Parse a filter into SQL with placeholders and its ordered parameter values."""
857+
params: list[Any] = []
858+
859+
def parse(node: ast.AST) -> sql.Composable:
860+
match node:
861+
case ast.Compare():
862+
if len(node.ops) > 1:
863+
comparisons = []
864+
for idx, op in enumerate(node.ops):
865+
left_node = node.left if idx == 0 else node.comparators[idx - 1]
866+
right_node = node.comparators[idx]
867+
comparisons.append(parse(ast.Compare(left=left_node, ops=[op], comparators=[right_node])))
868+
return sql.SQL("({})").format(sql.SQL(" AND ").join(comparisons))
869+
left_node = node.left
870+
right_node = node.comparators[0]
871+
op = node.ops[0]
872+
if isinstance(op, (ast.Eq, ast.NotEq)):
873+
if isinstance(left_node, ast.Constant) and left_node.value is None:
874+
left_node, right_node = right_node, left_node
875+
if isinstance(right_node, ast.Constant) and right_node.value is None:
876+
if isinstance(left_node, ast.Constant) and (
877+
left_node.value is None or isinstance(left_node.value, (str, int, float, bool))
878+
):
879+
matches = (left_node.value is None) == isinstance(op, ast.Eq)
880+
return sql.SQL("TRUE" if matches else "FALSE")
881+
return sql.SQL("{} IS NULL" if isinstance(op, ast.Eq) else "{} IS NOT NULL").format(
882+
parse(left_node)
883+
)
884+
left = parse(left_node)
885+
right = parse(right_node)
886+
match op:
887+
case ast.In():
888+
return sql.SQL("{} IN {}").format(left, right)
889+
case ast.NotIn():
890+
return sql.SQL("{} NOT IN {}").format(left, right)
891+
case ast.Eq():
892+
return sql.SQL("{} = {}").format(left, right)
893+
case ast.NotEq():
894+
return sql.SQL("{} <> {}").format(left, right)
895+
case ast.Gt():
896+
return sql.SQL("{} > {}").format(left, right)
897+
case ast.GtE():
898+
return sql.SQL("{} >= {}").format(left, right)
899+
case ast.Lt():
900+
return sql.SQL("{} < {}").format(left, right)
901+
case ast.LtE():
902+
return sql.SQL("{} <= {}").format(left, right)
903+
raise NotImplementedError(f"Unsupported operator: {type(op)}")
904+
case ast.BoolOp():
905+
values = [parse(value) for value in node.values]
906+
if isinstance(node.op, ast.And):
907+
return sql.SQL("({})").format(sql.SQL(" AND ").join(values))
908+
if isinstance(node.op, ast.Or):
909+
return sql.SQL("({})").format(sql.SQL(" OR ").join(values))
910+
raise NotImplementedError(f"Unsupported BoolOp: {type(node.op)}")
911+
case ast.UnaryOp():
912+
match node.op:
913+
case ast.Not():
914+
return sql.SQL("NOT ({})").format(parse(node.operand))
915+
case ast.UAdd() | ast.USub() | ast.Invert():
916+
raise NotImplementedError("Unary +, -, ~ are not supported in PostgreSQL filters.")
917+
case ast.Attribute() | ast.Name():
918+
name = node.attr if isinstance(node, ast.Attribute) else node.id
919+
if name not in self.definition.storage_names:
920+
raise VectorStoreOperationException(
921+
f"Field '{name}' not in data model (storage property names are used)."
922+
)
923+
return sql.Identifier(name)
924+
case ast.Constant():
925+
if not isinstance(node.value, (str, int, float, bool)) and node.value is not None:
926+
raise VectorStoreOperationException(f"Unsupported constant type: {type(node.value)}")
927+
params.append(node.value)
928+
return sql.Placeholder()
929+
case ast.List():
930+
if not node.elts:
931+
raise VectorStoreOperationException("Empty lists are not supported in PostgreSQL filters.")
932+
return sql.SQL("({})").format(sql.SQL(", ").join(parse(elt) for elt in node.elts))
933+
raise NotImplementedError(f"Unsupported AST node: {type(node)}")
934+
935+
return parse(node), params
927936

928937
@override
929938
def _get_record_from_result(self, result: dict[str, Any]) -> dict[str, Any]:

‎python/tests/integration/memory/postgres/test_postgres_int.py‎

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
from pydantic import BaseModel
1212

1313
from semantic_kernel.connectors.postgres import PostgresCollection, PostgresSettings, PostgresStore
14+
from semantic_kernel.data._shared import default_dynamic_filter_function
1415
from semantic_kernel.data.vector import (
1516
DistanceFunction,
1617
IndexKind,
@@ -22,6 +23,7 @@
2223
MemoryConnectorConnectionException,
2324
MemoryConnectorInitializationError,
2425
)
26+
from semantic_kernel.functions import KernelParameterMetadata
2527

2628
try:
2729
import psycopg # noqa: F401
@@ -257,3 +259,124 @@ async def test_search(vector_store: PostgresStore):
257259

258260
finally:
259261
await simple_collection.delete([r.id for r in records])
262+
263+
264+
@vectorstoremodel
265+
class FilteredDataModel(BaseModel):
266+
id: Annotated[int, VectorStoreField("key")]
267+
tenant: Annotated[str | None, VectorStoreField("data")]
268+
embedding: Annotated[
269+
list[float] | None,
270+
VectorStoreField(
271+
"vector", dimensions=3, index_kind=IndexKind.HNSW, distance_function=DistanceFunction.COSINE_DISTANCE
272+
),
273+
] = None
274+
275+
276+
@pytest_asyncio.fixture(params=["on", "off"])
277+
async def filtered_collection(request) -> AsyncGenerator[PostgresCollection[int, FilteredDataModel], None]:
278+
settings = pg_settings.model_copy(update={"min_pool": 1, "max_pool": 1})
279+
async with await settings.create_connection_pool(options=f"-c standard_conforming_strings={request.param}") as pool:
280+
async with pool.connection() as conn:
281+
cursor = await conn.execute("SHOW standard_conforming_strings")
282+
assert await cursor.fetchone() == (request.param,)
283+
284+
collection = PostgresCollection(
285+
record_type=FilteredDataModel,
286+
collection_name=f"test_filters_{uuid.uuid4().hex}",
287+
connection_pool=pool,
288+
)
289+
await collection.ensure_collection_exists()
290+
try:
291+
yield collection
292+
finally:
293+
await collection.ensure_collection_deleted()
294+
295+
296+
@pytest.mark.parametrize("include_total_count", [False, True])
297+
@pytest.mark.parametrize("filter_kind", ["single", "multiple", "membership", "dynamic"])
298+
@pytest.mark.parametrize("value", ["tenant_a", "O'Brien", "a\\' OR 1=1 --", "'; SELECT 1; --", "50%_%s\nnext"])
299+
async def test_search_filters(
300+
filtered_collection: PostgresCollection[int, FilteredDataModel],
301+
include_total_count: bool,
302+
filter_kind: str,
303+
value: str,
304+
):
305+
await filtered_collection.upsert([
306+
FilteredDataModel(id=1, tenant=value, embedding=[1.0, 0.0, 0.0]),
307+
FilteredDataModel(id=2, tenant="tenant_b", embedding=[1.0, 0.0, 0.0]),
308+
FilteredDataModel(id=3, tenant=value, embedding=[1.0, 0.0, 0.0]),
309+
])
310+
filter = f"lambda x: x.tenant == {value!r}"
311+
expected_ids = {1, 3}
312+
if filter_kind == "multiple":
313+
filter = [filter, "lambda x: x.id < 3"]
314+
expected_ids = {1}
315+
elif filter_kind == "membership":
316+
filter = f"lambda x: x.tenant in [{value!r}, 'missing'] and x.id < 3"
317+
expected_ids = {1}
318+
elif filter_kind == "dynamic":
319+
filter = default_dynamic_filter_function(
320+
filter="lambda x: x.id < 3",
321+
parameters=[KernelParameterMetadata(name="tenant")],
322+
tenant=value,
323+
)
324+
expected_ids = {1}
325+
326+
results = await filtered_collection.search(
327+
vector=[1.0, 0.0, 0.0], filter=filter, include_total_count=include_total_count, top=10
328+
)
329+
330+
assert {result.record.id async for result in results.results} == expected_ids
331+
assert results.total_count == (len(expected_ids) if include_total_count else None)
332+
333+
no_matches = await filtered_collection.search(
334+
vector=[1.0, 0.0, 0.0],
335+
filter="lambda x: x.tenant == 'missing'",
336+
include_total_count=include_total_count,
337+
)
338+
assert [result async for result in no_matches.results] == []
339+
assert no_matches.total_count == (0 if include_total_count else None)
340+
341+
342+
@pytest.mark.parametrize("include_total_count", [False, True])
343+
@pytest.mark.parametrize(
344+
"filter, expected_ids",
345+
[
346+
("lambda x: x.tenant == None", {1, 4}),
347+
("lambda x: x.tenant != None", {2, 3}),
348+
("lambda x: None == x.tenant", {1, 4}),
349+
("lambda x: None != x.tenant", {2, 3}),
350+
("lambda x: None == None", {1, 2, 3, 4}),
351+
("lambda x: None != None", set()),
352+
("lambda x: None != 'value'", {1, 2, 3, 4}),
353+
("lambda x: False == None", set()),
354+
("lambda x: not x.tenant == None", {2, 3}),
355+
("lambda x: not x.tenant != None", {1, 4}),
356+
("lambda x: x.tenant == None or x.tenant == 'value'", {1, 2, 4}),
357+
("lambda x: x.id > 1 and x.tenant == None", {4}),
358+
(["lambda x: x.id >= 3", "lambda x: x.tenant == None"], {4}),
359+
(["lambda x: x.id > 1", "lambda x: None != x.tenant"], {2, 3}),
360+
("lambda x: None == x.tenant == None", {1, 4}),
361+
("lambda x: x.tenant == ''", {3}),
362+
],
363+
)
364+
async def test_search_null_filters(
365+
filtered_collection: PostgresCollection[int, FilteredDataModel],
366+
include_total_count: bool,
367+
filter: str | list[str],
368+
expected_ids: set[int],
369+
):
370+
await filtered_collection.upsert([
371+
FilteredDataModel(id=1, tenant=None, embedding=[1.0, 0.0, 0.0]),
372+
FilteredDataModel(id=2, tenant="value", embedding=[1.0, 0.0, 0.0]),
373+
FilteredDataModel(id=3, tenant="", embedding=[1.0, 0.0, 0.0]),
374+
FilteredDataModel(id=4, tenant=None, embedding=[1.0, 0.0, 0.0]),
375+
])
376+
377+
results = await filtered_collection.search(
378+
vector=[1.0, 0.0, 0.0], filter=filter, include_total_count=include_total_count, top=10
379+
)
380+
381+
assert {result.record.id async for result in results.results} == expected_ids
382+
assert results.total_count == (len(expected_ids) if include_total_count else None)

0 commit comments

Comments
 (0)