diff --git a/src/documents/management/commands/document_index.py b/src/documents/management/commands/document_index.py index c4f72dd3a..10558c211 100644 --- a/src/documents/management/commands/document_index.py +++ b/src/documents/management/commands/document_index.py @@ -45,6 +45,18 @@ class Command(PaperlessCommand): "Safe to run on every startup or upgrade." ), ) + parser.add_argument( + "--heap-size-mb", + type=int, + default=None, + help=( + "Tantivy writer memory budget in MB for a full reindex " + "(split across writer threads). Defaults to 512MB; lower " + "this on memory-constrained hosts. Larger values buffer " + "more documents before flushing a segment, deferring " + "merge work rather than avoiding it." + ), + ) def handle(self, *args, **options): with transaction.atomic(): @@ -60,13 +72,26 @@ class Command(PaperlessCommand): "document_type", "storage_path", "owner", - ).prefetch_related("tags", "notes", "custom_fields", "versions") + ).prefetch_related( + "tags", + "notes__user", + "custom_fields__field", + "versions", + ) + total = documents.count() + rebuild_kwargs = {} + if options.get("heap_size_mb") is not None: + rebuild_kwargs["writer_heap_bytes"] = ( + options["heap_size_mb"] * 1_000_000 + ) get_backend().rebuild( documents, - iter_wrapper=lambda docs: self.track( - docs, + iter_wrapper=lambda pairs: self.track( + pairs, description="Indexing documents...", + total=total, ), + **rebuild_kwargs, ) reset_backend() diff --git a/src/documents/search/_backend.py b/src/documents/search/_backend.py index ea316920e..e56f27678 100644 --- a/src/documents/search/_backend.py +++ b/src/documents/search/_backend.py @@ -8,6 +8,7 @@ import time from datetime import UTC from datetime import datetime from enum import StrEnum +from itertools import islice from typing import TYPE_CHECKING from typing import Final from typing import Self @@ -16,13 +17,14 @@ from typing import TypeVar from typing import cast import filelock -import regex import tantivy from django.conf import settings +from django.contrib.contenttypes.models import ContentType from django.utils.timezone import get_current_timezone from guardian.shortcuts import get_users_with_perms from documents.search._query import build_permission_filter +from documents.search._query import extract_cjk_text from documents.search._query import parse_simple_text_highlight_query from documents.search._query import parse_simple_text_query from documents.search._query import parse_simple_title_query @@ -32,11 +34,14 @@ from documents.search._schema import build_schema from documents.search._schema import open_or_rebuild_index from documents.search._schema import wipe_index from documents.search._tokenizer import ascii_fold +from documents.search._tokenizer import autocomplete_tokens from documents.search._tokenizer import register_tokenizers from documents.utils import IterWrapper from documents.utils import identity if TYPE_CHECKING: + from collections.abc import Iterator + from collections.abc import Sequence from pathlib import Path from django.contrib.auth.models import AbstractUser @@ -51,9 +56,6 @@ _LOCK_RETRY_ATTEMPTS: Final[int] = 4 # total attempts (1 initial + 3 retries) _LOCK_BACKOFF_BASE: Final[float] = 1.0 # seconds _LOCK_BACKOFF_CAP: Final[float] = 10.0 # seconds -_WORD_RE = regex.compile(r"\w+") -_AUTOCOMPLETE_REGEX_TIMEOUT = 1.0 # seconds; guards against ReDoS on untrusted content - T = TypeVar("T") @@ -66,25 +68,16 @@ class SearchMode(StrEnum): def _extract_autocomplete_words(text_sources: list[str]) -> set[str]: """Extract and normalize words for autocomplete. - Splits on non-word characters (matching Tantivy's simple tokenizer), lowercases, - and ascii-folds each token. Uses the regex library with a timeout to guard against - ReDoS on untrusted document content. + Tokenizes with Tantivy's simple analyzer (simple -> lowercase -> ascii_fold) + so the extracted words match how document content is indexed, and runs the + whole pass in Rust. This replaces a Python regex scan plus per-token folding, + which dominated full-reindex CPU time. Tokenizing natively also removes the + ReDoS exposure of running a regex over untrusted content. """ words = set() for text in text_sources: - if not text: - continue - try: - tokens = _WORD_RE.findall(text, timeout=_AUTOCOMPLETE_REGEX_TIMEOUT) - except TimeoutError: # pragma: no cover - logger.warning( - "Autocomplete word extraction timed out for a text source; skipping.", - ) - continue - for token in tokens: - normalized = ascii_fold(token.lower()) - if normalized: - words.add(normalized) + if text: + words.update(autocomplete_tokens(text)) return words @@ -387,6 +380,7 @@ class TantivyBackend: self, document: Document, effective_content: str | None = None, + viewer_ids: list[int] | None = None, ) -> tantivy.Document: """Build a tantivy Document from a Django Document instance. @@ -405,10 +399,14 @@ class TantivyBackend: doc.add_text("title", document.title) doc.add_text("title_sort", document.title) doc.add_text("simple_title", document.title) - doc.add_text("bigram_title", document.title) doc.add_text("content", content) - doc.add_text("bigram_content", content) doc.add_text("simple_content", content) + # Bigram (character-ngram) fields exist for CJK substring search, + # no need to bloat the bigram index with latin characters. + if cjk_title := extract_cjk_text(document.title): + doc.add_text("bigram_title", cjk_title) + if content and (cjk_content := extract_cjk_text(content)): + doc.add_text("bigram_content", cjk_content) # Original filename - only add if not None/empty if document.original_filename: @@ -418,14 +416,16 @@ class TantivyBackend: if document.correspondent: doc.add_text("correspondent", document.correspondent.name) doc.add_text("correspondent_sort", document.correspondent.name) - doc.add_text("bigram_correspondent", document.correspondent.name) + if cjk_corr := extract_cjk_text(document.correspondent.name): + doc.add_text("bigram_correspondent", cjk_corr) doc.add_unsigned("correspondent_id", document.correspondent_id) # Document type if document.document_type: doc.add_text("document_type", document.document_type.name) doc.add_text("type_sort", document.document_type.name) - doc.add_text("bigram_document_type", document.document_type.name) + if cjk_type := extract_cjk_text(document.document_type.name): + doc.add_text("bigram_document_type", cjk_type) doc.add_unsigned("document_type_id", document.document_type_id) # Storage path @@ -437,7 +437,8 @@ class TantivyBackend: tag_names: list[str] = [] for tag in document.tags.all(): doc.add_text("tag", tag.name) - doc.add_text("bigram_tag", tag.name) + if cjk_tag := extract_cjk_text(tag.name): + doc.add_text("bigram_tag", cjk_tag) doc.add_unsigned("tag_id", tag.pk) tag_names.append(tag.name) @@ -492,12 +493,14 @@ class TantivyBackend: doc.add_unsigned("owner_id", document.owner_id) # Viewers with permission - users_with_perms = get_users_with_perms( - document, - only_with_perms_in=["view_document"], - ) - for user in users_with_perms: - doc.add_unsigned("viewer_id", user.pk) + if viewer_ids is None: + users_with_perms = get_users_with_perms( + document, + only_with_perms_in=["view_document"], + ) + viewer_ids = [int(u.id) for u in users_with_perms] + for viewer_id in viewer_ids: + doc.add_unsigned("viewer_id", viewer_id) # Autocomplete words text_sources = [document.title, content] @@ -912,7 +915,8 @@ class TantivyBackend: def rebuild( self, documents: QuerySet[Document], - iter_wrapper: IterWrapper[Document] = identity, + iter_wrapper: IterWrapper[tuple[Document, list[int]]] = identity, + writer_heap_bytes: int = 512_000_000, ) -> None: """ Rebuild the entire search index from scratch. @@ -923,7 +927,12 @@ class TantivyBackend: Args: documents: QuerySet of Document instances to index iter_wrapper: Optional wrapper function for progress tracking - (e.g., progress bar). Should yield each document unchanged. + (e.g., progress bar). Wraps an iterable of + ``(document, viewer_ids)`` pairs and should yield each pair + unchanged, advancing one step per document. + writer_heap_bytes: Tantivy writer memory budget (split across the + writer's threads). Larger values buffer more docs in RAM before + flushing a segment, deferring merge work; they do not avoid it. """ # Create new index (on-disk or in-memory) if self._path is not None: @@ -938,13 +947,17 @@ class TantivyBackend: old_index, old_schema = self._raw_index, self._raw_schema self._raw_index = new_index self._raw_schema = new_index.schema - + # Stream documents one-by-one (so the progress bar advances per + # document) while fetching viewer permissions one SQL query per chunk. + # The stream is Sized, so iter_wrapper can still discover the total. + documents_stream = _DocumentViewerStream(documents, chunk_size=1000) try: - writer = new_index.writer() - for document in iter_wrapper(documents): + writer = new_index.writer(heap_size=writer_heap_bytes) + for document, viewer_ids in iter_wrapper(documents_stream): doc = self._build_tantivy_doc( document, document.get_effective_content(), + viewer_ids=viewer_ids, ) writer.add_document(doc) writer.commit() @@ -959,6 +972,97 @@ class TantivyBackend: raise +def chunked(iterable, size): + iterator = iter(iterable) + while chunk := list(islice(iterator, size)): + yield chunk + + +class _DocumentViewerStream: + """Yield (document, viewer_ids) pairs while batch-loading viewer ids. + + Viewer permissions are fetched one SQL query per chunk (see + ``_bulk_get_viewer_ids``), but documents are yielded individually so a + progress bar wrapped around this stream advances per document rather than + jumping a whole chunk at a time. ``__len__`` lets the progress helper still + discover the total (it inspects ``QuerySet``/``Sized``). + + The viewer ids travel with each document in the yielded pair rather than + through a separate mutable attribute, so the pairing survives regardless + of how ``iter_wrapper`` consumes the stream (buffering, batching, etc.) — + there is no reliance on the caller advancing this generator in lock-step. + """ + + def __init__(self, documents: QuerySet[Document], *, chunk_size: int) -> None: + self._documents = documents + self._chunk_size = chunk_size + + def __len__(self) -> int: + return self._documents.count() + + def __iter__(self) -> Iterator[tuple[Document, list[int]]]: + # iterator(chunk_size=…) streams from a server-side cursor instead of + # materialising the whole queryset in memory; since Django 4.1 it still + # honours prefetch_related, running the prefetches one batch at a time. + documents = self._documents.iterator(chunk_size=self._chunk_size) + for chunk in chunked(documents, self._chunk_size): + viewer_ids_by_pk = _bulk_get_viewer_ids([doc.pk for doc in chunk]) + for doc in chunk: + yield doc, viewer_ids_by_pk.get(doc.pk, []) + + +def _bulk_get_viewer_ids(doc_pks: Sequence[int]) -> dict[int, list[int]]: + """Fetch all view_document permissions for a batch of documents in one query. + + Mirrors get_users_with_perms(doc, only_with_perms_in=["view_document"]) + (with_group_users defaults to True there): a user counts as a viewer if + they hold the permission directly, OR via membership in a group that + holds the permission. Missing the group case would silently drop search + access for group-only viewers. + """ + from collections import defaultdict + + from guardian.models import GroupObjectPermission + from guardian.models import UserObjectPermission + + from documents.models import Document + + # get_for_model is cached by Django, so this costs at most one query total. + ct = ContentType.objects.get_for_model(Document) + str_pks = [str(pk) for pk in doc_pks] + + viewer_map: dict[int, set[int]] = defaultdict(set) + + # Fold the permission lookup into the query via a join on codename instead + # of a separate Permission.objects.get(), which would otherwise run once per + # chunk during a full reindex. + user_qs = UserObjectPermission.objects.filter( + content_type=ct, + permission__content_type=ct, + permission__codename="view_document", + object_pk__in=str_pks, + ).values_list("object_pk", "user_id") + for object_pk, user_id in user_qs: + viewer_map[int(object_pk)].add(user_id) + + # User.groups has related_query_name="user", so group__user__id joins + # through the group membership m2m to the member users' ids in the same + # single-query fashion as user_qs above (values_list compiles to one SQL + # JOIN; no per-row Python-side lookups follow, so select_related / + # prefetch_related do not apply here). + group_qs = GroupObjectPermission.objects.filter( + content_type=ct, + permission__content_type=ct, + permission__codename="view_document", + object_pk__in=str_pks, + ).values_list("object_pk", "group__user__id") + for object_pk, user_id in group_qs: + if user_id is not None: + viewer_map[int(object_pk)].add(user_id) + + return {object_pk: list(user_ids) for object_pk, user_ids in viewer_map.items()} + + # Module-level singleton with proper thread safety _backend: TantivyBackend | None = None _backend_path: Path | None = None # tracks which INDEX_DIR the singleton uses diff --git a/src/documents/search/_query.py b/src/documents/search/_query.py index 749098447..534f32ad2 100644 --- a/src/documents/search/_query.py +++ b/src/documents/search/_query.py @@ -40,6 +40,17 @@ def _has_cjk(text: str) -> bool: return bool(_CJK_RE.search(text)) +def extract_cjk_text(text: str) -> str: + """Join the CJK runs in ``text`` for indexing into bigram (char-ngram) fields. + + Mirrors the query side (``_build_cjk_query``): only CJK runs are ever searched + against the bigram fields, so only CJK runs are worth indexing there. Latin + text fed to a character-bigram field is never matched and only bloats the + index and slows indexing/merge. Returns "" when there is no CJK text. + """ + return " ".join(_CJK_RE.findall(text)) + + def _build_cjk_query( index: tantivy.Index, raw_query: str, diff --git a/src/documents/search/_tokenizer.py b/src/documents/search/_tokenizer.py index 6879a66af..c84dd9093 100644 --- a/src/documents/search/_tokenizer.py +++ b/src/documents/search/_tokenizer.py @@ -154,6 +154,23 @@ def simple_search_tokens(text: str) -> list[str]: return _SIMPLE_SEARCH_ANALYZER.analyze(text) +# Autocomplete word extraction: tokenize -> lowercase -> ascii_fold in a single +# Rust pass. Uses the simple tokenizer so extracted words match how document +# content is actually indexed (the content tokenizer _paperless_text also uses +# simple()), replacing a Python regex scan plus per-token folding. +_AUTOCOMPLETE_ANALYZER: Final = ( + tantivy.TextAnalyzerBuilder(tantivy.Tokenizer.simple()) + .filter(tantivy.Filter.lowercase()) + .filter(tantivy.Filter.ascii_fold()) + .build() +) + + +def autocomplete_tokens(text: str) -> list[str]: + """Tokenize text into normalized autocomplete words (lowercased, ascii-folded).""" + return _AUTOCOMPLETE_ANALYZER.analyze(text) + + def ascii_fold(text: str) -> str: """Fold text to ASCII using the same mapping as the content tokenizers. diff --git a/src/documents/tests/search/test_backend.py b/src/documents/tests/search/test_backend.py index 5c804439b..f92f46700 100644 --- a/src/documents/tests/search/test_backend.py +++ b/src/documents/tests/search/test_backend.py @@ -1,5 +1,7 @@ import pytest +from django.contrib.auth.models import Group from django.contrib.auth.models import User +from guardian.shortcuts import assign_perm from pytest_mock import MockerFixture from documents.models import CustomField @@ -15,6 +17,7 @@ from documents.tests.factories import CorrespondentFactory from documents.tests.factories import DocumentFactory from documents.tests.factories import DocumentTypeFactory from documents.tests.factories import TagFactory +from documents.tests.factories import UserFactory pytestmark = [pytest.mark.search, pytest.mark.django_db] @@ -561,18 +564,48 @@ class TestRebuild: """Test index rebuilding functionality.""" def test_with_iter_wrapper_called(self, backend: TantivyBackend) -> None: - """Index rebuild must pass documents through iter_wrapper for progress tracking.""" + """Index rebuild must pass (document, viewer_ids) pairs through iter_wrapper.""" seen = [] - def wrapper(docs): - for doc in docs: + def wrapper(pairs): + for doc, viewer_ids in pairs: seen.append(doc.pk) - yield doc + yield doc, viewer_ids Document.objects.create(title="Tracked", content="x", checksum="TW1", pk=30) backend.rebuild(Document.objects.all(), iter_wrapper=wrapper) assert 30 in seen + def test_includes_group_granted_viewers(self, backend: TantivyBackend) -> None: + """Rebuild must index viewer ids for group-only grants, not just direct ones. + + The batched viewer-id lookup used during rebuild() must mirror + get_users_with_perms(with_group_users=True), which is the default + used by the non-batched per-document indexing path. Without it, a + user who can only see a document via group membership would lose + search access to it after any full reindex. + """ + owner = UserFactory() + group_member = UserFactory() + group = Group.objects.create(name="viewers") + group_member.groups.add(group) + + doc = DocumentFactory( + title="Group shared doc", + content="group secret keyword", + owner=owner, + ) + assign_perm("view_document", group, doc) + + backend.rebuild(Document.objects.all()) + + ids = backend.search_ids( + "group secret", + user=group_member, + search_mode=SearchMode.QUERY, + ) + assert ids == [doc.pk] + class TestAutocomplete: """Test autocomplete functionality.""" diff --git a/src/documents/tests/test_management.py b/src/documents/tests/test_management.py index 276da942d..aa6befb9b 100644 --- a/src/documents/tests/test_management.py +++ b/src/documents/tests/test_management.py @@ -176,6 +176,35 @@ class TestMakeIndex: call_command("document_index", "reindex", if_needed=True, skip_checks=True) mock_get_backend.return_value.rebuild.assert_called_once() + def test_reindex_default_heap_size_not_overridden( + self, + mocker: MockerFixture, + ) -> None: + """Without --heap-size-mb, rebuild must use its own default heap size.""" + mock_get_backend = mocker.patch( + "documents.management.commands.document_index.get_backend", + ) + call_command("document_index", "reindex", skip_checks=True) + _, kwargs = mock_get_backend.return_value.rebuild.call_args + assert "writer_heap_bytes" not in kwargs + + def test_reindex_heap_size_mb_passed_to_rebuild( + self, + mocker: MockerFixture, + ) -> None: + """--heap-size-mb must convert to bytes and pass through to rebuild.""" + mock_get_backend = mocker.patch( + "documents.management.commands.document_index.get_backend", + ) + call_command( + "document_index", + "reindex", + heap_size_mb=128, + skip_checks=True, + ) + _, kwargs = mock_get_backend.return_value.rebuild.call_args + assert kwargs["writer_heap_bytes"] == 128_000_000 + @pytest.mark.management class TestRenamer(DirectoriesMixin, FileSystemAssertsMixin, TestCase):