mirror of
https://github.com/paperless-ngx/paperless-ngx.git
synced 2026-09-07 10:17:59 +00:00
Feature: Replace Whoosh with tantivy search backend (#12471)
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Co-authored-by: Antoine Mérino <3023499+Merinorus@users.noreply.github.com>
This commit is contained in:
co-authored by
Claude Sonnet 4.6
Antoine Mérino
parent
e01a762e81
commit
aed9abe48c
@@ -0,0 +1,21 @@
|
||||
from documents.search._backend import SearchIndexLockError
|
||||
from documents.search._backend import SearchResults
|
||||
from documents.search._backend import TantivyBackend
|
||||
from documents.search._backend import TantivyRelevanceList
|
||||
from documents.search._backend import WriteBatch
|
||||
from documents.search._backend import get_backend
|
||||
from documents.search._backend import reset_backend
|
||||
from documents.search._schema import needs_rebuild
|
||||
from documents.search._schema import wipe_index
|
||||
|
||||
__all__ = [
|
||||
"SearchIndexLockError",
|
||||
"SearchResults",
|
||||
"TantivyBackend",
|
||||
"TantivyRelevanceList",
|
||||
"WriteBatch",
|
||||
"get_backend",
|
||||
"needs_rebuild",
|
||||
"reset_backend",
|
||||
"wipe_index",
|
||||
]
|
||||
@@ -0,0 +1,858 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import threading
|
||||
import unicodedata
|
||||
from collections import Counter
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING
|
||||
from typing import Self
|
||||
from typing import TypedDict
|
||||
from typing import TypeVar
|
||||
|
||||
import filelock
|
||||
import regex
|
||||
import tantivy
|
||||
from django.conf import settings
|
||||
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 parse_user_query
|
||||
from documents.search._schema import _write_sentinels
|
||||
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 register_tokenizers
|
||||
from documents.utils import IterWrapper
|
||||
from documents.utils import identity
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pathlib import Path
|
||||
|
||||
from django.contrib.auth.base_user import AbstractBaseUser
|
||||
from django.db.models import QuerySet
|
||||
|
||||
from documents.models import Document
|
||||
|
||||
logger = logging.getLogger("paperless.search")
|
||||
|
||||
_WORD_RE = regex.compile(r"\w+")
|
||||
_AUTOCOMPLETE_REGEX_TIMEOUT = 1.0 # seconds; guards against ReDoS on untrusted content
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
def _ascii_fold(s: str) -> str:
|
||||
"""
|
||||
Normalize unicode to ASCII equivalent characters for search consistency.
|
||||
|
||||
Converts accented characters (e.g., "café") to their ASCII base forms ("cafe")
|
||||
to enable cross-language searching without requiring exact diacritic matching.
|
||||
"""
|
||||
return unicodedata.normalize("NFD", s).encode("ascii", "ignore").decode()
|
||||
|
||||
|
||||
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.
|
||||
"""
|
||||
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)
|
||||
return words
|
||||
|
||||
|
||||
class SearchHit(TypedDict):
|
||||
"""Type definition for search result hits."""
|
||||
|
||||
id: int
|
||||
score: float
|
||||
rank: int
|
||||
highlights: dict[str, str]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SearchResults:
|
||||
"""
|
||||
Container for search results with pagination metadata.
|
||||
|
||||
Attributes:
|
||||
hits: List of search results with scores and highlights
|
||||
total: Total matching documents across all pages (for pagination)
|
||||
query: Preprocessed query string after date/syntax rewriting
|
||||
"""
|
||||
|
||||
hits: list[SearchHit]
|
||||
total: int # total matching documents (for pagination)
|
||||
query: str # preprocessed query string
|
||||
|
||||
|
||||
class TantivyRelevanceList:
|
||||
"""
|
||||
DRF-compatible list wrapper for Tantivy search hits.
|
||||
|
||||
Provides paginated access to search results while storing all hits in memory
|
||||
for efficient ID retrieval. Used by Django REST framework for pagination.
|
||||
|
||||
Methods:
|
||||
__len__: Returns total hit count for pagination calculations
|
||||
__getitem__: Slices the hit list for page-specific results
|
||||
|
||||
Note: Stores ALL post-filter hits so get_all_result_ids() can return
|
||||
every matching document ID without requiring a second search query.
|
||||
"""
|
||||
|
||||
def __init__(self, hits: list[SearchHit]) -> None:
|
||||
self._hits = hits
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self._hits)
|
||||
|
||||
def __getitem__(self, key: slice) -> list[SearchHit]:
|
||||
return self._hits[key]
|
||||
|
||||
|
||||
class SearchIndexLockError(Exception):
|
||||
"""Raised when the search index file lock cannot be acquired within the timeout."""
|
||||
|
||||
|
||||
class WriteBatch:
|
||||
"""
|
||||
Context manager for bulk index operations with file locking.
|
||||
|
||||
Provides transactional batch updates to the search index with proper
|
||||
concurrency control via file locking. All operations within the batch
|
||||
are committed atomically or rolled back on exception.
|
||||
|
||||
Usage:
|
||||
with backend.batch_update() as batch:
|
||||
batch.add_or_update(document)
|
||||
batch.remove(doc_id)
|
||||
"""
|
||||
|
||||
def __init__(self, backend: TantivyBackend, lock_timeout: float):
|
||||
self._backend = backend
|
||||
self._lock_timeout = lock_timeout
|
||||
self._writer = None
|
||||
self._lock = None
|
||||
|
||||
def __enter__(self) -> Self:
|
||||
if self._backend._path is not None:
|
||||
lock_path = self._backend._path / ".tantivy.lock"
|
||||
self._lock = filelock.FileLock(str(lock_path))
|
||||
try:
|
||||
self._lock.acquire(timeout=self._lock_timeout)
|
||||
except filelock.Timeout as e: # pragma: no cover
|
||||
raise SearchIndexLockError(
|
||||
f"Could not acquire index lock within {self._lock_timeout}s",
|
||||
) from e
|
||||
|
||||
self._writer = self._backend._index.writer()
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
try:
|
||||
if exc_type is None:
|
||||
self._writer.commit()
|
||||
self._backend._index.reload()
|
||||
# Explicitly delete writer to release tantivy's internal lock.
|
||||
# On exception the uncommitted writer is simply discarded.
|
||||
if self._writer is not None:
|
||||
del self._writer
|
||||
self._writer = None
|
||||
finally:
|
||||
if self._lock is not None:
|
||||
self._lock.release()
|
||||
|
||||
def add_or_update(
|
||||
self,
|
||||
document: Document,
|
||||
effective_content: str | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Add or update a document in the batch.
|
||||
|
||||
Implements upsert behavior by deleting any existing document with the same ID
|
||||
and adding the new version. This ensures stale document data (e.g., after
|
||||
permission changes) doesn't persist in the index.
|
||||
|
||||
Args:
|
||||
document: Django Document instance to index
|
||||
effective_content: Override document.content for indexing (used when
|
||||
re-indexing with newer OCR text from document versions)
|
||||
"""
|
||||
self.remove(document.pk)
|
||||
doc = self._backend._build_tantivy_doc(document, effective_content)
|
||||
self._writer.add_document(doc)
|
||||
|
||||
def remove(self, doc_id: int) -> None:
|
||||
"""
|
||||
Remove a document from the batch by its primary key.
|
||||
|
||||
Uses range query instead of term query to work around unsigned integer
|
||||
type detection bug in tantivy-py 0.25.
|
||||
"""
|
||||
# Use range query to work around u64 deletion bug
|
||||
self._writer.delete_documents_by_query(
|
||||
tantivy.Query.range_query(
|
||||
self._backend._schema,
|
||||
"id",
|
||||
tantivy.FieldType.Unsigned,
|
||||
doc_id,
|
||||
doc_id,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class TantivyBackend:
|
||||
"""
|
||||
Tantivy search backend with explicit lifecycle management.
|
||||
|
||||
Provides full-text search capabilities using the Tantivy search engine.
|
||||
Supports in-memory indexes (for testing) and persistent on-disk indexes
|
||||
(for production use). Handles document indexing, search queries, autocompletion,
|
||||
and "more like this" functionality.
|
||||
|
||||
The backend manages its own connection lifecycle and can be reset when
|
||||
the underlying index directory changes (e.g., during test isolation).
|
||||
"""
|
||||
|
||||
def __init__(self, path: Path | None = None):
|
||||
# path=None → in-memory index (for tests)
|
||||
# path=some_dir → on-disk index (for production)
|
||||
self._path = path
|
||||
self._index = None
|
||||
self._schema = None
|
||||
|
||||
def open(self) -> None:
|
||||
"""
|
||||
Open or rebuild the index as needed.
|
||||
|
||||
For disk-based indexes, checks if rebuilding is needed due to schema
|
||||
version or language changes. Registers custom tokenizers after opening.
|
||||
Safe to call multiple times - subsequent calls are no-ops.
|
||||
"""
|
||||
if self._index is not None:
|
||||
return # pragma: no cover
|
||||
if self._path is not None:
|
||||
self._index = open_or_rebuild_index(self._path)
|
||||
else:
|
||||
self._index = tantivy.Index(build_schema())
|
||||
register_tokenizers(self._index, settings.SEARCH_LANGUAGE)
|
||||
self._schema = self._index.schema
|
||||
|
||||
def close(self) -> None:
|
||||
"""
|
||||
Close the index and release resources.
|
||||
|
||||
Safe to call multiple times - subsequent calls are no-ops.
|
||||
"""
|
||||
self._index = None
|
||||
self._schema = None
|
||||
|
||||
def _ensure_open(self) -> None:
|
||||
"""Ensure the index is open before operations."""
|
||||
if self._index is None:
|
||||
self.open() # pragma: no cover
|
||||
|
||||
def _build_tantivy_doc(
|
||||
self,
|
||||
document: Document,
|
||||
effective_content: str | None = None,
|
||||
) -> tantivy.Document:
|
||||
"""Build a tantivy Document from a Django Document instance.
|
||||
|
||||
``effective_content`` overrides ``document.content`` for indexing —
|
||||
used when re-indexing a root document with a newer version's OCR text.
|
||||
"""
|
||||
content = (
|
||||
effective_content if effective_content is not None else document.content
|
||||
)
|
||||
|
||||
doc = tantivy.Document()
|
||||
|
||||
# Basic fields
|
||||
doc.add_unsigned("id", document.pk)
|
||||
doc.add_text("checksum", document.checksum)
|
||||
doc.add_text("title", document.title)
|
||||
doc.add_text("title_sort", document.title)
|
||||
doc.add_text("content", content)
|
||||
doc.add_text("bigram_content", content)
|
||||
|
||||
# Original filename - only add if not None/empty
|
||||
if document.original_filename:
|
||||
doc.add_text("original_filename", document.original_filename)
|
||||
|
||||
# Correspondent
|
||||
if document.correspondent:
|
||||
doc.add_text("correspondent", document.correspondent.name)
|
||||
doc.add_text("correspondent_sort", document.correspondent.name)
|
||||
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_unsigned("document_type_id", document.document_type_id)
|
||||
|
||||
# Storage path
|
||||
if document.storage_path:
|
||||
doc.add_text("storage_path", document.storage_path.name)
|
||||
doc.add_unsigned("storage_path_id", document.storage_path_id)
|
||||
|
||||
# Tags — collect names for autocomplete in the same pass
|
||||
tag_names: list[str] = []
|
||||
for tag in document.tags.all():
|
||||
doc.add_text("tag", tag.name)
|
||||
doc.add_unsigned("tag_id", tag.pk)
|
||||
tag_names.append(tag.name)
|
||||
|
||||
# Notes — JSON for structured queries (notes.user:alice, notes.note:text),
|
||||
# companion text field for default full-text search.
|
||||
num_notes = 0
|
||||
for note in document.notes.all():
|
||||
num_notes += 1
|
||||
doc.add_json("notes", {"note": note.note, "user": note.user.username})
|
||||
|
||||
# Custom fields — JSON for structured queries (custom_fields.name:x, custom_fields.value:y),
|
||||
# companion text field for default full-text search.
|
||||
for cfi in document.custom_fields.all():
|
||||
search_value = cfi.value_for_search
|
||||
# Skip fields where there is no value yet
|
||||
if search_value is None:
|
||||
continue
|
||||
doc.add_json(
|
||||
"custom_fields",
|
||||
{
|
||||
"name": cfi.field.name,
|
||||
"value": search_value,
|
||||
},
|
||||
)
|
||||
|
||||
# Dates
|
||||
created_date = datetime(
|
||||
document.created.year,
|
||||
document.created.month,
|
||||
document.created.day,
|
||||
tzinfo=UTC,
|
||||
)
|
||||
doc.add_date("created", created_date)
|
||||
doc.add_date("modified", document.modified)
|
||||
doc.add_date("added", document.added)
|
||||
|
||||
if document.archive_serial_number is not None:
|
||||
doc.add_unsigned("asn", document.archive_serial_number)
|
||||
|
||||
if document.page_count is not None:
|
||||
doc.add_unsigned("page_count", document.page_count)
|
||||
|
||||
doc.add_unsigned("num_notes", num_notes)
|
||||
|
||||
# Owner
|
||||
if document.owner_id:
|
||||
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)
|
||||
|
||||
# Autocomplete words
|
||||
text_sources = [document.title, content]
|
||||
if document.correspondent:
|
||||
text_sources.append(document.correspondent.name)
|
||||
if document.document_type:
|
||||
text_sources.append(document.document_type.name)
|
||||
text_sources.extend(tag_names)
|
||||
|
||||
for word in sorted(_extract_autocomplete_words(text_sources)):
|
||||
doc.add_text("autocomplete_word", word)
|
||||
|
||||
return doc
|
||||
|
||||
def add_or_update(
|
||||
self,
|
||||
document: Document,
|
||||
effective_content: str | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Add or update a single document with file locking.
|
||||
|
||||
Convenience method for single-document updates. For bulk operations,
|
||||
use batch_update() context manager for better performance.
|
||||
|
||||
Args:
|
||||
document: Django Document instance to index
|
||||
effective_content: Override document.content for indexing
|
||||
"""
|
||||
self._ensure_open()
|
||||
with self.batch_update(lock_timeout=5.0) as batch:
|
||||
batch.add_or_update(document, effective_content)
|
||||
|
||||
def remove(self, doc_id: int) -> None:
|
||||
"""
|
||||
Remove a single document from the index with file locking.
|
||||
|
||||
Convenience method for single-document removal. For bulk operations,
|
||||
use batch_update() context manager for better performance.
|
||||
|
||||
Args:
|
||||
doc_id: Primary key of the document to remove
|
||||
"""
|
||||
self._ensure_open()
|
||||
with self.batch_update(lock_timeout=5.0) as batch:
|
||||
batch.remove(doc_id)
|
||||
|
||||
def search(
|
||||
self,
|
||||
query: str,
|
||||
user: AbstractBaseUser | None,
|
||||
page: int,
|
||||
page_size: int,
|
||||
sort_field: str | None,
|
||||
*,
|
||||
sort_reverse: bool,
|
||||
) -> SearchResults:
|
||||
"""
|
||||
Execute a search query against the document index.
|
||||
|
||||
Processes the user query through date rewriting, normalization, and
|
||||
permission filtering before executing against Tantivy. Supports both
|
||||
relevance-based and field-based sorting.
|
||||
|
||||
Args:
|
||||
query: User's search query (supports natural date keywords, field filters)
|
||||
user: User for permission filtering (None for superuser/no filtering)
|
||||
page: Page number (1-indexed) for pagination
|
||||
page_size: Number of results per page
|
||||
sort_field: Field to sort by (None for relevance ranking)
|
||||
sort_reverse: Whether to reverse the sort order
|
||||
|
||||
Returns:
|
||||
SearchResults with hits, total count, and processed query
|
||||
"""
|
||||
self._ensure_open()
|
||||
tz = get_current_timezone()
|
||||
user_query = parse_user_query(self._index, query, tz)
|
||||
|
||||
# Apply permission filter if user is not None (not superuser)
|
||||
if user is not None:
|
||||
permission_filter = build_permission_filter(self._schema, user)
|
||||
final_query = tantivy.Query.boolean_query(
|
||||
[
|
||||
(tantivy.Occur.Must, user_query),
|
||||
(tantivy.Occur.Must, permission_filter),
|
||||
],
|
||||
)
|
||||
else:
|
||||
final_query = user_query
|
||||
|
||||
searcher = self._index.searcher()
|
||||
offset = (page - 1) * page_size
|
||||
|
||||
# Map sort fields
|
||||
sort_field_map = {
|
||||
"title": "title_sort",
|
||||
"correspondent__name": "correspondent_sort",
|
||||
"document_type__name": "type_sort",
|
||||
"created": "created",
|
||||
"added": "added",
|
||||
"modified": "modified",
|
||||
"archive_serial_number": "asn",
|
||||
"page_count": "page_count",
|
||||
"num_notes": "num_notes",
|
||||
}
|
||||
|
||||
# Perform search
|
||||
if sort_field and sort_field in sort_field_map:
|
||||
mapped_field = sort_field_map[sort_field]
|
||||
results = searcher.search(
|
||||
final_query,
|
||||
limit=offset + page_size,
|
||||
order_by_field=mapped_field,
|
||||
order=tantivy.Order.Desc if sort_reverse else tantivy.Order.Asc,
|
||||
)
|
||||
# Field sorting: hits are still (score, DocAddress) tuples; score unused
|
||||
all_hits = [(hit[1], 0.0) for hit in results.hits]
|
||||
else:
|
||||
# Score-based search: hits are (score, DocAddress) tuples
|
||||
results = searcher.search(final_query, limit=offset + page_size)
|
||||
all_hits = [(hit[1], hit[0]) for hit in results.hits]
|
||||
|
||||
total = results.count
|
||||
|
||||
# Normalize scores for score-based searches
|
||||
if not sort_field and all_hits:
|
||||
max_score = max(hit[1] for hit in all_hits) or 1.0
|
||||
all_hits = [(hit[0], hit[1] / max_score) for hit in all_hits]
|
||||
|
||||
# Apply threshold filter if configured (score-based search only)
|
||||
threshold = settings.ADVANCED_FUZZY_SEARCH_THRESHOLD
|
||||
if threshold is not None and not sort_field:
|
||||
all_hits = [hit for hit in all_hits if hit[1] >= threshold]
|
||||
|
||||
# Get the page's hits
|
||||
page_hits = all_hits[offset : offset + page_size]
|
||||
|
||||
# Build result hits with highlights
|
||||
hits: list[SearchHit] = []
|
||||
snippet_generator = None
|
||||
|
||||
for rank, (doc_address, score) in enumerate(page_hits, start=offset + 1):
|
||||
# Get the actual document from the searcher using the doc address
|
||||
actual_doc = searcher.doc(doc_address)
|
||||
doc_dict = actual_doc.to_dict()
|
||||
doc_id = doc_dict["id"][0]
|
||||
|
||||
highlights: dict[str, str] = {}
|
||||
|
||||
# Generate highlights if score > 0
|
||||
if score > 0:
|
||||
try:
|
||||
if snippet_generator is None:
|
||||
snippet_generator = tantivy.SnippetGenerator.create(
|
||||
searcher,
|
||||
final_query,
|
||||
self._schema,
|
||||
"content",
|
||||
)
|
||||
|
||||
content_snippet = snippet_generator.snippet_from_doc(actual_doc)
|
||||
if content_snippet:
|
||||
highlights["content"] = str(content_snippet)
|
||||
|
||||
# Try notes highlights
|
||||
if "notes" in doc_dict:
|
||||
notes_generator = tantivy.SnippetGenerator.create(
|
||||
searcher,
|
||||
final_query,
|
||||
self._schema,
|
||||
"notes",
|
||||
)
|
||||
notes_snippet = notes_generator.snippet_from_doc(actual_doc)
|
||||
if notes_snippet:
|
||||
highlights["notes"] = str(notes_snippet)
|
||||
|
||||
except Exception: # pragma: no cover
|
||||
logger.debug("Failed to generate highlights for doc %s", doc_id)
|
||||
|
||||
hits.append(
|
||||
SearchHit(
|
||||
id=doc_id,
|
||||
score=score,
|
||||
rank=rank,
|
||||
highlights=highlights,
|
||||
),
|
||||
)
|
||||
|
||||
return SearchResults(
|
||||
hits=hits,
|
||||
total=total,
|
||||
query=query,
|
||||
)
|
||||
|
||||
def autocomplete(
|
||||
self,
|
||||
term: str,
|
||||
limit: int,
|
||||
user: AbstractBaseUser | None = None,
|
||||
) -> list[str]:
|
||||
"""
|
||||
Get autocomplete suggestions for search queries.
|
||||
|
||||
Returns words that start with the given term prefix, ranked by document
|
||||
frequency (how many documents contain each word). Optionally filters
|
||||
results to only words from documents visible to the specified user.
|
||||
|
||||
Args:
|
||||
term: Prefix to match against autocomplete words
|
||||
limit: Maximum number of suggestions to return
|
||||
user: User for permission filtering (None for no filtering)
|
||||
|
||||
Returns:
|
||||
List of word suggestions ordered by frequency, then alphabetically
|
||||
"""
|
||||
self._ensure_open()
|
||||
normalized_term = _ascii_fold(term.lower())
|
||||
|
||||
searcher = self._index.searcher()
|
||||
|
||||
# Apply permission filter for non-superusers so autocomplete words
|
||||
# from invisible documents don't leak to other users.
|
||||
if user is not None and not user.is_superuser:
|
||||
base_query = build_permission_filter(self._schema, user)
|
||||
else:
|
||||
base_query = tantivy.Query.all_query()
|
||||
|
||||
results = searcher.search(base_query, limit=10000)
|
||||
|
||||
# Count how many visible documents each word appears in.
|
||||
# Using Counter (not set) preserves per-word document frequency so
|
||||
# we can rank suggestions by how commonly they occur — the same
|
||||
# signal Whoosh used for Tf/Idf-based autocomplete ordering.
|
||||
word_counts: Counter[str] = Counter()
|
||||
for _score, doc_address in results.hits:
|
||||
stored_doc = searcher.doc(doc_address)
|
||||
doc_dict = stored_doc.to_dict()
|
||||
if "autocomplete_word" in doc_dict:
|
||||
word_counts.update(doc_dict["autocomplete_word"])
|
||||
|
||||
# Filter to prefix matches, sort by document frequency descending;
|
||||
# break ties alphabetically for stable, deterministic output.
|
||||
matches = sorted(
|
||||
(w for w in word_counts if w.startswith(normalized_term)),
|
||||
key=lambda w: (-word_counts[w], w),
|
||||
)
|
||||
|
||||
return matches[:limit]
|
||||
|
||||
def more_like_this(
|
||||
self,
|
||||
doc_id: int,
|
||||
user: AbstractBaseUser | None,
|
||||
page: int,
|
||||
page_size: int,
|
||||
) -> SearchResults:
|
||||
"""
|
||||
Find documents similar to the given document using content analysis.
|
||||
|
||||
Uses Tantivy's "more like this" query to find documents with similar
|
||||
content patterns. The original document is excluded from results.
|
||||
|
||||
Args:
|
||||
doc_id: Primary key of the reference document
|
||||
user: User for permission filtering (None for no filtering)
|
||||
page: Page number (1-indexed) for pagination
|
||||
page_size: Number of results per page
|
||||
|
||||
Returns:
|
||||
SearchResults with similar documents (excluding the original)
|
||||
"""
|
||||
self._ensure_open()
|
||||
searcher = self._index.searcher()
|
||||
|
||||
# First find the document address
|
||||
id_query = tantivy.Query.range_query(
|
||||
self._schema,
|
||||
"id",
|
||||
tantivy.FieldType.Unsigned,
|
||||
doc_id,
|
||||
doc_id,
|
||||
)
|
||||
results = searcher.search(id_query, limit=1)
|
||||
|
||||
if not results.hits:
|
||||
# Document not found
|
||||
return SearchResults(hits=[], total=0, query=f"more_like:{doc_id}")
|
||||
|
||||
# Extract doc_address from (score, doc_address) tuple
|
||||
doc_address = results.hits[0][1]
|
||||
|
||||
# Build more like this query
|
||||
mlt_query = tantivy.Query.more_like_this_query(
|
||||
doc_address,
|
||||
min_doc_frequency=1,
|
||||
max_doc_frequency=None,
|
||||
min_term_frequency=1,
|
||||
max_query_terms=12,
|
||||
min_word_length=None,
|
||||
max_word_length=None,
|
||||
boost_factor=None,
|
||||
)
|
||||
|
||||
# Apply permission filter
|
||||
if user is not None:
|
||||
permission_filter = build_permission_filter(self._schema, user)
|
||||
final_query = tantivy.Query.boolean_query(
|
||||
[
|
||||
(tantivy.Occur.Must, mlt_query),
|
||||
(tantivy.Occur.Must, permission_filter),
|
||||
],
|
||||
)
|
||||
else:
|
||||
final_query = mlt_query
|
||||
|
||||
# Search
|
||||
offset = (page - 1) * page_size
|
||||
results = searcher.search(final_query, limit=offset + page_size)
|
||||
|
||||
total = results.count
|
||||
# Convert from (score, doc_address) to (doc_address, score)
|
||||
all_hits = [(hit[1], hit[0]) for hit in results.hits]
|
||||
|
||||
# Normalize scores
|
||||
if all_hits:
|
||||
max_score = max(hit[1] for hit in all_hits) or 1.0
|
||||
all_hits = [(hit[0], hit[1] / max_score) for hit in all_hits]
|
||||
|
||||
# Get page hits
|
||||
page_hits = all_hits[offset : offset + page_size]
|
||||
|
||||
# Build results
|
||||
hits: list[SearchHit] = []
|
||||
for rank, (doc_address, score) in enumerate(page_hits, start=offset + 1):
|
||||
actual_doc = searcher.doc(doc_address)
|
||||
doc_dict = actual_doc.to_dict()
|
||||
result_doc_id = doc_dict["id"][0]
|
||||
|
||||
# Skip the original document
|
||||
if result_doc_id == doc_id:
|
||||
continue
|
||||
|
||||
hits.append(
|
||||
SearchHit(
|
||||
id=result_doc_id,
|
||||
score=score,
|
||||
rank=rank,
|
||||
highlights={}, # MLT doesn't generate highlights
|
||||
),
|
||||
)
|
||||
|
||||
return SearchResults(
|
||||
hits=hits,
|
||||
total=max(0, total - 1), # Subtract 1 for the original document
|
||||
query=f"more_like:{doc_id}",
|
||||
)
|
||||
|
||||
def batch_update(self, lock_timeout: float = 30.0) -> WriteBatch:
|
||||
"""
|
||||
Get a batch context manager for bulk index operations.
|
||||
|
||||
Use this for efficient bulk document updates/deletions. All operations
|
||||
within the batch are committed atomically at the end of the context.
|
||||
|
||||
Args:
|
||||
lock_timeout: Seconds to wait for file lock acquisition
|
||||
|
||||
Returns:
|
||||
WriteBatch context manager
|
||||
|
||||
Raises:
|
||||
SearchIndexLockError: If lock cannot be acquired within timeout
|
||||
"""
|
||||
self._ensure_open()
|
||||
return WriteBatch(self, lock_timeout)
|
||||
|
||||
def rebuild(
|
||||
self,
|
||||
documents: QuerySet[Document],
|
||||
iter_wrapper: IterWrapper[Document] = identity,
|
||||
) -> None:
|
||||
"""
|
||||
Rebuild the entire search index from scratch.
|
||||
|
||||
Wipes the existing index and re-indexes all provided documents.
|
||||
On failure, restores the previous index state to keep the backend usable.
|
||||
|
||||
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.
|
||||
"""
|
||||
# Create new index (on-disk or in-memory)
|
||||
if self._path is not None:
|
||||
wipe_index(self._path)
|
||||
new_index = tantivy.Index(build_schema(), path=str(self._path))
|
||||
_write_sentinels(self._path)
|
||||
else:
|
||||
new_index = tantivy.Index(build_schema())
|
||||
register_tokenizers(new_index, settings.SEARCH_LANGUAGE)
|
||||
|
||||
# Point instance at the new index so _build_tantivy_doc uses it
|
||||
old_index, old_schema = self._index, self._schema
|
||||
self._index = new_index
|
||||
self._schema = new_index.schema
|
||||
|
||||
try:
|
||||
writer = new_index.writer()
|
||||
for document in iter_wrapper(documents):
|
||||
doc = self._build_tantivy_doc(
|
||||
document,
|
||||
document.get_effective_content(),
|
||||
)
|
||||
writer.add_document(doc)
|
||||
writer.commit()
|
||||
new_index.reload()
|
||||
except BaseException: # pragma: no cover
|
||||
# Restore old index on failure so the backend remains usable
|
||||
self._index = old_index
|
||||
self._schema = old_schema
|
||||
raise
|
||||
|
||||
|
||||
# Module-level singleton with proper thread safety
|
||||
_backend: TantivyBackend | None = None
|
||||
_backend_path: Path | None = None # tracks which INDEX_DIR the singleton uses
|
||||
_backend_lock = threading.RLock()
|
||||
|
||||
|
||||
def get_backend() -> TantivyBackend:
|
||||
"""
|
||||
Get the global backend instance with thread safety.
|
||||
|
||||
Returns a singleton TantivyBackend instance, automatically reinitializing
|
||||
when settings.INDEX_DIR changes. This ensures proper test isolation when
|
||||
using pytest-xdist or @override_settings that change the index directory.
|
||||
|
||||
Returns:
|
||||
Thread-safe singleton TantivyBackend instance
|
||||
"""
|
||||
global _backend, _backend_path
|
||||
|
||||
current_path: Path = settings.INDEX_DIR
|
||||
|
||||
# Fast path: backend is initialized and path hasn't changed (no lock needed)
|
||||
if _backend is not None and _backend_path == current_path:
|
||||
return _backend
|
||||
|
||||
# Slow path: first call, or INDEX_DIR changed between calls
|
||||
with _backend_lock:
|
||||
# Double-check after acquiring lock — another thread may have beaten us
|
||||
if _backend is not None and _backend_path == current_path:
|
||||
return _backend # pragma: no cover
|
||||
|
||||
if _backend is not None:
|
||||
_backend.close()
|
||||
|
||||
_backend = TantivyBackend(path=current_path)
|
||||
_backend.open()
|
||||
_backend_path = current_path
|
||||
|
||||
return _backend
|
||||
|
||||
|
||||
def reset_backend() -> None:
|
||||
"""
|
||||
Reset the global backend instance with thread safety.
|
||||
|
||||
Forces creation of a new backend instance on the next get_backend() call.
|
||||
Used for test isolation and when switching between different index directories.
|
||||
"""
|
||||
global _backend, _backend_path
|
||||
|
||||
with _backend_lock:
|
||||
if _backend is not None:
|
||||
_backend.close()
|
||||
_backend = None
|
||||
_backend_path = None
|
||||
@@ -0,0 +1,497 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import UTC
|
||||
from datetime import date
|
||||
from datetime import datetime
|
||||
from datetime import timedelta
|
||||
from typing import TYPE_CHECKING
|
||||
from typing import Final
|
||||
|
||||
import regex
|
||||
import tantivy
|
||||
from dateutil.relativedelta import relativedelta
|
||||
from django.conf import settings
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from datetime import tzinfo
|
||||
|
||||
from django.contrib.auth.base_user import AbstractBaseUser
|
||||
|
||||
# Maximum seconds any single regex substitution may run.
|
||||
# Prevents ReDoS on adversarial user-supplied query strings.
|
||||
_REGEX_TIMEOUT: Final[float] = 1.0
|
||||
|
||||
_DATE_ONLY_FIELDS = frozenset({"created"})
|
||||
|
||||
_DATE_KEYWORDS = frozenset(
|
||||
{
|
||||
"today",
|
||||
"yesterday",
|
||||
"this_week",
|
||||
"last_week",
|
||||
"this_month",
|
||||
"last_month",
|
||||
"this_year",
|
||||
"last_year",
|
||||
},
|
||||
)
|
||||
|
||||
_FIELD_DATE_RE = regex.compile(
|
||||
r"(\w+):(" + "|".join(_DATE_KEYWORDS) + r")\b",
|
||||
)
|
||||
_COMPACT_DATE_RE = regex.compile(r"\b(\d{14})\b")
|
||||
_RELATIVE_RANGE_RE = regex.compile(
|
||||
r"\[now([+-]\d+[dhm])?\s+TO\s+now([+-]\d+[dhm])?\]",
|
||||
regex.IGNORECASE,
|
||||
)
|
||||
# Whoosh-style relative date range: e.g. [-1 week to now], [-7 days to now]
|
||||
_WHOOSH_REL_RANGE_RE = regex.compile(
|
||||
r"\[-(?P<n>\d+)\s+(?P<unit>second|minute|hour|day|week|month|year)s?\s+to\s+now\]",
|
||||
regex.IGNORECASE,
|
||||
)
|
||||
# Whoosh-style 8-digit date: field:YYYYMMDD — field-aware so timezone can be applied correctly
|
||||
_DATE8_RE = regex.compile(r"(?P<field>\w+):(?P<date8>\d{8})\b")
|
||||
|
||||
|
||||
def _fmt(dt: datetime) -> str:
|
||||
"""Format a datetime as an ISO 8601 UTC string for use in Tantivy range queries."""
|
||||
return dt.astimezone(UTC).strftime("%Y-%m-%dT%H:%M:%SZ")
|
||||
|
||||
|
||||
def _iso_range(lo: datetime, hi: datetime) -> str:
|
||||
"""Format a [lo TO hi] range string in ISO 8601 for Tantivy query syntax."""
|
||||
return f"[{_fmt(lo)} TO {_fmt(hi)}]"
|
||||
|
||||
|
||||
def _date_only_range(keyword: str, tz: tzinfo) -> str:
|
||||
"""
|
||||
For `created` (DateField): use the local calendar date, converted to
|
||||
midnight UTC boundaries. No offset arithmetic — date only.
|
||||
"""
|
||||
|
||||
today = datetime.now(tz).date()
|
||||
|
||||
if keyword == "today":
|
||||
lo = datetime(today.year, today.month, today.day, tzinfo=UTC)
|
||||
return _iso_range(lo, lo + timedelta(days=1))
|
||||
if keyword == "yesterday":
|
||||
y = today - timedelta(days=1)
|
||||
lo = datetime(y.year, y.month, y.day, tzinfo=UTC)
|
||||
hi = datetime(today.year, today.month, today.day, tzinfo=UTC)
|
||||
return _iso_range(lo, hi)
|
||||
if keyword == "this_week":
|
||||
mon = today - timedelta(days=today.weekday())
|
||||
lo = datetime(mon.year, mon.month, mon.day, tzinfo=UTC)
|
||||
return _iso_range(lo, lo + timedelta(weeks=1))
|
||||
if keyword == "last_week":
|
||||
this_mon = today - timedelta(days=today.weekday())
|
||||
last_mon = this_mon - timedelta(weeks=1)
|
||||
lo = datetime(last_mon.year, last_mon.month, last_mon.day, tzinfo=UTC)
|
||||
hi = datetime(this_mon.year, this_mon.month, this_mon.day, tzinfo=UTC)
|
||||
return _iso_range(lo, hi)
|
||||
if keyword == "this_month":
|
||||
lo = datetime(today.year, today.month, 1, tzinfo=UTC)
|
||||
if today.month == 12:
|
||||
hi = datetime(today.year + 1, 1, 1, tzinfo=UTC)
|
||||
else:
|
||||
hi = datetime(today.year, today.month + 1, 1, tzinfo=UTC)
|
||||
return _iso_range(lo, hi)
|
||||
if keyword == "last_month":
|
||||
if today.month == 1:
|
||||
lo = datetime(today.year - 1, 12, 1, tzinfo=UTC)
|
||||
else:
|
||||
lo = datetime(today.year, today.month - 1, 1, tzinfo=UTC)
|
||||
hi = datetime(today.year, today.month, 1, tzinfo=UTC)
|
||||
return _iso_range(lo, hi)
|
||||
if keyword == "this_year":
|
||||
lo = datetime(today.year, 1, 1, tzinfo=UTC)
|
||||
return _iso_range(lo, datetime(today.year + 1, 1, 1, tzinfo=UTC))
|
||||
if keyword == "last_year":
|
||||
lo = datetime(today.year - 1, 1, 1, tzinfo=UTC)
|
||||
return _iso_range(lo, datetime(today.year, 1, 1, tzinfo=UTC))
|
||||
raise ValueError(f"Unknown keyword: {keyword}")
|
||||
|
||||
|
||||
def _datetime_range(keyword: str, tz: tzinfo) -> str:
|
||||
"""
|
||||
For `added` / `modified` (DateTimeField, stored as UTC): convert local day
|
||||
boundaries to UTC — full offset arithmetic required.
|
||||
"""
|
||||
|
||||
now_local = datetime.now(tz)
|
||||
today = now_local.date()
|
||||
|
||||
def _midnight(d: date) -> datetime:
|
||||
return datetime(d.year, d.month, d.day, tzinfo=tz).astimezone(UTC)
|
||||
|
||||
if keyword == "today":
|
||||
return _iso_range(_midnight(today), _midnight(today + timedelta(days=1)))
|
||||
if keyword == "yesterday":
|
||||
y = today - timedelta(days=1)
|
||||
return _iso_range(_midnight(y), _midnight(today))
|
||||
if keyword == "this_week":
|
||||
mon = today - timedelta(days=today.weekday())
|
||||
return _iso_range(_midnight(mon), _midnight(mon + timedelta(weeks=1)))
|
||||
if keyword == "last_week":
|
||||
this_mon = today - timedelta(days=today.weekday())
|
||||
last_mon = this_mon - timedelta(weeks=1)
|
||||
return _iso_range(_midnight(last_mon), _midnight(this_mon))
|
||||
if keyword == "this_month":
|
||||
first = today.replace(day=1)
|
||||
if today.month == 12:
|
||||
next_first = date(today.year + 1, 1, 1)
|
||||
else:
|
||||
next_first = date(today.year, today.month + 1, 1)
|
||||
return _iso_range(_midnight(first), _midnight(next_first))
|
||||
if keyword == "last_month":
|
||||
this_first = today.replace(day=1)
|
||||
if today.month == 1:
|
||||
last_first = date(today.year - 1, 12, 1)
|
||||
else:
|
||||
last_first = date(today.year, today.month - 1, 1)
|
||||
return _iso_range(_midnight(last_first), _midnight(this_first))
|
||||
if keyword == "this_year":
|
||||
return _iso_range(
|
||||
_midnight(date(today.year, 1, 1)),
|
||||
_midnight(date(today.year + 1, 1, 1)),
|
||||
)
|
||||
if keyword == "last_year":
|
||||
return _iso_range(
|
||||
_midnight(date(today.year - 1, 1, 1)),
|
||||
_midnight(date(today.year, 1, 1)),
|
||||
)
|
||||
raise ValueError(f"Unknown keyword: {keyword}")
|
||||
|
||||
|
||||
def _rewrite_compact_date(query: str) -> str:
|
||||
"""Rewrite Whoosh compact date tokens (14-digit YYYYMMDDHHmmss) to ISO 8601."""
|
||||
|
||||
def _sub(m: regex.Match[str]) -> str:
|
||||
raw = m.group(1)
|
||||
try:
|
||||
dt = datetime(
|
||||
int(raw[0:4]),
|
||||
int(raw[4:6]),
|
||||
int(raw[6:8]),
|
||||
int(raw[8:10]),
|
||||
int(raw[10:12]),
|
||||
int(raw[12:14]),
|
||||
tzinfo=UTC,
|
||||
)
|
||||
return dt.strftime("%Y-%m-%dT%H:%M:%SZ")
|
||||
except ValueError:
|
||||
return str(m.group(0))
|
||||
|
||||
try:
|
||||
return _COMPACT_DATE_RE.sub(_sub, query, timeout=_REGEX_TIMEOUT)
|
||||
except TimeoutError: # pragma: no cover
|
||||
raise ValueError(
|
||||
"Query too complex to process (compact date rewrite timed out)",
|
||||
)
|
||||
|
||||
|
||||
def _rewrite_relative_range(query: str) -> str:
|
||||
"""Rewrite Whoosh relative ranges ([now-7d TO now]) to concrete ISO 8601 UTC boundaries."""
|
||||
|
||||
def _sub(m: regex.Match[str]) -> str:
|
||||
now = datetime.now(UTC)
|
||||
|
||||
def _offset(s: str | None) -> timedelta:
|
||||
if not s:
|
||||
return timedelta(0)
|
||||
sign = 1 if s[0] == "+" else -1
|
||||
n, unit = int(s[1:-1]), s[-1]
|
||||
return (
|
||||
sign
|
||||
* {
|
||||
"d": timedelta(days=n),
|
||||
"h": timedelta(hours=n),
|
||||
"m": timedelta(minutes=n),
|
||||
}[unit]
|
||||
)
|
||||
|
||||
lo, hi = now + _offset(m.group(1)), now + _offset(m.group(2))
|
||||
if lo > hi:
|
||||
lo, hi = hi, lo
|
||||
return f"[{_fmt(lo)} TO {_fmt(hi)}]"
|
||||
|
||||
try:
|
||||
return _RELATIVE_RANGE_RE.sub(_sub, query, timeout=_REGEX_TIMEOUT)
|
||||
except TimeoutError: # pragma: no cover
|
||||
raise ValueError(
|
||||
"Query too complex to process (relative range rewrite timed out)",
|
||||
)
|
||||
|
||||
|
||||
def _rewrite_whoosh_relative_range(query: str) -> str:
|
||||
"""Rewrite Whoosh-style relative date ranges ([-N unit to now]) to ISO 8601.
|
||||
|
||||
Supports: second, minute, hour, day, week, month, year (singular and plural).
|
||||
Example: ``added:[-1 week to now]`` → ``added:[2025-01-01T… TO 2025-01-08T…]``
|
||||
"""
|
||||
now = datetime.now(UTC)
|
||||
|
||||
def _sub(m: regex.Match[str]) -> str:
|
||||
n = int(m.group("n"))
|
||||
unit = m.group("unit").lower()
|
||||
delta_map: dict[str, timedelta | relativedelta] = {
|
||||
"second": timedelta(seconds=n),
|
||||
"minute": timedelta(minutes=n),
|
||||
"hour": timedelta(hours=n),
|
||||
"day": timedelta(days=n),
|
||||
"week": timedelta(weeks=n),
|
||||
"month": relativedelta(months=n),
|
||||
"year": relativedelta(years=n),
|
||||
}
|
||||
lo = now - delta_map[unit]
|
||||
return f"[{_fmt(lo)} TO {_fmt(now)}]"
|
||||
|
||||
try:
|
||||
return _WHOOSH_REL_RANGE_RE.sub(_sub, query, timeout=_REGEX_TIMEOUT)
|
||||
except TimeoutError: # pragma: no cover
|
||||
raise ValueError(
|
||||
"Query too complex to process (Whoosh relative range rewrite timed out)",
|
||||
)
|
||||
|
||||
|
||||
def _rewrite_8digit_date(query: str, tz: tzinfo) -> str:
|
||||
"""Rewrite field:YYYYMMDD date tokens to an ISO 8601 day range.
|
||||
|
||||
Runs after ``_rewrite_compact_date`` so 14-digit timestamps are already
|
||||
converted and won't spuriously match here.
|
||||
|
||||
For DateField fields (e.g. ``created``) uses UTC midnight boundaries.
|
||||
For DateTimeField fields (e.g. ``added``, ``modified``) uses local TZ
|
||||
midnight boundaries converted to UTC — matching the ``_datetime_range``
|
||||
behaviour for keyword dates.
|
||||
"""
|
||||
|
||||
def _sub(m: regex.Match[str]) -> str:
|
||||
field = m.group("field")
|
||||
raw = m.group("date8")
|
||||
try:
|
||||
year, month, day = int(raw[0:4]), int(raw[4:6]), int(raw[6:8])
|
||||
d = date(year, month, day)
|
||||
if field in _DATE_ONLY_FIELDS:
|
||||
lo = datetime(d.year, d.month, d.day, tzinfo=UTC)
|
||||
hi = lo + timedelta(days=1)
|
||||
else:
|
||||
# DateTimeField: use local-timezone midnight → UTC
|
||||
lo = datetime(d.year, d.month, d.day, tzinfo=tz).astimezone(UTC)
|
||||
hi = datetime(
|
||||
(d + timedelta(days=1)).year,
|
||||
(d + timedelta(days=1)).month,
|
||||
(d + timedelta(days=1)).day,
|
||||
tzinfo=tz,
|
||||
).astimezone(UTC)
|
||||
return f"{field}:[{_fmt(lo)} TO {_fmt(hi)}]"
|
||||
except ValueError:
|
||||
return m.group(0)
|
||||
|
||||
try:
|
||||
return _DATE8_RE.sub(_sub, query, timeout=_REGEX_TIMEOUT)
|
||||
except TimeoutError: # pragma: no cover
|
||||
raise ValueError(
|
||||
"Query too complex to process (8-digit date rewrite timed out)",
|
||||
)
|
||||
|
||||
|
||||
def rewrite_natural_date_keywords(query: str, tz: tzinfo) -> str:
|
||||
"""
|
||||
Rewrite natural date syntax to ISO 8601 format for Tantivy compatibility.
|
||||
|
||||
Performs the first stage of query preprocessing, converting various date
|
||||
formats and keywords to ISO 8601 datetime ranges that Tantivy can parse:
|
||||
- Compact 14-digit dates (YYYYMMDDHHmmss)
|
||||
- Whoosh relative ranges ([-7 days to now], [now-1h TO now+2h])
|
||||
- 8-digit dates with field awareness (created:20240115)
|
||||
- Natural keywords (field:today, field:last_week, etc.)
|
||||
|
||||
Args:
|
||||
query: Raw user query string
|
||||
tz: Timezone for converting local date boundaries to UTC
|
||||
|
||||
Returns:
|
||||
Query with date syntax rewritten to ISO 8601 ranges
|
||||
|
||||
Note:
|
||||
Bare keywords without field prefixes pass through unchanged.
|
||||
"""
|
||||
query = _rewrite_compact_date(query)
|
||||
query = _rewrite_whoosh_relative_range(query)
|
||||
query = _rewrite_8digit_date(query, tz)
|
||||
query = _rewrite_relative_range(query)
|
||||
|
||||
def _replace(m: regex.Match[str]) -> str:
|
||||
field, keyword = m.group(1), m.group(2)
|
||||
if field in _DATE_ONLY_FIELDS:
|
||||
return f"{field}:{_date_only_range(keyword, tz)}"
|
||||
return f"{field}:{_datetime_range(keyword, tz)}"
|
||||
|
||||
try:
|
||||
return _FIELD_DATE_RE.sub(_replace, query, timeout=_REGEX_TIMEOUT)
|
||||
except TimeoutError: # pragma: no cover
|
||||
raise ValueError(
|
||||
"Query too complex to process (date keyword rewrite timed out)",
|
||||
)
|
||||
|
||||
|
||||
def normalize_query(query: str) -> str:
|
||||
"""
|
||||
Normalize query syntax for better search behavior.
|
||||
|
||||
Expands comma-separated field values to explicit AND clauses and
|
||||
collapses excessive whitespace for cleaner parsing:
|
||||
- tag:foo,bar → tag:foo AND tag:bar
|
||||
- multiple spaces → single spaces
|
||||
|
||||
Args:
|
||||
query: Query string after date rewriting
|
||||
|
||||
Returns:
|
||||
Normalized query string ready for Tantivy parsing
|
||||
"""
|
||||
|
||||
def _expand(m: regex.Match[str]) -> str:
|
||||
field = m.group(1)
|
||||
values = [v.strip() for v in m.group(2).split(",") if v.strip()]
|
||||
return " AND ".join(f"{field}:{v}" for v in values)
|
||||
|
||||
try:
|
||||
query = regex.sub(
|
||||
r"(\w+):([^\s\[\]]+(?:,[^\s\[\]]+)+)",
|
||||
_expand,
|
||||
query,
|
||||
timeout=_REGEX_TIMEOUT,
|
||||
)
|
||||
return regex.sub(r" {2,}", " ", query, timeout=_REGEX_TIMEOUT).strip()
|
||||
except TimeoutError: # pragma: no cover
|
||||
raise ValueError("Query too complex to process (normalization timed out)")
|
||||
|
||||
|
||||
_MAX_U64 = 2**64 - 1 # u64 max — used as inclusive upper bound for "any owner" range
|
||||
|
||||
|
||||
def build_permission_filter(
|
||||
schema: tantivy.Schema,
|
||||
user: AbstractBaseUser,
|
||||
) -> tantivy.Query:
|
||||
"""
|
||||
Build a query filter for user document permissions.
|
||||
|
||||
Creates a query that matches only documents visible to the specified user
|
||||
according to paperless-ngx permission rules:
|
||||
- Public documents (no owner) are visible to all users
|
||||
- Private documents are visible to their owner
|
||||
- Documents explicitly shared with the user are visible
|
||||
|
||||
Args:
|
||||
schema: Tantivy schema for field validation
|
||||
user: User to check permissions for
|
||||
|
||||
Returns:
|
||||
Tantivy query that filters results to visible documents
|
||||
|
||||
Implementation Notes:
|
||||
- Uses range_query instead of term_query to work around unsigned integer
|
||||
type detection bug in tantivy-py 0.25
|
||||
- Uses boolean_query for "no owner" check since exists_query is not
|
||||
available in tantivy-py 0.25.1 (available in master)
|
||||
- Uses disjunction_max_query to combine permission clauses with OR logic
|
||||
"""
|
||||
owner_any = tantivy.Query.range_query(
|
||||
schema,
|
||||
"owner_id",
|
||||
tantivy.FieldType.Unsigned,
|
||||
1,
|
||||
_MAX_U64,
|
||||
)
|
||||
no_owner = tantivy.Query.boolean_query(
|
||||
[
|
||||
(tantivy.Occur.Must, tantivy.Query.all_query()),
|
||||
(tantivy.Occur.MustNot, owner_any),
|
||||
],
|
||||
)
|
||||
owned = tantivy.Query.range_query(
|
||||
schema,
|
||||
"owner_id",
|
||||
tantivy.FieldType.Unsigned,
|
||||
user.pk,
|
||||
user.pk,
|
||||
)
|
||||
shared = tantivy.Query.range_query(
|
||||
schema,
|
||||
"viewer_id",
|
||||
tantivy.FieldType.Unsigned,
|
||||
user.pk,
|
||||
user.pk,
|
||||
)
|
||||
return tantivy.Query.disjunction_max_query([no_owner, owned, shared])
|
||||
|
||||
|
||||
DEFAULT_SEARCH_FIELDS = [
|
||||
"title",
|
||||
"content",
|
||||
"correspondent",
|
||||
"document_type",
|
||||
"tag",
|
||||
]
|
||||
_FIELD_BOOSTS = {"title": 2.0}
|
||||
|
||||
|
||||
def parse_user_query(
|
||||
index: tantivy.Index,
|
||||
raw_query: str,
|
||||
tz: tzinfo,
|
||||
) -> tantivy.Query:
|
||||
"""
|
||||
Parse user query through the complete preprocessing pipeline.
|
||||
|
||||
Transforms the raw user query through multiple stages:
|
||||
1. Date keyword rewriting (today → ISO 8601 ranges)
|
||||
2. Query normalization (comma expansion, whitespace cleanup)
|
||||
3. Tantivy parsing with field boosts
|
||||
4. Optional fuzzy query blending (if ADVANCED_FUZZY_SEARCH_THRESHOLD set)
|
||||
|
||||
Args:
|
||||
index: Tantivy index with registered tokenizers
|
||||
raw_query: Original user query string
|
||||
tz: Timezone for date boundary calculations
|
||||
|
||||
Returns:
|
||||
Parsed Tantivy query ready for execution
|
||||
|
||||
Note:
|
||||
When ADVANCED_FUZZY_SEARCH_THRESHOLD is configured, adds a low-priority
|
||||
fuzzy query as a Should clause (0.1 boost) to catch approximate matches
|
||||
while keeping exact matches ranked higher. The threshold value is applied
|
||||
as a post-search score filter, not during query construction.
|
||||
"""
|
||||
|
||||
query_str = rewrite_natural_date_keywords(raw_query, tz)
|
||||
query_str = normalize_query(query_str)
|
||||
|
||||
exact = index.parse_query(
|
||||
query_str,
|
||||
DEFAULT_SEARCH_FIELDS,
|
||||
field_boosts=_FIELD_BOOSTS,
|
||||
)
|
||||
|
||||
threshold = settings.ADVANCED_FUZZY_SEARCH_THRESHOLD
|
||||
if threshold is not None:
|
||||
fuzzy = index.parse_query(
|
||||
query_str,
|
||||
DEFAULT_SEARCH_FIELDS,
|
||||
field_boosts=_FIELD_BOOSTS,
|
||||
# (prefix=True, distance=1, transposition_cost_one=True) — edit-distance fuzziness
|
||||
fuzzy_fields={f: (True, 1, True) for f in DEFAULT_SEARCH_FIELDS},
|
||||
)
|
||||
return tantivy.Query.boolean_query(
|
||||
[
|
||||
(tantivy.Occur.Should, exact),
|
||||
# 0.1 boost keeps fuzzy hits ranked below exact matches (intentional)
|
||||
(tantivy.Occur.Should, tantivy.Query.boost_query(fuzzy, 0.1)),
|
||||
],
|
||||
)
|
||||
|
||||
return exact
|
||||
@@ -0,0 +1,165 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import shutil
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import tantivy
|
||||
from django.conf import settings
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pathlib import Path
|
||||
|
||||
logger = logging.getLogger("paperless.search")
|
||||
|
||||
SCHEMA_VERSION = 1
|
||||
|
||||
|
||||
def build_schema() -> tantivy.Schema:
|
||||
"""
|
||||
Build the Tantivy schema for the paperless document index.
|
||||
|
||||
Creates a comprehensive schema supporting full-text search, filtering,
|
||||
sorting, and autocomplete functionality. Includes fields for document
|
||||
content, metadata, permissions, custom fields, and notes.
|
||||
|
||||
Returns:
|
||||
Configured Tantivy schema ready for index creation
|
||||
"""
|
||||
sb = tantivy.SchemaBuilder()
|
||||
|
||||
sb.add_unsigned_field("id", stored=True, indexed=True, fast=True)
|
||||
sb.add_text_field("checksum", stored=True, tokenizer_name="raw")
|
||||
|
||||
for field in (
|
||||
"title",
|
||||
"correspondent",
|
||||
"document_type",
|
||||
"storage_path",
|
||||
"original_filename",
|
||||
"content",
|
||||
):
|
||||
sb.add_text_field(field, stored=True, tokenizer_name="paperless_text")
|
||||
|
||||
# Shadow sort fields - fast, not stored/indexed
|
||||
for field in ("title_sort", "correspondent_sort", "type_sort"):
|
||||
sb.add_text_field(
|
||||
field,
|
||||
stored=False,
|
||||
tokenizer_name="simple_analyzer",
|
||||
fast=True,
|
||||
)
|
||||
|
||||
# CJK support - not stored, indexed only
|
||||
sb.add_text_field("bigram_content", stored=False, tokenizer_name="bigram_analyzer")
|
||||
|
||||
# Autocomplete prefix scan - stored, not indexed
|
||||
sb.add_text_field("autocomplete_word", stored=True, tokenizer_name="raw")
|
||||
|
||||
sb.add_text_field("tag", stored=True, tokenizer_name="paperless_text")
|
||||
|
||||
# JSON fields — structured queries: notes.user:alice, custom_fields.name:invoice
|
||||
sb.add_json_field("notes", stored=True, tokenizer_name="paperless_text")
|
||||
sb.add_json_field("custom_fields", stored=True, tokenizer_name="paperless_text")
|
||||
|
||||
for field in (
|
||||
"correspondent_id",
|
||||
"document_type_id",
|
||||
"storage_path_id",
|
||||
"tag_id",
|
||||
"owner_id",
|
||||
"viewer_id",
|
||||
):
|
||||
sb.add_unsigned_field(field, stored=False, indexed=True, fast=True)
|
||||
|
||||
for field in ("created", "modified", "added"):
|
||||
sb.add_date_field(field, stored=True, indexed=True, fast=True)
|
||||
|
||||
for field in ("asn", "page_count", "num_notes"):
|
||||
sb.add_unsigned_field(field, stored=True, indexed=True, fast=True)
|
||||
|
||||
return sb.build()
|
||||
|
||||
|
||||
def needs_rebuild(index_dir: Path) -> bool:
|
||||
"""
|
||||
Check if the search index needs rebuilding.
|
||||
|
||||
Compares the current schema version and search language configuration
|
||||
against sentinel files to determine if the index is compatible with
|
||||
the current paperless-ngx version and settings.
|
||||
|
||||
Args:
|
||||
index_dir: Path to the search index directory
|
||||
|
||||
Returns:
|
||||
True if the index needs rebuilding, False if it's up to date
|
||||
"""
|
||||
version_file = index_dir / ".schema_version"
|
||||
if not version_file.exists():
|
||||
return True
|
||||
try:
|
||||
if int(version_file.read_text().strip()) != SCHEMA_VERSION:
|
||||
logger.info("Search index schema version mismatch - rebuilding.")
|
||||
return True
|
||||
except ValueError:
|
||||
return True
|
||||
|
||||
language_file = index_dir / ".schema_language"
|
||||
if not language_file.exists():
|
||||
logger.info("Search index language sentinel missing - rebuilding.")
|
||||
return True
|
||||
if language_file.read_text().strip() != (settings.SEARCH_LANGUAGE or ""):
|
||||
logger.info("Search index language changed - rebuilding.")
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def wipe_index(index_dir: Path) -> None:
|
||||
"""
|
||||
Delete all contents of the index directory to prepare for rebuild.
|
||||
|
||||
Recursively removes all files and subdirectories within the index
|
||||
directory while preserving the directory itself.
|
||||
|
||||
Args:
|
||||
index_dir: Path to the search index directory to clear
|
||||
"""
|
||||
for child in index_dir.iterdir():
|
||||
if child.is_dir():
|
||||
shutil.rmtree(child)
|
||||
else:
|
||||
child.unlink()
|
||||
|
||||
|
||||
def _write_sentinels(index_dir: Path) -> None:
|
||||
"""Write schema version and language sentinel files so the next index open can skip rebuilding."""
|
||||
(index_dir / ".schema_version").write_text(str(SCHEMA_VERSION))
|
||||
(index_dir / ".schema_language").write_text(settings.SEARCH_LANGUAGE or "")
|
||||
|
||||
|
||||
def open_or_rebuild_index(index_dir: Path | None = None) -> tantivy.Index:
|
||||
"""
|
||||
Open the Tantivy index, creating or rebuilding as needed.
|
||||
|
||||
Checks if the index needs rebuilding due to schema version or language
|
||||
changes. If rebuilding is needed, wipes the directory and creates a fresh
|
||||
index with the current schema and configuration.
|
||||
|
||||
Args:
|
||||
index_dir: Path to index directory (defaults to settings.INDEX_DIR)
|
||||
|
||||
Returns:
|
||||
Opened Tantivy index (caller must register custom tokenizers)
|
||||
"""
|
||||
if index_dir is None:
|
||||
index_dir = settings.INDEX_DIR
|
||||
if not index_dir.exists():
|
||||
return tantivy.Index(build_schema())
|
||||
if needs_rebuild(index_dir):
|
||||
wipe_index(index_dir)
|
||||
idx = tantivy.Index(build_schema(), path=str(index_dir))
|
||||
_write_sentinels(index_dir)
|
||||
return idx
|
||||
return tantivy.Index.open(str(index_dir))
|
||||
@@ -0,0 +1,116 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
import tantivy
|
||||
|
||||
logger = logging.getLogger("paperless.search")
|
||||
|
||||
# Mapping of ISO 639-1 codes (and common aliases) -> Tantivy Snowball name
|
||||
_LANGUAGE_MAP: dict[str, str] = {
|
||||
"ar": "Arabic",
|
||||
"arabic": "Arabic",
|
||||
"da": "Danish",
|
||||
"danish": "Danish",
|
||||
"nl": "Dutch",
|
||||
"dutch": "Dutch",
|
||||
"en": "English",
|
||||
"english": "English",
|
||||
"fi": "Finnish",
|
||||
"finnish": "Finnish",
|
||||
"fr": "French",
|
||||
"french": "French",
|
||||
"de": "German",
|
||||
"german": "German",
|
||||
"el": "Greek",
|
||||
"greek": "Greek",
|
||||
"hu": "Hungarian",
|
||||
"hungarian": "Hungarian",
|
||||
"it": "Italian",
|
||||
"italian": "Italian",
|
||||
"no": "Norwegian",
|
||||
"norwegian": "Norwegian",
|
||||
"pt": "Portuguese",
|
||||
"portuguese": "Portuguese",
|
||||
"ro": "Romanian",
|
||||
"romanian": "Romanian",
|
||||
"ru": "Russian",
|
||||
"russian": "Russian",
|
||||
"es": "Spanish",
|
||||
"spanish": "Spanish",
|
||||
"sv": "Swedish",
|
||||
"swedish": "Swedish",
|
||||
"ta": "Tamil",
|
||||
"tamil": "Tamil",
|
||||
"tr": "Turkish",
|
||||
"turkish": "Turkish",
|
||||
}
|
||||
|
||||
SUPPORTED_LANGUAGES: frozenset[str] = frozenset(_LANGUAGE_MAP)
|
||||
|
||||
|
||||
def register_tokenizers(index: tantivy.Index, language: str | None) -> None:
|
||||
"""
|
||||
Register all custom tokenizers required by the paperless schema.
|
||||
|
||||
Must be called on every Index instance since Tantivy requires tokenizer
|
||||
re-registration after each index open/creation. Registers tokenizers for
|
||||
full-text search, sorting, CJK language support, and fast-field indexing.
|
||||
|
||||
Args:
|
||||
index: Tantivy index instance to register tokenizers on
|
||||
language: ISO 639-1 language code for stemming (None to disable)
|
||||
|
||||
Note:
|
||||
simple_analyzer is registered as both a text and fast-field tokenizer
|
||||
since sort shadow fields (title_sort, correspondent_sort, type_sort)
|
||||
use fast=True and Tantivy requires fast-field tokenizers to exist
|
||||
even for documents that omit those fields.
|
||||
"""
|
||||
index.register_tokenizer("paperless_text", _paperless_text(language))
|
||||
index.register_tokenizer("simple_analyzer", _simple_analyzer())
|
||||
index.register_tokenizer("bigram_analyzer", _bigram_analyzer())
|
||||
# Fast-field tokenizer required for fast=True text fields in the schema
|
||||
index.register_fast_field_tokenizer("simple_analyzer", _simple_analyzer())
|
||||
|
||||
|
||||
def _paperless_text(language: str | None) -> tantivy.TextAnalyzer:
|
||||
"""Main full-text tokenizer for content, title, etc: simple -> remove_long(65) -> lowercase -> ascii_fold [-> stemmer]"""
|
||||
builder = (
|
||||
tantivy.TextAnalyzerBuilder(tantivy.Tokenizer.simple())
|
||||
.filter(tantivy.Filter.remove_long(65))
|
||||
.filter(tantivy.Filter.lowercase())
|
||||
.filter(tantivy.Filter.ascii_fold())
|
||||
)
|
||||
if language:
|
||||
tantivy_lang = _LANGUAGE_MAP.get(language.lower())
|
||||
if tantivy_lang:
|
||||
builder = builder.filter(tantivy.Filter.stemmer(tantivy_lang))
|
||||
else:
|
||||
logger.warning(
|
||||
"Unsupported search language '%s' - stemming disabled. Supported: %s",
|
||||
language,
|
||||
", ".join(sorted(SUPPORTED_LANGUAGES)),
|
||||
)
|
||||
return builder.build()
|
||||
|
||||
|
||||
def _simple_analyzer() -> tantivy.TextAnalyzer:
|
||||
"""Tokenizer for shadow sort fields (title_sort, correspondent_sort, type_sort): simple -> lowercase -> ascii_fold."""
|
||||
return (
|
||||
tantivy.TextAnalyzerBuilder(tantivy.Tokenizer.simple())
|
||||
.filter(tantivy.Filter.lowercase())
|
||||
.filter(tantivy.Filter.ascii_fold())
|
||||
.build()
|
||||
)
|
||||
|
||||
|
||||
def _bigram_analyzer() -> tantivy.TextAnalyzer:
|
||||
"""Enables substring search in CJK text: ngram(2,2) -> lowercase. CJK / no-whitespace language support."""
|
||||
return (
|
||||
tantivy.TextAnalyzerBuilder(
|
||||
tantivy.Tokenizer.ngram(min_gram=2, max_gram=2, prefix_only=False),
|
||||
)
|
||||
.filter(tantivy.Filter.lowercase())
|
||||
.build()
|
||||
)
|
||||
Reference in New Issue
Block a user