From af7a1fcdba3d93ff258e2138dd70f94442f48a1a Mon Sep 17 00:00:00 2001 From: Trenton Holmes <797416+stumpylog@users.noreply.github.com> Date: Sun, 26 Jul 2026 21:25:25 -0700 Subject: [PATCH] Simplify: dedupe row-copy/migration-test helpers, type fixture params Co-Authored-By: Claude Sonnet 5 --- src/paperless_ai/tests/test_vector_store.py | 464 +++++++++++--------- src/paperless_ai/vector_store.py | 137 +++--- 2 files changed, 323 insertions(+), 278 deletions(-) diff --git a/src/paperless_ai/tests/test_vector_store.py b/src/paperless_ai/tests/test_vector_store.py index 97c483633..7b42cf9e0 100644 --- a/src/paperless_ai/tests/test_vector_store.py +++ b/src/paperless_ai/tests/test_vector_store.py @@ -1,5 +1,7 @@ import sqlite3 from collections.abc import Generator +from collections.abc import Iterator +from contextlib import contextmanager from pathlib import Path import pytest @@ -9,6 +11,7 @@ from llama_index.core.vector_stores.types import MetadataFilter from llama_index.core.vector_stores.types import MetadataFilters from llama_index.core.vector_stores.types import VectorStoreQuery +from paperless_ai import vector_store as vs_mod from paperless_ai.vector_store import DB_FILENAME from paperless_ai.vector_store import DEFAULT_TABLE_NAME from paperless_ai.vector_store import MIGRATIONS @@ -89,8 +92,87 @@ def _ne_filter(document_id: str): ) +def _chunk_index_rows( + store: PaperlessSqliteVecVectorStore, + document_id: str | None = None, +) -> list[tuple[str, str]]: + """The document_chunks side table, as (chunk_id, document_id) pairs.""" + sql = "SELECT chunk_id, document_id FROM document_chunks" + params: list[str] = [] + if document_id is not None: + sql += " WHERE document_id = ?" + params.append(document_id) + rows = store.client.execute(sql + " ORDER BY chunk_id", params).fetchall() + return [(r["chunk_id"], r["document_id"]) for r in rows] + + +@contextmanager +def _pending_migrations(*migrations: Migration, schema_version: int) -> Iterator[None]: + """Register ``migrations`` and raise the module's SCHEMA_VERSION for the + duration of the block, so check_and_run_migrations() sees them as pending. + """ + original = vs_mod.SCHEMA_VERSION + MIGRATIONS.extend(migrations) + vs_mod.SCHEMA_VERSION = schema_version + try: + yield + finally: + for migration in migrations: + MIGRATIONS.remove(migration) + vs_mod.SCHEMA_VERSION = original + + +def _copying_apply(src: sqlite3.Connection, dst: sqlite3.Connection, dim: int) -> None: + """A structural migration apply() that just copies every row across. + + Deliberately spells its SQL out rather than reusing the production + helpers, so these tests exercise the migration machinery against a + realistic third-party apply() instead of co-drifting with it. + """ + dst.execute( # nosemgrep + f"CREATE VIRTUAL TABLE {DEFAULT_TABLE_NAME} USING vec0(" + "id TEXT PRIMARY KEY, document_id TEXT, modified TEXT," + f" +node_content TEXT, embedding float[{dim}] distance_metric=cosine" + ")", + ) + dst.execute( + "INSERT INTO index_meta (key, value) VALUES ('dim', ?) " + "ON CONFLICT(key) DO UPDATE SET value = excluded.value", + (str(dim),), + ) + rows = src.execute( + "SELECT id, document_id, modified, node_content, embedding " + f"FROM {DEFAULT_TABLE_NAME}", + ).fetchall() + dst.execute("BEGIN IMMEDIATE") + dst.executemany( + f"INSERT INTO {DEFAULT_TABLE_NAME} " + "(id, document_id, modified, node_content, embedding) " + "VALUES (?, ?, ?, ?, ?)", + [ + ( + r["id"], + r["document_id"], + r["modified"], + r["node_content"], + bytes(r["embedding"]), + ) + for r in rows + ], + ) + dst.execute( + "INSERT INTO index_meta (key, value) VALUES ('total_inserts', ?) " + "ON CONFLICT(key) DO UPDATE SET value = excluded.value", + (str(len(rows)),), + ) + dst.execute("COMMIT") + + class TestCrud: - def test_add_then_query_returns_node(self, store) -> None: + def test_add_then_query_returns_node( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: node = make_node("n1", "1") assert store.add([node]) == ["n1"] result = _query(store, node.embedding, top_k=1) @@ -99,21 +181,30 @@ class TestCrud: # cosine distance of the identical vector is 0 -> similarity 1 assert result.similarities[0] == pytest.approx(1.0) - def test_query_empty_store_returns_empty_no_raise(self, store) -> None: + def test_query_empty_store_returns_empty_no_raise( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: result = _query(store, [0.0] * DIM) assert result.ids == [] and result.nodes == [] and result.similarities == [] - def test_add_empty_list_is_noop(self, store) -> None: + def test_add_empty_list_is_noop(self, store: PaperlessSqliteVecVectorStore) -> None: assert store.add([]) == [] assert not store.table_exists() - def test_delete_removes_all_chunks_of_document(self, store) -> None: + def test_delete_removes_all_chunks_of_document( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: store.add([make_node("a1", "1"), make_node("a2", "1"), make_node("b1", "2")]) store.delete("1") result = _query(store, [0.0] * DIM, top_k=10) assert result.ids == ["b1"] - def test_query_with_in_filter_scopes_results(self, store) -> None: + def test_query_with_in_filter_scopes_results( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: store.add( [ make_node("a1", "1", seed=0.0), @@ -124,7 +215,10 @@ class TestCrud: result = _query(store, [0.0] * DIM, top_k=10, filters=_in_filter(["2", "3"])) assert sorted(result.ids) == ["b1", "c1"] - def test_query_respects_top_k_with_filter(self, store) -> None: + def test_query_respects_top_k_with_filter( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: # k semantics: global top-k even with IN filters (document_id is a # metadata column, not a partition key -- see design doc). store.add( @@ -139,7 +233,10 @@ class TestCrud: assert len(result.ids) == 3 assert result.similarities == sorted(result.similarities, reverse=True) - def test_get_nodes_filter_and_empty_paths(self, store) -> None: + def test_get_nodes_filter_and_empty_paths( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: assert store.get_nodes(filters=_in_filter(["1"])) == [] # no table yet store.add([make_node("a1", "1"), make_node("b1", "2")]) nodes = store.get_nodes(filters=_in_filter(["1"])) @@ -147,7 +244,10 @@ class TestCrud: assert nodes[0].embedding is not None assert store.get_nodes(filters=_in_filter(["999"])) == [] - def test_query_with_eq_filter_scopes_results(self, store) -> None: + def test_query_with_eq_filter_scopes_results( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: store.add( [ make_node("a1", "1", seed=0.0), @@ -163,18 +263,25 @@ class TestCrud: ) assert result.ids == ["b1"] - def test_get_nodes_node_ids_not_implemented(self, store) -> None: + def test_get_nodes_node_ids_not_implemented( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: with pytest.raises(NotImplementedError): store.get_nodes(node_ids=["x"]) - def test_fresh_instance_sees_existing_table(self, store, tmp_path: Path) -> None: + def test_fresh_instance_sees_existing_table( + self, + store: PaperlessSqliteVecVectorStore, + tmp_path: Path, + ) -> None: store.add([make_node("a1", "1")]) with PaperlessSqliteVecVectorStore(uri=str(tmp_path)) as reopened: assert reopened.table_exists() assert reopened.vector_dim() == DIM assert _query(reopened, [0.0] * DIM, top_k=1).ids == ["a1"] - def test_table_exists_and_drop(self, store) -> None: + def test_table_exists_and_drop(self, store: PaperlessSqliteVecVectorStore) -> None: assert not store.table_exists() store.add([make_node("a1", "1")]) assert store.table_exists() @@ -189,7 +296,10 @@ class TestBuildWhere: assert where == "(document_id != ?)" assert params == ["1"] - def test_query_with_ne_filter_excludes_matching_document(self, store) -> None: + def test_query_with_ne_filter_excludes_matching_document( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: store.add([make_node("a1", "1"), make_node("b1", "2")]) assert sorted( _query(store, [0.0] * DIM, top_k=5, filters=_ne_filter("1")).ids, @@ -214,7 +324,10 @@ class TestBuildWhere: assert where == "1 = 0" assert params == [] - def test_query_with_untranslatable_filter_returns_no_rows(self, store) -> None: + def test_query_with_untranslatable_filter_returns_no_rows( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: store.add([make_node("a1", "1"), make_node("b1", "2")]) nested = MetadataFilters( filters=[ @@ -232,7 +345,10 @@ class TestBuildWhere: class TestUpsert: - def test_upsert_replaces_and_prunes_stale_chunks(self, store) -> None: + def test_upsert_replaces_and_prunes_stale_chunks( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: store.add( [make_node("d1c1", "1"), make_node("d1c2", "1"), make_node("d2c1", "2")], ) @@ -240,18 +356,24 @@ class TestUpsert: result = _query(store, [0.0] * DIM, top_k=10) assert sorted(result.ids) == ["d1new", "d2c1"] - def test_upsert_creates_table_when_missing(self, store) -> None: + def test_upsert_creates_table_when_missing( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: store.upsert_document("1", [make_node("a1", "1")]) assert _query(store, [0.0] * DIM, top_k=1).ids == ["a1"] - def test_upsert_empty_nodes_removes_document(self, store) -> None: + def test_upsert_empty_nodes_removes_document( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: store.add([make_node("a1", "1"), make_node("b1", "2")]) store.upsert_document("1", []) assert _query(store, [0.0] * DIM, top_k=10).ids == ["b1"] def test_upsert_is_atomic_for_concurrent_readers( self, - store, + store: PaperlessSqliteVecVectorStore, tmp_path: Path, ) -> None: """A second connection must never observe document 1 half-replaced.""" @@ -267,38 +389,47 @@ class TestDocumentChunksIndex: ids without a vec0 full table scan on document_id -- see PaperlessSqliteVecVectorStore._delete_chunks_by_document_id.""" - def _chunk_index_rows(self, store) -> list[tuple[str, str]]: - rows = store.client.execute( - "SELECT chunk_id, document_id FROM document_chunks ORDER BY chunk_id", - ).fetchall() - return [(r["chunk_id"], r["document_id"]) for r in rows] - - def test_add_populates_document_chunks(self, store) -> None: + def test_add_populates_document_chunks( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: store.add([make_node("a1", "1"), make_node("a2", "1"), make_node("b1", "2")]) - assert self._chunk_index_rows(store) == [ + assert _chunk_index_rows(store) == [ ("a1", "1"), ("a2", "1"), ("b1", "2"), ] - def test_delete_clears_document_chunks(self, store) -> None: + def test_delete_clears_document_chunks( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: store.add([make_node("a1", "1"), make_node("a2", "1"), make_node("b1", "2")]) store.delete("1") - assert self._chunk_index_rows(store) == [("b1", "2")] + assert _chunk_index_rows(store) == [("b1", "2")] - def test_upsert_replaces_document_chunks(self, store) -> None: + def test_upsert_replaces_document_chunks( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: store.add([make_node("a1", "1"), make_node("b1", "2")]) store.upsert_document("1", [make_node("a2", "1")]) - assert self._chunk_index_rows(store) == [("a2", "1"), ("b1", "2")] + assert _chunk_index_rows(store) == [("a2", "1"), ("b1", "2")] - def test_drop_table_clears_document_chunks(self, store) -> None: + def test_drop_table_clears_document_chunks( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: store.add([make_node("a1", "1")]) store.drop_table() - assert self._chunk_index_rows(store) == [] + assert _chunk_index_rows(store) == [] class TestMetadataCoercion: - def test_none_metadata_values_become_empty_strings(self, store) -> None: + def test_none_metadata_values_become_empty_strings( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: node = make_node("a1", "1") node.metadata["modified"] = None store.add([node]) # must not raise (vec0 rejects NULL metadata) @@ -343,10 +474,16 @@ class TestModelNameTracking: class TestGetModifiedTimes: - def test_empty_store_returns_empty_dict(self, store) -> None: + def test_empty_store_returns_empty_dict( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: assert store.get_modified_times() == {} - def test_returns_one_entry_per_document(self, store) -> None: + def test_returns_one_entry_per_document( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: store.add( [ make_node("a1", "1", modified="2026-01-01T00:00:00"), @@ -361,7 +498,7 @@ class TestGetModifiedTimes: class TestCompact: - def _bloat_ratio(self, store) -> float: + def _bloat_ratio(self, store: PaperlessSqliteVecVectorStore) -> float: live = store.client.execute( "SELECT count(*) FROM documents", ).fetchone()[0] @@ -373,19 +510,25 @@ class TestCompact: total = int(row["value"]) if row else live return total / max(live, 1) - def _churn(self, store, cycles: int) -> None: + def _churn(self, store: PaperlessSqliteVecVectorStore, cycles: int) -> None: for i in range(cycles): store.upsert_document( "1", [make_node(f"gen{i}-{j}", "1", seed=float(j)) for j in range(20)], ) - def test_compact_noop_below_threshold(self, store) -> None: + def test_compact_noop_below_threshold( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: store.add([make_node("a1", "1")]) store.compact() assert _query(store, [0.0] * DIM, top_k=1).ids == ["a1"] - def test_force_compact_preserves_rows_and_metadata(self, store) -> None: + def test_force_compact_preserves_rows_and_metadata( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: store.add([make_node("a1", "1"), make_node("b1", "2", seed=3.0)]) self._churn(store, 5) before = { @@ -405,18 +548,27 @@ class TestCompact: store.upsert_document("3", [make_node("c1", "3", seed=100.0)]) assert "c1" in _query(store, [100.0] * DIM, top_k=1).ids - def test_auto_compact_triggers_on_churn(self, store) -> None: + def test_auto_compact_triggers_on_churn( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: store.add([make_node(f"s{j}", "1", seed=float(j)) for j in range(20)]) self._churn(store, 5) assert self._bloat_ratio(store) > 2 store.compact() assert self._bloat_ratio(store) == pytest.approx(1.0) - def test_compact_on_missing_table_is_noop(self, store) -> None: + def test_compact_on_missing_table_is_noop( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: store.compact() store.compact(force=True) - def test_compact_preserves_document_chunks(self, store) -> None: + def test_compact_preserves_document_chunks( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: """document_chunks must survive the file-swap rebuild, or delete() would silently stop finding chunk ids for anything indexed before a compaction ran.""" @@ -424,19 +576,16 @@ class TestCompact: self._churn(store, 5) store.compact(force=True) - rows = store.client.execute( - "SELECT chunk_id, document_id FROM document_chunks WHERE document_id = '2'", - ).fetchall() - assert [(r["chunk_id"], r["document_id"]) for r in rows] == [("b1", "2")] + assert _chunk_index_rows(store, "2") == [("b1", "2")] store.delete("2") assert "b1" not in _query(store, [0.0] * DIM, top_k=10).ids def test_failed_compact_removes_temp_wal_and_shm( self, - store, + store: PaperlessSqliteVecVectorStore, tmp_path: Path, - monkeypatch, + monkeypatch: pytest.MonkeyPatch, ) -> None: """A compact() that raises mid-rebuild must leave no .compact* files. @@ -471,8 +620,8 @@ class TestCompact: def test_force_compact_streams_rows_across_batches( self, - store, - monkeypatch, + store: PaperlessSqliteVecVectorStore, + monkeypatch: pytest.MonkeyPatch, ) -> None: """Rebuild must preserve every row when rows span multiple batches. @@ -488,11 +637,15 @@ class TestCompact: class TestDbFile: - def test_single_db_file_in_index_dir(self, store, tmp_path: Path) -> None: + def test_single_db_file_in_index_dir( + self, + store: PaperlessSqliteVecVectorStore, + tmp_path: Path, + ) -> None: store.add([make_node("a1", "1")]) assert (tmp_path / DB_FILENAME).exists() - def test_wal_mode_enabled(self, store) -> None: + def test_wal_mode_enabled(self, store: PaperlessSqliteVecVectorStore) -> None: assert ( store.client.execute("PRAGMA journal_mode").fetchone()[0].lower() == "wal" ) @@ -507,121 +660,76 @@ class TestMigrations: ).fetchone() return int(row[0]) if row else None - def test_new_table_records_schema_version(self, store) -> None: + def test_new_table_records_schema_version( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: store.add([make_node("a1", "1")]) assert self._schema_version(store) == SCHEMA_VERSION - def test_check_migrations_no_table_returns_false(self, store) -> None: + def test_check_migrations_no_table_returns_false( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: assert store.check_and_run_migrations() is False - def test_check_migrations_current_version_returns_false(self, store) -> None: + def test_check_migrations_current_version_returns_false( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: store.add([make_node("a1", "1")]) assert store.check_and_run_migrations() is False - def test_reembed_migration_returns_true(self, store, tmp_path: Path) -> None: + def test_reembed_migration_returns_true( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: store.add([make_node("a1", "1")]) # store was just created at the real (current) SCHEMA_VERSION, so the # simulated pending migration must target one version past that -- # not a hardcoded 1 -> 2, which would collide with the real # document_chunks migration already registered in MIGRATIONS. - from paperless_ai import vector_store as vs_mod - - original = vs_mod.SCHEMA_VERSION migration = Migration( - from_version=original, - to_version=original + 1, + from_version=SCHEMA_VERSION, + to_version=SCHEMA_VERSION + 1, kind="re-embed", description="test re-embed", ) - MIGRATIONS.append(migration) - try: - vs_mod.SCHEMA_VERSION = original + 1 - result = store.check_and_run_migrations() - finally: - MIGRATIONS.remove(migration) - vs_mod.SCHEMA_VERSION = original - assert result is True + with _pending_migrations(migration, schema_version=SCHEMA_VERSION + 1): + assert store.check_and_run_migrations() is True def test_structural_migration_copies_rows_and_updates_version( self, - store, - tmp_path: Path, + store: PaperlessSqliteVecVectorStore, ) -> None: store.add([make_node("a1", "1"), make_node("b1", "2")]) - - def apply( - src: sqlite3.Connection, - dst: sqlite3.Connection, - dim: int, - ) -> None: - dst.execute( # nosemgrep - f"CREATE VIRTUAL TABLE {DEFAULT_TABLE_NAME} USING vec0(" - "id TEXT PRIMARY KEY, document_id TEXT, modified TEXT," - f" +node_content TEXT, embedding float[{dim}] distance_metric=cosine" - ")", - ) - dst.execute( - "INSERT INTO index_meta (key, value) VALUES ('dim', ?) " - "ON CONFLICT(key) DO UPDATE SET value = excluded.value", - (str(dim),), - ) - rows = src.execute( - "SELECT id, document_id, modified, node_content, embedding " - f"FROM {DEFAULT_TABLE_NAME}", - ).fetchall() - dst.execute("BEGIN IMMEDIATE") - dst.executemany( - f"INSERT INTO {DEFAULT_TABLE_NAME} " - "(id, document_id, modified, node_content, embedding) " - "VALUES (?, ?, ?, ?, ?)", - [ - ( - r["id"], - r["document_id"], - r["modified"], - r["node_content"], - bytes(r["embedding"]), - ) - for r in rows - ], - ) - dst.execute( - "INSERT INTO index_meta (key, value) VALUES ('total_inserts', ?) " - "ON CONFLICT(key) DO UPDATE SET value = excluded.value", - (str(len(rows)),), - ) - dst.execute("COMMIT") - - from paperless_ai import vector_store as vs_mod - - original = vs_mod.SCHEMA_VERSION migration = Migration( - from_version=original, - to_version=original + 1, + from_version=SCHEMA_VERSION, + to_version=SCHEMA_VERSION + 1, kind="structural", description="test structural", - apply=apply, + apply=_copying_apply, ) - MIGRATIONS.append(migration) - try: - vs_mod.SCHEMA_VERSION = original + 1 - result = store.check_and_run_migrations() - finally: - MIGRATIONS.remove(migration) - vs_mod.SCHEMA_VERSION = original + with _pending_migrations(migration, schema_version=SCHEMA_VERSION + 1): + assert store.check_and_run_migrations() is False - assert result is False - assert self._schema_version(store) == original + 1 + assert self._schema_version(store) == SCHEMA_VERSION + 1 ids = {n.node_id for n in store.get_nodes()} assert ids == {"a1", "b1"} - def test_compact_preserves_schema_version(self, store) -> None: + def test_compact_preserves_schema_version( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: store.add([make_node("a1", "1")]) assert self._schema_version(store) == SCHEMA_VERSION store.compact(force=True) assert self._schema_version(store) == SCHEMA_VERSION - def test_v1_to_v2_migration_backfills_document_chunks(self, store) -> None: + def test_v1_to_v2_migration_backfills_document_chunks( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: """The real v1 -> v2 migration (registered in MIGRATIONS) must backfill document_chunks for rows written before it existed, or delete()/upsert_document() would find zero chunk ids for them and @@ -639,10 +747,7 @@ class TestMigrations: assert store.check_and_run_migrations() is False # structural, not re-embed assert self._schema_version(store) == SCHEMA_VERSION - rows = store.client.execute( - "SELECT chunk_id, document_id FROM document_chunks ORDER BY chunk_id", - ).fetchall() - assert [(r["chunk_id"], r["document_id"]) for r in rows] == [ + assert _chunk_index_rows(store) == [ ("a1", "1"), ("a2", "1"), ("b1", "2"), @@ -653,83 +758,38 @@ class TestMigrations: store.delete("1") assert sorted(_query(store, [0.0] * DIM, top_k=10).ids) == ["b1"] - def test_stop_at_reembed_boundary(self, store) -> None: - # Registry: structural v2, re-embed v3, structural v4. - # Only v2 should apply; the re-embed boundary must stop execution - # before v4 runs, and the stored version must stay at 2. + def test_stop_at_reembed_boundary( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: + # Registry, relative to the store's current (real) SCHEMA_VERSION N: + # structural N+1, re-embed N+2, structural N+3. Only N+1 should apply; + # the re-embed boundary must stop execution before N+3 runs, and the + # stored version must stay at N+1. store.add([make_node("a1", "1"), make_node("b1", "2")]) - - def copy_apply( - src: sqlite3.Connection, - dst: sqlite3.Connection, - dim: int, - ) -> None: - dst.execute( # nosemgrep - f"CREATE VIRTUAL TABLE {DEFAULT_TABLE_NAME} USING vec0(" - "id TEXT PRIMARY KEY, document_id TEXT, modified TEXT," - f" +node_content TEXT, embedding float[{dim}] distance_metric=cosine" - ")", - ) - dst.execute( - "INSERT INTO index_meta (key, value) VALUES ('dim', ?) " - "ON CONFLICT(key) DO UPDATE SET value = excluded.value", - (str(dim),), - ) - rows = src.execute( - "SELECT id, document_id, modified, node_content, embedding " - f"FROM {DEFAULT_TABLE_NAME}", - ).fetchall() - dst.execute("BEGIN IMMEDIATE") - dst.executemany( - f"INSERT INTO {DEFAULT_TABLE_NAME} " - "(id, document_id, modified, node_content, embedding) " - "VALUES (?, ?, ?, ?, ?)", - [ - ( - r["id"], - r["document_id"], - r["modified"], - r["node_content"], - bytes(r["embedding"]), - ) - for r in rows - ], - ) - dst.execute("COMMIT") - migrations = [ Migration( - from_version=1, - to_version=2, + from_version=SCHEMA_VERSION, + to_version=SCHEMA_VERSION + 1, kind="structural", - description="v2 structural", - apply=copy_apply, + description="structural", + apply=_copying_apply, ), Migration( - from_version=2, - to_version=3, + from_version=SCHEMA_VERSION + 1, + to_version=SCHEMA_VERSION + 2, kind="re-embed", - description="v3 re-embed boundary", + description="re-embed boundary", ), Migration( - from_version=3, - to_version=4, + from_version=SCHEMA_VERSION + 2, + to_version=SCHEMA_VERSION + 3, kind="structural", - description="v4 structural - must not run", - apply=copy_apply, + description="structural - must not run", + apply=_copying_apply, ), ] - MIGRATIONS.extend(migrations) - try: - from paperless_ai import vector_store as vs_mod + with _pending_migrations(*migrations, schema_version=SCHEMA_VERSION + 3): + assert store.check_and_run_migrations() is True - original = vs_mod.SCHEMA_VERSION - vs_mod.SCHEMA_VERSION = 4 - result = store.check_and_run_migrations() - finally: - for m in migrations: - MIGRATIONS.remove(m) - vs_mod.SCHEMA_VERSION = original - - assert result is True - assert self._schema_version(store) == 2 + assert self._schema_version(store) == SCHEMA_VERSION + 1 diff --git a/src/paperless_ai/vector_store.py b/src/paperless_ai/vector_store.py index 7b61e45e2..fc94fe6f2 100644 --- a/src/paperless_ai/vector_store.py +++ b/src/paperless_ai/vector_store.py @@ -105,6 +105,42 @@ def _unpack(blob: bytes) -> list[float]: return list(struct.unpack(f"{len(blob) // 4}f", blob)) +def _copy_rows(src_conn: sqlite3.Connection, dst_conn: sqlite3.Connection) -> int: + """Copy every live vec0 row from ``src_conn`` into ``dst_conn``, recording + each one in ``dst_conn``'s document_chunks side table. Returns the number + of rows copied. The caller owns ``dst_conn``'s transaction. + + Rows are streamed from the source cursor in batches instead of being + materialized all at once, so a large index does not cause an OOM during a + routine compaction or migration. + """ + src_cursor = src_conn.execute( + "SELECT id, document_id, modified, node_content, embedding FROM " + + DEFAULT_TABLE_NAME, + ) + copied = 0 + while batch := src_cursor.fetchmany(COMPACT_BATCH_SIZE): + dst_conn.executemany( + _INSERT, + [ + ( + r["id"], + r["document_id"], + r["modified"], + r["node_content"], + bytes(r["embedding"]), + ) + for r in batch + ], + ) + dst_conn.executemany( + _INSERT_CHUNK_INDEX, + [(r["id"], r["document_id"]) for r in batch], + ) + copied += len(batch) + return copied + + def _build_where(filters: MetadataFilters | None) -> tuple[str, list[str]]: """Translate the EQ / IN / NE filters we use into a parameterized SQL clause on vec0 metadata columns. Returns ("", []) when there is nothing to filter. @@ -247,13 +283,17 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): else: self._conn.execute("COMMIT") - def _meta_get(self, key: str) -> str | None: - row = self._conn.execute( + @staticmethod + def _meta_get_on(conn: sqlite3.Connection, key: str) -> str | None: + row = conn.execute( "SELECT value FROM index_meta WHERE key = ?", (key,), ).fetchone() return row["value"] if row else None + def _meta_get(self, key: str) -> str | None: + return self._meta_get_on(self._conn, key) + @staticmethod def _meta_set_on(conn: sqlite3.Connection, key: str, value: str) -> None: conn.execute( @@ -355,7 +395,7 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): side table, kept in lockstep with every insert into the vec0 table.""" self._conn.executemany( _INSERT_CHUNK_INDEX, - [(row[0], row[1]) for row in rows], + [(chunk_id, document_id) for chunk_id, document_id, *_ in rows], ) def _delete_chunks_by_document_id(self, document_id: str) -> None: @@ -368,21 +408,18 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): indexed lookup) and deleting each by its `id` primary key instead turns that scan into a handful of O(1) point deletes. """ - chunk_ids = [ - row[0] - for row in self._conn.execute( - "SELECT chunk_id FROM document_chunks WHERE document_id = ?", - (str(document_id),), - ) - ] - for chunk_id in chunk_ids: - self._conn.execute( - "DELETE FROM " + DEFAULT_TABLE_NAME + " WHERE id = ?", - (chunk_id,), - ) + doc_id = str(document_id) + chunk_rows = self._conn.execute( + "SELECT chunk_id FROM document_chunks WHERE document_id = ?", + (doc_id,), + ).fetchall() + self._conn.executemany( + "DELETE FROM " + DEFAULT_TABLE_NAME + " WHERE id = ?", + [(row["chunk_id"],) for row in chunk_rows], + ) self._conn.execute( "DELETE FROM document_chunks WHERE document_id = ?", - (str(document_id),), + (doc_id,), ) def _increment_total_inserts(self, count: int) -> None: @@ -558,32 +595,8 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): value = self._meta_get(key) if value is not None: self._meta_set_on(new_conn, key, value) - src_cursor = self._conn.execute( - "SELECT id, document_id, modified, node_content, embedding " - "FROM " + DEFAULT_TABLE_NAME, - ) new_conn.execute("BEGIN IMMEDIATE") - # Stream rows from the source cursor in batches instead of - # materializing the whole table in memory, so a large index does - # not cause an OOM during routine maintenance compactions. - while batch := src_cursor.fetchmany(COMPACT_BATCH_SIZE): - new_conn.executemany( - _INSERT, - [ - ( - r["id"], - r["document_id"], - r["modified"], - r["node_content"], - bytes(r["embedding"]), - ) - for r in batch - ], - ) - new_conn.executemany( - _INSERT_CHUNK_INDEX, - [(r["id"], r["document_id"]) for r in batch], - ) + _copy_rows(self._conn, new_conn) # Reset the cumulative counter: after compact, total_inserts == live. self._meta_set_on(new_conn, "total_inserts", str(live)) new_conn.execute("COMMIT") @@ -682,46 +695,18 @@ def _migrate_v1_to_v2_add_document_chunks( full table scan on the document_id metadata column. Every row written before this migration predates that table, so without backfilling, deleting a pre-migration document would find zero chunk ids and leave its - vec0 rows permanently orphaned. This copies all live rows into the new - file -- the same approach compact() uses -- indexing each as it goes. + vec0 rows permanently orphaned. Backfilling is a plain row copy into the + new-schema file (``_copy_rows``, the same helper compact() uses), which + records every copied row in document_chunks as it goes. """ PaperlessSqliteVecVectorStore._create_vec_table(dst_conn, dim) PaperlessSqliteVecVectorStore._meta_set_on(dst_conn, "dim", str(dim)) - embed_model_row = src_conn.execute( - "SELECT value FROM index_meta WHERE key = 'embed_model'", - ).fetchone() - if embed_model_row is not None: - PaperlessSqliteVecVectorStore._meta_set_on( - dst_conn, - "embed_model", - embed_model_row["value"], - ) + embed_model = PaperlessSqliteVecVectorStore._meta_get_on(src_conn, "embed_model") + if embed_model is not None: + PaperlessSqliteVecVectorStore._meta_set_on(dst_conn, "embed_model", embed_model) - src_cursor = src_conn.execute( - "SELECT id, document_id, modified, node_content, embedding FROM " - + DEFAULT_TABLE_NAME, - ) dst_conn.execute("BEGIN IMMEDIATE") - live = 0 - while batch := src_cursor.fetchmany(COMPACT_BATCH_SIZE): - dst_conn.executemany( - _INSERT, - [ - ( - r["id"], - r["document_id"], - r["modified"], - r["node_content"], - bytes(r["embedding"]), - ) - for r in batch - ], - ) - dst_conn.executemany( - _INSERT_CHUNK_INDEX, - [(r["id"], r["document_id"]) for r in batch], - ) - live += len(batch) + live = _copy_rows(src_conn, dst_conn) # This migration only ever copies live rows (like compact()), so the # cumulative counter resets to match -- the new file has no bloat yet. PaperlessSqliteVecVectorStore._meta_set_on(dst_conn, "total_inserts", str(live))