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

we need to be able to specify time bounds on the bulk index
to allow recovering from checkpoint after crash

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

CHANGELOG.md

@@ -14,6 +14,7 @@ and this project adheres to
 
 ### Changed
 
+- ✨(backend) improve indexing command with checkpoint recovery and asynchronicity 
 - ♿️(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

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

@@ -3,11 +3,12 @@ Handle search setup that needs to be done at bootstrap time.
 """
 
 import logging
-import time
+from datetime import datetime
 
 from django.core.management.base import BaseCommand, CommandError
 
 from core.services.search_indexers import get_document_indexer
+from core.tasks.indexing import batch_document_indexer_task
 
 logger = logging.getLogger("docs.search.bootstrap_search")
 
@@ -27,6 +28,31 @@ class Command(BaseCommand):
             default=50,
             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(
+            "--crash-safe-mode",
+            action="store_true",
+            dest="crash_safe_mode",
+            default=True,
+            help="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.",
+        )
 
     def handle(self, *args, **options):
         """Launch and log search index generation."""
@@ -35,18 +61,18 @@ class Command(BaseCommand):
         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"]
-
         try:
-            count = indexer.index(batch_size=batch_size, crash_safe_mode=True)
+            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": options["crash_safe_mode"],
+                },
+            )
         except Exception as err:
-            raise CommandError("Unable to regenerate index") from err
+            raise CommandError("Unable send indexing task to worker") from err
 
-        duration = time.perf_counter() - start
         logger.info(
-            "Search index regenerated from %d document(s) in %.2f seconds.",
-            count,
-            duration,
+            "Document indexing task sent to worker",
         )

src/backend/core/tasks/search.py

@@ -4,7 +4,6 @@ from logging import getLogger
 
 from django.conf import settings
 from django.core.cache import cache
-from django.db.models import Q
 
 from django_redis.cache import RedisCache
 
@@ -12,6 +11,7 @@ from core import models
 from core.services.search_indexers import (
     get_document_indexer,
 )
+from core.utils import build_indexable_documents_queryset
 
 from impress.celery_app import app
 
@@ -22,7 +22,7 @@ logger = getLogger(__file__)
 def document_indexer_task(document_id):
     """
     Celery Task: Indexes a single document by its ID.
-    
+
     Args:
         document_id: Primary key of the document to index.
     """
@@ -36,14 +36,14 @@ def document_indexer_task(document_id):
 def batch_indexer_throttle_acquire(timeout: int = 0, atomic: bool = True):
     """
     Acquire a throttle lock to prevent multiple batch indexation tasks during countdown.
-    
+
     implements a debouncing pattern: only the first call during the timeout period
     will succeed, subsequent calls are skipped until the timeout expires.
 
     Args:
         timeout (int): Lock duration in seconds (countdown period).
         atomic (bool): Use Redis locks for atomic operations if available.
-    
+
     Returns:
         bool: True if lock acquired (first call), False if already held (subsequent calls).
     """
@@ -64,31 +64,35 @@ def batch_indexer_throttle_acquire(timeout: int = 0, atomic: bool = True):
 
 
 @app.task
-def batch_document_indexer_task(timestamp):
+def batch_document_indexer_task(lower_time_bound=None, upper_time_bound=None, **kwargs):
     """
     Celery Task: Batch indexes all documents modified since timestamp.
-    
+
     Args:
-        timestamp: ISO timestamp to filter documents by updated_at, deleted_at, 
-                   or ancestors_deleted_at.
+        lower_time_bound (datetime|isoformat, optional):
+            indexes documents updated or deleted after this timestamp.
+        upper_time_bound (datetime|isoformat, optional):
+            indexes documents updated or deleted before this timestamp.
     """
     indexer = get_document_indexer()
 
-    if indexer:
-        queryset = models.Document.objects.filter(
-            Q(updated_at__gte=timestamp)
-            | Q(deleted_at__gte=timestamp)
-            | Q(ancestors_deleted_at__gte=timestamp)
-        )
+    if not indexer:
+        logger.warning("Indexing task triggered but no indexer configured: skipping")
+        return
 
-        count = indexer.index(queryset)
-        logger.info("Indexed %d documents", count)
+    count = indexer.index(
+        queryset=build_indexable_documents_queryset(
+            lower_time_bound=lower_time_bound, upper_time_bound=upper_time_bound
+        ),
+        **kwargs,
+    )
+    logger.info("Indexed %d documents", count)
 
 
 def trigger_batch_document_indexer(document):
     """
     Trigger document indexation with optional debounce mechanism.
-    
+
     behavior depends on SEARCH_INDEXER_COUNTDOWN setting:
     - if countdown > 0 sec (async batch mode):
       * schedules a batch indexation task after countdown in seconds
@@ -97,7 +101,7 @@ def trigger_batch_document_indexer(document):
     - if countdown == 0 sec (sync mode):
       * executes indexation synchronously in the current thread
       * no batching, no throttling, no Celery task queuing
-    
+
     Args:
         document (Document): the document instance that triggered the indexation.
     """
@@ -118,7 +122,7 @@ def trigger_batch_document_indexer(document):
             )
 
             batch_document_indexer_task.apply_async(
-                args=[document.updated_at], countdown=countdown
+                kwargs={"lower_time_bound": document.updated_at}, countdown=countdown
             )
         else:
             logger.info("Skip task for batch document %s indexation", document.pk)

src/backend/core/tests/commands/test_index.py

@@ -2,21 +2,24 @@
 Unit test for `index` command.
 """
 
+from datetime import datetime, timedelta, timezone
 from operator import itemgetter
 from unittest import mock
 
+from django.core.cache import cache
 from django.core.management import CommandError, call_command
 from django.db import transaction
 
 import pytest
 
 from core import factories
-from core.services.search_indexers import FindDocumentIndexer
+from core.services.search_indexers import BULK_INDEXER_CHECKPOINT, FindDocumentIndexer
+from core.tests.conftest import create_document_with_updated_at
 
 
 @pytest.mark.django_db
 @pytest.mark.usefixtures("indexer_settings")
-def test_index():
+def test_index_without_bound_success():
     """Test the command `index` that run the Find app indexer for all the available documents."""
     user = factories.UserFactory()
     indexer = FindDocumentIndexer()
@@ -37,20 +40,152 @@ def test_index():
     }
 
     with mock.patch.object(FindDocumentIndexer, "push") as mock_push:
