diff --git a/src/paperless_ai/migrations/__init__.py b/src/paperless_ai/migrations/__init__.py new file mode 100644 index 000000000..af63f3400 --- /dev/null +++ b/src/paperless_ai/migrations/__init__.py @@ -0,0 +1,60 @@ +"""Schema migrations for the sqlite-vec vector store. + +Each migration lives in its own module here, named ``mNNNN_description.py`` +(e.g. ``m0001_v1_to_v2.py`` -- a leading digit isn't a valid Python +identifier, hence the ``m`` prefix, unlike Django's own numbered migrations, +which load via a dynamic ``importlib.import_module()`` call rather than a +static import statement), and registers itself into ``MIGRATIONS`` at import +time. ``vector_store.py`` imports those modules at the bottom of the file, +purely for that registration side effect, after ``PaperlessSqliteVecVectorStore`` +is fully defined -- migrations need it to implement ``apply()`` (see +``Migration`` below). + +To add a new migration: add a new ``mNNNN_description.py`` module here that +imports ``PaperlessSqliteVecVectorStore`` from ``paperless_ai.vector_store``, +defines its ``apply()``, and appends a ``Migration`` to ``MIGRATIONS``; then +import that module at the bottom of ``vector_store.py`` and bump +``SCHEMA_VERSION`` there. A migration must freeze its own historical DDL for +any side table its target version depends on (``DROP TABLE IF EXISTS`` + +its own literal ``CREATE TABLE``/``CREATE INDEX`` statements) rather than +delegating to any "current schema" helper -- see ``m0001_v1_to_v2.py`` for +why and the worked example. +""" + +import sqlite3 +from collections.abc import Callable +from dataclasses import dataclass +from dataclasses import field +from typing import Literal + + +@dataclass +class Migration: + """A schema migration for the sqlite-vec vector store. + + kind="structural": rows are copied into a new-schema file with no + re-embedding needed. Supply ``apply(src_conn, dst_conn, dim)``, which + must create every table its target schema needs in ``dst_conn`` and copy + ``src_conn``'s rows and relevant ``index_meta`` keys into it. + ``schema_version`` is written by the migration runner after ``apply`` + returns, not by ``apply`` itself. + + kind="re-embed": the new schema requires fresh embeddings. + ``check_and_run_migrations()`` returns True when it encounters one of + these so the caller can force a full rebuild (which recreates the table + at the current SCHEMA_VERSION). + """ + + from_version: int + to_version: int + kind: Literal["structural", "re-embed"] + description: str + apply: Callable[[sqlite3.Connection, sqlite3.Connection, int], None] | None = field( + default=None, + repr=False, + ) + + +# Registry of all schema migrations in order, populated by each migration +# module's import-time registration (see the module docstring above). +MIGRATIONS: list[Migration] = [] diff --git a/src/paperless_ai/tests/test_vector_store.py b/src/paperless_ai/tests/test_vector_store.py index 75cc0d17e..908627d27 100644 --- a/src/paperless_ai/tests/test_vector_store.py +++ b/src/paperless_ai/tests/test_vector_store.py @@ -9,11 +9,11 @@ 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.migrations import MIGRATIONS +from paperless_ai.migrations import Migration from paperless_ai.vector_store import DB_FILENAME from paperless_ai.vector_store import DEFAULT_TABLE_NAME -from paperless_ai.vector_store import MIGRATIONS from paperless_ai.vector_store import SCHEMA_VERSION -from paperless_ai.vector_store import Migration from paperless_ai.vector_store import PaperlessSqliteVecVectorStore from paperless_ai.vector_store import _build_where @@ -646,3 +646,50 @@ class TestMigrations: assert result is True assert self._schema_version(store) == 2 + + def test_has_pending_migration_false_when_no_table( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: + """ + GIVEN: + - A vector store with no table created yet + WHEN: + - has_pending_migration() is checked + THEN: + - False is returned (nothing to migrate before anything exists) + """ + assert store.has_pending_migration() is False + + def test_has_pending_migration_false_at_current_version( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: + """ + GIVEN: + - A store at the current SCHEMA_VERSION + WHEN: + - has_pending_migration() is checked + THEN: + - False is returned + """ + store.add([make_node("a1", "1")]) + assert store.has_pending_migration() is False + + def test_has_pending_migration_true_when_behind( + self, + store: PaperlessSqliteVecVectorStore, + ) -> None: + """ + GIVEN: + - A store whose schema_version has been forced behind SCHEMA_VERSION + WHEN: + - has_pending_migration() is checked + THEN: + - True is returned + """ + 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 diff --git a/src/paperless_ai/vector_store.py b/src/paperless_ai/vector_store.py index f6f333576..4c4d0b7a2 100644 --- a/src/paperless_ai/vector_store.py +++ b/src/paperless_ai/vector_store.py @@ -2,16 +2,12 @@ import json import logging import sqlite3 import struct -from collections.abc import Callable from collections.abc import Iterator from collections.abc import Sequence from contextlib import contextmanager -from dataclasses import dataclass -from dataclasses import field from pathlib import Path from types import TracebackType from typing import Any -from typing import Literal import sqlite_vec from llama_index.core.bridge.pydantic import PrivateAttr @@ -26,6 +22,9 @@ from llama_index.core.vector_stores.types import VectorStoreQueryResult from llama_index.core.vector_stores.utils import metadata_dict_to_node from llama_index.core.vector_stores.utils import node_to_metadata_dict +from paperless_ai.migrations import MIGRATIONS +from paperless_ai.migrations import Migration + logger = logging.getLogger("paperless_ai.vector_store") DB_FILENAME = "llmindex.db" @@ -53,38 +52,6 @@ COMPACT_BATCH_SIZE = 500 _FILTER_COLUMNS = frozenset({"document_id", "modified"}) -@dataclass -class Migration: - """A schema migration for the sqlite-vec vector store. - - kind="structural": rows are copied into a new-schema file with no - re-embedding needed. Supply ``apply(src_conn, dst_conn, dim)`` which - must create the vec0 table in ``dst_conn``, copy all rows from - ``src_conn``, and write ``dim`` / ``embed_model`` / ``total_inserts`` to - ``dst_conn``'s ``index_meta``. ``schema_version`` is written by the - migration runner after ``apply`` returns. - - kind="re-embed": the new schema requires fresh embeddings. - ``check_and_run_migrations()`` returns True when it encounters one of - these so the caller can force a full rebuild (which recreates the table - at the current SCHEMA_VERSION). - """ - - from_version: int - to_version: int - kind: Literal["structural", "re-embed"] - description: str - apply: Callable[[sqlite3.Connection, sqlite3.Connection, int], None] | None = field( - default=None, - repr=False, - ) - - -# Registry of all schema migrations in order. Empty at v1 -- this is the -# baseline. Add entries here (and bump SCHEMA_VERSION) when the schema changes. -MIGRATIONS: list[Migration] = [] - - def _pack(embedding: Sequence[float]) -> bytes: return struct.pack(f"{len(embedding)}f", *embedding) @@ -551,6 +518,31 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): Path(compact_path).replace(db_path) self._conn = self._open_connection(db_path) + def _stored_schema_version(self) -> int | None: + """The schema_version recorded in index_meta, or None if no table + exists. A missing key (a store predating version tracking) is + treated as SCHEMA_VERSION -- i.e. already current -- since no + migration in MIGRATIONS targets a version before tracking began. + """ + if not self.table_exists(): + return None + raw = self._meta_get("schema_version") + return int(raw) if raw is not None else SCHEMA_VERSION + + def has_pending_migration(self) -> bool: + """Cheaply check whether a migration is pending, with no exclusive + access needed -- just a metadata read under the connection callers + already hold via the write FileLock. + + Callers should only pay for check_and_run_migrations()'s exclusive + access (a structural migration's file swap must not run while + readers are active) when this returns True, so that the common + case -- already at SCHEMA_VERSION -- never contends with readers + or a concurrent compaction. + """ + current = self._stored_schema_version() + return current is not None and current < SCHEMA_VERSION + def check_and_run_migrations(self) -> bool: """Apply any pending schema migrations to the store. @@ -559,15 +551,13 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): this method returns True when one is encountered so the caller can force a full rebuild (which recreates the table at SCHEMA_VERSION). - Must be called under the write FileLock. No-op when the table does - not exist or is already at SCHEMA_VERSION. + Must be called under the write FileLock, with readers excluded (see + has_pending_migration() for a cheap pre-check that avoids paying for + that exclusion in the common case). No-op when the table does not + exist or is already at SCHEMA_VERSION. """ - if not self.table_exists(): - return False - - raw = self._meta_get("schema_version") - current = int(raw) if raw is not None else SCHEMA_VERSION - if current >= SCHEMA_VERSION: + current = self._stored_schema_version() + if current is None or current >= SCHEMA_VERSION: return False pending = sorted( @@ -579,7 +569,7 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore): if migration.kind == "re-embed": logger.warning( "LLM index schema v%d -> v%d requires re-embedding (%s); " - "forcing full rebuild.", + "the caller must force a rebuild.", migration.from_version, migration.to_version, migration.description,