pgvector-haystack 8.1.2__tar.gz → 8.2.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-8.1.2 → pgvector_haystack-8.2.0}/CHANGELOG.md +7 -0
  2. {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/PKG-INFO +1 -1
  3. {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/src/haystack_integrations/document_stores/pgvector/document_store.py +171 -0
  4. pgvector_haystack-8.2.0/src/haystack_integrations/document_stores/pgvector/table_replacement.py +440 -0
  5. {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/tests/conftest.py +39 -0
  6. {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/tests/test_document_store.py +471 -1
  7. {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/tests/test_document_store_async.py +176 -1
  8. {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/.gitignore +0 -0
  9. {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/LICENSE.txt +0 -0
  10. {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/README.md +0 -0
  11. {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/examples/embedding_retrieval.py +0 -0
  12. {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/examples/hybrid_retrieval.py +0 -0
  13. {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/pydoc/config_docusaurus.yml +0 -0
  14. {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/pyproject.toml +0 -0
  15. {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/src/haystack_integrations/components/retrievers/pgvector/__init__.py +0 -0
  16. {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/src/haystack_integrations/components/retrievers/pgvector/embedding_retriever.py +0 -0
  17. {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/src/haystack_integrations/components/retrievers/pgvector/keyword_retriever.py +0 -0
  18. {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/src/haystack_integrations/components/retrievers/py.typed +0 -0
  19. {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/src/haystack_integrations/document_stores/pgvector/__init__.py +0 -0
  20. {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/src/haystack_integrations/document_stores/pgvector/converters.py +0 -0
  21. {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/src/haystack_integrations/document_stores/pgvector/filters.py +0 -0
  22. {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/src/haystack_integrations/document_stores/py.typed +0 -0
  23. {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/tests/__init__.py +0 -0
  24. {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/tests/test_converters.py +0 -0
  25. {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/tests/test_filters.py +0 -0
  26. {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/tests/test_retrieval.py +0 -0
  27. {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/tests/test_retrieval_async.py +0 -0
  28. {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/tests/test_retrievers.py +0 -0
@@ -1,5 +1,12 @@
1
1
  # Changelog
2
2
 
3
+ ## [integrations/pgvector-v8.1.2] - 2026-10-21
4
+
5
+ ### 🚀 Features
6
+
7
+ - *(pgvector)* Ensure safe SQL handling for HNSW index creation and embedding queries (#4139)
8
+
9
+
3
10
  ## [integrations/pgvector-v8.1.1] - 2026-10-21
4
11
 
5
12
  ### 🐛 Bug Fixes
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: pgvector-haystack
3
- Version: 8.1.2
3
+ Version: 8.2.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
@@ -22,6 +22,7 @@ from pgvector.psycopg import register_vector, register_vector_async
22
22
 
23
23
  from .converters import _from_haystack_to_pg_documents, _from_pg_to_haystack_documents
24
24
  from .filters import _convert_filters_to_where_clause_and_params, _validate_filters
25
+ from .table_replacement import replace_table, replace_table_async
25
26
 
26
27
  logger = logging.getLogger(__name__)
27
28
 
@@ -635,6 +636,176 @@ class PgvectorDocumentStore:
635
636
  error_msg=f"Could not delete table {self.schema_name}.{self.table_name} in PgvectorDocumentStore",
636
637
  )
637
638
 
639
+ def replace_table(
640
+ self,
641
+ staging_store: "PgvectorDocumentStore",
642
+ *,
643
+ backup_table_name: str | None = None,
644
+ lock_timeout_ms: int = 5000,
645
+ ) -> str:
646
+ """
647
+ Atomically replace this store's table with the table of a populated staging store.
648
+
649
+ Use this to rebuild a document store without downtime: write the new documents into a staging store, then
650
+ swap it in. Queries see either the old or the new documents, never a partially filled table. The former
651
+ table is kept as a backup until you drop it.
652
+
653
+ Requirements:
654
+ - `staging_store` must connect to the same PostgreSQL database and use the same `schema_name`,
655
+ `embedding_dimension`, `vector_type`, `vector_function`, `language`, and `search_strategy`, plus the same
656
+ `hnsw_index_creation_kwargs` with the `hnsw` search strategy.
657
+ - Index names are unique per schema, so `staging_store` needs its own `keyword_index_name` and, with the
658
+ `hnsw` search strategy, its own `hnsw_index_name`. A staging store with the default index names fails
659
+ when it creates its table.
660
+ - The staging table must already exist. This method never creates it.
661
+ - Pause writes to this store during the swap. A write that is in flight when the swap starts lands in the
662
+ backup table, not in the promoted one.
663
+
664
+ Locking: the swap takes an `ACCESS EXCLUSIVE` lock on both tables, so it first waits for running queries on
665
+ them to finish. While it waits, new reads and writes on both tables wait too. If the locks are not acquired
666
+ within `lock_timeout_ms`, nothing is changed. The renames themselves are fast.
667
+
668
+ After a successful swap, this store serves the promoted documents under its own table and index names, and
669
+ `staging_store` is repointed to the backup: its `table_name`, `keyword_index_name`, and, with the `hnsw`
670
+ search strategy, `hnsw_index_name` are set to the backup names. Call `staging_store.delete_table()` to
671
+ drop the former documents once you no longer need them. Because these attributes change,
672
+ `staging_store.to_dict()` also returns the backup names after the swap.
673
+
674
+ Don't create stores that serve queries with `recreate_table=True`. Any new instance with that setting, for
675
+ example in another process or a pipeline loaded from YAML, drops the live table when it first connects.
676
+
677
+ ### Usage example
678
+
679
+ ```python
680
+ from haystack import Document
681
+ from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore
682
+
683
+ live_store = PgvectorDocumentStore(table_name="docs", embedding_dimension=768)
684
+ live_store.write_documents([Document(content="The current corpus")]) # already serving queries
685
+
686
+ # The staging store can reuse the same names on every rebuild: after a swap they are free again.
687
+ staging_store = PgvectorDocumentStore(
688
+ table_name="docs_staging",
689
+ embedding_dimension=768,
690
+ keyword_index_name="docs_staging_keyword_index",
691
+ )
692
+ staging_store.write_documents([Document(content="The rebuilt corpus")])
693
+
694
+ # Pause writes to live_store, then swap.
695
+ backup_table_name = live_store.replace_table(staging_store)
696
+
697
+ # live_store now serves the rebuilt corpus. staging_store points at the backup table.
698
+ staging_store.delete_table()
699
+ ```
700
+
701
+ :param staging_store: A populated store whose table replaces this store's table. Must meet the
702
+ requirements above.
703
+ :param backup_table_name: Name to keep the former table under. Must not exist yet. If not set, a unique
704
+ name of the form `<table_name>_backup_<suffix>` is generated.
705
+ :param lock_timeout_ms: How long to wait for the table locks, in milliseconds, before giving up. Reads and
706
+ writes on both tables also wait during this time. Must be a positive integer.
707
+ :returns: The name of the backup table that now holds the former documents.
708
+ :raises TypeError: If `lock_timeout_ms` is not an `int`.
709
+ :raises ValueError: If `staging_store` is configured incompatibly with this store as described in the
710
+ requirements above, connects to a different database, or if `backup_table_name` is invalid or
711
+ `lock_timeout_ms` is not positive.
712
+ :raises DocumentStoreError: If either table doesn't exist, the tables have different columns or
713
+ constraints, an expected index is missing, `backup_table_name` or a generated backup index name is
714
+ already taken, the locks are not acquired within `lock_timeout_ms`, or the database reports another
715
+ error. In every case the swap is rolled back and nothing is changed.
716
+ """
717
+ return replace_table(
718
+ self,
719
+ staging_store,
720
+ backup_table_name=backup_table_name,
721
+ lock_timeout_ms=lock_timeout_ms,
722
+ )
723
+
724
+ async def replace_table_async(
725
+ self,
726
+ staging_store: "PgvectorDocumentStore",
727
+ *,
728
+ backup_table_name: str | None = None,
729
+ lock_timeout_ms: int = 5000,
730
+ ) -> str:
731
+ """
732
+ Asynchronously and atomically replace this store's table with the table of a populated staging store.
733
+
734
+ Use this to rebuild a document store without downtime: write the new documents into a staging store, then
735
+ swap it in. Queries see either the old or the new documents, never a partially filled table. The former
736
+ table is kept as a backup until you drop it.
737
+
738
+ Requirements:
739
+ - `staging_store` must connect to the same PostgreSQL database and use the same `schema_name`,
740
+ `embedding_dimension`, `vector_type`, `vector_function`, `language`, and `search_strategy`, plus the same
741
+ `hnsw_index_creation_kwargs` with the `hnsw` search strategy.
742
+ - Index names are unique per schema, so `staging_store` needs its own `keyword_index_name` and, with the
743
+ `hnsw` search strategy, its own `hnsw_index_name`. A staging store with the default index names fails
744
+ when it creates its table.
745
+ - The staging table must already exist. This method never creates it.
746
+ - Pause writes to this store during the swap. A write that is in flight when the swap starts lands in the
747
+ backup table, not in the promoted one.
748
+
749
+ Locking: the swap takes an `ACCESS EXCLUSIVE` lock on both tables, so it first waits for running queries on
750
+ them to finish. While it waits, new reads and writes on both tables wait too. If the locks are not acquired
751
+ within `lock_timeout_ms`, nothing is changed. The renames themselves are fast.
752
+
753
+ After a successful swap, this store serves the promoted documents under its own table and index names, and
754
+ `staging_store` is repointed to the backup: its `table_name`, `keyword_index_name`, and, with the `hnsw`
755
+ search strategy, `hnsw_index_name` are set to the backup names. Call `staging_store.delete_table_async()` to
756
+ drop the former documents once you no longer need them. Because these attributes change,
757
+ `staging_store.to_dict()` also returns the backup names after the swap.
758
+
759
+ Don't create stores that serve queries with `recreate_table=True`. Any new instance with that setting, for
760
+ example in another process or a pipeline loaded from YAML, drops the live table when it first connects.
761
+
762
+ ### Usage example
763
+
764
+ ```python
765
+ from haystack import Document
766
+ from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore
767
+
768
+ live_store = PgvectorDocumentStore(table_name="docs", embedding_dimension=768)
769
+ await live_store.write_documents_async([Document(content="The current corpus")]) # already serving queries
770
+
771
+ # The staging store can reuse the same names on every rebuild: after a swap they are free again.
772
+ staging_store = PgvectorDocumentStore(
773
+ table_name="docs_staging",
774
+ embedding_dimension=768,
775
+ keyword_index_name="docs_staging_keyword_index",
776
+ )
777
+ await staging_store.write_documents_async([Document(content="The rebuilt corpus")])
778
+
779
+ # Pause writes to live_store, then swap.
780
+ backup_table_name = await live_store.replace_table_async(staging_store)
781
+
782
+ # live_store now serves the rebuilt corpus. staging_store points at the backup table.
783
+ await staging_store.delete_table_async()
784
+ ```
785
+
786
+ :param staging_store: A populated store whose table replaces this store's table. Must meet the
787
+ requirements above.
788
+ :param backup_table_name: Name to keep the former table under. Must not exist yet. If not set, a unique
789
+ name of the form `<table_name>_backup_<suffix>` is generated.
790
+ :param lock_timeout_ms: How long to wait for the table locks, in milliseconds, before giving up. Reads and
791
+ writes on both tables also wait during this time. Must be a positive integer.
792
+ :returns: The name of the backup table that now holds the former documents.
793
+ :raises TypeError: If `lock_timeout_ms` is not an `int`.
794
+ :raises ValueError: If `staging_store` is configured incompatibly with this store as described in the
795
+ requirements above, connects to a different database, or if `backup_table_name` is invalid or
796
+ `lock_timeout_ms` is not positive.
797
+ :raises DocumentStoreError: If either table doesn't exist, the tables have different columns or
798
+ constraints, an expected index is missing, `backup_table_name` or a generated backup index name is
799
+ already taken, the locks are not acquired within `lock_timeout_ms`, or the database reports another
800
+ error. In every case the swap is rolled back and nothing is changed.
801
+ """
802
+ return await replace_table_async(
803
+ self,
804
+ staging_store,
805
+ backup_table_name=backup_table_name,
806
+ lock_timeout_ms=lock_timeout_ms,
807
+ )
808
+
638
809
  def _build_hnsw_queries(self) -> tuple[Composed | None, SQL, Composed, Composed]:
639
810
  """Common method to build all HNSW-related SQL queries"""
640
811
 
@@ -0,0 +1,440 @@
1
+ # SPDX-FileCopyrightText: 2023-present deepset GmbH <info@deepset.ai>
2
+ #
3
+ # SPDX-License-Identifier: Apache-2.0
4
+
5
+ """Safe transactional table replacement helpers for PgvectorDocumentStore."""
6
+
7
+ from __future__ import annotations
8
+
9
+ from collections.abc import Iterator
10
+ from contextlib import contextmanager
11
+ from dataclasses import dataclass
12
+ from typing import TYPE_CHECKING, Any
13
+ from uuid import uuid4
14
+
15
+ from haystack.document_stores.errors import DocumentStoreError
16
+ from psycopg import AsyncCursor, Cursor, Error
17
+ from psycopg.sql import SQL, Composed, Identifier
18
+
19
+ if TYPE_CHECKING:
20
+ from .document_store import PgvectorDocumentStore
21
+
22
+ _POSTGRES_IDENTIFIER_MAX_BYTES = 63
23
+ # Identifies "same database on the same running server" from the server side only, so it does not depend on
24
+ # how the client spelled the host (localhost vs 127.0.0.1, IPv4 vs IPv6, TCP vs Unix socket).
25
+ # pg_postmaster_start_time() is callable by any role, unlike pg_control_system().
26
+ _DATABASE_IDENTITY_QUERY = SQL(
27
+ "SELECT current_database(), (SELECT oid FROM pg_database WHERE datname = current_database()), "
28
+ "pg_postmaster_start_time()"
29
+ )
30
+ _TABLE_KINDS_QUERY = SQL(
31
+ "SELECT c.relname, c.relkind FROM pg_class AS c "
32
+ "JOIN pg_namespace AS n ON n.oid = c.relnamespace "
33
+ "WHERE n.nspname = %s AND c.relname = ANY(%s::name[])"
34
+ )
35
+ _COLUMNS_QUERY = SQL(
36
+ "SELECT a.attname, pg_catalog.format_type(a.atttypid, a.atttypmod), a.attnotnull, "
37
+ "a.attidentity, a.attgenerated "
38
+ "FROM pg_attribute AS a "
39
+ "JOIN pg_class AS c ON c.oid = a.attrelid "
40
+ "JOIN pg_namespace AS n ON n.oid = c.relnamespace "
41
+ "WHERE n.nspname = %s AND c.relname = %s AND a.attnum > 0 AND NOT a.attisdropped "
42
+ "ORDER BY a.attnum"
43
+ )
44
+ _CONSTRAINTS_QUERY = SQL(
45
+ "SELECT con.contype, pg_catalog.pg_get_constraintdef(con.oid) "
46
+ "FROM pg_constraint AS con "
47
+ "JOIN pg_class AS c ON c.oid = con.conrelid "
48
+ "JOIN pg_namespace AS n ON n.oid = c.relnamespace "
49
+ "WHERE n.nspname = %s AND c.relname = %s "
50
+ "ORDER BY con.contype, pg_catalog.pg_get_constraintdef(con.oid)"
51
+ )
52
+ _INDEXES_QUERY = SQL(
53
+ "SELECT t.relname, i.relname, x.indisprimary FROM pg_index AS x "
54
+ "JOIN pg_class AS t ON t.oid = x.indrelid "
55
+ "JOIN pg_class AS i ON i.oid = x.indexrelid "
56
+ "JOIN pg_namespace AS n ON n.oid = t.relnamespace "
57
+ "WHERE n.nspname = %s AND t.relname = ANY(%s::name[])"
58
+ )
59
+ _EXISTING_RELATIONS_QUERY = SQL(
60
+ "SELECT c.relname FROM pg_class AS c "
61
+ "JOIN pg_namespace AS n ON n.oid = c.relnamespace "
62
+ "WHERE n.nspname = %s AND c.relname = ANY(%s::name[])"
63
+ )
64
+ _RENAME_INDEX = "ALTER INDEX {schema}.{old_name} RENAME TO {new_name}"
65
+ _RENAME_TABLE = "ALTER TABLE {schema}.{old_name} RENAME TO {new_name}"
66
+
67
+
68
+ def _truncate_identifier(identifier: str) -> str:
69
+ """Keep a generated PostgreSQL identifier within PostgreSQL's 63-byte limit."""
70
+ encoded = identifier.encode("utf-8")
71
+ if len(encoded) <= _POSTGRES_IDENTIFIER_MAX_BYTES:
72
+ return identifier
73
+
74
+ suffix = f"_{uuid4().hex[:8]}"
75
+ prefix_bytes = encoded[: _POSTGRES_IDENTIFIER_MAX_BYTES - len(suffix)]
76
+ return prefix_bytes.decode("utf-8", errors="ignore") + suffix
77
+
78
+
79
+ def _validate_identifier(identifier: str, label: str) -> None:
80
+ if not identifier:
81
+ msg = f"{label} must not be empty"
82
+ raise ValueError(msg)
83
+ if len(identifier.encode("utf-8")) > _POSTGRES_IDENTIFIER_MAX_BYTES:
84
+ msg = f"{label} must be at most {_POSTGRES_IDENTIFIER_MAX_BYTES} UTF-8 bytes"
85
+ raise ValueError(msg)
86
+
87
+
88
+ def _validate_stores(
89
+ active_store: PgvectorDocumentStore,
90
+ staging_store: PgvectorDocumentStore,
91
+ backup_table_name: str | None,
92
+ lock_timeout_ms: int,
93
+ ) -> tuple[str, list[tuple[str, str, str]]]:
94
+ if staging_store is active_store:
95
+ msg = "staging_store must be a different PgvectorDocumentStore instance"
96
+ raise ValueError(msg)
97
+ if active_store.schema_name != staging_store.schema_name:
98
+ msg = "Active and staging document stores must use the same schema"
99
+ raise ValueError(msg)
100
+ if active_store.table_name == staging_store.table_name:
101
+ msg = "Active and staging document stores must use different table names"
102
+ raise ValueError(msg)
103
+ if isinstance(lock_timeout_ms, bool) or not isinstance(lock_timeout_ms, int):
104
+ msg = f"lock_timeout_ms must be an int, got {type(lock_timeout_ms).__name__}"
105
+ raise TypeError(msg)
106
+ if lock_timeout_ms <= 0:
107
+ msg = "lock_timeout_ms must be a positive integer"
108
+ raise ValueError(msg)
109
+
110
+ compatible_options = (
111
+ "embedding_dimension",
112
+ "vector_type",
113
+ "vector_function",
114
+ "language",
115
+ "search_strategy",
116
+ )
117
+ for option in compatible_options:
118
+ if getattr(active_store, option) != getattr(staging_store, option):
119
+ msg = f"Active and staging document stores must use the same {option}"
120
+ raise ValueError(msg)
121
+ if active_store.search_strategy == "hnsw" and (
122
+ active_store.hnsw_index_creation_kwargs != staging_store.hnsw_index_creation_kwargs
123
+ ):
124
+ msg = "Active and staging document stores must use the same hnsw_index_creation_kwargs"
125
+ raise ValueError(msg)
126
+ if active_store.keyword_index_name == staging_store.keyword_index_name:
127
+ msg = "Active and staging keyword indexes must have different names"
128
+ raise ValueError(msg)
129
+ if active_store.search_strategy == "hnsw" and active_store.hnsw_index_name == staging_store.hnsw_index_name:
130
+ msg = "Active and staging HNSW indexes must have different names"
131
+ raise ValueError(msg)
132
+
133
+ for name, label in (
134
+ (active_store.schema_name, "schema_name"),
135
+ (active_store.table_name, "active table_name"),
136
+ (staging_store.table_name, "staging table_name"),
137
+ (active_store.keyword_index_name, "active keyword_index_name"),
138
+ (staging_store.keyword_index_name, "staging keyword_index_name"),
139
+ ):
140
+ _validate_identifier(name, label)
141
+ if active_store.search_strategy == "hnsw":
142
+ _validate_identifier(active_store.hnsw_index_name, "active hnsw_index_name")
143
+ _validate_identifier(staging_store.hnsw_index_name, "staging hnsw_index_name")
144
+
145
+ if backup_table_name is None:
146
+ backup_table_name = _truncate_identifier(f"{active_store.table_name}_backup_{uuid4().hex[:8]}")
147
+ else:
148
+ _validate_identifier(backup_table_name, "backup_table_name")
149
+ if backup_table_name in {active_store.table_name, staging_store.table_name}:
150
+ msg = "backup_table_name must differ from the active and staging table names"
151
+ raise ValueError(msg)
152
+
153
+ token = uuid4().hex[:8]
154
+ backup_keyword_index_name = _truncate_identifier(f"{backup_table_name}_kw_{token}")
155
+ index_renames = [(active_store.keyword_index_name, staging_store.keyword_index_name, backup_keyword_index_name)]
156
+ if active_store.search_strategy == "hnsw":
157
+ backup_hnsw_index_name = _truncate_identifier(f"{backup_table_name}_hnsw_{token}")
158
+ index_renames.append((active_store.hnsw_index_name, staging_store.hnsw_index_name, backup_hnsw_index_name))
159
+ return backup_table_name, index_renames
160
+
161
+
162
+ @contextmanager
163
+ def _without_table_setup(store: PgvectorDocumentStore) -> Iterator[None]:
164
+ """
165
+ Let `store` connect without its first-connect table setup, then restore its settings.
166
+
167
+ Setup creates a missing table and, with `recreate_table=True`, drops an existing one first. A replacement must
168
+ neither promote a silently created empty table nor drop live data, so setup is skipped while connecting.
169
+ """
170
+ settings = (store.recreate_table, store.hnsw_recreate_index_if_exists, store._table_initialized)
171
+ store.recreate_table = False
172
+ store.hnsw_recreate_index_if_exists = False
173
+ store._table_initialized = True
174
+ try:
175
+ yield
176
+ finally:
177
+ store.recreate_table, store.hnsw_recreate_index_if_exists, store._table_initialized = settings
178
+
179
+
180
+ def _prepare_sync_store(store: PgvectorDocumentStore) -> None:
181
+ """Open a connection without initializing or recreating a table."""
182
+ with _without_table_setup(store):
183
+ store._ensure_db_setup()
184
+
185
+
186
+ async def _prepare_async_store(store: PgvectorDocumentStore) -> None:
187
+ """Open an async connection without initializing or recreating a table."""
188
+ with _without_table_setup(store):
189
+ await store._ensure_db_setup_async()
190
+
191
+
192
+ _Row = tuple[Any, ...]
193
+
194
+
195
+ def _database_identity(cursor: Cursor) -> list[_Row]:
196
+ return cursor.execute(_DATABASE_IDENTITY_QUERY).fetchall()
197
+
198
+
199
+ async def _database_identity_async(cursor: AsyncCursor) -> list[_Row]:
200
+ await cursor.execute(_DATABASE_IDENTITY_QUERY)
201
+ return await cursor.fetchall()
202
+
203
+
204
+ @dataclass(frozen=True)
205
+ class _Swap:
206
+ """
207
+ Names involved in one table replacement, with the SQL and checks to perform it.
208
+
209
+ Holds no connection and performs no I/O, so the sync and async paths share every query, check, and
210
+ statement and only differ in how they execute them.
211
+ """
212
+
213
+ schema: str
214
+ active_table: str
215
+ staging_table: str
216
+ backup_table: str
217
+ # (active name, staging name, backup name) for the keyword index and, with HNSW, the HNSW index.
218
+ index_renames: tuple[tuple[str, str, str], ...]
219
+ backup_pk: str
220
+ lock_timeout_ms: int
221
+
222
+ @classmethod
223
+ def create(
224
+ cls,
225
+ active_store: PgvectorDocumentStore,
226
+ staging_store: PgvectorDocumentStore,
227
+ backup_table_name: str | None,
228
+ lock_timeout_ms: int,
229
+ ) -> _Swap:
230
+ backup_table, index_renames = _validate_stores(active_store, staging_store, backup_table_name, lock_timeout_ms)
231
+ return cls(
232
+ schema=active_store.schema_name,
233
+ active_table=active_store.table_name,
234
+ staging_table=staging_store.table_name,
235
+ backup_table=backup_table,
236
+ index_renames=tuple(index_renames),
237
+ backup_pk=_truncate_identifier(f"{backup_table}_pkey"),
238
+ lock_timeout_ms=lock_timeout_ms,
239
+ )
240
+
241
+ def lock_timeout_query(self) -> tuple[SQL, tuple[str]]:
242
+ return SQL("SELECT set_config('lock_timeout', %s, true)"), (f"{self.lock_timeout_ms}ms",)
243
+
244
+ def table_kinds_query(self) -> tuple[SQL, tuple[str, list[str]]]:
245
+ return _TABLE_KINDS_QUERY, (self.schema, [self.active_table, self.staging_table])
246
+
247
+ def check_tables_exist(self, rows: list[_Row]) -> None:
248
+ kinds: dict[str, str] = dict(rows)
249
+ for label, table in (("Active", self.active_table), ("Staging", self.staging_table)):
250
+ if kinds.get(table) != "r":
251
+ msg = f"{label} table {self.schema}.{table} does not exist as a regular table"
252
+ raise DocumentStoreError(msg)
253
+
254
+ def lock_statement(self) -> Composed:
255
+ return SQL("LOCK TABLE {schema}.{active}, {schema}.{staging} IN ACCESS EXCLUSIVE MODE").format(
256
+ schema=Identifier(self.schema),
257
+ active=Identifier(self.active_table),
258
+ staging=Identifier(self.staging_table),
259
+ )
260
+
261
+ def catalog_queries(self) -> list[tuple[SQL, tuple[Any, ...]]]:
262
+ """Queries whose results `plan` expects, in this order. Run them after taking the locks."""
263
+ backup_relations = [self.backup_table, self.backup_pk, *(backup for _, _, backup in self.index_renames)]
264
+ return [
265
+ (_COLUMNS_QUERY, (self.schema, self.active_table)),
266
+ (_COLUMNS_QUERY, (self.schema, self.staging_table)),
267
+ (_CONSTRAINTS_QUERY, (self.schema, self.active_table)),
268
+ (_CONSTRAINTS_QUERY, (self.schema, self.staging_table)),
269
+ (_INDEXES_QUERY, (self.schema, [self.active_table, self.staging_table])),
270
+ (_EXISTING_RELATIONS_QUERY, (self.schema, backup_relations)),
271
+ ]
272
+
273
+ def plan(self, results: list[list[_Row]]) -> list[Composed]:
274
+ """
275
+ Check the catalog results of `catalog_queries` and return the rename statements to execute, in order.
276
+
277
+ :raises DocumentStoreError: If the tables are incompatible or a backup name is already taken.
278
+ """
279
+ active_columns, staging_columns, active_constraints, staging_constraints, index_rows, existing_rows = results
280
+ if active_columns != staging_columns:
281
+ msg = "Active and staging tables have incompatible column definitions"
282
+ raise DocumentStoreError(msg)
283
+ if active_constraints != staging_constraints:
284
+ msg = "Active and staging tables have incompatible constraints"
285
+ raise DocumentStoreError(msg)
286
+ existing = {name for (name,) in existing_rows}
287
+ if self.backup_table in existing:
288
+ msg = f"Backup relation {self.schema}.{self.backup_table} already exists"
289
+ raise DocumentStoreError(msg)
290
+
291
+ indexes = {(table, index) for table, index, _ in index_rows}
292
+ primary_keys = {table: index for table, index, is_primary in index_rows if is_primary}
293
+ renames = list(self.index_renames)
294
+ # PostgreSQL does not rename a table's primary key with the table, so without this the promoted table
295
+ # keeps the staging name and names drift on every replacement. Renaming the index also renames its
296
+ # constraint. Matching constraints guarantee both tables have a primary key or neither.
297
+ if self.active_table in primary_keys and self.staging_table in primary_keys:
298
+ renames.append((primary_keys[self.active_table], primary_keys[self.staging_table], self.backup_pk))
299
+
300
+ for active_index, staging_index, backup_index in renames:
301
+ if (self.active_table, active_index) not in indexes:
302
+ msg = f"Expected index {active_index!r} was not found on active table {self.active_table!r}"
303
+ raise DocumentStoreError(msg)
304
+ if (self.staging_table, staging_index) not in indexes:
305
+ msg = f"Expected index {staging_index!r} was not found on staging table {self.staging_table!r}"
306
+ raise DocumentStoreError(msg)
307
+ if backup_index in existing:
308
+ msg = f"Backup index relation {self.schema}.{backup_index} already exists"
309
+ raise DocumentStoreError(msg)
310
+
311
+ # Active names are freed first, so the staging objects can take them over.
312
+ return [
313
+ *(self._rename(_RENAME_INDEX, active, backup) for active, _, backup in renames),
314
+ *(self._rename(_RENAME_INDEX, staging, active) for active, staging, _ in renames),
315
+ self._rename(_RENAME_TABLE, self.active_table, self.backup_table),
316
+ self._rename(_RENAME_TABLE, self.staging_table, self.active_table),
317
+ ]
318
+
319
+ def _rename(self, template: str, old_name: str, new_name: str) -> Composed:
320
+ return SQL(template).format(
321
+ schema=Identifier(self.schema), old_name=Identifier(old_name), new_name=Identifier(new_name)
322
+ )
323
+
324
+ def error(self, error: Error) -> DocumentStoreError:
325
+ msg = (
326
+ f"Could not atomically replace {self.schema}.{self.active_table} "
327
+ f"with staging table {self.schema}.{self.staging_table}: {error}"
328
+ )
329
+ return DocumentStoreError(msg)
330
+
331
+
332
+ def _point_stores_after_swap(
333
+ active_store: PgvectorDocumentStore, staging_store: PgvectorDocumentStore, swap: _Swap
334
+ ) -> None:
335
+ """Repoint the staging instance at the backup and mark both tables as initialized."""
336
+ staging_store.table_name = swap.backup_table
337
+ staging_store.keyword_index_name = swap.index_renames[0][2]
338
+ if active_store.search_strategy == "hnsw":
339
+ staging_store.hnsw_index_name = swap.index_renames[1][2]
340
+ # Both tables already exist. Without this, an instance that had not initialized its table before the swap
341
+ # would run table initialization on its next reconnect, and `recreate_table=True` would drop the promoted table.
342
+ active_store._table_initialized = True
343
+ staging_store._table_initialized = True
344
+
345
+
346
+ def replace_table(
347
+ active_store: PgvectorDocumentStore,
348
+ staging_store: PgvectorDocumentStore,
349
+ *,
350
+ backup_table_name: str | None = None,
351
+ lock_timeout_ms: int = 5000,
352
+ ) -> str:
353
+ """
354
+ Replace the table of `active_store` with the table of `staging_store`.
355
+
356
+ See `PgvectorDocumentStore.replace_table` for requirements, locking, and side effects.
357
+
358
+ :param active_store: The document store that currently serves queries.
359
+ :param staging_store: A populated document store whose table replaces the active table.
360
+ :param backup_table_name: Name to keep the former table under, or `None` to generate one.
361
+ :param lock_timeout_ms: How long to wait for the table locks, in milliseconds.
362
+ :returns: The name of the backup table that now holds the former documents.
363
+ """
364
+ swap = _Swap.create(active_store, staging_store, backup_table_name, lock_timeout_ms)
365
+ _prepare_sync_store(active_store)
366
+ _prepare_sync_store(staging_store)
367
+ assert active_store._connection is not None
368
+ assert staging_store._connection is not None
369
+
370
+ try:
371
+ with active_store._connection.cursor() as active_cursor, staging_store._connection.cursor() as staging_cursor:
372
+ if _database_identity(active_cursor) != _database_identity(staging_cursor):
373
+ msg = "Active and staging document stores must connect to the same PostgreSQL database"
374
+ raise ValueError(msg)
375
+
376
+ connection = active_store._connection
377
+ with connection.transaction(), connection.cursor() as cursor:
378
+ cursor.execute(*swap.lock_timeout_query())
379
+ swap.check_tables_exist(cursor.execute(*swap.table_kinds_query()).fetchall())
380
+ cursor.execute(swap.lock_statement())
381
+ results = [cursor.execute(query, params).fetchall() for query, params in swap.catalog_queries()]
382
+ for statement in swap.plan(results):
383
+ cursor.execute(statement)
384
+ except Error as error:
385
+ raise swap.error(error) from error
386
+
387
+ _point_stores_after_swap(active_store, staging_store, swap)
388
+ return swap.backup_table
389
+
390
+
391
+ async def replace_table_async(
392
+ active_store: PgvectorDocumentStore,
393
+ staging_store: PgvectorDocumentStore,
394
+ *,
395
+ backup_table_name: str | None = None,
396
+ lock_timeout_ms: int = 5000,
397
+ ) -> str:
398
+ """
399
+ Asynchronously replace the table of `active_store` with the table of `staging_store`.
400
+
401
+ See `PgvectorDocumentStore.replace_table_async` for requirements, locking, and side effects.
402
+
403
+ :param active_store: The document store that currently serves queries.
404
+ :param staging_store: A populated document store whose table replaces the active table.
405
+ :param backup_table_name: Name to keep the former table under, or `None` to generate one.
406
+ :param lock_timeout_ms: How long to wait for the table locks, in milliseconds.
407
+ :returns: The name of the backup table that now holds the former documents.
408
+ """
409
+ swap = _Swap.create(active_store, staging_store, backup_table_name, lock_timeout_ms)
410
+ await _prepare_async_store(active_store)
411
+ await _prepare_async_store(staging_store)
412
+ assert active_store._async_connection is not None
413
+ assert staging_store._async_connection is not None
414
+
415
+ try:
416
+ async with (
417
+ active_store._async_connection.cursor() as active_cursor,
418
+ staging_store._async_connection.cursor() as staging_cursor,
419
+ ):
420
+ if await _database_identity_async(active_cursor) != await _database_identity_async(staging_cursor):
421
+ msg = "Active and staging document stores must connect to the same PostgreSQL database"
422
+ raise ValueError(msg)
423
+
424
+ connection = active_store._async_connection
425
+ async with connection.transaction(), connection.cursor() as cursor:
426
+ await cursor.execute(*swap.lock_timeout_query())
427
+ await cursor.execute(*swap.table_kinds_query())
428
+ swap.check_tables_exist(await cursor.fetchall())
429
+ await cursor.execute(swap.lock_statement())
430
+ results = []
431
+ for query, params in swap.catalog_queries():
432
+ await cursor.execute(query, params)
433
+ results.append(await cursor.fetchall())
434
+ for statement in swap.plan(results):
435
+ await cursor.execute(statement)
436
+ except Error as error:
437
+ raise swap.error(error) from error
438
+
439
+ _point_stores_after_swap(active_store, staging_store, swap)
440
+ return swap.backup_table