pgvector-haystack 5.4.0__tar.gz → 5.5.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/CHANGELOG.md +19 -0
- {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/PKG-INFO +1 -1
- {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/src/haystack_integrations/document_stores/pgvector/document_store.py +186 -0
- {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/tests/test_document_store.py +218 -0
- pgvector_haystack-5.5.0/tests/test_document_store_async.py +450 -0
- pgvector_haystack-5.4.0/tests/test_document_store_async.py +0 -232
- {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/.gitignore +0 -0
- {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/LICENSE.txt +0 -0
- {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/README.md +0 -0
- {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/examples/embedding_retrieval.py +0 -0
- {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/examples/hybrid_retrieval.py +0 -0
- {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/pydoc/config_docusaurus.yml +0 -0
- {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/pyproject.toml +0 -0
- {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/src/haystack_integrations/components/retrievers/pgvector/__init__.py +0 -0
- {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/src/haystack_integrations/components/retrievers/pgvector/embedding_retriever.py +0 -0
- {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/src/haystack_integrations/components/retrievers/pgvector/keyword_retriever.py +0 -0
- {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/src/haystack_integrations/components/retrievers/py.typed +0 -0
- {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/src/haystack_integrations/document_stores/pgvector/__init__.py +0 -0
- {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/src/haystack_integrations/document_stores/pgvector/converters.py +0 -0
- {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/src/haystack_integrations/document_stores/pgvector/filters.py +0 -0
- {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/src/haystack_integrations/document_stores/py.typed +0 -0
- {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/tests/__init__.py +0 -0
- {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/tests/conftest.py +0 -0
- {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/tests/test_converters.py +0 -0
- {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/tests/test_filters.py +0 -0
- {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/tests/test_retrieval.py +0 -0
- {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/tests/test_retrieval_async.py +0 -0
- {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/tests/test_retrievers.py +0 -0
|
@@ -1,5 +1,24 @@
|
|
|
1
1
|
# Changelog
|
|
2
2
|
|
|
3
|
+
## [integrations/pgvector-v5.4.0] - 2025-12-05
|
|
4
|
+
|
|
5
|
+
### 📚 Documentation
|
|
6
|
+
|
|
7
|
+
- Add pydoc configurations for Docusaurus (#2411)
|
|
8
|
+
|
|
9
|
+
### ⚙️ CI
|
|
10
|
+
|
|
11
|
+
- Change pytest command (#2475)
|
|
12
|
+
|
|
13
|
+
### 🧹 Chores
|
|
14
|
+
|
|
15
|
+
- Remove Readme API CI workflow and configs (#2573)
|
|
16
|
+
- Pgvector - document expected connection string and raise informative error (#2583)
|
|
17
|
+
|
|
18
|
+
### 🌀 Miscellaneous
|
|
19
|
+
|
|
20
|
+
- Enhancement: Adopt PEP 585 type hinting (part 5) (#2528)
|
|
21
|
+
|
|
3
22
|
## [integrations/pgvector-v5.3.0] - 2025-10-17
|
|
4
23
|
|
|
5
24
|
### 🚀 Features
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: pgvector-haystack
|
|
3
|
-
Version: 5.
|
|
3
|
+
Version: 5.5.0
|
|
4
4
|
Summary: An integration of pgvector (vector search extension for Postgres) with Haystack
|
|
5
5
|
Project-URL: Source, https://github.com/deepset-ai/haystack-core-integrations
|
|
6
6
|
Project-URL: Documentation, https://github.com/deepset-ai/haystack-core-integrations/blob/main/integrations/pgvector/README.md
|
|
@@ -1,6 +1,7 @@
|
|
|
1
1
|
# SPDX-FileCopyrightText: 2023-present deepset GmbH <info@deepset.ai>
|
|
2
2
|
#
|
|
3
3
|
# SPDX-License-Identifier: Apache-2.0
|
|
4
|
+
|
|
4
5
|
from typing import Any, Literal, Optional, Union, overload
|
|
5
6
|
|
|
6
7
|
from haystack import default_from_dict, default_to_dict, logging
|
|
@@ -13,6 +14,7 @@ from psycopg.cursor_async import AsyncCursor
|
|
|
13
14
|
from psycopg.rows import DictRow, dict_row
|
|
14
15
|
from psycopg.sql import SQL, Composed, Identifier
|
|
15
16
|
from psycopg.sql import Literal as SQLLiteral
|
|
17
|
+
from psycopg.types.json import Jsonb
|
|
16
18
|
|
|
17
19
|
from pgvector.psycopg import register_vector, register_vector_async
|
|
18
20
|
|
|
@@ -1036,6 +1038,190 @@ class PgvectorDocumentStore:
|
|
|
1036
1038
|
error_msg="Could not delete all documents from PgvectorDocumentStore",
|
|
1037
1039
|
)
|
|
1038
1040
|
|
|
1041
|
+
def delete_by_filter(self, filters: dict[str, Any]) -> int:
|
|
1042
|
+
"""
|
|
1043
|
+
Deletes all documents that match the provided filters.
|
|
1044
|
+
|
|
1045
|
+
:param filters: The filters to apply to select documents for deletion.
|
|
1046
|
+
For filter syntax, see [Haystack metadata filtering](https://docs.haystack.deepset.ai/docs/metadata-filtering)
|
|
1047
|
+
:returns: The number of documents deleted.
|
|
1048
|
+
"""
|
|
1049
|
+
_validate_filters(filters)
|
|
1050
|
+
|
|
1051
|
+
delete_sql = SQL("DELETE FROM {schema_name}.{table_name}").format(
|
|
1052
|
+
schema_name=Identifier(self.schema_name),
|
|
1053
|
+
table_name=Identifier(self.table_name),
|
|
1054
|
+
)
|
|
1055
|
+
|
|
1056
|
+
params = ()
|
|
1057
|
+
if filters:
|
|
1058
|
+
sql_where_clause, params = _convert_filters_to_where_clause_and_params(filters)
|
|
1059
|
+
delete_sql += sql_where_clause
|
|
1060
|
+
|
|
1061
|
+
self._ensure_db_setup()
|
|
1062
|
+
assert self._cursor is not None
|
|
1063
|
+
|
|
1064
|
+
try:
|
|
1065
|
+
self._execute_sql(
|
|
1066
|
+
cursor=self._cursor,
|
|
1067
|
+
sql_query=delete_sql,
|
|
1068
|
+
params=params,
|
|
1069
|
+
error_msg="Could not delete documents by filter from PgvectorDocumentStore",
|
|
1070
|
+
)
|
|
1071
|
+
deleted_count = self._cursor.rowcount
|
|
1072
|
+
logger.info(
|
|
1073
|
+
"Deleted {n_docs} documents from table '{schema}.{table}' using filters.",
|
|
1074
|
+
n_docs=deleted_count,
|
|
1075
|
+
schema=self.schema_name,
|
|
1076
|
+
table=self.table_name,
|
|
1077
|
+
)
|
|
1078
|
+
return deleted_count
|
|
1079
|
+
except Error as e:
|
|
1080
|
+
msg = f"Failed to delete documents by filter from PgvectorDocumentStore: {e!s}"
|
|
1081
|
+
raise DocumentStoreError(msg) from e
|
|
1082
|
+
|
|
1083
|
+
async def delete_by_filter_async(self, filters: dict[str, Any]) -> int:
|
|
1084
|
+
"""
|
|
1085
|
+
Asynchronously deletes all documents that match the provided filters.
|
|
1086
|
+
|
|
1087
|
+
:param filters: The filters to apply to select documents for deletion.
|
|
1088
|
+
For filter syntax, see [Haystack metadata filtering](https://docs.haystack.deepset.ai/docs/metadata-filtering)
|
|
1089
|
+
:returns: The number of documents deleted.
|
|
1090
|
+
"""
|
|
1091
|
+
_validate_filters(filters)
|
|
1092
|
+
|
|
1093
|
+
delete_sql = SQL("DELETE FROM {schema_name}.{table_name}").format(
|
|
1094
|
+
schema_name=Identifier(self.schema_name),
|
|
1095
|
+
table_name=Identifier(self.table_name),
|
|
1096
|
+
)
|
|
1097
|
+
|
|
1098
|
+
params = ()
|
|
1099
|
+
if filters:
|
|
1100
|
+
sql_where_clause, params = _convert_filters_to_where_clause_and_params(filters)
|
|
1101
|
+
delete_sql += sql_where_clause
|
|
1102
|
+
|
|
1103
|
+
await self._ensure_db_setup_async()
|
|
1104
|
+
assert self._async_cursor is not None
|
|
1105
|
+
|
|
1106
|
+
try:
|
|
1107
|
+
await self._execute_sql_async(
|
|
1108
|
+
cursor=self._async_cursor,
|
|
1109
|
+
sql_query=delete_sql,
|
|
1110
|
+
params=params,
|
|
1111
|
+
error_msg="Could not delete documents by filter from PgvectorDocumentStore",
|
|
1112
|
+
)
|
|
1113
|
+
deleted_count = self._async_cursor.rowcount
|
|
1114
|
+
logger.info(
|
|
1115
|
+
"Deleted {n_docs} documents from table '{schema}.{table}' using filters.",
|
|
1116
|
+
n_docs=deleted_count,
|
|
1117
|
+
schema=self.schema_name,
|
|
1118
|
+
table=self.table_name,
|
|
1119
|
+
)
|
|
1120
|
+
return deleted_count
|
|
1121
|
+
except Error as e:
|
|
1122
|
+
msg = f"Failed to delete documents by filter from PgvectorDocumentStore: {e!s}"
|
|
1123
|
+
raise DocumentStoreError(msg) from e
|
|
1124
|
+
|
|
1125
|
+
def update_by_filter(self, filters: dict[str, Any], meta: dict[str, Any]) -> int:
|
|
1126
|
+
"""
|
|
1127
|
+
Updates the metadata of all documents that match the provided filters.
|
|
1128
|
+
|
|
1129
|
+
:param filters: The filters to apply to select documents for updating.
|
|
1130
|
+
For filter syntax, see [Haystack metadata filtering](https://docs.haystack.deepset.ai/docs/metadata-filtering)
|
|
1131
|
+
:param meta: The metadata fields to update.
|
|
1132
|
+
:returns: The number of documents updated.
|
|
1133
|
+
"""
|
|
1134
|
+
_validate_filters(filters)
|
|
1135
|
+
|
|
1136
|
+
if not meta:
|
|
1137
|
+
msg = "meta must be a non-empty dictionary"
|
|
1138
|
+
raise ValueError(msg)
|
|
1139
|
+
|
|
1140
|
+
update_sql = SQL(
|
|
1141
|
+
"UPDATE {schema_name}.{table_name} SET meta = COALESCE(meta, '{{}}'::jsonb) || %s::jsonb"
|
|
1142
|
+
).format(
|
|
1143
|
+
schema_name=Identifier(self.schema_name),
|
|
1144
|
+
table_name=Identifier(self.table_name),
|
|
1145
|
+
)
|
|
1146
|
+
|
|
1147
|
+
params: tuple[Any, ...] = (Jsonb(meta),)
|
|
1148
|
+
if filters:
|
|
1149
|
+
sql_where_clause, where_params = _convert_filters_to_where_clause_and_params(filters)
|
|
1150
|
+
update_sql += sql_where_clause
|
|
1151
|
+
params = params + where_params
|
|
1152
|
+
|
|
1153
|
+
self._ensure_db_setup()
|
|
1154
|
+
assert self._cursor is not None
|
|
1155
|
+
|
|
1156
|
+
try:
|
|
1157
|
+
self._execute_sql(
|
|
1158
|
+
cursor=self._cursor,
|
|
1159
|
+
sql_query=update_sql,
|
|
1160
|
+
params=params,
|
|
1161
|
+
error_msg="Could not update documents by filter from PgvectorDocumentStore",
|
|
1162
|
+
)
|
|
1163
|
+
updated_count = self._cursor.rowcount
|
|
1164
|
+
logger.info(
|
|
1165
|
+
"Updated {n_docs} documents in table '{schema}.{table}' using filters.",
|
|
1166
|
+
n_docs=updated_count,
|
|
1167
|
+
schema=self.schema_name,
|
|
1168
|
+
table=self.table_name,
|
|
1169
|
+
)
|
|
1170
|
+
return updated_count
|
|
1171
|
+
except Error as e:
|
|
1172
|
+
msg = f"Failed to update documents by filter in PgvectorDocumentStore: {e!s}"
|
|
1173
|
+
raise DocumentStoreError(msg) from e
|
|
1174
|
+
|
|
1175
|
+
async def update_by_filter_async(self, filters: dict[str, Any], meta: dict[str, Any]) -> int:
|
|
1176
|
+
"""
|
|
1177
|
+
Asynchronously updates the metadata of all documents that match the provided filters.
|
|
1178
|
+
|
|
1179
|
+
:param filters: The filters to apply to select documents for updating.
|
|
1180
|
+
For filter syntax, see [Haystack metadata filtering](https://docs.haystack.deepset.ai/docs/metadata-filtering)
|
|
1181
|
+
:param meta: The metadata fields to update.
|
|
1182
|
+
:returns: The number of documents updated.
|
|
1183
|
+
"""
|
|
1184
|
+
_validate_filters(filters)
|
|
1185
|
+
|
|
1186
|
+
if not meta:
|
|
1187
|
+
msg = "meta must be a non-empty dictionary"
|
|
1188
|
+
raise ValueError(msg)
|
|
1189
|
+
|
|
1190
|
+
update_sql = SQL(
|
|
1191
|
+
"UPDATE {schema_name}.{table_name} SET meta = COALESCE(meta, '{{}}'::jsonb) || %s::jsonb"
|
|
1192
|
+
).format(
|
|
1193
|
+
schema_name=Identifier(self.schema_name),
|
|
1194
|
+
table_name=Identifier(self.table_name),
|
|
1195
|
+
)
|
|
1196
|
+
|
|
1197
|
+
params: tuple[Any, ...] = (Jsonb(meta),)
|
|
1198
|
+
if filters:
|
|
1199
|
+
sql_where_clause, where_params = _convert_filters_to_where_clause_and_params(filters)
|
|
1200
|
+
update_sql += sql_where_clause
|
|
1201
|
+
params = params + where_params
|
|
1202
|
+
|
|
1203
|
+
await self._ensure_db_setup_async()
|
|
1204
|
+
assert self._async_cursor is not None
|
|
1205
|
+
|
|
1206
|
+
try:
|
|
1207
|
+
await self._execute_sql_async(
|
|
1208
|
+
cursor=self._async_cursor,
|
|
1209
|
+
sql_query=update_sql,
|
|
1210
|
+
params=params,
|
|
1211
|
+
error_msg="Could not update documents by filter from PgvectorDocumentStore",
|
|
1212
|
+
)
|
|
1213
|
+
updated_count = self._async_cursor.rowcount
|
|
1214
|
+
logger.info(
|
|
1215
|
+
"Updated {n_docs} documents in table '{schema}.{table}' using filters.",
|
|
1216
|
+
n_docs=updated_count,
|
|
1217
|
+
schema=self.schema_name,
|
|
1218
|
+
table=self.table_name,
|
|
1219
|
+
)
|
|
1220
|
+
return updated_count
|
|
1221
|
+
except Error as e:
|
|
1222
|
+
msg = f"Failed to update documents by filter in PgvectorDocumentStore: {e!s}"
|
|
1223
|
+
raise DocumentStoreError(msg) from e
|
|
1224
|
+
|
|
1039
1225
|
def _build_keyword_retrieval_query(
|
|
1040
1226
|
self, query: str, top_k: int, filters: Optional[dict[str, Any]] = None
|
|
1041
1227
|
) -> tuple[Composed, tuple]:
|
|
@@ -260,3 +260,221 @@ def test_delete_table_first_call(document_store):
|
|
|
260
260
|
without triggering errors due to an uninitialized state.
|
|
261
261
|
"""
|
|
262
262
|
document_store.delete_table() # if throw error, test fails
|
|
263
|
+
|
|
264
|
+
|
|
265
|
+
@pytest.mark.integration
|
|
266
|
+
def test_delete_by_filter(document_store: PgvectorDocumentStore):
|
|
267
|
+
docs = [
|
|
268
|
+
Document(content="Doc 1", meta={"category": "A", "year": 2023}),
|
|
269
|
+
Document(content="Doc 2", meta={"category": "B", "year": 2023}),
|
|
270
|
+
Document(content="Doc 3", meta={"category": "A", "year": 2024}),
|
|
271
|
+
]
|
|
272
|
+
document_store.write_documents(docs)
|
|
273
|
+
assert document_store.count_documents() == 3
|
|
274
|
+
|
|
275
|
+
deleted_count = document_store.delete_by_filter(filters={"field": "meta.category", "operator": "==", "value": "A"})
|
|
276
|
+
assert deleted_count == 2
|
|
277
|
+
assert document_store.count_documents() == 1
|
|
278
|
+
|
|
279
|
+
remaining_docs = document_store.filter_documents()
|
|
280
|
+
assert len(remaining_docs) == 1
|
|
281
|
+
assert remaining_docs[0].meta["category"] == "B"
|
|
282
|
+
|
|
283
|
+
deleted_count = document_store.delete_by_filter(filters={"field": "meta.year", "operator": "==", "value": 2023})
|
|
284
|
+
assert deleted_count == 1
|
|
285
|
+
assert document_store.count_documents() == 0
|
|
286
|
+
|
|
287
|
+
|
|
288
|
+
@pytest.mark.integration
|
|
289
|
+
def test_delete_by_filter_no_matches(document_store: PgvectorDocumentStore):
|
|
290
|
+
docs = [
|
|
291
|
+
Document(content="Doc 1", meta={"category": "A"}),
|
|
292
|
+
Document(content="Doc 2", meta={"category": "B"}),
|
|
293
|
+
]
|
|
294
|
+
document_store.write_documents(docs)
|
|
295
|
+
assert document_store.count_documents() == 2
|
|
296
|
+
|
|
297
|
+
deleted_count = document_store.delete_by_filter(filters={"field": "meta.category", "operator": "==", "value": "C"})
|
|
298
|
+
assert deleted_count == 0
|
|
299
|
+
assert document_store.count_documents() == 2
|
|
300
|
+
|
|
301
|
+
|
|
302
|
+
@pytest.mark.integration
|
|
303
|
+
def test_delete_by_filter_advanced_filters(document_store: PgvectorDocumentStore):
|
|
304
|
+
docs = [
|
|
305
|
+
Document(content="Doc 1", meta={"category": "A", "year": 2023, "status": "draft"}),
|
|
306
|
+
Document(content="Doc 2", meta={"category": "A", "year": 2024, "status": "published"}),
|
|
307
|
+
Document(content="Doc 3", meta={"category": "B", "year": 2023, "status": "draft"}),
|
|
308
|
+
]
|
|
309
|
+
document_store.write_documents(docs)
|
|
310
|
+
assert document_store.count_documents() == 3
|
|
311
|
+
|
|
312
|
+
# AND condition
|
|
313
|
+
deleted_count = document_store.delete_by_filter(
|
|
314
|
+
filters={
|
|
315
|
+
"operator": "AND",
|
|
316
|
+
"conditions": [
|
|
317
|
+
{"field": "meta.category", "operator": "==", "value": "A"},
|
|
318
|
+
{"field": "meta.year", "operator": "==", "value": 2023},
|
|
319
|
+
],
|
|
320
|
+
}
|
|
321
|
+
)
|
|
322
|
+
assert deleted_count == 1
|
|
323
|
+
assert document_store.count_documents() == 2
|
|
324
|
+
|
|
325
|
+
# OR condition
|
|
326
|
+
deleted_count = document_store.delete_by_filter(
|
|
327
|
+
filters={
|
|
328
|
+
"operator": "OR",
|
|
329
|
+
"conditions": [
|
|
330
|
+
{"field": "meta.category", "operator": "==", "value": "B"},
|
|
331
|
+
{"field": "meta.status", "operator": "==", "value": "published"},
|
|
332
|
+
],
|
|
333
|
+
}
|
|
334
|
+
)
|
|
335
|
+
assert deleted_count == 2
|
|
336
|
+
assert document_store.count_documents() == 0
|
|
337
|
+
|
|
338
|
+
|
|
339
|
+
@pytest.mark.integration
|
|
340
|
+
def test_update_by_filter(document_store: PgvectorDocumentStore):
|
|
341
|
+
docs = [
|
|
342
|
+
Document(content="Doc 1", meta={"category": "A", "status": "draft"}),
|
|
343
|
+
Document(content="Doc 2", meta={"category": "B", "status": "draft"}),
|
|
344
|
+
Document(content="Doc 3", meta={"category": "A", "status": "draft"}),
|
|
345
|
+
]
|
|
346
|
+
document_store.write_documents(docs)
|
|
347
|
+
assert document_store.count_documents() == 3
|
|
348
|
+
|
|
349
|
+
# update status for category="A" documents
|
|
350
|
+
updated_count = document_store.update_by_filter(
|
|
351
|
+
filters={"field": "meta.category", "operator": "==", "value": "A"}, meta={"status": "published"}
|
|
352
|
+
)
|
|
353
|
+
assert updated_count == 2
|
|
354
|
+
|
|
355
|
+
# verify
|
|
356
|
+
published_docs = document_store.filter_documents(
|
|
357
|
+
filters={"field": "meta.status", "operator": "==", "value": "published"}
|
|
358
|
+
)
|
|
359
|
+
assert len(published_docs) == 2
|
|
360
|
+
for doc in published_docs:
|
|
361
|
+
assert doc.meta["category"] == "A"
|
|
362
|
+
assert doc.meta["status"] == "published"
|
|
363
|
+
|
|
364
|
+
# Verify category B still has draft status
|
|
365
|
+
draft_docs = document_store.filter_documents(filters={"field": "meta.status", "operator": "==", "value": "draft"})
|
|
366
|
+
assert len(draft_docs) == 1
|
|
367
|
+
assert draft_docs[0].meta["category"] == "B"
|
|
368
|
+
|
|
369
|
+
|
|
370
|
+
@pytest.mark.integration
|
|
371
|
+
def test_update_by_filter_multiple_fields(document_store: PgvectorDocumentStore):
|
|
372
|
+
docs = [
|
|
373
|
+
Document(content="Doc 1", meta={"category": "A", "year": 2023}),
|
|
374
|
+
Document(content="Doc 2", meta={"category": "A", "year": 2023}),
|
|
375
|
+
Document(content="Doc 3", meta={"category": "B", "year": 2024}),
|
|
376
|
+
]
|
|
377
|
+
document_store.write_documents(docs)
|
|
378
|
+
assert document_store.count_documents() == 3
|
|
379
|
+
|
|
380
|
+
# update multiple fields for category="A" documents
|
|
381
|
+
updated_count = document_store.update_by_filter(
|
|
382
|
+
filters={"field": "meta.category", "operator": "==", "value": "A"},
|
|
383
|
+
meta={"status": "published", "priority": "high", "reviewed": True},
|
|
384
|
+
)
|
|
385
|
+
assert updated_count == 2
|
|
386
|
+
|
|
387
|
+
# verify
|
|
388
|
+
published_docs = document_store.filter_documents(filters={"field": "meta.category", "operator": "==", "value": "A"})
|
|
389
|
+
assert len(published_docs) == 2
|
|
390
|
+
for doc in published_docs:
|
|
391
|
+
assert doc.meta["status"] == "published"
|
|
392
|
+
assert doc.meta["priority"] == "high"
|
|
393
|
+
assert doc.meta["reviewed"] is True
|
|
394
|
+
assert doc.meta["year"] == 2023 # Original field should still be present
|
|
395
|
+
|
|
396
|
+
# verify category B was not updated
|
|
397
|
+
b_docs = document_store.filter_documents(filters={"field": "meta.category", "operator": "==", "value": "B"})
|
|
398
|
+
assert len(b_docs) == 1
|
|
399
|
+
assert "status" not in b_docs[0].meta
|
|
400
|
+
assert "priority" not in b_docs[0].meta
|
|
401
|
+
|
|
402
|
+
|
|
403
|
+
@pytest.mark.integration
|
|
404
|
+
def test_update_by_filter_no_matches(document_store: PgvectorDocumentStore):
|
|
405
|
+
docs = [
|
|
406
|
+
Document(content="Doc 1", meta={"category": "A"}),
|
|
407
|
+
Document(content="Doc 2", meta={"category": "B"}),
|
|
408
|
+
]
|
|
409
|
+
document_store.write_documents(docs)
|
|
410
|
+
assert document_store.count_documents() == 2
|
|
411
|
+
|
|
412
|
+
# update documents with category="C" (no matches)
|
|
413
|
+
updated_count = document_store.update_by_filter(
|
|
414
|
+
filters={"field": "meta.category", "operator": "==", "value": "C"}, meta={"status": "published"}
|
|
415
|
+
)
|
|
416
|
+
assert updated_count == 0
|
|
417
|
+
assert document_store.count_documents() == 2
|
|
418
|
+
|
|
419
|
+
# verify no documents were updated
|
|
420
|
+
published_docs = document_store.filter_documents(
|
|
421
|
+
filters={"field": "meta.status", "operator": "==", "value": "published"}
|
|
422
|
+
)
|
|
423
|
+
assert len(published_docs) == 0
|
|
424
|
+
|
|
425
|
+
|
|
426
|
+
@pytest.mark.integration
|
|
427
|
+
def test_update_by_filter_advanced_filters(document_store: PgvectorDocumentStore):
|
|
428
|
+
docs = [
|
|
429
|
+
Document(content="Doc 1", meta={"category": "A", "year": 2023, "status": "draft"}),
|
|
430
|
+
Document(content="Doc 2", meta={"category": "A", "year": 2024, "status": "draft"}),
|
|
431
|
+
Document(content="Doc 3", meta={"category": "B", "year": 2023, "status": "draft"}),
|
|
432
|
+
]
|
|
433
|
+
document_store.write_documents(docs)
|
|
434
|
+
assert document_store.count_documents() == 3
|
|
435
|
+
|
|
436
|
+
# AND condition
|
|
437
|
+
updated_count = document_store.update_by_filter(
|
|
438
|
+
filters={
|
|
439
|
+
"operator": "AND",
|
|
440
|
+
"conditions": [
|
|
441
|
+
{"field": "meta.category", "operator": "==", "value": "A"},
|
|
442
|
+
{"field": "meta.year", "operator": "==", "value": 2023},
|
|
443
|
+
],
|
|
444
|
+
},
|
|
445
|
+
meta={"status": "published"},
|
|
446
|
+
)
|
|
447
|
+
assert updated_count == 1
|
|
448
|
+
|
|
449
|
+
# verify only one document was updated
|
|
450
|
+
published_docs = document_store.filter_documents(
|
|
451
|
+
filters={"field": "meta.status", "operator": "==", "value": "published"}
|
|
452
|
+
)
|
|
453
|
+
assert len(published_docs) == 1
|
|
454
|
+
assert published_docs[0].meta["category"] == "A"
|
|
455
|
+
assert published_docs[0].meta["year"] == 2023
|
|
456
|
+
|
|
457
|
+
# OR condition
|
|
458
|
+
updated_count = document_store.update_by_filter(
|
|
459
|
+
filters={
|
|
460
|
+
"operator": "OR",
|
|
461
|
+
"conditions": [
|
|
462
|
+
{"field": "meta.category", "operator": "==", "value": "B"},
|
|
463
|
+
{"field": "meta.year", "operator": "==", "value": 2024},
|
|
464
|
+
],
|
|
465
|
+
},
|
|
466
|
+
meta={"featured": True},
|
|
467
|
+
)
|
|
468
|
+
assert updated_count == 2
|
|
469
|
+
|
|
470
|
+
featured_docs = document_store.filter_documents(filters={"field": "meta.featured", "operator": "==", "value": True})
|
|
471
|
+
assert len(featured_docs) == 2
|
|
472
|
+
|
|
473
|
+
|
|
474
|
+
@pytest.mark.integration
|
|
475
|
+
def test_update_by_filter_empty_meta_raises_error(document_store: PgvectorDocumentStore):
|
|
476
|
+
docs = [Document(content="Doc 1", meta={"category": "A"})]
|
|
477
|
+
document_store.write_documents(docs)
|
|
478
|
+
|
|
479
|
+
with pytest.raises(ValueError, match="meta must be a non-empty dictionary"):
|
|
480
|
+
document_store.update_by_filter(filters={"field": "meta.category", "operator": "==", "value": "A"}, meta={})
|
|
@@ -0,0 +1,450 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: 2023-present deepset GmbH <info@deepset.ai>
|
|
2
|
+
#
|
|
3
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
4
|
+
|
|
5
|
+
from unittest.mock import patch
|
|
6
|
+
|
|
7
|
+
import psycopg
|
|
8
|
+
import pytest
|
|
9
|
+
from haystack.dataclasses.document import ByteStream, Document
|
|
10
|
+
from haystack.document_stores.errors import DocumentStoreError, DuplicateDocumentError
|
|
11
|
+
from haystack.document_stores.types import DuplicatePolicy
|
|
12
|
+
from haystack.utils import Secret
|
|
13
|
+
from psycopg.sql import SQL
|
|
14
|
+
|
|
15
|
+
from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
@pytest.mark.integration
|
|
19
|
+
@pytest.mark.asyncio
|
|
20
|
+
class TestDocumentStoreAsync:
|
|
21
|
+
async def test_write_documents(self, document_store: PgvectorDocumentStore):
|
|
22
|
+
docs = [Document(id="1")]
|
|
23
|
+
assert await document_store.write_documents_async(docs) == 1
|
|
24
|
+
with pytest.raises(DuplicateDocumentError):
|
|
25
|
+
await document_store.write_documents_async(docs, DuplicatePolicy.FAIL)
|
|
26
|
+
|
|
27
|
+
async def test_write_blob(self, document_store: PgvectorDocumentStore):
|
|
28
|
+
bytestream = ByteStream(b"test", meta={"meta_key": "meta_value"}, mime_type="mime_type")
|
|
29
|
+
docs = [Document(id="1", blob=bytestream)]
|
|
30
|
+
await document_store.write_documents_async(docs)
|
|
31
|
+
|
|
32
|
+
retrieved_docs = await document_store.filter_documents_async()
|
|
33
|
+
assert retrieved_docs == docs
|
|
34
|
+
|
|
35
|
+
async def test_count_documents(self, document_store: PgvectorDocumentStore):
|
|
36
|
+
await document_store.write_documents_async(
|
|
37
|
+
[
|
|
38
|
+
Document(content="test doc 1"),
|
|
39
|
+
Document(content="test doc 2"),
|
|
40
|
+
Document(content="test doc 3"),
|
|
41
|
+
]
|
|
42
|
+
)
|
|
43
|
+
assert await document_store.count_documents_async() == 3
|
|
44
|
+
|
|
45
|
+
async def test_filter_documents(self, document_store: PgvectorDocumentStore):
|
|
46
|
+
filterable_docs = [
|
|
47
|
+
Document(
|
|
48
|
+
content="1",
|
|
49
|
+
meta={
|
|
50
|
+
"number": -10,
|
|
51
|
+
},
|
|
52
|
+
),
|
|
53
|
+
Document(
|
|
54
|
+
content="2",
|
|
55
|
+
meta={
|
|
56
|
+
"number": 100,
|
|
57
|
+
},
|
|
58
|
+
),
|
|
59
|
+
]
|
|
60
|
+
await document_store.write_documents_async(filterable_docs)
|
|
61
|
+
result = await document_store.filter_documents_async(
|
|
62
|
+
filters={"field": "meta.number", "operator": "==", "value": 100}
|
|
63
|
+
)
|
|
64
|
+
|
|
65
|
+
assert result == [d for d in filterable_docs if d.meta.get("number") == 100]
|
|
66
|
+
|
|
67
|
+
async def test_delete_documents(self, document_store: PgvectorDocumentStore):
|
|
68
|
+
doc = Document(content="test doc")
|
|
69
|
+
await document_store.write_documents_async([doc])
|
|
70
|
+
assert await document_store.count_documents_async() == 1
|
|
71
|
+
|
|
72
|
+
await document_store.delete_documents_async([doc.id])
|
|
73
|
+
assert await document_store.count_documents_async() == 0
|
|
74
|
+
|
|
75
|
+
async def test_connection_check_and_recreation(self, document_store: PgvectorDocumentStore):
|
|
76
|
+
await document_store._ensure_db_setup_async()
|
|
77
|
+
original_connection = document_store._async_connection
|
|
78
|
+
|
|
79
|
+
with patch.object(PgvectorDocumentStore, "_connection_is_valid_async", return_value=False):
|
|
80
|
+
await document_store._ensure_db_setup_async()
|
|
81
|
+
new_connection = document_store._async_connection
|
|
82
|
+
|
|
83
|
+
# verify that a new connection is created
|
|
84
|
+
assert new_connection is not original_connection
|
|
85
|
+
assert document_store._async_connection == new_connection
|
|
86
|
+
assert original_connection.closed
|
|
87
|
+
|
|
88
|
+
assert document_store._async_cursor is not None
|
|
89
|
+
assert document_store._async_dict_cursor is not None
|
|
90
|
+
|
|
91
|
+
# test with new connection
|
|
92
|
+
with patch.object(PgvectorDocumentStore, "_connection_is_valid_async", return_value=True):
|
|
93
|
+
await document_store._ensure_db_setup_async()
|
|
94
|
+
same_connection = document_store._async_connection
|
|
95
|
+
assert same_connection is document_store._async_connection
|
|
96
|
+
|
|
97
|
+
async def test_delete_all_documents_async(self, document_store: PgvectorDocumentStore) -> None:
|
|
98
|
+
document_store.write_documents([Document(id=str(i)) for i in range(10)])
|
|
99
|
+
await document_store.delete_all_documents_async()
|
|
100
|
+
assert document_store.count_documents() == 0
|
|
101
|
+
document_store.write_documents([Document(id="1")])
|
|
102
|
+
assert document_store.count_documents() == 1
|
|
103
|
+
|
|
104
|
+
async def test_invalid_connection_string(self, monkeypatch):
|
|
105
|
+
monkeypatch.setenv("PG_CONN_STR", "invalid_connection_string")
|
|
106
|
+
document_store = PgvectorDocumentStore()
|
|
107
|
+
with pytest.raises(DocumentStoreError) as e:
|
|
108
|
+
await document_store._ensure_db_setup_async()
|
|
109
|
+
assert "Failed to connect to PostgreSQL database" in str(e)
|
|
110
|
+
|
|
111
|
+
async def test_delete_by_filter_async(self, document_store: PgvectorDocumentStore):
|
|
112
|
+
docs = [
|
|
113
|
+
Document(content="Doc 1", meta={"category": "A", "year": 2023}),
|
|
114
|
+
Document(content="Doc 2", meta={"category": "B", "year": 2023}),
|
|
115
|
+
Document(content="Doc 3", meta={"category": "A", "year": 2024}),
|
|
116
|
+
]
|
|
117
|
+
await document_store.write_documents_async(docs)
|
|
118
|
+
assert await document_store.count_documents_async() == 3
|
|
119
|
+
|
|
120
|
+
deleted_count = await document_store.delete_by_filter_async(
|
|
121
|
+
filters={"field": "meta.category", "operator": "==", "value": "A"}
|
|
122
|
+
)
|
|
123
|
+
assert deleted_count == 2
|
|
124
|
+
assert await document_store.count_documents_async() == 1
|
|
125
|
+
|
|
126
|
+
remaining_docs = await document_store.filter_documents_async()
|
|
127
|
+
assert len(remaining_docs) == 1
|
|
128
|
+
assert remaining_docs[0].meta["category"] == "B"
|
|
129
|
+
|
|
130
|
+
deleted_count = await document_store.delete_by_filter_async(
|
|
131
|
+
filters={"field": "meta.year", "operator": "==", "value": 2023}
|
|
132
|
+
)
|
|
133
|
+
assert deleted_count == 1
|
|
134
|
+
assert await document_store.count_documents_async() == 0
|
|
135
|
+
|
|
136
|
+
async def test_delete_by_filter_async_no_matches(self, document_store: PgvectorDocumentStore):
|
|
137
|
+
docs = [
|
|
138
|
+
Document(content="Doc 1", meta={"category": "A"}),
|
|
139
|
+
Document(content="Doc 2", meta={"category": "B"}),
|
|
140
|
+
]
|
|
141
|
+
await document_store.write_documents_async(docs)
|
|
142
|
+
assert await document_store.count_documents_async() == 2
|
|
143
|
+
|
|
144
|
+
deleted_count = await document_store.delete_by_filter_async(
|
|
145
|
+
filters={"field": "meta.category", "operator": "==", "value": "C"}
|
|
146
|
+
)
|
|
147
|
+
assert deleted_count == 0
|
|
148
|
+
assert await document_store.count_documents_async() == 2
|
|
149
|
+
|
|
150
|
+
async def test_delete_by_filter_async_advanced_filters(self, document_store: PgvectorDocumentStore):
|
|
151
|
+
docs = [
|
|
152
|
+
Document(content="Doc 1", meta={"category": "A", "year": 2023, "status": "draft"}),
|
|
153
|
+
Document(content="Doc 2", meta={"category": "A", "year": 2024, "status": "published"}),
|
|
154
|
+
Document(content="Doc 3", meta={"category": "B", "year": 2023, "status": "draft"}),
|
|
155
|
+
]
|
|
156
|
+
await document_store.write_documents_async(docs)
|
|
157
|
+
assert await document_store.count_documents_async() == 3
|
|
158
|
+
|
|
159
|
+
# AND condition
|
|
160
|
+
deleted_count = await document_store.delete_by_filter_async(
|
|
161
|
+
filters={
|
|
162
|
+
"operator": "AND",
|
|
163
|
+
"conditions": [
|
|
164
|
+
{"field": "meta.category", "operator": "==", "value": "A"},
|
|
165
|
+
{"field": "meta.year", "operator": "==", "value": 2023},
|
|
166
|
+
],
|
|
167
|
+
}
|
|
168
|
+
)
|
|
169
|
+
assert deleted_count == 1
|
|
170
|
+
assert await document_store.count_documents_async() == 2
|
|
171
|
+
|
|
172
|
+
# OR condition
|
|
173
|
+
deleted_count = await document_store.delete_by_filter_async(
|
|
174
|
+
filters={
|
|
175
|
+
"operator": "OR",
|
|
176
|
+
"conditions": [
|
|
177
|
+
{"field": "meta.category", "operator": "==", "value": "B"},
|
|
178
|
+
{"field": "meta.status", "operator": "==", "value": "published"},
|
|
179
|
+
],
|
|
180
|
+
}
|
|
181
|
+
)
|
|
182
|
+
assert deleted_count == 2
|
|
183
|
+
assert await document_store.count_documents_async() == 0
|
|
184
|
+
|
|
185
|
+
async def test_update_by_filter_async(self, document_store: PgvectorDocumentStore):
|
|
186
|
+
docs = [
|
|
187
|
+
Document(content="Doc 1", meta={"category": "A", "status": "draft"}),
|
|
188
|
+
Document(content="Doc 2", meta={"category": "B", "status": "draft"}),
|
|
189
|
+
Document(content="Doc 3", meta={"category": "A", "status": "draft"}),
|
|
190
|
+
]
|
|
191
|
+
await document_store.write_documents_async(docs)
|
|
192
|
+
assert await document_store.count_documents_async() == 3
|
|
193
|
+
|
|
194
|
+
# Update status for category="A" documents
|
|
195
|
+
updated_count = await document_store.update_by_filter_async(
|
|
196
|
+
filters={"field": "meta.category", "operator": "==", "value": "A"}, meta={"status": "published"}
|
|
197
|
+
)
|
|
198
|
+
assert updated_count == 2
|
|
199
|
+
|
|
200
|
+
# Verify the updates
|
|
201
|
+
published_docs = await document_store.filter_documents_async(
|
|
202
|
+
filters={"field": "meta.status", "operator": "==", "value": "published"}
|
|
203
|
+
)
|
|
204
|
+
assert len(published_docs) == 2
|
|
205
|
+
for doc in published_docs:
|
|
206
|
+
assert doc.meta["category"] == "A"
|
|
207
|
+
assert doc.meta["status"] == "published"
|
|
208
|
+
|
|
209
|
+
# Verify category B still has draft status
|
|
210
|
+
draft_docs = await document_store.filter_documents_async(
|
|
211
|
+
filters={"field": "meta.status", "operator": "==", "value": "draft"}
|
|
212
|
+
)
|
|
213
|
+
assert len(draft_docs) == 1
|
|
214
|
+
assert draft_docs[0].meta["category"] == "B"
|
|
215
|
+
|
|
216
|
+
async def test_update_by_filter_async_multiple_fields(self, document_store: PgvectorDocumentStore):
|
|
217
|
+
docs = [
|
|
218
|
+
Document(content="Doc 1", meta={"category": "A", "year": 2023}),
|
|
219
|
+
Document(content="Doc 2", meta={"category": "A", "year": 2023}),
|
|
220
|
+
Document(content="Doc 3", meta={"category": "B", "year": 2024}),
|
|
221
|
+
]
|
|
222
|
+
await document_store.write_documents_async(docs)
|
|
223
|
+
assert await document_store.count_documents_async() == 3
|
|
224
|
+
|
|
225
|
+
# update multiple fields for category="A" documents
|
|
226
|
+
updated_count = await document_store.update_by_filter_async(
|
|
227
|
+
filters={"field": "meta.category", "operator": "==", "value": "A"},
|
|
228
|
+
meta={"status": "published", "priority": "high", "reviewed": True},
|
|
229
|
+
)
|
|
230
|
+
assert updated_count == 2
|
|
231
|
+
|
|
232
|
+
# verify
|
|
233
|
+
published_docs = await document_store.filter_documents_async(
|
|
234
|
+
filters={"field": "meta.category", "operator": "==", "value": "A"}
|
|
235
|
+
)
|
|
236
|
+
assert len(published_docs) == 2
|
|
237
|
+
for doc in published_docs:
|
|
238
|
+
assert doc.meta["status"] == "published"
|
|
239
|
+
assert doc.meta["priority"] == "high"
|
|
240
|
+
assert doc.meta["reviewed"] is True
|
|
241
|
+
assert doc.meta["year"] == 2023 # Original field should still be present
|
|
242
|
+
|
|
243
|
+
# verify category B was not updated
|
|
244
|
+
b_docs = await document_store.filter_documents_async(
|
|
245
|
+
filters={"field": "meta.category", "operator": "==", "value": "B"}
|
|
246
|
+
)
|
|
247
|
+
assert len(b_docs) == 1
|
|
248
|
+
assert "status" not in b_docs[0].meta
|
|
249
|
+
assert "priority" not in b_docs[0].meta
|
|
250
|
+
|
|
251
|
+
async def test_update_by_filter_async_no_matches(self, document_store: PgvectorDocumentStore):
|
|
252
|
+
docs = [
|
|
253
|
+
Document(content="Doc 1", meta={"category": "A"}),
|
|
254
|
+
Document(content="Doc 2", meta={"category": "B"}),
|
|
255
|
+
]
|
|
256
|
+
await document_store.write_documents_async(docs)
|
|
257
|
+
assert await document_store.count_documents_async() == 2
|
|
258
|
+
|
|
259
|
+
# update documents with category="C" (no matches)
|
|
260
|
+
updated_count = await document_store.update_by_filter_async(
|
|
261
|
+
filters={"field": "meta.category", "operator": "==", "value": "C"}, meta={"status": "published"}
|
|
262
|
+
)
|
|
263
|
+
assert updated_count == 0
|
|
264
|
+
assert await document_store.count_documents_async() == 2
|
|
265
|
+
|
|
266
|
+
# verify no documents were updated
|
|
267
|
+
published_docs = await document_store.filter_documents_async(
|
|
268
|
+
filters={"field": "meta.status", "operator": "==", "value": "published"}
|
|
269
|
+
)
|
|
270
|
+
assert len(published_docs) == 0
|
|
271
|
+
|
|
272
|
+
async def test_update_by_filter_async_advanced_filters(self, document_store: PgvectorDocumentStore):
|
|
273
|
+
docs = [
|
|
274
|
+
Document(content="Doc 1", meta={"category": "A", "year": 2023, "status": "draft"}),
|
|
275
|
+
Document(content="Doc 2", meta={"category": "A", "year": 2024, "status": "draft"}),
|
|
276
|
+
Document(content="Doc 3", meta={"category": "B", "year": 2023, "status": "draft"}),
|
|
277
|
+
]
|
|
278
|
+
await document_store.write_documents_async(docs)
|
|
279
|
+
assert await document_store.count_documents_async() == 3
|
|
280
|
+
|
|
281
|
+
# AND condition
|
|
282
|
+
updated_count = await document_store.update_by_filter_async(
|
|
283
|
+
filters={
|
|
284
|
+
"operator": "AND",
|
|
285
|
+
"conditions": [
|
|
286
|
+
{"field": "meta.category", "operator": "==", "value": "A"},
|
|
287
|
+
{"field": "meta.year", "operator": "==", "value": 2023},
|
|
288
|
+
],
|
|
289
|
+
},
|
|
290
|
+
meta={"status": "published"},
|
|
291
|
+
)
|
|
292
|
+
assert updated_count == 1
|
|
293
|
+
|
|
294
|
+
# verify
|
|
295
|
+
published_docs = await document_store.filter_documents_async(
|
|
296
|
+
filters={"field": "meta.status", "operator": "==", "value": "published"}
|
|
297
|
+
)
|
|
298
|
+
assert len(published_docs) == 1
|
|
299
|
+
assert published_docs[0].meta["category"] == "A"
|
|
300
|
+
assert published_docs[0].meta["year"] == 2023
|
|
301
|
+
|
|
302
|
+
# OR condition
|
|
303
|
+
updated_count = await document_store.update_by_filter_async(
|
|
304
|
+
filters={
|
|
305
|
+
"operator": "OR",
|
|
306
|
+
"conditions": [
|
|
307
|
+
{"field": "meta.category", "operator": "==", "value": "B"},
|
|
308
|
+
{"field": "meta.year", "operator": "==", "value": 2024},
|
|
309
|
+
],
|
|
310
|
+
},
|
|
311
|
+
meta={"featured": True},
|
|
312
|
+
)
|
|
313
|
+
assert updated_count == 2
|
|
314
|
+
|
|
315
|
+
featured_docs = await document_store.filter_documents_async(
|
|
316
|
+
filters={"field": "meta.featured", "operator": "==", "value": True}
|
|
317
|
+
)
|
|
318
|
+
assert len(featured_docs) == 2
|
|
319
|
+
|
|
320
|
+
async def test_update_by_filter_async_empty_meta_raises_error(self, document_store: PgvectorDocumentStore):
|
|
321
|
+
docs = [Document(content="Doc 1", meta={"category": "A"})]
|
|
322
|
+
await document_store.write_documents_async(docs)
|
|
323
|
+
|
|
324
|
+
with pytest.raises(ValueError, match="meta must be a non-empty dictionary"):
|
|
325
|
+
await document_store.update_by_filter_async(
|
|
326
|
+
filters={"field": "meta.category", "operator": "==", "value": "A"}, meta={}
|
|
327
|
+
)
|
|
328
|
+
|
|
329
|
+
|
|
330
|
+
@pytest.mark.integration
|
|
331
|
+
@pytest.mark.asyncio
|
|
332
|
+
async def test_hnsw_index_recreation():
|
|
333
|
+
async def get_index_oid(cursor, document_store, schema_name, index_name):
|
|
334
|
+
sql_get_index_oid = SQL("""
|
|
335
|
+
SELECT c.oid
|
|
336
|
+
FROM pg_class c
|
|
337
|
+
JOIN pg_namespace n ON n.oid = c.relnamespace
|
|
338
|
+
WHERE c.relkind = 'i'
|
|
339
|
+
AND n.nspname = %s
|
|
340
|
+
AND c.relname = %s;
|
|
341
|
+
""")
|
|
342
|
+
result = await document_store._execute_sql_async(
|
|
343
|
+
cursor=cursor, sql_query=sql_get_index_oid, params=(schema_name, index_name)
|
|
344
|
+
)
|
|
345
|
+
return (await result.fetchone())[0]
|
|
346
|
+
|
|
347
|
+
# create a new schema
|
|
348
|
+
connection_string = "postgresql://postgres:postgres@localhost:5432/postgres"
|
|
349
|
+
schema_name = "test_schema"
|
|
350
|
+
async with await psycopg.AsyncConnection.connect(connection_string, autocommit=True) as conn:
|
|
351
|
+
await conn.execute(f"CREATE SCHEMA IF NOT EXISTS {schema_name}")
|
|
352
|
+
|
|
353
|
+
# create a first document store and trigger the creation of the hnsw index
|
|
354
|
+
params = {
|
|
355
|
+
"connection_string": Secret.from_token(connection_string),
|
|
356
|
+
"schema_name": schema_name,
|
|
357
|
+
"table_name": "haystack_test_hnsw_index_recreation",
|
|
358
|
+
"search_strategy": "hnsw",
|
|
359
|
+
}
|
|
360
|
+
ds1 = PgvectorDocumentStore(**params)
|
|
361
|
+
await ds1._ensure_db_setup_async()
|
|
362
|
+
|
|
363
|
+
# get the hnsw index oid
|
|
364
|
+
hnws_index_name = "haystack_hnsw_index"
|
|
365
|
+
first_oid = await get_index_oid(
|
|
366
|
+
cursor=ds1._async_cursor, document_store=ds1, schema_name=ds1.schema_name, index_name=hnws_index_name
|
|
367
|
+
)
|
|
368
|
+
|
|
369
|
+
# create second document store with recreation enabled
|
|
370
|
+
ds2 = PgvectorDocumentStore(**params, hnsw_recreate_index_if_exists=True)
|
|
371
|
+
await ds2._ensure_db_setup_async()
|
|
372
|
+
|
|
373
|
+
# get the index oid
|
|
374
|
+
second_oid = await get_index_oid(
|
|
375
|
+
cursor=ds2._async_cursor, document_store=ds2, schema_name=ds2.schema_name, index_name=hnws_index_name
|
|
376
|
+
)
|
|
377
|
+
|
|
378
|
+
# verify that oids differ
|
|
379
|
+
assert second_oid != first_oid, "Index was not recreated (OID remained the same)"
|
|
380
|
+
|
|
381
|
+
# Clean up: drop the schema after the test
|
|
382
|
+
async with await psycopg.AsyncConnection.connect(connection_string, autocommit=True) as conn:
|
|
383
|
+
await conn.execute(f"DROP SCHEMA IF EXISTS {schema_name} CASCADE")
|
|
384
|
+
|
|
385
|
+
|
|
386
|
+
@pytest.mark.integration
|
|
387
|
+
@pytest.mark.asyncio
|
|
388
|
+
async def test_create_table_if_not_exists():
|
|
389
|
+
async def get_table_oid(cursor, document_store, schema_name, table_name):
|
|
390
|
+
sql_get_table_oid = SQL("""
|
|
391
|
+
SELECT c.oid
|
|
392
|
+
FROM pg_class c
|
|
393
|
+
JOIN pg_namespace n ON n.oid = c.relnamespace
|
|
394
|
+
WHERE c.relkind = 'r'
|
|
395
|
+
AND n.nspname = %s
|
|
396
|
+
AND c.relname = %s;
|
|
397
|
+
""")
|
|
398
|
+
result = await document_store._execute_sql_async(
|
|
399
|
+
cursor=cursor, sql_query=sql_get_table_oid, params=(schema_name, table_name)
|
|
400
|
+
)
|
|
401
|
+
return (await result.fetchone())[0]
|
|
402
|
+
|
|
403
|
+
connection_string = "postgresql://postgres:postgres@localhost:5432/postgres"
|
|
404
|
+
schema_name = "test_schema"
|
|
405
|
+
table_name = "test_table"
|
|
406
|
+
|
|
407
|
+
# Create a new schema
|
|
408
|
+
async with await psycopg.AsyncConnection.connect(connection_string, autocommit=True) as conn:
|
|
409
|
+
await conn.execute(f"CREATE SCHEMA {schema_name}")
|
|
410
|
+
|
|
411
|
+
document_store = PgvectorDocumentStore(
|
|
412
|
+
connection_string=Secret.from_token(connection_string),
|
|
413
|
+
schema_name=schema_name,
|
|
414
|
+
table_name=table_name,
|
|
415
|
+
)
|
|
416
|
+
|
|
417
|
+
await document_store._ensure_db_setup_async()
|
|
418
|
+
await document_store._initialize_table_async()
|
|
419
|
+
|
|
420
|
+
first_table_oid = await get_table_oid(
|
|
421
|
+
cursor=document_store._async_cursor,
|
|
422
|
+
document_store=document_store,
|
|
423
|
+
schema_name=schema_name,
|
|
424
|
+
table_name=table_name,
|
|
425
|
+
)
|
|
426
|
+
assert first_table_oid is not None, "Table was not created"
|
|
427
|
+
|
|
428
|
+
await document_store._initialize_table_async()
|
|
429
|
+
second_table_oid = await get_table_oid(
|
|
430
|
+
cursor=document_store._async_cursor,
|
|
431
|
+
document_store=document_store,
|
|
432
|
+
schema_name=schema_name,
|
|
433
|
+
table_name=table_name,
|
|
434
|
+
)
|
|
435
|
+
|
|
436
|
+
assert first_table_oid == second_table_oid, "Table was recreated"
|
|
437
|
+
|
|
438
|
+
# Clean up: drop the schema after the test
|
|
439
|
+
async with await psycopg.AsyncConnection.connect(connection_string, autocommit=True) as conn:
|
|
440
|
+
await conn.execute(f"DROP SCHEMA IF EXISTS {schema_name} CASCADE")
|
|
441
|
+
|
|
442
|
+
|
|
443
|
+
@pytest.mark.integration
|
|
444
|
+
@pytest.mark.asyncio
|
|
445
|
+
async def test_delete_table_async_first_call(document_store):
|
|
446
|
+
"""
|
|
447
|
+
Test that delete_table_async can be executed as the initial operation on the Document Store
|
|
448
|
+
without triggering errors due to an uninitialized state.
|
|
449
|
+
"""
|
|
450
|
+
await document_store.delete_table_async() # if throw error, test fails
|
|
@@ -1,232 +0,0 @@
|
|
|
1
|
-
# SPDX-FileCopyrightText: 2023-present deepset GmbH <info@deepset.ai>
|
|
2
|
-
#
|
|
3
|
-
# SPDX-License-Identifier: Apache-2.0
|
|
4
|
-
|
|
5
|
-
from unittest.mock import patch
|
|
6
|
-
|
|
7
|
-
import psycopg
|
|
8
|
-
import pytest
|
|
9
|
-
from haystack.dataclasses.document import ByteStream, Document
|
|
10
|
-
from haystack.document_stores.errors import DocumentStoreError, DuplicateDocumentError
|
|
11
|
-
from haystack.document_stores.types import DuplicatePolicy
|
|
12
|
-
from haystack.utils import Secret
|
|
13
|
-
from psycopg.sql import SQL
|
|
14
|
-
|
|
15
|
-
from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore
|
|
16
|
-
|
|
17
|
-
|
|
18
|
-
@pytest.mark.integration
|
|
19
|
-
@pytest.mark.asyncio
|
|
20
|
-
class TestDocumentStoreAsync:
|
|
21
|
-
async def test_write_documents(self, document_store: PgvectorDocumentStore):
|
|
22
|
-
docs = [Document(id="1")]
|
|
23
|
-
assert await document_store.write_documents_async(docs) == 1
|
|
24
|
-
with pytest.raises(DuplicateDocumentError):
|
|
25
|
-
await document_store.write_documents_async(docs, DuplicatePolicy.FAIL)
|
|
26
|
-
|
|
27
|
-
async def test_write_blob(self, document_store: PgvectorDocumentStore):
|
|
28
|
-
bytestream = ByteStream(b"test", meta={"meta_key": "meta_value"}, mime_type="mime_type")
|
|
29
|
-
docs = [Document(id="1", blob=bytestream)]
|
|
30
|
-
await document_store.write_documents_async(docs)
|
|
31
|
-
|
|
32
|
-
retrieved_docs = await document_store.filter_documents_async()
|
|
33
|
-
assert retrieved_docs == docs
|
|
34
|
-
|
|
35
|
-
async def test_count_documents(self, document_store: PgvectorDocumentStore):
|
|
36
|
-
await document_store.write_documents_async(
|
|
37
|
-
[
|
|
38
|
-
Document(content="test doc 1"),
|
|
39
|
-
Document(content="test doc 2"),
|
|
40
|
-
Document(content="test doc 3"),
|
|
41
|
-
]
|
|
42
|
-
)
|
|
43
|
-
assert await document_store.count_documents_async() == 3
|
|
44
|
-
|
|
45
|
-
async def test_filter_documents(self, document_store: PgvectorDocumentStore):
|
|
46
|
-
filterable_docs = [
|
|
47
|
-
Document(
|
|
48
|
-
content="1",
|
|
49
|
-
meta={
|
|
50
|
-
"number": -10,
|
|
51
|
-
},
|
|
52
|
-
),
|
|
53
|
-
Document(
|
|
54
|
-
content="2",
|
|
55
|
-
meta={
|
|
56
|
-
"number": 100,
|
|
57
|
-
},
|
|
58
|
-
),
|
|
59
|
-
]
|
|
60
|
-
await document_store.write_documents_async(filterable_docs)
|
|
61
|
-
result = await document_store.filter_documents_async(
|
|
62
|
-
filters={"field": "meta.number", "operator": "==", "value": 100}
|
|
63
|
-
)
|
|
64
|
-
|
|
65
|
-
assert result == [d for d in filterable_docs if d.meta.get("number") == 100]
|
|
66
|
-
|
|
67
|
-
async def test_delete_documents(self, document_store: PgvectorDocumentStore):
|
|
68
|
-
doc = Document(content="test doc")
|
|
69
|
-
await document_store.write_documents_async([doc])
|
|
70
|
-
assert await document_store.count_documents_async() == 1
|
|
71
|
-
|
|
72
|
-
await document_store.delete_documents_async([doc.id])
|
|
73
|
-
assert await document_store.count_documents_async() == 0
|
|
74
|
-
|
|
75
|
-
async def test_connection_check_and_recreation(self, document_store: PgvectorDocumentStore):
|
|
76
|
-
await document_store._ensure_db_setup_async()
|
|
77
|
-
original_connection = document_store._async_connection
|
|
78
|
-
|
|
79
|
-
with patch.object(PgvectorDocumentStore, "_connection_is_valid_async", return_value=False):
|
|
80
|
-
await document_store._ensure_db_setup_async()
|
|
81
|
-
new_connection = document_store._async_connection
|
|
82
|
-
|
|
83
|
-
# verify that a new connection is created
|
|
84
|
-
assert new_connection is not original_connection
|
|
85
|
-
assert document_store._async_connection == new_connection
|
|
86
|
-
assert original_connection.closed
|
|
87
|
-
|
|
88
|
-
assert document_store._async_cursor is not None
|
|
89
|
-
assert document_store._async_dict_cursor is not None
|
|
90
|
-
|
|
91
|
-
# test with new connection
|
|
92
|
-
with patch.object(PgvectorDocumentStore, "_connection_is_valid_async", return_value=True):
|
|
93
|
-
await document_store._ensure_db_setup_async()
|
|
94
|
-
same_connection = document_store._async_connection
|
|
95
|
-
assert same_connection is document_store._async_connection
|
|
96
|
-
|
|
97
|
-
async def test_delete_all_documents_async(self, document_store: PgvectorDocumentStore) -> None:
|
|
98
|
-
document_store.write_documents([Document(id=str(i)) for i in range(10)])
|
|
99
|
-
await document_store.delete_all_documents_async()
|
|
100
|
-
assert document_store.count_documents() == 0
|
|
101
|
-
document_store.write_documents([Document(id="1")])
|
|
102
|
-
assert document_store.count_documents() == 1
|
|
103
|
-
|
|
104
|
-
async def test_invalid_connection_string(self, monkeypatch):
|
|
105
|
-
monkeypatch.setenv("PG_CONN_STR", "invalid_connection_string")
|
|
106
|
-
document_store = PgvectorDocumentStore()
|
|
107
|
-
with pytest.raises(DocumentStoreError) as e:
|
|
108
|
-
await document_store._ensure_db_setup_async()
|
|
109
|
-
assert "Failed to connect to PostgreSQL database" in str(e)
|
|
110
|
-
|
|
111
|
-
|
|
112
|
-
@pytest.mark.integration
|
|
113
|
-
@pytest.mark.asyncio
|
|
114
|
-
async def test_hnsw_index_recreation():
|
|
115
|
-
async def get_index_oid(cursor, document_store, schema_name, index_name):
|
|
116
|
-
sql_get_index_oid = SQL("""
|
|
117
|
-
SELECT c.oid
|
|
118
|
-
FROM pg_class c
|
|
119
|
-
JOIN pg_namespace n ON n.oid = c.relnamespace
|
|
120
|
-
WHERE c.relkind = 'i'
|
|
121
|
-
AND n.nspname = %s
|
|
122
|
-
AND c.relname = %s;
|
|
123
|
-
""")
|
|
124
|
-
result = await document_store._execute_sql_async(
|
|
125
|
-
cursor=cursor, sql_query=sql_get_index_oid, params=(schema_name, index_name)
|
|
126
|
-
)
|
|
127
|
-
return (await result.fetchone())[0]
|
|
128
|
-
|
|
129
|
-
# create a new schema
|
|
130
|
-
connection_string = "postgresql://postgres:postgres@localhost:5432/postgres"
|
|
131
|
-
schema_name = "test_schema"
|
|
132
|
-
async with await psycopg.AsyncConnection.connect(connection_string, autocommit=True) as conn:
|
|
133
|
-
await conn.execute(f"CREATE SCHEMA IF NOT EXISTS {schema_name}")
|
|
134
|
-
|
|
135
|
-
# create a first document store and trigger the creation of the hnsw index
|
|
136
|
-
params = {
|
|
137
|
-
"connection_string": Secret.from_token(connection_string),
|
|
138
|
-
"schema_name": schema_name,
|
|
139
|
-
"table_name": "haystack_test_hnsw_index_recreation",
|
|
140
|
-
"search_strategy": "hnsw",
|
|
141
|
-
}
|
|
142
|
-
ds1 = PgvectorDocumentStore(**params)
|
|
143
|
-
await ds1._ensure_db_setup_async()
|
|
144
|
-
|
|
145
|
-
# get the hnsw index oid
|
|
146
|
-
hnws_index_name = "haystack_hnsw_index"
|
|
147
|
-
first_oid = await get_index_oid(
|
|
148
|
-
cursor=ds1._async_cursor, document_store=ds1, schema_name=ds1.schema_name, index_name=hnws_index_name
|
|
149
|
-
)
|
|
150
|
-
|
|
151
|
-
# create second document store with recreation enabled
|
|
152
|
-
ds2 = PgvectorDocumentStore(**params, hnsw_recreate_index_if_exists=True)
|
|
153
|
-
await ds2._ensure_db_setup_async()
|
|
154
|
-
|
|
155
|
-
# get the index oid
|
|
156
|
-
second_oid = await get_index_oid(
|
|
157
|
-
cursor=ds2._async_cursor, document_store=ds2, schema_name=ds2.schema_name, index_name=hnws_index_name
|
|
158
|
-
)
|
|
159
|
-
|
|
160
|
-
# verify that oids differ
|
|
161
|
-
assert second_oid != first_oid, "Index was not recreated (OID remained the same)"
|
|
162
|
-
|
|
163
|
-
# Clean up: drop the schema after the test
|
|
164
|
-
async with await psycopg.AsyncConnection.connect(connection_string, autocommit=True) as conn:
|
|
165
|
-
await conn.execute(f"DROP SCHEMA IF EXISTS {schema_name} CASCADE")
|
|
166
|
-
|
|
167
|
-
|
|
168
|
-
@pytest.mark.integration
|
|
169
|
-
@pytest.mark.asyncio
|
|
170
|
-
async def test_create_table_if_not_exists():
|
|
171
|
-
async def get_table_oid(cursor, document_store, schema_name, table_name):
|
|
172
|
-
sql_get_table_oid = SQL("""
|
|
173
|
-
SELECT c.oid
|
|
174
|
-
FROM pg_class c
|
|
175
|
-
JOIN pg_namespace n ON n.oid = c.relnamespace
|
|
176
|
-
WHERE c.relkind = 'r'
|
|
177
|
-
AND n.nspname = %s
|
|
178
|
-
AND c.relname = %s;
|
|
179
|
-
""")
|
|
180
|
-
result = await document_store._execute_sql_async(
|
|
181
|
-
cursor=cursor, sql_query=sql_get_table_oid, params=(schema_name, table_name)
|
|
182
|
-
)
|
|
183
|
-
return (await result.fetchone())[0]
|
|
184
|
-
|
|
185
|
-
connection_string = "postgresql://postgres:postgres@localhost:5432/postgres"
|
|
186
|
-
schema_name = "test_schema"
|
|
187
|
-
table_name = "test_table"
|
|
188
|
-
|
|
189
|
-
# Create a new schema
|
|
190
|
-
async with await psycopg.AsyncConnection.connect(connection_string, autocommit=True) as conn:
|
|
191
|
-
await conn.execute(f"CREATE SCHEMA {schema_name}")
|
|
192
|
-
|
|
193
|
-
document_store = PgvectorDocumentStore(
|
|
194
|
-
connection_string=Secret.from_token(connection_string),
|
|
195
|
-
schema_name=schema_name,
|
|
196
|
-
table_name=table_name,
|
|
197
|
-
)
|
|
198
|
-
|
|
199
|
-
await document_store._ensure_db_setup_async()
|
|
200
|
-
await document_store._initialize_table_async()
|
|
201
|
-
|
|
202
|
-
first_table_oid = await get_table_oid(
|
|
203
|
-
cursor=document_store._async_cursor,
|
|
204
|
-
document_store=document_store,
|
|
205
|
-
schema_name=schema_name,
|
|
206
|
-
table_name=table_name,
|
|
207
|
-
)
|
|
208
|
-
assert first_table_oid is not None, "Table was not created"
|
|
209
|
-
|
|
210
|
-
await document_store._initialize_table_async()
|
|
211
|
-
second_table_oid = await get_table_oid(
|
|
212
|
-
cursor=document_store._async_cursor,
|
|
213
|
-
document_store=document_store,
|
|
214
|
-
schema_name=schema_name,
|
|
215
|
-
table_name=table_name,
|
|
216
|
-
)
|
|
217
|
-
|
|
218
|
-
assert first_table_oid == second_table_oid, "Table was recreated"
|
|
219
|
-
|
|
220
|
-
# Clean up: drop the schema after the test
|
|
221
|
-
async with await psycopg.AsyncConnection.connect(connection_string, autocommit=True) as conn:
|
|
222
|
-
await conn.execute(f"DROP SCHEMA IF EXISTS {schema_name} CASCADE")
|
|
223
|
-
|
|
224
|
-
|
|
225
|
-
@pytest.mark.integration
|
|
226
|
-
@pytest.mark.asyncio
|
|
227
|
-
async def test_delete_table_async_first_call(document_store):
|
|
228
|
-
"""
|
|
229
|
-
Test that delete_table_async can be executed as the initial operation on the Document Store
|
|
230
|
-
without triggering errors due to an uninitialized state.
|
|
231
|
-
"""
|
|
232
|
-
await document_store.delete_table_async() # if throw error, test fails
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|