Fix: make migration-check result tri-state to avoid deferred-vs-current ambiguity

This commit is contained in:
stumpylog
2026-07-30 07:56:04 -07:00
parent 3bc8636e62
commit 3d2977f272
2 changed files with 175 additions and 14 deletions
+77 -14
View File
@@ -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))
@@ -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,
@@ -1017,6 +1084,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: