Improve indexing command

CHANGELOG.md

@@ -14,6 +14,10 @@ and this project adheres to
 
 ### Changed
 
+- ✨(backend) improve indexing command
+  - checkpoint recovery 
+  - asynchronicity 
+  - admin command trigger
 - ♿️(frontend) improve version history list accessibility #2033
 - ♿(frontend) focus skip link on headings and skip grid dropzone #1983
 - ♿️(frontend) add sr-only format to export download button #2088

Makefile

@@ -260,7 +260,7 @@ demo: ## flush db then create a demo for load testing purpose
 .PHONY: demo
 
 index: ## index all documents to remote search
-	@$(MANAGE) index
+	@$(MANAGE) index $(args)
 .PHONY: index
 
 # Nota bene: Black should come after isort just in case they don't agree...

docs/commands/index.md

@@ -0,0 +1,65 @@
+# Index Command
+
+The `index` management command is used to index documents to the remote search indexer.
+
+It sends an asynchronous task to the Celery worker.
+
+## Usage
+
+### Make Command
+
+```bash
+# Basic usage with defaults
+make index
+
+# With custom parameters
+make index args="--batch-size 100 --lower-time-bound 2024-01-01T00:00:00 --upper-time-bound 2026-01-01T00:00:00"
+
+```
+
+### Command line
+
+```bash
+python manage.py index \
+  --lower-time-bound "2024-01-01T00:00:00" \
+  --upper-time-bound "2024-01-31T23:59:59" \
+  --batch-size 200 \
+  --async_mode
+```
+
+### Django Admin
+
+The command is available in the Django admin interface:
+
+1. Go to `/admin/core/run-indexing/`, you arrive at the "Run Indexing Command" page
+3. Fill in the form with the desired parameters
+4. Click **"Run Indexing Command"**
+
+## Parameters
+
+### `--batch-size`
+- **type:** Integer
+- **default:** `settings.SEARCH_INDEXER_BATCH_SIZE`
+- **description:** Number of documents to process per batch. Higher values may improve performance but use more memory.
+
+### `--lower-time-bound`
+- **optional**: true
+- **type:** ISO 8601 datetime string
+- **default:** `None`
+- **description:** Only documents updated after this date will be indexed.
+
+### `--upper-time-bound`
+- **optional**: true
+- **type:** ISO 8601 datetime string
+- **default:** `None`
+- **description:** Only documents updated before this date will be indexed.
+
+## `--async_mode`
+- **type:** Boolean flag
+- **default:** `False`
+- **description:** Runs asynchronously is async_mode==True. 
+
+## Crash Safe Mode
+
+The command saves the updated.at of the last document of each successful batch into the `bulk-indexer-checkpoint` cache variable.
+If the process crashes, this value can be used as `lower-time-bound` to resume from the last successfully indexed document.

src/backend/core/admin.py

@@ -1,16 +1,55 @@
 """Admin classes and registrations for core app."""
 
+from datetime import datetime
+
 from django.contrib import admin, messages
 from django.contrib.auth import admin as auth_admin
-from django.shortcuts import redirect
+from django.core.management import call_command
+from django.http import HttpRequest
+from django.shortcuts import redirect, render
 from django.utils.translation import gettext_lazy as _
 
 from treebeard.admin import TreeAdmin
 
 from core import models
+from core.forms import RunIndexingForm
 from core.tasks.user_reconciliation import user_reconciliation_csv_import_job
 
 
+# Customize the default admin site's get_app_list method
+_original_get_app_list = admin.site.get_app_list
+
+
+def custom_get_app_list(self, request, app_label=None):
+    """Add custom commands to the app list."""
+    app_list = _original_get_app_list(request, app_label)
+
+    # Add Commands app with Run Indexing command
+    commands_app = {
+        "name": _("Commands"),
+        "app_label": "commands",
+        "app_url": "#",
+        "has_module_perms": True,
+        "models": [
+            {
+                "name": _("Run indexing"),
+                "object_name": "RunIndexing",
+                "admin_url": "/admin/core/run-indexing/",
+                "view_only": False,
+                "add_url": None,
+                "change_url": None,
+            }
+        ],
+    }
+
+    app_list.append(commands_app)
+    return app_list
+
+
+# Monkey-patch the admin site
+admin.site.get_app_list = custom_get_app_list.__get__(admin.site, admin.AdminSite)
+
+
 @admin.register(models.User)
 class UserAdmin(auth_admin.UserAdmin):
     """Admin class for the User model"""
