mirror of
https://github.com/paperless-ngx/paperless-ngx.git
synced 2026-07-31 16:15:58 +00:00
Fix: make migration-check result tri-state to avoid deferred-vs-current ambiguity
This commit is contained in:
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user