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.
- {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/CHANGELOG.md +7 -0
- {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/PKG-INFO +1 -1
- {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/src/haystack_integrations/document_stores/pgvector/document_store.py +171 -0
- pgvector_haystack-8.2.0/src/haystack_integrations/document_stores/pgvector/table_replacement.py +440 -0
- {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/tests/conftest.py +39 -0
- {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/tests/test_document_store.py +471 -1
- {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/tests/test_document_store_async.py +176 -1
- {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/.gitignore +0 -0
- {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/LICENSE.txt +0 -0
- {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/README.md +0 -0
- {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/examples/embedding_retrieval.py +0 -0
- {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/examples/hybrid_retrieval.py +0 -0
- {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/pydoc/config_docusaurus.yml +0 -0
- {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/pyproject.toml +0 -0
- {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/src/haystack_integrations/components/retrievers/pgvector/__init__.py +0 -0
- {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/src/haystack_integrations/components/retrievers/pgvector/embedding_retriever.py +0 -0
- {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/src/haystack_integrations/components/retrievers/pgvector/keyword_retriever.py +0 -0
- {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/src/haystack_integrations/components/retrievers/py.typed +0 -0
- {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/src/haystack_integrations/document_stores/pgvector/__init__.py +0 -0
- {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/src/haystack_integrations/document_stores/pgvector/converters.py +0 -0
- {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/src/haystack_integrations/document_stores/pgvector/filters.py +0 -0
- {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/src/haystack_integrations/document_stores/py.typed +0 -0
- {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/tests/__init__.py +0 -0
- {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/tests/test_converters.py +0 -0
- {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/tests/test_filters.py +0 -0
- {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/tests/test_retrieval.py +0 -0
- {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/tests/test_retrieval_async.py +0 -0
- {pgvector_haystack-8.1.2 → pgvector_haystack-8.2.0}/tests/test_retrievers.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.5
|
|
2
2
|
Name: pgvector-haystack
|
|
3
|
-
Version: 8.
|
|
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
|
|
pgvector_haystack-8.2.0/src/haystack_integrations/document_stores/pgvector/table_replacement.py
ADDED
|
@@ -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
|