@@ -227,3 +266,43 @@ class InvitationAdmin(admin.ModelAdmin):
     def save_model(self, request, obj, form, change):
         obj.issuer = request.user
         obj.save()
+
+
+def run_indexing_view(request: HttpRequest):
+    """Custom admin view for running indexing commands."""
+    if request.method == "POST":
+        form = RunIndexingForm(request.POST)
+        if form.is_valid():
+            call_command(
+                "index",
+                batch_size=int(request.POST.get("batch_size")),
+                lower_time_bound=convert_to_isoformat(
+                    request.POST.get("lower_time_bound")
+                ),
+                upper_time_bound=convert_to_isoformat(
+                    request.POST.get("upper_time_bound")
+                ),
+                async_mode=True,
+            )
+            messages.success(request, _("Indexing triggered!"))
+        else:
+            messages.error(request, _("Please correct the errors below."))
+    else:
+        form = RunIndexingForm()
+
+    return render(
+        request=request,
+        template_name="runindexing.html",
+        context={
+            **admin.site.each_context(request),
+            "title": "Run Indexing Command",
+            "form": form,
+        },
+    )
+
+
+def convert_to_isoformat(value: str) -> str | None:
+    """Convert datetime-local input to ISO format."""
+    if value:
+        return datetime.fromisoformat(value).isoformat()
+    return None

src/backend/core/factories.py

@@ -6,6 +6,7 @@
 from django.contrib.auth.hashers import make_password
 
 import factory.fuzzy
+from factory import post_generation
 from faker import Faker
 
 from core import models
@@ -159,6 +160,20 @@ def masked_by(self, create, extracted, **kwargs):
                     document=self, user=item, defaults={"is_masked": True}
                 )
 
+    @post_generation
+    def updated_at(self, create, extracted, **kwargs):
+        """
+        the BaseModel.updated_at has auto_now=True.
+        This prevents setting a specific updated_at value with the factory.
+
+        This post_generation method bypasses this behavior.
+        """
+        if not create or not extracted:
+            return
+
+        self.__class__.objects.filter(pk=self.pk).update(updated_at=extracted)
+        self.refresh_from_db()
+
 
 class UserDocumentAccessFactory(factory.django.DjangoModelFactory):
     """Create fake document user accesses for testing."""

src/backend/core/forms.py

@@ -0,0 +1,40 @@
+"""Forms for the core app."""
+
+from django import forms
+from django.conf import settings
+
+
+class RunIndexingForm(forms.Form):
+    """
+    Form for running the indexing process.
+    """
+
+    batch_size = forms.IntegerField(
+        min_value=1,
+        initial=settings.SEARCH_INDEXER_BATCH_SIZE,
+    )
+    lower_time_bound = forms.DateTimeField(
+        required=False, widget=forms.TextInput(attrs={"type": "datetime-local"})
+    )
+    upper_time_bound = forms.DateTimeField(
+        required=False, widget=forms.TextInput(attrs={"type": "datetime-local"})
+    )
+
+    def clean(self):
+        """Override clean to validate time bounds."""
+        cleaned_data = super().clean()
+        self.check_time_bounds()
+        return cleaned_data
+
+    def check_time_bounds(self):
+        """Validate that lower_time_bound is before upper_time_bound."""
+        lower_time_bound = self.cleaned_data.get("lower_time_bound")
+        upper_time_bound = self.cleaned_data.get("upper_time_bound")
+        if (
+            lower_time_bound
+            and upper_time_bound
+            and lower_time_bound > upper_time_bound
+        ):
+            raise forms.ValidationError(
+                "Lower time bound must be before upper time bound."
+            )

src/backend/core/management/commands/index.py

@@ -4,12 +4,16 @@
 
 import logging
 import time
+from datetime import datetime
 
+from django.conf import settings
 from django.core.management.base import BaseCommand, CommandError
 
+from core import models
 from core.services.search_indexers import get_document_indexer
+from core.tasks.search import batch_document_indexer_task
 
-logger = logging.getLogger("docs.search.bootstrap_search")
+logger = logging.getLogger(__name__)
 
 
 class Command(BaseCommand):
