diff --git a/src/llama_index_cloud_sql_pg/async_vector_store.py b/src/llama_index_cloud_sql_pg/async_vector_store.py index c2c75a2..64b0910 100644 --- a/src/llama_index_cloud_sql_pg/async_vector_store.py +++ b/src/llama_index_cloud_sql_pg/async_vector_store.py @@ -598,7 +598,7 @@ def __parse_metadata_filters_recursively( def __parse_metadata_filter( self, filter: MetadataFilter, bind_params: dict[str, Any] ) -> str: - key = self.__to_postgres_key(filter.key) + key = self.__to_postgres_key(filter.key, bind_params) op = self.__to_postgres_operator(filter.operator) param_name = f"param_{len(bind_params)}" if filter.operator == FilterOperator.IS_EMPTY: @@ -617,7 +617,7 @@ def __parse_metadata_filter( Value -> '{filter.value}'""" ) return "" - bind_params[param_name] = f'["{filter.value}"]' + bind_params[param_name] = json.dumps([filter.value]) return f"({key})::jsonb {op} :{param_name} " if filter.operator == FilterOperator.TEXT_MATCH: bind_params[param_name] = f"%{filter.value}%" @@ -698,12 +698,14 @@ def __to_postgres_operator(self, operator: FilterOperator) -> str: warnings.warn(f"Unknown operator: {operator}, fallback to '='") return "=" - def __to_postgres_key(self, key: str) -> str: + def __to_postgres_key(self, key: str, bind_params: dict[str, Any]) -> str: if key in [ *self._metadata_columns, self._id_column, self._ref_doc_id_column, self._text_column, ]: - return key - return f"{self._metadata_json_column}->>'{key}'" + return f'"{key}"' + key_param = f"key_{len(bind_params)}" + bind_params[key_param] = key + return f"{self._metadata_json_column}->>:{key_param}" diff --git a/tests/test_async_vector_store.py b/tests/test_async_vector_store.py index 9a07d51..b5bc845 100644 --- a/tests/test_async_vector_store.py +++ b/tests/test_async_vector_store.py @@ -467,3 +467,54 @@ async def test_delete_nodes(self, vs): async def test_clear(self, vs): with pytest.raises(Exception, match=sync_method_exception_str): vs.clear() + + async def test_parse_metadata_filters_parameterizes_key_and_value( + self, engine, custom_vs + ): + injected_key = "k' IS NULL OR 1=1 OR li_metadata->>'k" + injected_value = "x') OR 1=1;--" + + bind_params: dict[str, Any] = {} + filters = MetadataFilters( + filters=[ + MetadataFilter( + key=injected_key, + value=injected_value, + operator=FilterOperator.EQ, + ), + MetadataFilter( + key="len", + value=3, + operator=FilterOperator.EQ, + ), + MetadataFilter( + key="tags", + value='a"b', + operator=FilterOperator.CONTAINS, + ), + ], + condition=FilterCondition.AND, + ) + + clause = ( + custom_vs._AsyncPostgresVectorStore__parse_metadata_filters_recursively( + filters, bind_params + ) + ) + + assert injected_key not in clause + assert injected_value not in clause + assert ( + clause + == 'li_metadata->>:key_0 = :param_1 and "len"::float = :param_2 and (li_metadata->>:key_3)::jsonb @> :param_4 ' + ) + assert bind_params == { + "key_0": injected_key, + "param_1": injected_value, + "param_2": 3, + "key_3": "tags", + "param_4": '["a\\"b"]', + } + + results = await run_on_background(engine, custom_vs.aget_nodes(filters=filters)) + assert len(results) == 0