-        call_command("index")
+        call_command("index", crash_safe_mode=False)
 
-        push_call_args = [call.args[0] for call in mock_push.call_args_list]
+    push_call_args = [call.args[0] for call in mock_push.call_args_list]
 
-        # called once but with a batch of docs
-        mock_push.assert_called_once()
+    # called once but with a batch of docs
+    mock_push.assert_called_once()
 
-        assert sorted(push_call_args[0], key=itemgetter("id")) == sorted(
-            [
-                indexer.serialize_document(doc, accesses),
-                indexer.serialize_document(no_title_doc, accesses),
-            ],
-            key=itemgetter("id"),
+    assert sorted(push_call_args[0], key=itemgetter("id")) == sorted(
+        [
+            indexer.serialize_document(doc, accesses),
+            indexer.serialize_document(no_title_doc, accesses),
+        ],
+        key=itemgetter("id"),
+    )
+
+    # crash_safe_mode deactivated -> no checkpoint stored
+    assert cache.get(BULK_INDEXER_CHECKPOINT) is None
+
+
+@pytest.mark.django_db
+@pytest.mark.usefixtures("indexer_settings")
+def test_index_with_both_bounds_success():
+    """Test the command `index` for all documents within time bound."""
+    cache.clear()
+    lower_time_bound = datetime(2024, 2, 1, tzinfo=timezone.utc)
+    upper_time_bound = lower_time_bound + timedelta(days=30)
+
+    document_too_early = create_document_with_updated_at(
+        lower_time_bound - timedelta(days=10)
+    )
+    document_in_window_1 = create_document_with_updated_at(
+        lower_time_bound + timedelta(days=5)
+    )
+    document_in_window_2 = create_document_with_updated_at(
+        lower_time_bound + timedelta(days=15)
+    )
+    document_too_late = create_document_with_updated_at(
+        upper_time_bound + timedelta(days=10)
+    )
+
+    with mock.patch.object(FindDocumentIndexer, "push") as mock_push:
+        call_command(
+            "index",
+            lower_time_bound=lower_time_bound.isoformat(),
+            upper_time_bound=upper_time_bound.isoformat(),
         )
+    all_push_call_args = [
+        document["id"]
+        for call_arg_list in mock_push.call_args_list
+        for document in call_arg_list.args[0]
+    ]
+
+    # Only documents in window should be indexed
+    assert str(document_too_early.id) not in all_push_call_args
+    assert str(document_in_window_1.id) in all_push_call_args
+    assert str(document_in_window_2.id) in all_push_call_args
+    assert str(document_too_late.id) not in all_push_call_args
+
+    # Checkpoint should be set to last indexed document's updated_at
+    assert (
+        cache.get(BULK_INDEXER_CHECKPOINT)
+        == document_in_window_2.updated_at.isoformat()
+    )
+
+
+@pytest.mark.django_db
+@pytest.mark.usefixtures("indexer_settings")
+def test_index_with_crash_recovery():
+    """Test resuming indexing from checkpoint after a crash."""
+    cache.clear()
+    lower_time_bound = datetime(2024, 2, 1, tzinfo=timezone.utc)
+    upper_time_bound = lower_time_bound + timedelta(days=60)
+
+    batch_size = 2
+    documents = [
+        # batch 0
+        create_document_with_updated_at(lower_time_bound + timedelta(days=5)),
+        create_document_with_updated_at(lower_time_bound + timedelta(days=10)),
+        # batch 1
+        create_document_with_updated_at(lower_time_bound + timedelta(days=20)),
+        create_document_with_updated_at(lower_time_bound + timedelta(days=25)),
+        # batch 2 - will crash here
+        create_document_with_updated_at(lower_time_bound + timedelta(days=30)),
+        create_document_with_updated_at(lower_time_bound + timedelta(days=35)),
+        # batch 3
+        create_document_with_updated_at(lower_time_bound + timedelta(days=40)),
+        create_document_with_updated_at(lower_time_bound + timedelta(days=45)),
+    ]
+
+    # First run: simulate crash on batch 3
+    with mock.patch.object(FindDocumentIndexer, "push") as mock_push:
+
+        def push_with_failure(data):
+            # Crash when we encounter document at index 4 (batch 2 with batch_size=2)
+            if str(documents[4].id) in [document["id"] for document in data]:
+                raise ConnectionError("Simulated indexing error")
+
+        mock_push.side_effect = push_with_failure
+
+        call_command(
+            "index",
+            batch_size=batch_size,
+            lower_time_bound=lower_time_bound.isoformat(),
+            upper_time_bound=upper_time_bound.isoformat(),
+        )
+    all_push_call_args = [
+        document["id"]
+        for call_arg_list in mock_push.call_args_list
+        for document in call_arg_list.args[0]
+    ]
+
+    # First 2 batches should have been indexed and checkpoint saved
+    checkpoint = cache.get(BULK_INDEXER_CHECKPOINT)
+    assert checkpoint == documents[3].updated_at.isoformat()
+    # first 2 batches should be indexed successfully
+    for i in range(0, 4):
+        assert str(documents[i].id) in all_push_call_args
+    # next batch should have been attempted but failed
+    for i in range(4, 6):
+        assert str(documents[i].id) in all_push_call_args
+    # last batches indexing should not have been attempted
+    for i in range(6, 8):
+        assert str(documents[i].id) not in all_push_call_args
+
+    # Second run: resume from checkpoint
+    with mock.patch.object(FindDocumentIndexer, "push") as mock_push:
+        call_command(
+            "index",
+            batch_size=batch_size,
+            lower_time_bound=checkpoint,
+            upper_time_bound=upper_time_bound.isoformat(),
+        )
+        all_push_call_args = [
+            document["id"]
+            for call_arg_list in mock_push.call_args_list
+            for document in call_arg_list.args[0]
+        ]
+
+        # first 2 batches should NOT be re-indexed
+        # except the last document of the last batch which is on the checkpoint boundary
+        for i in range(0, 3):
+            assert str(documents[i].id) not in all_push_call_args
+        # next batches should be indexed including the document at the checkpoint boundary
+        # which has already been indexed and is re-indexed
+        for i in range(4, 8):
+            assert str(documents[i].id) in all_push_call_args
 
 
 @pytest.mark.django_db

src/backend/core/tests/conftest.py

@@ -8,7 +8,7 @@ from django.core.cache import cache
 import pytest
 import responses
 
-from core import factories
+from core import factories, models
 from core.tests.utils.urls import reload_urls
 
 USER = "user"
@@ -143,3 +143,10 @@ def user_token():
     A fixture to create a user token for testing.
     """
     return build_authorization_bearer("some_token")
+
+
+def create_document_with_updated_at(updated_at, **kwargs):
+    """Helper to create a document with a specific updated_at, bypassing auto_now=True."""
+    document = factories.DocumentFactory(**kwargs)
+    models.Document.objects.filter(pk=document.pk).update(updated_at=updated_at)
+    return models.Document.objects.get(pk=document.pk)

src/backend/core/tests/test_utils.py

@@ -2,6 +2,7 @@
 
 import base64
 import uuid
+from datetime import datetime, timedelta, timezone
 
 from django.core.cache import cache
 
@@ -9,6 +10,7 @@ import pycrdt
 import pytest
 
 from core import factories, utils
+from core.tests.conftest import create_document_with_updated_at
 
 pytestmark = pytest.mark.django_db
 
@@ -240,3 +242,127 @@ def test_utils_get_value_by_pattern_no_match():
     result = utils.get_value_by_pattern(data, r"^title\.")
 
     assert result == []
+
+
+def test_utils_build_indexable_documents_queryset_no_filters():
+    """Test build_indexable_documents_queryset with no filters returns all documents."""
+    factories.DocumentFactory.create_batch(3)
+
+    queryset = utils.build_indexable_documents_queryset()
+
+    assert queryset.count() == 3
+
+
+def test_utils_build_indexable_documents_queryset_oldest_updated_at():
+    """Test filtering by oldest_updated_at includes documents updated after timestamp."""
+    lower_time_bound = datetime(2024, 2, 1, tzinfo=timezone.utc)
+
+    document_too_old = create_document_with_updated_at(
+        updated_at=lower_time_bound - timedelta(days=30)
+    )
+    document_at_boundary = create_document_with_updated_at(lower_time_bound)
+    document_recent = create_document_with_updated_at(
+        updated_at=lower_time_bound + timedelta(days=15)
+    )
+
+    queryset = utils.build_indexable_documents_queryset(
+        lower_time_bound=lower_time_bound
+    )
+
+    assert queryset.count() == 2
+    assert document_too_old not in queryset
+    assert document_at_boundary in queryset
+    assert document_recent in queryset
+
+
+def test_utils_build_indexable_documents_queryset_newest_updated_at():
+    """Test filtering by newest_updated_at includes documents updated before timestamp."""
+    upper_time_bound = datetime(2024, 2, 1, tzinfo=timezone.utc)
+
+    document_old = create_document_with_updated_at(
+        updated_at=upper_time_bound - timedelta(days=30)
+    )
+    document_at_boundary = create_document_with_updated_at(upper_time_bound)
+    document_too_recent = create_document_with_updated_at(
+        updated_at=upper_time_bound + timedelta(days=15)
+    )
+
+    queryset = utils.build_indexable_documents_queryset(
+        upper_time_bound=upper_time_bound
+    )
+
+    assert queryset.count() == 2
+    assert document_old in queryset
+    assert document_at_boundary in queryset
+    assert document_too_recent not in queryset
+
+
+def test_utils_build_indexable_documents_queryset_both_bounds():
+    """Test filtering with both oldest and newest bounds."""
+    lower_time_bound = datetime(2024, 2, 1, tzinfo=timezone.utc)
+    upper_time_bound = lower_time_bound + timedelta(days=30)
+
+    document_too_early = create_document_with_updated_at(
+        updated_at=lower_time_bound - timedelta(days=10)
+    )
+    document_in_window = create_document_with_updated_at(
+        updated_at=lower_time_bound + timedelta(days=5)
+    )
+    other_document_in_window = create_document_with_updated_at(
+        updated_at=lower_time_bound + timedelta(days=15)
+    )
+    document_too_late = create_document_with_updated_at(
+        updated_at=upper_time_bound + timedelta(days=10)
+    )
+
+    queryset = utils.build_indexable_documents_queryset(
+        lower_time_bound=lower_time_bound, upper_time_bound=upper_time_bound
+    )
+
+    assert queryset.count() == 2
+    assert document_too_early not in queryset
+    assert document_in_window in queryset
+    assert other_document_in_window in queryset
+    assert document_too_late not in queryset
+
+
+def test_utils_build_indexable_documents_queryset_deleted_at():
+    """Test filtering includes documents with deleted_at in range."""
+    lower_time_bound = datetime(2024, 2, 1, tzinfo=timezone.utc)
+
+    document_deleted_in_range = create_document_with_updated_at(
+        updated_at=lower_time_bound - timedelta(days=30),
+        deleted_at=lower_time_bound + timedelta(days=15),
+    )
+    document_updated_before_range = create_document_with_updated_at(
+        updated_at=lower_time_bound - timedelta(days=30)
+    )
+
+    queryset = utils.build_indexable_documents_queryset(
+        lower_time_bound=lower_time_bound
+    )
+
+    assert queryset.count() == 1
+    assert document_deleted_in_range in queryset
+    assert document_updated_before_range not in queryset
+
+
+def test_utils_build_indexable_documents_queryset_ancestors_deleted_at():
+    """Test filtering includes documents with ancestors_deleted_at in range."""
+    lower_time_bound = datetime(2024, 2, 1, tzinfo=timezone.utc)
+
+    document_ancestor_deleted_in_range = create_document_with_updated_at(
+        updated_at=lower_time_bound - timedelta(days=30),
+        ancestors_deleted_at=lower_time_bound + timedelta(days=15),
+    )
+    document_updated_before_range = create_document_with_updated_at(
+        updated_at=lower_time_bound - timedelta(days=30)
+    )
+
+    queryset = utils.build_indexable_documents_queryset(
+        lower_time_bound=lower_time_bound
+    )
+
+    assert queryset.count() == 1
+    assert document_ancestor_deleted_in_range in queryset
+    assert document_updated_before_range not in queryset

src/backend/core/utils.py

@@ -8,7 +8,7 @@ from collections import defaultdict
 
 from django.core.cache import cache
 from django.db import models as db
-from django.db.models import Subquery
+from django.db.models import Q, Subquery
 
 import pycrdt
 from bs4 import BeautifulSoup
@@ -168,3 +168,47 @@ def users_sharing_documents_with(user):
         elapsed,
     )
     return result
+
+
+def build_indexable_documents_queryset(lower_time_bound=None, upper_time_bound=None):
+    """
+    Build a queryset of documents to index filtered by timestamp range.
+
+    filters out documents that have been updated, or deleted outside the time range.
+
+    Args:
+        lower_time_bound (datetime, optional):
+            keeps documents updated or deleted after this timestamp.
+        upper_time_bound (datetime, optional):
+            keeps documents updated or deleted before this timestamp.
+
+    Returns:
+        QuerySet: Filtered Document queryset ready for indexation.
+    """
+
+    queryset = models.Document.objects.all()
+
+    conditions = Q()
+    if lower_time_bound and upper_time_bound:
+        conditions = (
+            Q(updated_at__gte=lower_time_bound, updated_at__lte=upper_time_bound)
+            | Q(deleted_at__gte=lower_time_bound, deleted_at__lte=upper_time_bound)
+            | Q(
+                ancestors_deleted_at__gte=lower_time_bound,
+                ancestors_deleted_at__lte=upper_time_bound,
+            )
+        )
+    elif lower_time_bound:
+        conditions = (
+            Q(updated_at__gte=lower_time_bound)
+            | Q(deleted_at__gte=lower_time_bound)
+            | Q(ancestors_deleted_at__gte=lower_time_bound)
+        )
+    elif upper_time_bound:
+        conditions = (
+            Q(updated_at__lte=upper_time_bound)
+            | Q(deleted_at__lte=upper_time_bound)
+            | Q(ancestors_deleted_at__lte=upper_time_bound)
+        )
+
+    return queryset.filter(conditions)