Refactor: extract migration infrastructure, add has_pending_migration()

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

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] = []

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

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,