@@ -24,9 +28,32 @@ def add_arguments(self, parser):
             action="store",
             dest="batch_size",
             type=int,
-            default=50,
+            default=settings.SEARCH_INDEXER_BATCH_SIZE,
             help="Indexation query batch size",
         )
+        parser.add_argument(
+            "--lower-time-bound",
+            action="store",
+            dest="lower_time_bound",
+            type=datetime.fromisoformat,
+            default=None,
+            help="DateTime in ISO format. Only documents updated after this date will be indexed",
+        )
+        parser.add_argument(
+            "--upper-time-bound",
+            action="store",
+            dest="upper_time_bound",
+            type=datetime.fromisoformat,
+            default=None,
+            help="DateTime in ISO format. Only documents updated before this date will be indexed",
+        )
+        parser.add_argument(
+            "--async",
+            action="store_true",
+            dest="async_mode",
+            default=False,
+            help="Whether to execute indexing asynchronously in a Celery task (default: False)",
+        )
 
     def handle(self, *args, **options):
         """Launch and log search index generation."""
@@ -35,18 +62,37 @@ def handle(self, *args, **options):
         if not indexer:
             raise CommandError("The indexer is not enabled or properly configured.")
 
-        logger.info("Starting to regenerate Find index...")
-        start = time.perf_counter()
-        batch_size = options["batch_size"]
+        if options["async_mode"]:
+            batch_document_indexer_task.apply_async(
+                kwargs={
+                    "lower_time_bound": options["lower_time_bound"],
+                    "upper_time_bound": options["upper_time_bound"],
+                    "batch_size": options["batch_size"],
+                    "crash_safe_mode": True,
+                },
+            )
+            logger.info(
+                "Document indexing task sent to worker",
+            )
+        else:
+            logger.info("Starting to regenerate Find index...")
+            start = time.perf_counter()
 
-        try:
-            count = indexer.index(batch_size=batch_size)
-        except Exception as err:
-            raise CommandError("Unable to regenerate index") from err
+            try:
+                count = indexer.index(
+                    queryset=models.Document.objects.filter_updated_at(
+                        lower_time_bound=options["lower_time_bound"],
+                        upper_time_bound=options["upper_time_bound"],
+                    ),
+                    batch_size=options["batch_size"],
+                    crash_safe_mode=True,
+                )
+            except Exception as err:
+                raise CommandError("Unable to regenerate index") from err
 
-        duration = time.perf_counter() - start
-        logger.info(
-            "Search index regenerated from %d document(s) in %.2f seconds.",
-            count,
-            duration,
-        )
+            duration = time.perf_counter() - start
+            logger.info(
+                "Search index regenerated from %d document(s) in %.2f seconds.",
+                count,
+                duration,
+            )

src/backend/core/models.py

@@ -859,6 +859,38 @@ def annotate_user_roles(self, user):
             user_roles=models.Value([], output_field=output_field),
         )
 
+    def filter_updated_at(self, lower_time_bound=None, upper_time_bound=None):
+        """
+        Filter documents by update_at.
+
+        Args:
+            lower_time_bound (datetime, optional):
+                Keep documents updated after this timestamp.
+            upper_time_bound (datetime, optional):
+                Keep documents updated before this timestamp.
+
+        Returns:
+            QuerySet: Filtered queryset ready for indexation.
+        """
+        conditions = models.Q()
+        if lower_time_bound and upper_time_bound:
+            conditions = models.Q(
+                updated_at__gte=lower_time_bound, updated_at__lte=upper_time_bound
+            ) | models.Q(
+                ancestors_deleted_at__gte=lower_time_bound,
+                ancestors_deleted_at__lte=upper_time_bound,
+            )
+        elif lower_time_bound:
+            conditions = models.Q(updated_at__gte=lower_time_bound) | models.Q(
+                ancestors_deleted_at__gte=lower_time_bound
+            )
+        elif upper_time_bound:
+            conditions = models.Q(updated_at__lte=upper_time_bound) | models.Q(
+                ancestors_deleted_at__lte=upper_time_bound
+            )
+
+        return self.filter(conditions)
+
 
 class DocumentManager(MP_NodeManager.from_queryset(DocumentQuerySet)):
     """