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.
Files changed (28) hide show
  1. {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/CHANGELOG.md +19 -0
  2. {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/PKG-INFO +1 -1
  3. {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/src/haystack_integrations/document_stores/pgvector/document_store.py +186 -0
  4. {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/tests/test_document_store.py +218 -0
  5. pgvector_haystack-5.5.0/tests/test_document_store_async.py +450 -0
  6. pgvector_haystack-5.4.0/tests/test_document_store_async.py +0 -232
  7. {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/.gitignore +0 -0
  8. {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/LICENSE.txt +0 -0
  9. {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/README.md +0 -0
  10. {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/examples/embedding_retrieval.py +0 -0
  11. {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/examples/hybrid_retrieval.py +0 -0
  12. {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/pydoc/config_docusaurus.yml +0 -0
  13. {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/pyproject.toml +0 -0
  14. {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/src/haystack_integrations/components/retrievers/pgvector/__init__.py +0 -0
  15. {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/src/haystack_integrations/components/retrievers/pgvector/embedding_retriever.py +0 -0
  16. {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/src/haystack_integrations/components/retrievers/pgvector/keyword_retriever.py +0 -0
  17. {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/src/haystack_integrations/components/retrievers/py.typed +0 -0
  18. {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/src/haystack_integrations/document_stores/pgvector/__init__.py +0 -0
  19. {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/src/haystack_integrations/document_stores/pgvector/converters.py +0 -0
  20. {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/src/haystack_integrations/document_stores/pgvector/filters.py +0 -0
  21. {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/src/haystack_integrations/document_stores/py.typed +0 -0
  22. {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/tests/__init__.py +0 -0
  23. {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/tests/conftest.py +0 -0
  24. {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/tests/test_converters.py +0 -0
  25. {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/tests/test_filters.py +0 -0
  26. {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/tests/test_retrieval.py +0 -0
  27. {pgvector_haystack-5.4.0 → pgvector_haystack-5.5.0}/tests/test_retrieval_async.py +0 -0
  28. {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.4.0
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