✨(backend): add crash-safe mode to index command

crash-save mode consist in indexing documents in ascending
updated_at order and save the last document.update_at.
This allows resuming indexing from the last successful batch
in case of a crash.

Signed-off-by: charles <[email protected]>

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
 

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