fixup! ✨(backend) add time-bound filtering options on index command

fixup! fixup! ✨(backend) add time-bound filtering options on index command

Signed-off-by: charles <charles.englebert@protonmail.com>

Author: mascarpon3

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/management/commands/index.py

@@ -3,15 +3,17 @@
 """
 
 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):
@@ -45,6 +47,13 @@ def add_arguments(self, parser):
             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."""
@@ -53,15 +62,37 @@ def handle(self, *args, **options):
         if not indexer:
             raise CommandError("The indexer is not enabled or properly configured.")
 
-        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,
-            },
-        )
+        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()
 
-        logger.info(
-            "Document indexing task sent to worker",
-        )
+            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,
+            )

src/backend/core/models.py

@@ -859,6 +859,31 @@ 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
+            )
+        elif lower_time_bound:
+            conditions = models.Q(updated_at__gte=lower_time_bound)
+        elif upper_time_bound:
+            conditions = models.Q(updated_at__lte=upper_time_bound)
+
+        return self.filter(conditions)
+
 
 class DocumentManager(MP_NodeManager.from_queryset(DocumentQuerySet)):
     """

src/backend/core/tasks/search.py

@@ -11,7 +11,6 @@
 from core.services.search_indexers import (
     get_document_indexer,
 )
-from core.utils import build_indexable_documents_queryset
 
 from impress.celery_app import app
 
@@ -81,7 +80,7 @@ def batch_document_indexer_task(lower_time_bound=None, upper_time_bound=None, **
         return
 
     count = indexer.index(
-        queryset=build_indexable_documents_queryset(
+        queryset=models.Document.objects.filter_updated_at(
             lower_time_bound=lower_time_bound, upper_time_bound=upper_time_bound
         ),
         **kwargs,