diff --git a/src/backend/core/management/commands/index.py b/src/backend/core/management/commands/index.py index af046e06..effc23d4 100644 --- a/src/backend/core/management/commands/index.py +++ b/src/backend/core/management/commands/index.py @@ -40,7 +40,7 @@ class Command(BaseCommand): batch_size = options["batch_size"] try: - count = indexer.index(batch_size=batch_size) + count = indexer.index(batch_size=batch_size, crash_safe_mode=True) except Exception as err: raise CommandError("Unable to regenerate index") from err diff --git a/src/backend/core/services/search_indexers.py b/src/backend/core/services/search_indexers.py index d2bf9ed0..5e4d1ac5 100644 --- a/src/backend/core/services/search_indexers.py +++ b/src/backend/core/services/search_indexers.py @@ -1,5 +1,6 @@ """Document search index management utilities and indexers""" +import itertools import logging from abc import ABC, abstractmethod from collections import defaultdict @@ -7,6 +8,7 @@ from functools import cache from django.conf import settings from django.contrib.auth.models import AnonymousUser +from django.core.cache import cache as django_cache from django.core.exceptions import ImproperlyConfigured from django.utils.module_loading import import_string @@ -18,6 +20,9 @@ from core.enums import SearchType logger = logging.getLogger(__name__) +BULK_INDEXER_CHECKPOINT = "bulk-indexer-checkpoint" + + @cache def get_document_indexer(): """Returns an instance of indexer service if enabled and properly configured.""" @@ -125,7 +130,7 @@ class BaseDocumentIndexer(ABC): if not self.search_url: raise ImproperlyConfigured("SEARCH_URL must be set in Django settings.") - def index(self, queryset=None, batch_size=None): + def index(self, queryset=None, batch_size=None, crash_safe_mode=False): """ Fetch documents in batches, serialize them, and push to the search backend. @@ -134,35 +139,37 @@ class BaseDocumentIndexer(ABC): Defaults to all documents without filter. batch_size (int, optional): Number of documents per batch. Defaults to settings.SEARCH_INDEXER_BATCH_SIZE. + crash_safe_mode (bool, optional): If True, order documents by updated_at + and store the last indexed document.updated_at in cache. + This allows resuming indexing from the last successful batch in case of a crash + but is more computationally expensive due to sorting. """ - last_id = 0 count = 0 queryset = queryset or models.Document.objects.all() batch_size = batch_size or self.batch_size - while True: - documents_batch = list( - queryset.filter( - id__gt=last_id, - ).order_by("id")[:batch_size] - ) - - if not documents_batch: - break + if crash_safe_mode: + queryset = queryset.order_by("updated_at") + for documents_batch in itertools.batched(queryset.iterator(), batch_size): doc_paths = [doc.path for doc in documents_batch] - last_id = documents_batch[-1].id accesses_by_document_path = get_batch_accesses_by_users_and_teams(doc_paths) - serialized_batch = [ self.serialize_document(document, accesses_by_document_path) for document in documents_batch if document.content or document.title ] - if serialized_batch: - self.push(serialized_batch) - count += len(serialized_batch) + if not serialized_batch: + continue + + self.push(serialized_batch) + count += len(serialized_batch) + + if crash_safe_mode: + django_cache.set( + BULK_INDEXER_CHECKPOINT, serialized_batch[-1]["updated_at"] + ) return count