Refactor: rename COMPACT_BATCH_SIZE to BATCH_SIZE, no longer compact()-specific

src/paperless_ai/indexing.py

@@ -1,3 +1,4 @@
+import enum
 import logging
 from collections.abc import Iterable
 from contextlib import contextmanager
@@ -186,18 +187,49 @@ def write_store(embed_model_name: str | None = None):
         yield store
 
 
-def _check_and_run_migrations(store: "PaperlessSqliteVecVectorStore") -> bool:
-    """Run any pending structural migrations, returning True if a pending
-    re-embed migration needs the caller to force a rebuild -- never
-    triggered automatically here. Safe to call before any write, including
+class MigrationCheckResult(enum.Enum):
+    """Outcome of _check_and_run_migrations().
+
+    CURRENT: no migration was pending, or a pending structural migration
+    was applied successfully -- safe to write.
+
+    REEMBED_REQUIRED: a pending migration needs fresh embeddings, which is
+    never triggered automatically -- the caller must force a rebuild.
+
+    DEFERRED: a migration was pending but could not run because active
+    index readers did not drain within LLM_INDEX_COMPACTION_LOCK_TIMEOUT --
+    the store is still on its old schema. Callers must NOT proceed to
+    write: collapsing this into the same falsy value as CURRENT (as a
+    plain bool return once did) would let a write proceed against an
+    unmigrated schema.
+    """
+
+    CURRENT = "current"
+    REEMBED_REQUIRED = "reembed_required"
+    DEFERRED = "deferred"
+
+
+def _check_and_run_migrations(
+    store: "PaperlessSqliteVecVectorStore",
+) -> MigrationCheckResult:
+    """Run any pending structural migrations, reporting the outcome as a
+    tri-state result. Safe to call before any write, including
     delete()/upsert_document(): has_pending_migration() (see its docstring)
     keeps this a no-op, with no exclusive access taken, once the store is
     current.
     """
     if not store.has_pending_migration():
-        return False
-    return bool(
-        _with_exclusive_access("migration check", store.check_and_run_migrations),
+        return MigrationCheckResult.CURRENT
+    result = _with_exclusive_access(
+        "migration check",
+        store.check_and_run_migrations,
+    )
+    if result is None:
+        return MigrationCheckResult.DEFERRED
+    return (
+        MigrationCheckResult.REEMBED_REQUIRED
+        if result
+        else MigrationCheckResult.CURRENT
     )
 
 
@@ -372,12 +404,21 @@ def update_llm_index(
     happens, since a rebuild always covers the whole library regardless.
     """
     with write_store() as store:
-        needs_reembed = _check_and_run_migrations(store)
-        if needs_reembed:
+        migration_result = _check_and_run_migrations(store)
+        if migration_result is MigrationCheckResult.REEMBED_REQUIRED:
             logger.warning(
                 "LLM index migration requires re-embedding; forcing rebuild.",
             )
             rebuild = True
