Compare commits

..
Author SHA1 Message Date
stumpylogandClaude Sonnet 5 744e073ee6 Perf/fix: stop building a full-library IN filter for unrestricted RAG context
get_context_for_document() always materialized every document id a user
can see into a Python list, even for a superuser (or no user at all),
passing it through as a SQL IN filter. For a superuser, that's the whole
library:

- Past ~32,763 documents, this crashes outright:
  sqlite3.OperationalError: too many SQL variables (SQLite's
  SQLITE_MAX_VARIABLE_NUMBER is 32766 by default, and the query already
  binds embedding + k + a NE self-exclusion clause alongside the ids).
- Below that cliff, vec0's IN-list evaluation is a nested loop (strncmp
  per row per allowed id), so it's quadratic in library size for no
  reason -- the filter was never going to exclude anything.

get_objects_for_user_owner_aware() already returns every Document for a
superuser (guardian's own with_superuser shortcut), so skipping straight
to document_ids=None changes nothing about which documents are
considered, only how we get there.

Also:
- Drop the pointless sorted() in _document_id_filters(): a SQL IN clause
  doesn't care about order, so it was pure overhead on every call.
- Add a hard guard in _build_where() so a future regression (or a
  legitimately huge permission-restricted user) fails closed -- no rows,
  a logged warning -- instead of a cryptic OperationalError. Since this
  filter scopes document access, failing closed rather than skipping the
  filter is the only safe way to handle an oversized list.

First item from VECTOR_STORE_PERF_BACKLOG.md (an audit done alongside
perf/13314-vecstore-point-delete, deferred to its own branch since that
one was already large).

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-07-30 14:47:15 -07:00
stumpylog 40dbe3d903 Refactor: rename COMPACT_BATCH_SIZE to BATCH_SIZE, no longer compact()-specific 2026-07-30 14:29:36 -07:00
stumpylog f77ace3d13 Test: cover update_llm_index()'s migration-check-deferred branch 2026-07-30 14:29:36 -07:00
stumpylog f6cb5e3ef1 Fix: make migration-check result tri-state to avoid deferred-vs-current ambiguity 2026-07-30 14:29:36 -07:00
8 changed files with 407 additions and 39 deletions
+19 -12
View File
@@ -96,19 +96,26 @@ def get_context_for_document(
user: User | None = None,
max_docs: int = 5,
) -> str:
visible_documents = (
get_objects_for_user_owner_aware(
user,
"view_document",
Document,
)
if user
else None
)
# None means "no restriction" to query_similar_documents. A superuser
# (like no user at all) can see every document, so skip materializing
# every visible pk into a Python list and passing it through as a SQL
# IN filter: for a large library that is a wasted quadratic scan in the
# vector store at best, and past ~32,763 documents a hard
# sqlite3.OperationalError (SQLite's bound-parameter limit) at worst.
# get_objects_for_user_owner_aware() would return every Document for a
# superuser anyway (guardian's own with_superuser shortcut), so this
# changes nothing about which documents are considered -- only how we
# get there.
visible_document_ids = (
list(visible_documents.values_list("pk", flat=True))
if visible_documents is not None
else None
None
if user is None or user.is_superuser
else list(
get_objects_for_user_owner_aware(
user,
"view_document",
Document,
).values_list("pk", flat=True),
)
)
similar_docs = query_similar_documents(
document=doc,
+78 -15
View File
@@ -1,3 +1,4 @@
import enum
import logging
from collections.abc import Iterable
from collections.abc import Iterator
@@ -217,18 +218,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
)
@@ -365,7 +397,7 @@ def _document_id_filters(doc_ids):
MetadataFilter(
key="document_id",
operator=FilterOperator.IN,
value=sorted(doc_ids),
value=list(doc_ids),
),
],
)
@@ -403,12 +435,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",
@@ -482,8 +523,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 "
@@ -491,6 +532,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)
@@ -509,14 +558,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:
@@ -528,7 +582,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 "
@@ -536,6 +591,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))
@@ -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] = {}
+1 -1
View File
@@ -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",
)
@@ -3,6 +3,8 @@ from unittest.mock import MagicMock
from unittest.mock import patch
import pytest
import pytest_mock
from django.contrib.auth.models import User
from django.test import override_settings
from documents.models import Document
@@ -278,3 +280,104 @@ def test_get_context_for_document_no_similar_docs(mock_document):
with patch("paperless_ai.ai_classifier.query_similar_documents", return_value=[]):
result = get_context_for_document(mock_document)
assert result == ""
class TestGetContextForDocumentVisibility:
"""get_context_for_document must not materialize every visible document
id for a user who can already see the whole library: a superuser (like
no user at all) gets document_ids=None (no restriction) straight
through to query_similar_documents(), instead of a full-library IN
filter that is wasteful at best and, past ~32,763 documents, a hard
sqlite3.OperationalError at worst (SQLite's bound-parameter limit).
"""
def test_skips_permission_lookup_for_superuser(
self,
mock_document: MagicMock,
mock_similar_documents: list[MagicMock],
mocker: pytest_mock.MockerFixture,
) -> None:
"""
GIVEN:
- A superuser
WHEN:
- get_context_for_document() is called
THEN:
- get_objects_for_user_owner_aware() is never called, and
query_similar_documents() is called with document_ids=None
"""
mock_query = mocker.patch(
"paperless_ai.ai_classifier.query_similar_documents",
return_value=mock_similar_documents,
)
mock_get_objects = mocker.patch(
"paperless_ai.ai_classifier.get_objects_for_user_owner_aware",
)
user = mocker.MagicMock(spec=User)
user.is_superuser = True
get_context_for_document(mock_document, user, max_docs=2)
mock_get_objects.assert_not_called()
assert mock_query.call_args.kwargs["document_ids"] is None
def test_skips_permission_lookup_when_no_user(
self,
mock_document: MagicMock,
mock_similar_documents: list[MagicMock],
mocker: pytest_mock.MockerFixture,
) -> None:
"""
GIVEN:
- No user (user=None)
WHEN:
- get_context_for_document() is called
THEN:
- get_objects_for_user_owner_aware() is never called, and
query_similar_documents() is called with document_ids=None
"""
mock_query = mocker.patch(
"paperless_ai.ai_classifier.query_similar_documents",
return_value=mock_similar_documents,
)
mock_get_objects = mocker.patch(
"paperless_ai.ai_classifier.get_objects_for_user_owner_aware",
)
get_context_for_document(mock_document, None, max_docs=2)
mock_get_objects.assert_not_called()
assert mock_query.call_args.kwargs["document_ids"] is None
def test_restricts_to_visible_documents_for_non_superuser(
self,
mock_document: MagicMock,
mock_similar_documents: list[MagicMock],
mocker: pytest_mock.MockerFixture,
) -> None:
"""
GIVEN:
- A non-superuser with a specific set of visible documents
WHEN:
- get_context_for_document() is called
THEN:
- query_similar_documents() is called with exactly that user's
visible document ids, unchanged from before this optimization
"""
mock_query = mocker.patch(
"paperless_ai.ai_classifier.query_similar_documents",
return_value=mock_similar_documents,
)
mock_queryset = mocker.MagicMock()
mock_queryset.values_list.return_value = [1, 2, 3]
mock_get_objects = mocker.patch(
"paperless_ai.ai_classifier.get_objects_for_user_owner_aware",
return_value=mock_queryset,
)
user = mocker.MagicMock(spec=User)
user.is_superuser = False
get_context_for_document(mock_document, user, max_docs=2)
mock_get_objects.assert_called_once_with(user, "view_document", Document)
assert mock_query.call_args.kwargs["document_ids"] == [1, 2, 3]
+127
View File
@@ -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
@@ -869,6 +870,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,
@@ -915,6 +954,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,
@@ -938,6 +1005,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
@@ -1056,6 +1152,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:
+43 -3
View File
@@ -17,6 +17,7 @@ from paperless_ai.migrations import Migration
from paperless_ai.migrations import m0001_v1_to_v2
from paperless_ai.tables import DocumentChunksTable
from paperless_ai.tables import DocumentMetaTable
from paperless_ai.vector_store import _MAX_IN_VALUES
from paperless_ai.vector_store import DB_FILENAME
from paperless_ai.vector_store import DEFAULT_TABLE_NAME
from paperless_ai.vector_store import SCHEMA_VERSION
@@ -296,8 +297,47 @@ class TestBuildWhere:
assert where == "1 = 0"
assert params == []
def test_query_with_untranslatable_filter_returns_no_rows(self, store) -> None:
store.add([make_node("a1", 1), make_node("b1", 2)])
def test_fails_closed_when_in_filter_exceeds_max_values(
self,
caplog: pytest.LogCaptureFixture,
) -> None:
"""
GIVEN:
- An IN filter with more values than _MAX_IN_VALUES (SQLite's
own bound-parameter limit is 32766; this guard sits below
that with headroom for the query's other bound parameters)
WHEN:
- _build_where() translates it to SQL
THEN:
- It fails closed ("1 = 0", no params) instead of building an
IN clause SQLite would reject, and logs a warning -- this
filter scopes document access, so refusing to build it must
never widen the scope to "everything" by accident
"""
oversized = _in_filter([str(i) for i in range(_MAX_IN_VALUES + 1)])
with caplog.at_level("WARNING"):
where, params = _build_where(oversized)
assert where == "(1 = 0)"
assert params == []
assert "document_id" in caplog.text
def test_query_with_untranslatable_filter_returns_no_rows(
self,
store: PaperlessSqliteVecVectorStore,
) -> None:
"""
GIVEN:
- A store with one chunk each for documents 1 and 2
WHEN:
- A query and a get_nodes() call are issued with a filter whose
only clause is an untranslatable nested MetadataFilters
THEN:
- Neither call raises, and both fail closed by returning no rows
(rather than emitting invalid or unfiltered SQL)
"""
store.add([make_node("a1", "1"), make_node("b1", "2")])
nested = MetadataFilters(
filters=[
MetadataFilter(
@@ -508,7 +548,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]))}
+34 -6
View File
@@ -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
@@ -73,6 +75,17 @@ class _Row(NamedTuple):
embedding: bytes
# _build_where(): the largest IN value list translated into bound SQL
# parameters. SQLite's own hard limit (SQLITE_MAX_VARIABLE_NUMBER) is 32766
# by default; this leaves headroom below that for the query's other bound
# parameters (the embedding blob, k, and any NE clause) and for the limit
# itself to move. An IN filter this large should not happen in practice --
# callers are expected to pass None (no filter) rather than every id when
# the filter would not actually narrow anything -- so this is a guard
# against a future regression, not a normal code path.
_MAX_IN_VALUES = 32700
def _pack(embedding: Sequence[float]) -> bytes:
return struct.pack(f"{len(embedding)}f", *embedding)
@@ -117,6 +130,21 @@ def _build_where(filters: MetadataFilters | None) -> tuple[str, list[int]]:
if not values: # pragma: no cover
clauses.append("1 = 0")
continue
if len(values) > _MAX_IN_VALUES:
# Fail closed (see the empty-clauses case below) rather than
# let SQLite raise "too many SQL variables" past its own
# limit: this filter scopes document access, so an IN list
# too large to safely bind must match no rows, never widen
# the scope to "everything" by accident.
logger.warning(
"Refusing to build an IN filter on %r with %d values "
"(over the %d-value safety limit); returning no rows.",
f.key,
len(values),
_MAX_IN_VALUES,
)
clauses.append("1 = 0")
continue
placeholders = ",".join("?" for _ in values)
clauses.append(f"{f.key} IN ({placeholders})")
params.extend(values)
@@ -605,7 +633,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 +651,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")