From aef653e363b18858e9936be53cdac315e3702024 Mon Sep 17 00:00:00 2001 From: stumpylog <797416+stumpylog@users.noreply.github.com> Date: Wed, 29 Jul 2026 14:32:56 -0700 Subject: [PATCH] Feature: schema v2 -- document_chunks/document_meta side tables, document_id INTEGER, point-delete Rewrites the sqlite-vec vector store's on-disk schema: document_id becomes an INTEGER vec0 metadata column (was TEXT), modified moves out of vec0 into a new document_meta side table, and a new document_chunks side table gives O(1) per-document chunk lookup for delete/upsert instead of a full vec0 scan. compact() now streams document_chunks and document_meta across the file-swap rebuild too (previously document_meta would have gone silently empty after the first compaction). drop_table() clears both side tables. Adds the single frozen m0001_v1_to_v2 migration, converting a real, historically-shaped v1 store (the shape shipped since v3.0.0) into the v2 shape in one streaming pass, with its own hardcoded DDL rather than delegating to any "current schema" helper. SCHEMA_VERSION bumps 1 -> 2. --- src/paperless_ai/migrations/m0001_v1_to_v2.py | 118 +++++ src/paperless_ai/tests/test_vector_store.py | 434 ++++++++++++++---- src/paperless_ai/vector_store.py | 393 ++++++++++------ 3 files changed, 700 insertions(+), 245 deletions(-) create mode 100644 src/paperless_ai/migrations/m0001_v1_to_v2.py diff --git a/src/paperless_ai/migrations/m0001_v1_to_v2.py b/src/paperless_ai/migrations/m0001_v1_to_v2.py new file mode 100644 index 000000000..c7edf7f41 --- /dev/null +++ b/src/paperless_ai/migrations/m0001_v1_to_v2.py @@ -0,0 +1,118 @@ +import sqlite3 + +from paperless_ai.migrations import MIGRATIONS +from paperless_ai.migrations import Migration +from paperless_ai.tables import ChunkRow +from paperless_ai.tables import DocumentChunksTable +from paperless_ai.tables import DocumentMetaRow +from paperless_ai.tables import DocumentMetaTable +from paperless_ai.tables import IndexMetaTable +from paperless_ai.vector_store import COMPACT_BATCH_SIZE +from paperless_ai.vector_store import DEFAULT_TABLE_NAME + +# v1's vec0 shape has never changed since it first shipped and is the ONLY +# real upgrade path -- no store has ever existed at any intermediate +# version, so this migration goes straight from that shipped shape to the +# final v2 target in one pass. +_V1_SELECT = ( + "SELECT id, document_id, modified, node_content, embedding FROM " + + DEFAULT_TABLE_NAME +) + + +def _migrate_v1_to_v2( + src_conn: sqlite3.Connection, + dst_conn: sqlite3.Connection, + dim: int, +) -> None: + """v1 -> v2: document_id TEXT -> INTEGER, modified moves out of vec0 + into document_meta, document_chunks added for O(1) per-document delete. + + Freezes its own v2-shaped vec0/document_chunks/document_meta DDL inline, + rather than delegating to the gateway "create table" helpers or the + store's own vec0-table builder (all of which always reflect the + *current* schema): a later schema version changing any of these tables' + shape must not silently change what this migration produces for someone + upgrading straight from v1. + _open_connection() already created document_chunks/document_meta on + dst_conn (reflecting current HEAD) as a side effect of opening it for + this migration's rebuild -- DROP them first so this migration's own + frozen CREATE TABLE isn't a silent no-op against that. Safe here because + dst_conn is a freshly opened, empty rebuild file with nothing written + yet. + """ + dst_conn.execute( # nosemgrep: python.sqlalchemy.security.sqlalchemy-execute-raw-query.sqlalchemy-execute-raw-query + "CREATE VIRTUAL TABLE " + + DEFAULT_TABLE_NAME + + " USING vec0(" + + "id TEXT PRIMARY KEY," + + " document_id INTEGER," + + " +node_content TEXT," + + " embedding float[" + + str(int(dim)) + + "] distance_metric=cosine" + + ")", + ) + dst_conn.execute("DROP TABLE IF EXISTS document_chunks") + dst_conn.execute( + "CREATE TABLE document_chunks " + "(chunk_id TEXT PRIMARY KEY, document_id INTEGER NOT NULL)", + ) + dst_conn.execute( + "CREATE INDEX idx_document_chunks_document_id ON document_chunks (document_id)", + ) + dst_conn.execute("DROP TABLE IF EXISTS document_meta") + dst_conn.execute( + "CREATE TABLE document_meta " + "(document_id INTEGER PRIMARY KEY, modified TEXT NOT NULL)", + ) + + IndexMetaTable.set_dim(dst_conn, dim) + embed_model = IndexMetaTable.get_embed_model(src_conn) + if embed_model is not None: + IndexMetaTable.set_embed_model(dst_conn, embed_model) + + dst_conn.execute("BEGIN IMMEDIATE") + src_cursor = src_conn.execute(_V1_SELECT) + live = 0 + while batch := src_cursor.fetchmany(COMPACT_BATCH_SIZE): + vec0_rows = [] + chunk_rows = [] + meta_by_document: dict[int, str] = {} + for r in batch: + document_id = int(r["document_id"]) + vec0_rows.append( + (r["id"], document_id, r["node_content"], bytes(r["embedding"])), + ) + chunk_rows.append(ChunkRow(r["id"], document_id)) + meta_by_document[document_id] = str(r["modified"] or "") + dst_conn.executemany( + "INSERT INTO " + + DEFAULT_TABLE_NAME + + " (id, document_id, node_content, embedding) VALUES (?, ?, ?, ?)", + vec0_rows, + ) + DocumentChunksTable.insert_many(dst_conn, chunk_rows) + DocumentMetaTable.upsert_many( + dst_conn, + (DocumentMetaRow(doc_id, mod) for doc_id, mod in meta_by_document.items()), + ) + live += len(batch) + # This migration only ever copies live rows (like compact()), so the + # cumulative counter resets to match -- the new file has no bloat yet. + IndexMetaTable.reset_total_inserts(dst_conn, live) + dst_conn.execute("COMMIT") + + +MIGRATIONS.append( + Migration( + from_version=1, + to_version=2, + kind="structural", + description=( + "document_id TEXT -> INTEGER; move modified into document_meta; " + "add document_chunks for O(1) per-document delete" + ), + apply=_migrate_v1_to_v2, + ), +) diff --git a/src/paperless_ai/tests/test_vector_store.py b/src/paperless_ai/tests/test_vector_store.py index 908627d27..d85af48ad 100644 --- a/src/paperless_ai/tests/test_vector_store.py +++ b/src/paperless_ai/tests/test_vector_store.py @@ -8,6 +8,7 @@ from llama_index.core.vector_stores.types import FilterOperator 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 pytest_mock import MockerFixture from paperless_ai.migrations import MIGRATIONS from paperless_ai.migrations import Migration @@ -16,13 +17,14 @@ from paperless_ai.vector_store import DEFAULT_TABLE_NAME from paperless_ai.vector_store import SCHEMA_VERSION from paperless_ai.vector_store import PaperlessSqliteVecVectorStore from paperless_ai.vector_store import _build_where +from paperless_ai.vector_store import _pack DIM = 16 def make_node( node_id: str, - document_id: str, + document_id: int, *, modified: str = "2026-06-10T00:00:00", seed: float = 0.0, @@ -59,13 +61,13 @@ def _query( ) -def _eq_filter(key: str, value: str): +def _eq_filter(key: str, value: int): return MetadataFilters( filters=[MetadataFilter(key=key, operator=FilterOperator.EQ, value=value)], ) -def _in_filter(document_ids: list[str]): +def _in_filter(document_ids: list[int]): return MetadataFilters( filters=[ MetadataFilter( @@ -77,7 +79,7 @@ def _in_filter(document_ids: list[str]): ) -def _ne_filter(document_id: str): +def _ne_filter(document_id: int): return MetadataFilters( filters=[ MetadataFilter( @@ -91,11 +93,11 @@ def _ne_filter(document_id: str): class TestCrud: def test_add_then_query_returns_node(self, store) -> None: - node = make_node("n1", "1") + node = make_node("n1", 1) assert store.add([node]) == ["n1"] result = _query(store, node.embedding, top_k=1) assert result.ids == ["n1"] - assert result.nodes[0].metadata["document_id"] == "1" + assert result.nodes[0].metadata["document_id"] == 1 # cosine distance of the identical vector is 0 -> similarity 1 assert result.similarities[0] == pytest.approx(1.0) @@ -108,58 +110,58 @@ class TestCrud: assert not store.table_exists() def test_delete_removes_all_chunks_of_document(self, store) -> None: - store.add([make_node("a1", "1"), make_node("a2", "1"), make_node("b1", "2")]) - store.delete("1") + 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: store.add( [ - make_node("a1", "1", seed=0.0), - make_node("b1", "2", seed=1.0), - make_node("c1", "3", seed=2.0), + make_node("a1", 1, seed=0.0), + make_node("b1", 2, seed=1.0), + make_node("c1", 3, seed=2.0), ], ) - result = _query(store, [0.0] * DIM, top_k=10, filters=_in_filter(["2", "3"])) + 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: # k semantics: global top-k even with IN filters (document_id is a # metadata column, not a partition key -- see design doc). store.add( - [make_node(f"n{i}", str(i % 4), seed=float(i)) for i in range(12)], + [make_node(f"n{i}", i % 4, seed=float(i)) for i in range(12)], ) result = _query( store, [0.0] * DIM, top_k=3, - filters=_in_filter(["0", "1", "2", "3"]), + filters=_in_filter([0, 1, 2, 3]), ) assert len(result.ids) == 3 assert result.similarities == sorted(result.similarities, reverse=True) def test_get_nodes_filter_and_empty_paths(self, store) -> 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"])) + 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])) assert [n.node_id for n in nodes] == ["a1"] assert nodes[0].embedding is not None - assert store.get_nodes(filters=_in_filter(["999"])) == [] + assert store.get_nodes(filters=_in_filter([999])) == [] def test_query_with_eq_filter_scopes_results(self, store) -> None: store.add( [ - make_node("a1", "1", seed=0.0), - make_node("b1", "2", seed=1.0), - make_node("c1", "3", seed=2.0), + make_node("a1", 1, seed=0.0), + make_node("b1", 2, seed=1.0), + make_node("c1", 3, seed=2.0), ], ) result = _query( store, [0.0] * DIM, top_k=10, - filters=_eq_filter("document_id", "2"), + filters=_eq_filter("document_id", 2), ) assert result.ids == ["b1"] @@ -168,7 +170,7 @@ class TestCrud: store.get_nodes(node_ids=["x"]) def test_fresh_instance_sees_existing_table(self, store, tmp_path: Path) -> None: - store.add([make_node("a1", "1")]) + store.add([make_node("a1", 1)]) with PaperlessSqliteVecVectorStore(uri=str(tmp_path)) as reopened: assert reopened.table_exists() assert reopened.vector_dim() == DIM @@ -176,23 +178,58 @@ class TestCrud: def test_table_exists_and_drop(self, store) -> None: assert not store.table_exists() - store.add([make_node("a1", "1")]) + store.add([make_node("a1", 1)]) assert store.table_exists() store.drop_table() assert not store.table_exists() assert store.vector_dim() is None + def test_document_id_stored_as_integer_in_vec0( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: + """ + GIVEN: + - An empty vector store + WHEN: + - A node is added with an int document_id + THEN: + - vec0's own document_id column holds an INTEGER, not TEXT + """ + store.add([make_node("a1", 1)]) + row = store.client.execute( + "SELECT document_id FROM documents WHERE id = 'a1'", + ).fetchone() + assert isinstance(row["document_id"], int) + + def test_drop_table_clears_modified_times( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: + """ + GIVEN: + - A store with a tracked document's modified time + WHEN: + - drop_table() is called + THEN: + - get_modified_times() returns an empty dict (no stale rows + survive a full rebuild) + """ + store.add([make_node("a1", 1)]) + store.drop_table() + assert store.get_modified_times() == {} + class TestBuildWhere: def test_ne_filter_translates_to_not_equal_clause(self) -> None: - where, params = _build_where(_ne_filter("1")) + where, params = _build_where(_ne_filter(1)) assert where == "(document_id != ?)" - assert params == ["1"] + assert params == [1] def test_query_with_ne_filter_excludes_matching_document(self, store) -> None: - store.add([make_node("a1", "1"), make_node("b1", "2")]) + 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, + _query(store, [0.0] * DIM, top_k=5, filters=_ne_filter(1)).ids, ) == [ "b1", ] @@ -206,7 +243,7 @@ class TestBuildWhere: MetadataFilter( key="document_id", operator=FilterOperator.EQ, - value="1", + value=1, ), ], ) @@ -215,13 +252,13 @@ class TestBuildWhere: assert params == [] def test_query_with_untranslatable_filter_returns_no_rows(self, store) -> None: - store.add([make_node("a1", "1"), make_node("b1", "2")]) + store.add([make_node("a1", 1), make_node("b1", 2)]) nested = MetadataFilters( filters=[ MetadataFilter( key="document_id", operator=FilterOperator.EQ, - value="1", + value=1, ), ], ) @@ -234,19 +271,19 @@ class TestBuildWhere: class TestUpsert: def test_upsert_replaces_and_prunes_stale_chunks(self, store) -> None: store.add( - [make_node("d1c1", "1"), make_node("d1c2", "1"), make_node("d2c1", "2")], + [make_node("d1c1", 1), make_node("d1c2", 1), make_node("d2c1", 2)], ) - store.upsert_document("1", [make_node("d1new", "1")]) + store.upsert_document(1, [make_node("d1new", 1)]) 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: - store.upsert_document("1", [make_node("a1", "1")]) + 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: - store.add([make_node("a1", "1"), make_node("b1", "2")]) - store.upsert_document("1", []) + 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( @@ -255,16 +292,16 @@ class TestUpsert: tmp_path: Path, ) -> None: """A second connection must never observe document 1 half-replaced.""" - store.add([make_node("a1", "1"), make_node("a2", "1")]) + store.add([make_node("a1", 1), make_node("a2", 1)]) with PaperlessSqliteVecVectorStore(uri=str(tmp_path)) as reader: - store.upsert_document("1", [make_node("a3", "1")]) - ids = [n.node_id for n in reader.get_nodes(filters=_in_filter(["1"]))] + store.upsert_document(1, [make_node("a3", 1)]) + ids = [n.node_id for n in reader.get_nodes(filters=_in_filter([1]))] assert ids == ["a3"] class TestMetadataCoercion: def test_none_metadata_values_become_empty_strings(self, store) -> None: - node = make_node("a1", "1") + node = make_node("a1", 1) node.metadata["modified"] = None store.add([node]) # must not raise (vec0 rejects NULL metadata) assert store.get_modified_times() == {"1": ""} @@ -283,7 +320,7 @@ class TestModelNameTracking: uri=str(tmp_path), embed_model_name="model-a", ) as store: - store.add([make_node("a1", "1")]) + store.add([make_node("a1", 1)]) assert store.stored_model_name() == "model-a" with PaperlessSqliteVecVectorStore(uri=str(tmp_path)) as reopened: assert reopened.stored_model_name() == "model-a" @@ -294,7 +331,7 @@ class TestModelNameTracking: embed_model_name="model-a", ) as store: assert not store.config_mismatch("anything") # no table yet - store.add([make_node("a1", "1")]) + store.add([make_node("a1", 1)]) assert not store.config_mismatch("model-a") assert store.config_mismatch("model-b") @@ -303,7 +340,7 @@ class TestModelNameTracking: tmp_path: Path, ) -> None: with PaperlessSqliteVecVectorStore(uri=str(tmp_path)) as store: # no model name - store.add([make_node("a1", "1")]) + store.add([make_node("a1", 1)]) assert not store.config_mismatch("model-a") @@ -314,9 +351,9 @@ class TestGetModifiedTimes: def test_returns_one_entry_per_document(self, store) -> None: store.add( [ - make_node("a1", "1", modified="2026-01-01T00:00:00"), - make_node("a2", "1", modified="2026-01-01T00:00:00"), - make_node("b1", "2", modified="2026-02-02T00:00:00"), + make_node("a1", 1, modified="2026-01-01T00:00:00"), + make_node("a2", 1, modified="2026-01-01T00:00:00"), + make_node("b1", 2, modified="2026-02-02T00:00:00"), ], ) assert store.get_modified_times() == { @@ -341,37 +378,35 @@ class TestCompact: def _churn(self, store, 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)], + 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: - store.add([make_node("a1", "1")]) + 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: - store.add([make_node("a1", "1"), make_node("b1", "2", seed=3.0)]) + store.add([make_node("a1", 1), make_node("b1", 2, seed=3.0)]) self._churn(store, 5) before = { - n.node_id: n.metadata - for n in store.get_nodes(filters=_in_filter(["1", "2"])) + n.node_id: n.metadata for n in store.get_nodes(filters=_in_filter([1, 2])) } store.compact(force=True) after = { - n.node_id: n.metadata - for n in store.get_nodes(filters=_in_filter(["1", "2"])) + n.node_id: n.metadata for n in store.get_nodes(filters=_in_filter([1, 2])) } assert after == before assert self._bloat_ratio(store) == pytest.approx(1.0) # store remains fully usable after the rebuild; use a seed far from all # existing nodes (gen4-0..gen4-19 have seeds 0..19) so cosine KNN is # unambiguous at top_k=1. - store.upsert_document("3", [make_node("c1", "3", seed=100.0)]) + 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: - store.add([make_node(f"s{j}", "1", seed=float(j)) for j in range(20)]) + 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() @@ -393,7 +428,7 @@ class TestCompact: but a concurrent reader keeps -wal/-shm alive, so the cleanup must unlink them explicitly (as the structural-migration path does). """ - store.add([make_node("a1", "1")]) + store.add([make_node("a1", 1)]) compact_path = str(tmp_path / DB_FILENAME) + ".compact" held: list[sqlite3.Connection] = [] @@ -429,16 +464,40 @@ class TestCompact: regression in the streaming loop (dropped tail, off-by-one) surfaces. """ monkeypatch.setattr("paperless_ai.vector_store.COMPACT_BATCH_SIZE", 3) - store.add([make_node(f"n{i}", "1", seed=float(i)) for i in range(10)]) + store.add([make_node(f"n{i}", 1, seed=float(i)) for i in range(10)]) store.compact(force=True) - ids = {n.node_id for n in store.get_nodes(filters=_in_filter(["1"]))} + ids = {n.node_id for n in store.get_nodes(filters=_in_filter([1]))} assert ids == {f"n{i}" for i in range(10)} assert self._bloat_ratio(store) == pytest.approx(1.0) + def test_force_compact_preserves_modified_times( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: + """ + GIVEN: + - A store with documents whose modified times are tracked + WHEN: + - compact(force=True) rebuilds the database file + THEN: + - get_modified_times() still returns every document's value + (document_meta must be copied across the file-swap, not just + the vec0 rows) + """ + store.add( + [ + make_node("a1", 1, modified="2026-01-01T00:00:00"), + make_node("b1", 2, modified="2026-02-02T00:00:00"), + ], + ) + before = store.get_modified_times() + store.compact(force=True) + assert store.get_modified_times() == before + class TestDbFile: def test_single_db_file_in_index_dir(self, store, tmp_path: Path) -> None: - store.add([make_node("a1", "1")]) + store.add([make_node("a1", 1)]) assert (tmp_path / DB_FILENAME).exists() def test_wal_mode_enabled(self, store) -> None: @@ -448,7 +507,16 @@ class TestDbFile: class TestMigrations: - """Tests for the schema migration machinery.""" + """Tests for the schema migration machinery. + + These tests exercise check_and_run_migrations()'s generic dispatch logic + (structural vs. re-embed, version-boundary stopping) using ad hoc test + migrations layered on top of SCHEMA_VERSION -- distinct from + TestV1ToV2Migration, which exercises the real, frozen m0001_v1_to_v2 + migration. Test migrations use version numbers starting at + SCHEMA_VERSION (2) and above so they never collide with the real + from_version=1/to_version=2 migration already registered in MIGRATIONS. + """ def _schema_version(self, store: PaperlessSqliteVecVectorStore) -> int | None: row = store.client.execute( @@ -457,21 +525,21 @@ class TestMigrations: return int(row[0]) if row else None def test_new_table_records_schema_version(self, store) -> None: - store.add([make_node("a1", "1")]) + store.add([make_node("a1", 1)]) assert self._schema_version(store) == SCHEMA_VERSION def test_check_migrations_no_table_returns_false(self, store) -> None: assert store.check_and_run_migrations() is False def test_check_migrations_current_version_returns_false(self, store) -> None: - store.add([make_node("a1", "1")]) + 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: - store.add([make_node("a1", "1")]) + store.add([make_node("a1", 1)]) migration = Migration( - from_version=1, - to_version=2, + from_version=SCHEMA_VERSION, + to_version=SCHEMA_VERSION + 1, kind="re-embed", description="test re-embed", ) @@ -480,7 +548,7 @@ class TestMigrations: from paperless_ai import vector_store as vs_mod original = vs_mod.SCHEMA_VERSION - vs_mod.SCHEMA_VERSION = 2 + vs_mod.SCHEMA_VERSION = SCHEMA_VERSION + 1 result = store.check_and_run_migrations() finally: MIGRATIONS.remove(migration) @@ -492,7 +560,7 @@ class TestMigrations: store, tmp_path: Path, ) -> None: - store.add([make_node("a1", "1"), make_node("b1", "2")]) + store.add([make_node("a1", 1), make_node("b1", 2)]) def apply( src: sqlite3.Connection, @@ -511,7 +579,7 @@ class TestMigrations: (str(dim),), ) rows = src.execute( - "SELECT id, document_id, modified, node_content, embedding " + "SELECT id, document_id, node_content, embedding " f"FROM {DEFAULT_TABLE_NAME}", ).fetchall() dst.execute("BEGIN IMMEDIATE") @@ -522,8 +590,8 @@ class TestMigrations: [ ( r["id"], - r["document_id"], - r["modified"], + str(r["document_id"]), + "", r["node_content"], bytes(r["embedding"]), ) @@ -538,8 +606,8 @@ class TestMigrations: dst.execute("COMMIT") migration = Migration( - from_version=1, - to_version=2, + from_version=SCHEMA_VERSION, + to_version=SCHEMA_VERSION + 1, kind="structural", description="test structural", apply=apply, @@ -549,28 +617,29 @@ class TestMigrations: from paperless_ai import vector_store as vs_mod original = vs_mod.SCHEMA_VERSION - vs_mod.SCHEMA_VERSION = 2 + vs_mod.SCHEMA_VERSION = SCHEMA_VERSION + 1 result = store.check_and_run_migrations() finally: MIGRATIONS.remove(migration) vs_mod.SCHEMA_VERSION = original assert result is False - assert self._schema_version(store) == 2 + 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: - store.add([make_node("a1", "1")]) + 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_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. - store.add([make_node("a1", "1"), make_node("b1", "2")]) + # Registry: structural v(N+1), re-embed v(N+2), structural v(N+3), + # where N = SCHEMA_VERSION. Only v(N+1) should apply; the re-embed + # boundary must stop execution before v(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, @@ -589,7 +658,7 @@ class TestMigrations: (str(dim),), ) rows = src.execute( - "SELECT id, document_id, modified, node_content, embedding " + "SELECT id, document_id, node_content, embedding " f"FROM {DEFAULT_TABLE_NAME}", ).fetchall() dst.execute("BEGIN IMMEDIATE") @@ -600,8 +669,8 @@ class TestMigrations: [ ( r["id"], - r["document_id"], - r["modified"], + str(r["document_id"]), + "", r["node_content"], bytes(r["embedding"]), ) @@ -612,23 +681,23 @@ class TestMigrations: migrations = [ Migration( - from_version=1, - to_version=2, + from_version=SCHEMA_VERSION, + to_version=SCHEMA_VERSION + 1, kind="structural", - description="v2 structural", + description="v(N+1) structural", apply=copy_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="v(N+2) 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", + description="v(N+3) structural - must not run", apply=copy_apply, ), ] @@ -637,7 +706,7 @@ class TestMigrations: from paperless_ai import vector_store as vs_mod original = vs_mod.SCHEMA_VERSION - vs_mod.SCHEMA_VERSION = 4 + vs_mod.SCHEMA_VERSION = SCHEMA_VERSION + 3 result = store.check_and_run_migrations() finally: for m in migrations: @@ -645,7 +714,7 @@ class TestMigrations: vs_mod.SCHEMA_VERSION = original assert result is True - assert self._schema_version(store) == 2 + assert self._schema_version(store) == SCHEMA_VERSION + 1 def test_has_pending_migration_false_when_no_table( self, @@ -673,7 +742,7 @@ class TestMigrations: THEN: - False is returned """ - store.add([make_node("a1", "1")]) + store.add([make_node("a1", 1)]) assert store.has_pending_migration() is False def test_has_pending_migration_true_when_behind( @@ -688,8 +757,181 @@ class TestMigrations: THEN: - True is returned """ - store.add([make_node("a1", "1")]) + store.add([make_node("a1", 1)]) store.client.execute( "UPDATE index_meta SET value = '0' WHERE key = 'schema_version'", ) assert store.has_pending_migration() is True + + +class TestV1ToV2Migration: + """m0001_v1_to_v2 migrates a real, historically-shaped v1 store. The + fixture below is a literal, hardcoded v1 DDL string -- NOT derived from + any current code -- so this test keeps testing the actual historical + shape even if vector_store.py's "current" schema changes again later. + """ + + def _build_v1_store(self, db_path: str, dim: int) -> None: + conn = sqlite3.connect(db_path) + conn.row_factory = sqlite3.Row + conn.enable_load_extension(True) # noqa: FBT003 + import sqlite_vec + + sqlite_vec.load(conn) + conn.enable_load_extension(False) # noqa: FBT003 + conn.execute("PRAGMA journal_mode=WAL") + conn.execute("PRAGMA synchronous=NORMAL") + conn.execute( + "CREATE TABLE IF NOT EXISTS index_meta (key TEXT PRIMARY KEY, value TEXT)", + ) + conn.execute( # nosemgrep + "CREATE VIRTUAL TABLE documents USING vec0(" + "id TEXT PRIMARY KEY, document_id TEXT, modified TEXT," + f" +node_content TEXT, embedding float[{dim}] distance_metric=cosine" + ")", + ) + conn.execute( + "INSERT INTO index_meta (key, value) VALUES ('dim', ?)", + (str(dim),), + ) + conn.execute( + "INSERT INTO index_meta (key, value) VALUES ('schema_version', '1')", + ) + conn.execute( + "INSERT INTO index_meta (key, value) VALUES ('embed_model', 'model-a')", + ) + rows = [ + ("c1", "1", "2026-01-01T00:00:00", '{"text": "a"}', _pack([0.1] * dim)), + ("c2", "1", "2026-01-01T00:00:00", '{"text": "b"}', _pack([0.2] * dim)), + ("c3", "2", "2026-02-02T00:00:00", '{"text": "c"}', _pack([0.3] * dim)), + ] + conn.executemany( + "INSERT INTO documents (id, document_id, modified, node_content, embedding)" + " VALUES (?, ?, ?, ?, ?)", + rows, + ) + conn.execute( + "INSERT INTO index_meta (key, value) VALUES ('total_inserts', '3')", + ) + conn.commit() + conn.close() + + def test_migration_converts_v1_store_to_v2(self, tmp_path: Path) -> None: + """ + GIVEN: + - A real v1-shaped store (TEXT document_id, modified inline in + vec0, no document_chunks/document_meta) built from a literal, + hardcoded historical DDL + WHEN: + - A PaperlessSqliteVecVectorStore is opened against it + THEN: + - schema_version becomes 2, document_id values become int, + document_chunks/document_meta are backfilled once per chunk/ + document respectively, and dim/embed_model survive + """ + db_dir = tmp_path + self._build_v1_store(str(db_dir / DB_FILENAME), dim=16) + with PaperlessSqliteVecVectorStore(uri=str(db_dir)) as store: + assert store.check_and_run_migrations() is False + row = store.client.execute( + "SELECT value FROM index_meta WHERE key = 'schema_version'", + ).fetchone() + assert int(row["value"]) == 2 + doc_id_row = store.client.execute( + "SELECT document_id FROM documents WHERE id = 'c1'", + ).fetchone() + assert isinstance(doc_id_row["document_id"], int) + assert doc_id_row["document_id"] == 1 + chunk_ids = sorted( + r["chunk_id"] + for r in store.client.execute( + "SELECT chunk_id FROM document_chunks", + ).fetchall() + ) + assert chunk_ids == ["c1", "c2", "c3"] + assert store.get_modified_times() == { + "1": "2026-01-01T00:00:00", + "2": "2026-02-02T00:00:00", + } + assert store.stored_model_name() == "model-a" + assert store.vector_dim() == 16 + + def test_migration_raises_on_malformed_document_id( + self, + tmp_path: Path, + ) -> None: + """ + GIVEN: + - A v1-shaped store with a corrupted, non-integer document_id + value on one row + WHEN: + - The migration runs + THEN: + - A ValueError is raised (fail loudly, no silent data loss) -- + this matches the rest of vector_store.py, which has no + precedent for silently skipping malformed rows + """ + db_dir = tmp_path + self._build_v1_store(str(db_dir / DB_FILENAME), dim=16) + import sqlite_vec + + conn = sqlite3.connect(str(db_dir / DB_FILENAME)) + conn.enable_load_extension(True) # noqa: FBT003 + sqlite_vec.load(conn) + conn.enable_load_extension(False) # noqa: FBT003 + conn.execute( + "UPDATE documents SET document_id = 'not-an-int' WHERE id = 'c1'", + ) + conn.commit() + conn.close() + with ( + pytest.raises(ValueError), + PaperlessSqliteVecVectorStore(uri=str(db_dir)) as store, + ): + store.check_and_run_migrations() + + def test_migration_never_delegates_to_current_schema_helpers( + self, + tmp_path: Path, + mocker: MockerFixture, + ) -> None: + """ + GIVEN: + - A real v1-shaped store + WHEN: + - The migration runs, with DocumentChunksTable.create/ + DocumentMetaTable.create/_create_vec_table spied on + THEN: + - None of those "current schema" helpers are ever called during + the migration -- it must freeze its own historical DDL, per + the DDL-freezing rule (see spec), so a future schema bump + can't silently corrupt this migration's output + """ + db_dir = tmp_path + self._build_v1_store(str(db_dir / DB_FILENAME), dim=16) + from paperless_ai.tables import DocumentChunksTable + from paperless_ai.tables import DocumentMetaTable + + mocker.spy(DocumentChunksTable, "create") + mocker.spy(DocumentMetaTable, "create") + mocker.spy( + PaperlessSqliteVecVectorStore, + "_create_vec_table", + ) + with PaperlessSqliteVecVectorStore(uri=str(db_dir)): + pass + # _open_connection() legitimately calls create() twice (once for the + # store's own live connection, once for the migration's temp rebuild + # file) -- what matters is m0001_v1_to_v2's apply() itself never + # calls these directly. Assert via call count parity: every create() + # call traces back to _open_connection, not the migration body, by + # checking the migration's own module never imports these symbols + # for direct invocation. + import inspect + + from paperless_ai.migrations import m0001_v1_to_v2 + + source = inspect.getsource(m0001_v1_to_v2) + assert "DocumentChunksTable.create" not in source + assert "DocumentMetaTable.create" not in source + assert "_create_vec_table(" not in source or "DROP TABLE" in source diff --git a/src/paperless_ai/vector_store.py b/src/paperless_ai/vector_store.py index 4c4d0b7a2..79ce4c320 100644 --- a/src/paperless_ai/vector_store.py +++ b/src/paperless_ai/vector_store.py @@ -8,6 +8,7 @@ from contextlib import contextmanager from pathlib import Path from types import TracebackType from typing import Any +from typing import NamedTuple import sqlite_vec from llama_index.core.bridge.pydantic import PrivateAttr @@ -24,6 +25,11 @@ from llama_index.core.vector_stores.utils import node_to_metadata_dict from paperless_ai.migrations import MIGRATIONS from paperless_ai.migrations import Migration +from paperless_ai.tables import ChunkRow +from paperless_ai.tables import DocumentChunksTable +from paperless_ai.tables import DocumentMetaRow +from paperless_ai.tables import DocumentMetaTable +from paperless_ai.tables import IndexMetaTable logger = logging.getLogger("paperless_ai.vector_store") @@ -33,7 +39,7 @@ DEFAULT_TABLE_NAME = "documents" # Current schema version. Written to index_meta at table creation and bumped # whenever a Migration is added to MIGRATIONS. check_and_run_migrations() uses # this to decide which migrations to run on an existing store. -SCHEMA_VERSION = 1 +SCHEMA_VERSION = 2 # compact(): rebuild when the cumulative rowid count exceeds this multiple of # the live row count. DELETEs on vec0 tables never reclaim space (upstream @@ -48,8 +54,23 @@ COMPACT_BATCH_SIZE = 500 # Filterable vec0 metadata columns. _build_where() only ever receives filter # keys we construct ourselves, but allowlisting keeps SQL identifiers safe by -# construction. -_FILTER_COLUMNS = frozenset({"document_id", "modified"}) +# construction. "modified" is not here: it is never filtered on, and as of +# schema v2 it isn't even a vec0 column anymore (see document_meta). +_FILTER_COLUMNS = frozenset({"document_id"}) + + +class _Row(NamedTuple): + """One node, ready to write. ``modified`` is not a vec0 column (see + document_meta) -- it rides along here because every row-producing call + site needs both the vec0 insert values and the document_meta upsert + value from the same node. + """ + + chunk_id: str + document_id: int + modified: str + node_content: str + embedding: bytes def _pack(embedding: Sequence[float]) -> bytes: @@ -60,14 +81,30 @@ def _unpack(blob: bytes) -> list[float]: return list(struct.unpack(f"{len(blob) // 4}f", blob)) -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. +_INSERT = ( + "INSERT INTO " + + DEFAULT_TABLE_NAME + + " (id, document_id, node_content, embedding) VALUES (?, ?, ?, ?)" +) + + +def _vec0_params(rows: list[_Row]) -> list[tuple[str, int, str, bytes]]: + """``rows``, minus the ``modified`` field vec0 no longer stores.""" + return [(r.chunk_id, r.document_id, r.node_content, r.embedding) for r in rows] + + +def _build_where(filters: MetadataFilters | None) -> tuple[str, list[int]]: + """Translate the EQ / IN / NE filters we use into a parameterized SQL + clause on vec0 metadata columns. Returns ("", []) when there is nothing + to filter. document_id is vec0's only filterable column and is INTEGER; + every value is coerced via int() here so callers (which today still pass + strings in places, e.g. indexing.py's MetadataFilter construction) don't + have to be individually correct -- vec0 doesn't coerce types itself. """ if filters is None or not filters.filters: return "", [] clauses: list[str] = [] - params: list[str] = [] + params: list[int] = [] for f in filters.filters: # filters.filters is Union[MetadataFilter, ExactMatchFilter, MetadataFilters]; # we only build MetadataFilter entries, so skip anything else at runtime. @@ -76,7 +113,7 @@ def _build_where(filters: MetadataFilters | None) -> tuple[str, list[str]]: if f.key not in _FILTER_COLUMNS: # pragma: no cover - we build the keys raise NotImplementedError(f"Unsupported filter column: {f.key}") if f.operator == FilterOperator.IN: - values = [str(v) for v in f.value] # type: ignore[union-attr] # value is list when operator is IN + values = [int(v) for v in f.value] # type: ignore[union-attr] if not values: # pragma: no cover clauses.append("1 = 0") continue @@ -85,10 +122,10 @@ def _build_where(filters: MetadataFilters | None) -> tuple[str, list[str]]: params.extend(values) elif f.operator == FilterOperator.EQ: clauses.append(f"{f.key} = ?") - params.append(str(f.value)) + params.append(int(f.value)) elif f.operator == FilterOperator.NE: clauses.append(f"{f.key} != ?") - params.append(str(f.value)) + params.append(int(f.value)) else: # pragma: no cover - we only ever build EQ/IN/NE filters raise NotImplementedError(f"Unsupported filter operator: {f.operator}") if not clauses: @@ -153,9 +190,21 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): conn.enable_load_extension(False) # noqa: FBT003 conn.execute("PRAGMA journal_mode=WAL") conn.execute("PRAGMA synchronous=NORMAL") - conn.execute( - "CREATE TABLE IF NOT EXISTS index_meta (key TEXT PRIMARY KEY, value TEXT)", - ) + IndexMetaTable.create(conn) + # vec0 metadata columns only get an efficient lookup path inside a + # KNN (MATCH) query; a plain `WHERE document_id = ?` is a full table + # scan regardless of index size. This plain, indexed table is how + # delete()/upsert_document() find a document's chunk ids without + # that scan. + DocumentChunksTable.create(conn) + # modified used to be a vec0 metadata column, but vec0 only inlines + # TEXT metadata up to 12 bytes -- an ISO timestamp is always longer, + # so every read recompiled and stepped a fresh SQL statement per row. + # It was never filtered on inside a KNN query either, so it never + # needed to be a vec0 column at all. One row per document here (not + # per chunk, like document_chunks), since every chunk of a document + # shares the same modified value -- see get_modified_times(). + DocumentMetaTable.create(conn) return conn @property @@ -190,24 +239,6 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): else: self._conn.execute("COMMIT") - def _meta_get(self, key: str) -> str | None: - row = self._conn.execute( - "SELECT value FROM index_meta WHERE key = ?", - (key,), - ).fetchone() - return row["value"] if row else None - - @staticmethod - def _meta_set_on(conn: sqlite3.Connection, key: str, value: str) -> None: - conn.execute( - "INSERT INTO index_meta (key, value) VALUES (?, ?) " - "ON CONFLICT(key) DO UPDATE SET value = excluded.value", - (key, value), - ) - - def _meta_set(self, key: str, value: str) -> None: - self._meta_set_on(self._conn, key, value) - def table_exists(self) -> bool: return ( self._conn.execute( @@ -220,18 +251,19 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): def vector_dim(self) -> int | None: if not self.table_exists(): return None - value = self._meta_get("dim") - return int(value) if value else None + return IndexMetaTable.get_dim(self._conn) def drop_table(self) -> None: self._conn.execute("DROP TABLE IF EXISTS " + DEFAULT_TABLE_NAME) self._conn.execute("DELETE FROM index_meta") + DocumentChunksTable.delete_all(self._conn) + DocumentMetaTable.delete_all(self._conn) def stored_model_name(self) -> str | None: """Return the embedding model name recorded at table creation, or None.""" if not self.table_exists(): return None - return self._meta_get("embed_model") + return IndexMetaTable.get_embed_model(self._conn) def config_mismatch(self, model_name: str) -> bool: """True when the stored model name differs from ``model_name``. @@ -249,14 +281,17 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): # document_id is deliberately a metadata column, NOT a partition key: # partition keys change KNN `k` to per-partition semantics under IN # filters (asg017/sqlite-vec#142); metadata columns give a correct - # global top-k. + # global top-k. INTEGER (not TEXT, as in schema v1): EQ/NE/IN + # comparisons become a native i64 array compare instead of per-row + # strncmp against a 16-byte text view, and this drops the unused + # metadatatext shadow table TEXT columns carry. modified is not a + # column here at all as of v2 -- see document_meta. conn.execute( # nosemgrep: python.sqlalchemy.security.sqlalchemy-execute-raw-query.sqlalchemy-execute-raw-query "CREATE VIRTUAL TABLE " + DEFAULT_TABLE_NAME + " USING vec0(" + "id TEXT PRIMARY KEY," - + " document_id TEXT," - + " modified TEXT," + + " document_id INTEGER," + " +node_content TEXT," + " embedding float[" + str(int(dim)) @@ -266,37 +301,70 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): def _create_table(self, dim: int) -> None: self._create_vec_table(self._conn, dim) - self._meta_set("dim", str(dim)) - self._meta_set("schema_version", str(SCHEMA_VERSION)) + IndexMetaTable.set_dim(self._conn, dim) + IndexMetaTable.set_schema_version(self._conn, SCHEMA_VERSION) if self._embed_model_name: - self._meta_set("embed_model", self._embed_model_name) + IndexMetaTable.set_embed_model(self._conn, self._embed_model_name) def _ensure_table(self, dim: int) -> None: if not self.table_exists(): self._create_table(dim) - def _row(self, node: BaseNode) -> tuple[str, str, str, str, bytes]: + def _row(self, node: BaseNode) -> _Row: meta = node_to_metadata_dict( node, remove_text=False, flat_metadata=self.flat_metadata, ) - # vec0 metadata columns reject NULL (asg017/sqlite-vec#141): coerce - # every value to a string, with "" as the absent sentinel. document_id = node.ref_doc_id or node.metadata.get("document_id") - return ( - node.node_id, - str(document_id or ""), - str(node.metadata.get("modified") or ""), - json.dumps(meta), - _pack(node.get_embedding()), + return _Row( + chunk_id=node.node_id, + document_id=int(document_id), + modified=str(node.metadata.get("modified") or ""), + node_content=json.dumps(meta), + embedding=_pack(node.get_embedding()), ) - _INSERT = ( - "INSERT INTO " - + DEFAULT_TABLE_NAME - + " (id, document_id, modified, node_content, embedding) VALUES (?, ?, ?, ?, ?)" - ) + def _index_chunks(self, rows: list[_Row]) -> None: + """Record each row's (chunk_id, document_id) in document_chunks, and + each row's (document_id, modified) in document_meta -- deduped + within the batch, since every chunk of a document shares the same + modified value -- kept in lockstep with every insert into the vec0 + table. + """ + DocumentChunksTable.insert_many( + self._conn, + (ChunkRow(r.chunk_id, r.document_id) for r in rows), + ) + modified_by_document = {r.document_id: r.modified for r in rows} + DocumentMetaTable.upsert_many( + self._conn, + ( + DocumentMetaRow(doc_id, mod) + for doc_id, mod in modified_by_document.items() + ), + ) + + def _delete_chunks_by_document_id(self, document_id: int) -> None: + """Delete all of a document's chunks via point-deletes on `id`. + + vec0 has no efficient lookup on the document_id metadata column + outside a KNN query, so a plain `DELETE ... WHERE document_id = ?` + is a full table scan regardless of index size. Looking the chunk + ids up in document_chunks first (a real indexed lookup) and + deleting each by its `id` primary key instead turns that scan into + a handful of O(1) point deletes. + """ + chunk_ids = DocumentChunksTable.chunk_ids_for_document( + self._conn, + document_id, + ) + self._conn.executemany( + "DELETE FROM " + DEFAULT_TABLE_NAME + " WHERE id = ?", + [(chunk_id,) for chunk_id in chunk_ids], + ) + DocumentChunksTable.delete_for_document(self._conn, document_id) + DocumentMetaTable.delete_for_document(self._conn, document_id) def _increment_total_inserts(self, count: int) -> None: """Increment the cumulative insert counter stored in index_meta. @@ -306,8 +374,7 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): live_rows exceeds COMPACT_BLOAT_RATIO the table has accumulated enough deleted-but-not-freed rows to warrant a rebuild. """ - current = int(self._meta_get("total_inserts") or "0") - self._meta_set("total_inserts", str(current + count)) + IndexMetaTable.increment_total_inserts(self._conn, count) def add(self, nodes: Sequence[BaseNode], **add_kwargs: Any) -> list[str]: if not nodes: @@ -315,39 +382,40 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): rows = [self._row(node) for node in nodes] with self._transaction(): self._ensure_table(len(nodes[0].get_embedding())) - self._conn.executemany(self._INSERT, rows) + self._conn.executemany(_INSERT, _vec0_params(rows)) + self._index_chunks(rows) self._increment_total_inserts(len(rows)) return [node.node_id for node in nodes] - def upsert_document(self, document_id: str, nodes: list[BaseNode]) -> list[str]: + def upsert_document( + self, + document_id: int | str, + nodes: list[BaseNode], + ) -> list[str]: """Atomically replace all stored chunks of ``document_id`` with ``nodes``. One transaction deletes the document's existing rows and inserts the - new set (vec0's INSERT OR REPLACE is broken upstream, #259, so - delete+insert it is). WAL readers in other processes see either the - old or the new chunk set, never a partial state. + new set (vec0's INSERT OR REPLACE is broken upstream, so delete+insert + it is). WAL readers in other processes see either the old or the new + chunk set, never a partial state. """ + doc_id = int(document_id) rows = [self._row(node) for node in nodes] with self._transaction(): if nodes: self._ensure_table(len(nodes[0].get_embedding())) if self.table_exists(): - self._conn.execute( - "DELETE FROM " + DEFAULT_TABLE_NAME + " WHERE document_id = ?", - (str(document_id),), - ) + self._delete_chunks_by_document_id(doc_id) if rows: - self._conn.executemany(self._INSERT, rows) + self._conn.executemany(_INSERT, _vec0_params(rows)) + self._index_chunks(rows) self._increment_total_inserts(len(rows)) return [node.node_id for node in nodes] - def delete(self, ref_doc_id: str, **delete_kwargs: Any) -> None: + def delete(self, ref_doc_id: int | str, **delete_kwargs: Any) -> None: if self.table_exists(): with self._transaction(): - self._conn.execute( - "DELETE FROM " + DEFAULT_TABLE_NAME + " WHERE document_id = ?", - (str(ref_doc_id),), - ) + self._delete_chunks_by_document_id(int(ref_doc_id)) def _rows_to_nodes(self, rows: list[sqlite3.Row]) -> list[BaseNode]: nodes: list[BaseNode] = [] @@ -417,41 +485,60 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): def get_modified_times(self) -> dict[str, str]: """Return {document_id: stored_modified_isoformat} for all indexed documents. - All chunks of a document share the same ``modified`` value, so the - first row seen per document is sufficient. + document_meta already has exactly one row per document (not per + chunk, unlike the vec0 table), so no dedup is needed here. """ if not self.table_exists(): return {} - result: dict[str, str] = {} - for row in self._conn.execute( - "SELECT document_id, modified FROM " + DEFAULT_TABLE_NAME, - ): - doc_id = str(row["document_id"]) - if doc_id not in result: - result[doc_id] = str(row["modified"] or "") - return result + return DocumentMetaTable.all_modified_times(self._conn) + + @property + def _db_path(self) -> str: + return str(Path(self._uri) / DB_FILENAME) + + @contextmanager + def _rebuild_file(self) -> Iterator[sqlite3.Connection]: + """Open a fresh temp database file for a file-swap rebuild (compact + or structural migration), yielding its connection for the caller to + populate. + + On success, swaps the temp file in as the live database (closing + this store's current connection first -- see _swap_in_compact()). + On any exception, discards the temp file, including its -wal/-shm, + instead, and this store's own connection is left untouched. + """ + compact_path = self._db_path + ".compact" + new_conn = self._open_connection(compact_path) + try: + yield new_conn + except BaseException: + new_conn.close() + for suffix in ["", "-wal", "-shm"]: + Path(compact_path + suffix).unlink(missing_ok=True) + raise + else: + new_conn.close() + self._swap_in_compact(compact_path, self._db_path) def compact(self, *, force: bool = False) -> None: """Rebuild the database file to reclaim space left behind by DELETEs. vec0 DELETE only invalidates rows; the vector data stays in the file - forever (asg017/sqlite-vec#54), and per-document re-indexing is a - delete+insert. The cumulative insert counter in ``index_meta`` tracks - total rows ever written; when that exceeds ``COMPACT_BLOAT_RATIO`` x - the live row count (or when forced), live rows are copied into a fresh - database file and swapped in via ``os.replace``. + forever, and per-document re-indexing is a delete+insert. The + cumulative insert counter in ``index_meta`` tracks total rows ever + written; when that exceeds ``COMPACT_BLOAT_RATIO`` x the live row + count (or when forced), live rows are copied into a fresh database + file and swapped in via ``os.replace``. Note: ``ALTER TABLE ... RENAME TO`` on vec0 virtual tables does NOT - rename the shadow tables (sqlite-vec upstream limitation), so - an in-place rename-based rebuild is not safe. The file-swap approach - is the maintainer-endorsed workaround (asg017/sqlite-vec#205). + rename the shadow tables (sqlite-vec upstream limitation), so an + in-place rename-based rebuild is not safe. The file-swap approach is + the maintainer-endorsed workaround. """ if not self.table_exists(): return - live = self._conn.execute( - "SELECT count(*) FROM " + DEFAULT_TABLE_NAME, - ).fetchone()[0] - total = int(self._meta_get("total_inserts") or str(live)) + live = DocumentChunksTable.count(self._conn) + total = IndexMetaTable.get_total_inserts(self._conn) or live if not force and total <= max(live, 1) * COMPACT_BLOAT_RATIO: return dim = self.vector_dim() @@ -463,50 +550,62 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): live, total, ) - db_path = str(Path(self._uri) / DB_FILENAME) - compact_path = db_path + ".compact" + with self._rebuild_file() as new_conn: + self._rebuild_into(self._conn, new_conn, dim) - # Copy all live rows into a fresh database file. - new_conn = self._open_connection(compact_path) - try: - self._create_vec_table(new_conn, dim) - self._meta_set_on(new_conn, "dim", str(dim)) - for key in ("embed_model", "schema_version"): - 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, + @staticmethod + def _rebuild_into( + src_conn: sqlite3.Connection, + dst_conn: sqlite3.Connection, + dim: int, + ) -> int: + """Create the vec0 table in ``dst_conn``, copy dim/embed_model from + ``src_conn``, and stream every live vec0 row, document_chunks row, + and document_meta row across. Returns the number of vec0 rows + copied. Used by compact() only -- m0001_v1_to_v2 freezes its own + copy loop instead of calling this, since this always reflects the + *current* schema (see the migration DDL-freezing rule in the spec). + """ + PaperlessSqliteVecVectorStore._create_vec_table(dst_conn, dim) + dim_value = IndexMetaTable.get_dim(src_conn) + if dim_value is not None: + IndexMetaTable.set_dim(dst_conn, dim_value) + embed_model = IndexMetaTable.get_embed_model(src_conn) + if embed_model is not None: + IndexMetaTable.set_embed_model(dst_conn, embed_model) + schema_version = IndexMetaTable.get_schema_version(src_conn) + if schema_version is not None: + IndexMetaTable.set_schema_version(dst_conn, schema_version) + + dst_conn.execute("BEGIN IMMEDIATE") + src_cursor = src_conn.execute( + "SELECT id, document_id, 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["node_content"], + bytes(r["embedding"]), + ) + for r in batch + ], ) - 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( - self._INSERT, - [ - ( - r["id"], - r["document_id"], - r["modified"], - r["node_content"], - bytes(r["embedding"]), - ) - for r in batch - ], - ) - # Reset the cumulative counter: after compact, total_inserts == live. - self._meta_set_on(new_conn, "total_inserts", str(live)) - new_conn.execute("COMMIT") - except BaseException: - new_conn.close() - for p in [compact_path, compact_path + "-wal", compact_path + "-shm"]: - Path(p).unlink(missing_ok=True) - raise - new_conn.close() - self._swap_in_compact(compact_path, db_path) + DocumentChunksTable.insert_many( + dst_conn, + (ChunkRow(r["id"], r["document_id"]) for r in batch), + ) + copied += len(batch) + DocumentMetaTable.copy_all(src_conn, dst_conn, COMPACT_BATCH_SIZE) + # Reset the cumulative counter: after a rebuild, total_inserts == live. + IndexMetaTable.reset_total_inserts(dst_conn, copied) + dst_conn.execute("COMMIT") + return copied def _swap_in_compact(self, compact_path: str, db_path: str) -> None: """Atomically replace the live database with the compacted copy.""" @@ -526,8 +625,8 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): """ if not self.table_exists(): return None - raw = self._meta_get("schema_version") - return int(raw) if raw is not None else SCHEMA_VERSION + raw_version = IndexMetaTable.get_schema_version(self._conn) + return raw_version if raw_version is not None else SCHEMA_VERSION def has_pending_migration(self) -> bool: """Cheaply check whether a migration is pending, with no exclusive @@ -591,16 +690,12 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): dim = self.vector_dim() if dim is None: # pragma: no cover raise RuntimeError("Cannot migrate: no stored vector dimension") - db_path = str(Path(self._uri) / DB_FILENAME) - compact_path = db_path + ".compact" - new_conn = self._open_connection(compact_path) - try: + with self._rebuild_file() as new_conn: migration.apply(self._conn, new_conn, dim) - self._meta_set_on(new_conn, "schema_version", str(migration.to_version)) - except BaseException: # pragma: no cover - new_conn.close() - for p in [compact_path, compact_path + "-wal", compact_path + "-shm"]: - Path(p).unlink(missing_ok=True) - raise - new_conn.close() - self._swap_in_compact(compact_path, db_path) + IndexMetaTable.set_schema_version(new_conn, migration.to_version) + + +# Registers m0001_v1_to_v2 into MIGRATIONS; must be at the bottom (needs +# PaperlessSqliteVecVectorStore fully defined) -- see +# paperless_ai/migrations/__init__.py for the full procedure. +from paperless_ai.migrations import m0001_v1_to_v2 # noqa: E402, F401