+        elif migration_result is MigrationCheckResult.DEFERRED:
+            logger.info(
+                "Skipping LLM index update: migration check deferred while "
+                "index readers are active; will retry next run.",
+            )
+            return (
+                "Skipping LLM index update: migration check deferred; "
+                "will retry next run."
+            )
     documents = Document.objects.select_related(
         "correspondent",
         "document_type",
@@ -451,8 +492,8 @@ def llm_index_add_or_update_document(document: Document):
         _embed_nodes(new_nodes, get_embedding_model(config))
 
     with write_store(embed_model_name=get_configured_model_name(config)) as store:
-        needs_reembed = _check_and_run_migrations(store)
-        if needs_reembed:
+        migration_result = _check_and_run_migrations(store)
+        if migration_result is MigrationCheckResult.REEMBED_REQUIRED:
             logger.warning(
                 "Skipping incremental LLM index update for document %s: the "
                 "index requires re-embedding first. Run 'document_llmindex "
@@ -460,6 +501,14 @@ def llm_index_add_or_update_document(document: Document):
                 document.id,
             )
             return
+        if migration_result is MigrationCheckResult.DEFERRED:
+            logger.info(
+                "Skipping incremental LLM index update for document %s: "
+                "migration check deferred while index readers are active; "
+                "will retry on the next write.",
+                document.id,
+            )
+            return
         store.upsert_document(str(document.id), new_nodes)
 
 
@@ -478,14 +527,19 @@ def llm_index_migrate() -> None:
     if not AIConfig().llm_index_enabled:
         return
     with write_store() as store:
-        needs_reembed = _check_and_run_migrations(store)
-    if needs_reembed:
+        migration_result = _check_and_run_migrations(store)
+    if migration_result is MigrationCheckResult.REEMBED_REQUIRED:
         logger.warning(
             "LLM index requires re-embedding, which this automatic migration "
             "check will not do on its own -- it can be slow and, for a "
             "metered embedding backend, cost money. Run "
             "'document_llmindex rebuild' manually when ready.",
         )
+    elif migration_result is MigrationCheckResult.DEFERRED:
+        logger.info(
+            "LLM index migration check deferred while index readers are "
+            "active; will retry next run.",
+        )
 
 
 def llm_index_compact() -> None:
@@ -497,7 +551,8 @@ def llm_index_compact() -> None:
 def llm_index_remove_document(document: Document):
     """Remove a document's chunks from the LLM index."""
     with write_store() as store:
-        if _check_and_run_migrations(store):
+        migration_result = _check_and_run_migrations(store)
+        if migration_result is MigrationCheckResult.REEMBED_REQUIRED:
             logger.warning(
                 "Skipping removal of document %s from the LLM index: the "
                 "index requires re-embedding first. Run 'document_llmindex "
@@ -505,6 +560,14 @@ def llm_index_remove_document(document: Document):
                 document.id,
             )
             return
+        if migration_result is MigrationCheckResult.DEFERRED:
+            logger.info(
+                "Skipping removal of document %s from the LLM index: "
+                "migration check deferred while index readers are active; "
+                "will retry on the next write.",
+                document.id,
+            )
+            return
         store.delete(str(document.id))
 
 

src/paperless_ai/migrations/m0001_v1_to_v2.py

@@ -7,7 +7,7 @@ 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 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
@@ -75,7 +75,7 @@ def _migrate_v1_to_v2(
     dst_conn.execute("BEGIN IMMEDIATE")
     src_cursor = src_conn.execute(_V1_SELECT)
     live = 0
-    while batch := src_cursor.fetchmany(COMPACT_BATCH_SIZE):
+    while batch := src_cursor.fetchmany(BATCH_SIZE):
         vec0_rows = []
         chunk_rows = []
         meta_by_document: dict[int, str] = {}

src/paperless_ai/tables.py

@@ -131,7 +131,7 @@ class DocumentMetaTable:
         connections (compact()/migrations) -- an unbounded fetchall here
         would defeat the same OOM-avoidance the vec0 row copy already relies
         on. batch_size has no default: forces the call site to think about
-        it (pass COMPACT_BATCH_SIZE)."""
+        it (pass BATCH_SIZE)."""
         cursor = src_conn.execute(
             "SELECT document_id, modified FROM document_meta",
         )

src/paperless_ai/tests/test_ai_indexing.py

@@ -9,6 +9,7 @@ from django.db import connection
 from django.test import override_settings
 from django.test.utils import CaptureQueriesContext
 from django.utils import timezone
+from filelock import Timeout
 from llama_index.core.schema import MetadataMode
 
 from documents.models import Correspondent
@@ -830,6 +831,44 @@ class TestLlmIndexLocking:
 
         mock_store.upsert_document.assert_not_called()
 
+    def test_add_or_update_document_skips_write_when_migration_check_deferred(
+        self,
+        temp_llm_index_dir: Path,
+        mock_embed_model: FakeEmbedding,
+        mocker: pytest_mock.MockerFixture,
+    ) -> None:
+        """A migration check that times out waiting for readers to drain
+        must be treated the same as a pending migration -- proceeding to
+        write would target a store still on its old schema. Regression
+        test for the tri-state fix: a bare bool collapsed this outcome
+        into the same falsy value as "already current".
+        """
+        mock_store = MagicMock()
+        mock_store.has_pending_migration.return_value = True
+        mocker.patch(
+            "paperless_ai.indexing.write_store",
+            return_value=mocker.MagicMock(
+                __enter__=mocker.MagicMock(return_value=mock_store),
+                __exit__=mocker.MagicMock(return_value=False),
+            ),
+        )
+        mocker.patch(
+            "paperless_ai.indexing._exclude_readers",
+            side_effect=Timeout("test"),
+        )
+        mock_node = MagicMock()
+        mock_node.get_content.return_value = "fake node text"
+        mocker.patch(
+            "paperless_ai.indexing.build_document_node",
+            return_value=[mock_node],
+        )
+
+        doc = MagicMock(spec=Document)
+        doc.id = 1
+        indexing.llm_index_add_or_update_document(doc)
+
+        mock_store.upsert_document.assert_not_called()
+
     def test_remove_document_uses_write_store(
         self,
         temp_llm_index_dir: Path,
@@ -876,6 +915,34 @@ class TestLlmIndexLocking:
 
         mock_store.delete.assert_not_called()
 
+    def test_remove_document_skips_write_when_migration_check_deferred(
+        self,
+        temp_llm_index_dir: Path,
+        mocker: pytest_mock.MockerFixture,
+    ) -> None:
+        """A migration check deferred by a reader-lock timeout must block
+        the delete too, for the same reason as the incremental-update path.
+        """
+        mock_store = MagicMock()
+        mock_store.has_pending_migration.return_value = True
+        mocker.patch(
+            "paperless_ai.indexing.write_store",
+            return_value=mocker.MagicMock(
+                __enter__=mocker.MagicMock(return_value=mock_store),
+                __exit__=mocker.MagicMock(return_value=False),
+            ),
+        )
+        mocker.patch(
+            "paperless_ai.indexing._exclude_readers",
+            side_effect=Timeout("test"),
+        )
+
+        doc = MagicMock(spec=Document)
+        doc.id = 1
+        indexing.llm_index_remove_document(doc)
+
+        mock_store.delete.assert_not_called()
+
     def test_update_llm_index_rebuild_uses_write_store(
         self,
         temp_llm_index_dir: Path,
@@ -899,6 +966,35 @@ class TestLlmIndexLocking:
 
         mock_store.drop_table.assert_called_once()
 
+    def test_update_llm_index_skips_when_migration_check_deferred(
+        self,
+        temp_llm_index_dir: Path,
+        mocker: pytest_mock.MockerFixture,
+    ) -> None:
+        """A migration check deferred by a reader-lock timeout must short-
+        circuit before the second write_store() block (document scanning,
+        add/upsert, compaction) ever runs -- that block would otherwise
+        write against a store still on its old schema.
+        """
+        mock_store = MagicMock()
+        mock_store.has_pending_migration.return_value = True
+        write_store_mock = mocker.patch(
+            "paperless_ai.indexing.write_store",
+            return_value=mocker.MagicMock(
+                __enter__=mocker.MagicMock(return_value=mock_store),
+                __exit__=mocker.MagicMock(return_value=False),
+            ),
+        )
+        mocker.patch(
+            "paperless_ai.indexing._exclude_readers",
+            side_effect=Timeout("test"),
+        )
+
+        result = indexing.update_llm_index(rebuild=False)
+
+        assert "deferred" in result
+        write_store_mock.assert_called_once()
+
 
 @pytest.mark.django_db
 @pytest.mark.django_db
@@ -1017,6 +1113,37 @@ class TestLlmIndexMigrate:
             indexing.llm_index_migrate()
         assert "requires re-embedding" in caplog.text
 
+    def test_logs_info_when_migration_check_deferred(
+        self,
+        mocker: pytest_mock.MockerFixture,
+        caplog: pytest.LogCaptureFixture,
+    ) -> None:
+        """
+        GIVEN:
+            - AI/LLM index support is enabled
+            - A pending migration cannot run because readers are active
+        WHEN:
+            - llm_index_migrate() is called
+        THEN:
+            - An info line notes the deferral, not the re-embed warning
+        """
+        mocker.patch(
+            "paperless_ai.indexing.AIConfig",
+            return_value=mocker.Mock(llm_index_enabled=True),
+        )
+        store_mock = mocker.MagicMock()
+        store_mock.has_pending_migration.return_value = True
+        write_store_cm = mocker.patch("paperless_ai.indexing.write_store")
+        write_store_cm.return_value.__enter__.return_value = store_mock
+        mocker.patch(
+            "paperless_ai.indexing._exclude_readers",
+            side_effect=Timeout("test"),
+        )
+        with caplog.at_level(logging.INFO, logger="paperless_ai.indexing"):
+            indexing.llm_index_migrate()
+        assert "deferred" in caplog.text
+        assert "requires re-embedding" not in caplog.text
+
 
 @pytest.mark.django_db
 class TestQuerySimilarDocuments:

src/paperless_ai/tests/test_vector_store.py

@@ -508,7 +508,7 @@ class TestCompact:
         A tiny batch size forces several fetchmany()/executemany() cycles so a
         regression in the streaming loop (dropped tail, off-by-one) surfaces.
         """
-        monkeypatch.setattr("paperless_ai.vector_store.COMPACT_BATCH_SIZE", 3)
+        monkeypatch.setattr("paperless_ai.vector_store.BATCH_SIZE", 3)
         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]))}

src/paperless_ai/vector_store.py

@@ -47,10 +47,12 @@ SCHEMA_VERSION = 2
 # a rebuild copies the live rows into a fresh table.
 COMPACT_BLOAT_RATIO = 2.0
 
-# compact(): number of rows copied per executemany() when rebuilding the file.
-# Rows are streamed from the source cursor in batches of this size rather than
-# materialized all at once, keeping memory bounded regardless of index size.
-COMPACT_BATCH_SIZE = 500
+# Number of rows fetched/copied per batch whenever this module streams rows
+# instead of materializing them all at once, keeping memory bounded regardless
+# of index size -- used by compact()'s rebuild, m0001_v1_to_v2's migration
+# copy, and DocumentMetaTable.copy_all(). No longer compact()-specific, hence
+# the plain name.
+BATCH_SIZE = 500
 
 # Filterable vec0 metadata columns. _build_where() only ever receives filter
 # keys we construct ourselves, but allowlisting keeps SQL identifiers safe by
@@ -605,7 +607,7 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore):
             + DEFAULT_TABLE_NAME,
         )
         copied = 0
-        while batch := src_cursor.fetchmany(COMPACT_BATCH_SIZE):
+        while batch := src_cursor.fetchmany(BATCH_SIZE):
             dst_conn.executemany(
                 _INSERT,
                 [
@@ -623,7 +625,7 @@ class PaperlessSqliteVecVectorStore(BasePydanticVectorStore):
                 (ChunkRow(r["id"], r["document_id"]) for r in batch),
             )
             copied += len(batch)
-        DocumentMetaTable.copy_all(src_conn, dst_conn, COMPACT_BATCH_SIZE)
+        DocumentMetaTable.copy_all(src_conn, dst_conn, BATCH_SIZE)
         # Reset the cumulative counter: after a rebuild, total_inserts == live.
         IndexMetaTable.reset_total_inserts(dst_conn, copied)
         dst_conn.execute("COMMIT")