✨ (backend) Add async triggers to enable document indexation with find.

Signed-off-by: Fabre Florian <[email protected]>

Author: joehybird

src/backend/core/models.py

@@ -20,7 +20,9 @@
 from django.core.files.storage import default_storage
 from django.core.mail import send_mail
 from django.db import models, transaction
+from django.db.models import signals
 from django.db.models.functions import Left, Length
+from django.dispatch import receiver
 from django.template.loader import render_to_string
 from django.utils import timezone
 from django.utils.functional import cached_property
@@ -39,6 +41,7 @@
     RoleChoices,
     get_equivalent_link_definition,
 )
+from .tasks.find import trigger_document_indexer
 
 logger = getLogger(__name__)
 
@@ -949,6 +952,16 @@ def restore(self):
             )
 
 
+@receiver(signals.post_save, sender=Document)
+def document_post_save(sender, instance, **kwargs):
+    """
+    Asynchronous call to the document indexer at the end of the transaction.
+    Note : Within the transaction we can have an empty content and a serialization
+    error.
+    """
+    trigger_document_indexer(instance, on_commit=True)
+
+
 class LinkTrace(BaseModel):
     """
     Relation model to trace accesses to a document via a link by a logged-in user.
@@ -1174,6 +1187,15 @@ def get_abilities(self, user):
         }
 
 
+@receiver(signals.post_save, sender=DocumentAccess)
+def document_access_post_save(sender, instance, created, **kwargs):
+    """
+    Asynchronous call to the document indexer at the end of the transaction.
+    """
+    if not created:
+        trigger_document_indexer(instance.document, on_commit=True)
+
+
 class DocumentAskForAccess(BaseModel):
     """Relation model to ask for access to a document."""
 

src/backend/core/services/search_indexers.py

@@ -19,7 +19,9 @@ def get_batch_accesses_by_users_and_teams(paths):
     grouped by users and teams, including all ancestor paths.
     """
     print("paths: ", paths)
-    ancestor_map = utils.get_ancestor_to_descendants_map(paths, steplen=models.Document.steplen)
+    ancestor_map = utils.get_ancestor_to_descendants_map(
+        paths, steplen=models.Document.steplen
+    )
     ancestor_paths = list(ancestor_map.keys())
     print("ancestor map: ", ancestor_map)
     print("ancestor paths: ", ancestor_paths)

src/backend/core/tasks/find.py

@@ -0,0 +1,82 @@
+"""Trigger document indexation using celery task."""
+
+from django.conf import settings
+from django.core.cache import cache
+from django.db import transaction
+
+from core import models
+from core.services.search_indexers import (
+    FindDocumentIndexer,
+    get_batch_accesses_by_users_and_teams,
+)
+
+from impress.celery_app import app
+
+
+def document_indexer_debounce_key(document_id):
+    return f"doc-indexer-debounce-{document_id}"
+
+
+def incr_counter(key):
+    try:
+        return cache.incr(key)
+    except ValueError:
+        cache.set(key, 1)
+        return 1
+
+
+def decr_counter(key):
+    try:
+        return cache.decr(key)
+    except ValueError:
+        cache.set(key, 0)
+        return 0
+
+
+@app.task
+def document_indexer_task(document_id):
+    """Send indexation query for a document using celery task."""
+    key = document_indexer_debounce_key(document_id)
+
+    # check if the counter : if still up, skip the task. only the last one
+    # within the countdown delay will do the query.
+    if decr_counter(key) > 0:
+        return
+
+    doc = models.Document.objects.get(pk=document_id)
+    indexer = FindDocumentIndexer()
+    accesses = get_batch_accesses_by_users_and_teams((doc.path,))
+
+    data = indexer.serialize_document(document=doc, accesses=accesses)
+
+    indexer.push(data)
+
+
+def trigger_document_indexer(document, on_commit=False):
+    """
+    Trigger indexation task with debounce a delay set by the SEARCH_INDEXER_COUNTDOWN setting.
+
+    Args:
+        document (Document): The document instance.
+        on_commit (bool): Wait for the end of the transaction before starting the task
+            (some fields may be in wrong state within the transaction)
+    """
+
+    if document.deleted_at or document.ancestors_deleted_at:
+        pass
+
+    if on_commit:
+
+        def _aux():
+            trigger_document_indexer(document, on_commit=False)
+
+        transaction.on_commit(_aux)
+    else:
+        key = document_indexer_debounce_key(document.pk)
+        countdown = getattr(settings, "SEARCH_INDEXER_COUNTDOWN", 1)
+
+        # Each time this method is called during the countdown, we increment the
+        # counter and each task decrease it, so the index be run only once.
+        incr_counter(key)
+
+        document_indexer_task.apply_async(args=[document.pk], countdown=countdown)

src/backend/core/tests/test_models_documents.py

@@ -5,6 +5,7 @@
 
 import random
 import smtplib
+import time
 from logging import Logger
 from unittest import mock
 
@@ -13,12 +14,15 @@
 from django.core.cache import cache
 from django.core.exceptions import ValidationError
 from django.core.files.storage import default_storage
+from django.db import transaction
 from django.test.utils import override_settings
 from django.utils import timezone
 
 import pytest
 
 from core import factories, models
+from core.services.search_indexers import FindDocumentIndexer
+from core.tasks.find import document_indexer_debounce_key
 
 pytestmark = pytest.mark.django_db
 
@@ -1323,3 +1327,99 @@ def test_models_documents_compute_ancestors_links_paths_mapping_structure(
                 {"link_reach": sibling.link_reach, "link_role": sibling.link_role},
             ],
         }
+
+
+@mock.patch.object(FindDocumentIndexer, "push")
+@pytest.mark.django_db(transaction=True)
+def test_models_documents_post_save_indexer(mock_push, settings):
+    settings.SEARCH_INDEXER_COUNTDOWN = 0
+
+    user = factories.UserFactory()
+
+    with transaction.atomic():
+        doc1, doc2, doc3 = factories.DocumentFactory.create_batch(3)
+
+        factories.UserDocumentAccessFactory(document=doc1, user=user)
+        factories.UserDocumentAccessFactory(document=doc2, user=user)
+        factories.UserDocumentAccessFactory(document=doc3, user=user)
+
+    time.sleep(0.1)  # waits for the end of the tasks
+
+    accesses = {
+        str(doc1.path): {"users": [user.sub]},
+        str(doc2.path): {"users": [user.sub]},
+        str(doc3.path): {"users": [user.sub]},
+    }
+
+    data = [call.args[0] for call in mock_push.call_args_list]
+
+    indexer = FindDocumentIndexer()
+
+    def sortkey(d):
+        return d["id"]
+
+    assert sorted(data, key=sortkey) == sorted(
+        [
+            indexer.serialize_document(doc1, accesses),
+            indexer.serialize_document(doc2, accesses),
+            indexer.serialize_document(doc3, accesses),
+        ],
+        key=sortkey,
+    )
+
+
+@mock.patch.object(FindDocumentIndexer, "push")
+@pytest.mark.django_db(transaction=True)
+def test_models_documents_post_save_indexer_debounce(mock_push, settings):
+    settings.SEARCH_INDEXER_COUNTDOWN = 0
+
+    indexer = FindDocumentIndexer()
+    user = factories.UserFactory()
+
+    with transaction.atomic():
+        doc = factories.DocumentFactory()
+        factories.UserDocumentAccessFactory(document=doc, user=user)
+
+    accesses = {
+        str(doc.path): {"users": [user.sub]},
+    }
+
+    time.sleep(0.1)  # waits for the end of the tasks
+
+    # One save from the factory
+    assert [call.args[0] for call in mock_push.call_args_list] == [
+        indexer.serialize_document(doc, accesses),
+    ]
+
+    # The debounce counter should be reset
+    assert cache.get(document_indexer_debounce_key(doc.pk)) == 0
+
+    # Now simulate 1 waiting task
+    cache.set(document_indexer_debounce_key(doc.pk), 1)
+
+    # save it again to trigger the indexer, but nothing should be done since
+    # the counter is over 0
+    with transaction.atomic():
+        updated_doc = models.Document.objects.get(pk=doc.pk)
+        updated_doc.save()
+
+    time.sleep(0.1)  # waits for the end of the tasks
+
+    # Still on the first indexation.
+    assert [call.args[0] for call in mock_push.call_args_list] == [
+        indexer.serialize_document(doc, accesses),
+    ]
+
+    # Reset the counter
+    cache.set(document_indexer_debounce_key(doc.pk), 0)
+
+    with transaction.atomic():
+        updated_doc = models.Document.objects.get(pk=doc.pk)
+        updated_doc.save()
+
+    time.sleep(0.1)  # waits for the end of the tasks
+
+    assert [call.args[0] for call in mock_push.call_args_list] == [
+        indexer.serialize_document(doc, accesses),
+        indexer.serialize_document(updated_doc, accesses),
+    ]

src/backend/core/tests/test_services_search_indexers.py

@@ -107,7 +107,7 @@ def test_services_search_indexers_serialize_document_returns_expected_json():
         "created_at": document.created_at.isoformat(),
         "updated_at": document.updated_at.isoformat(),
         "reach": document.link_reach,
-        'size': 13,
+        "size": 13,
         "is_active": True,
     }
 
@@ -168,7 +168,9 @@ def test_services_search_indexers_batches_pass_only_batch_accesses(mock_push, se
 def test_services_search_indexers_ancestors_link_reach(mock_push):
     """Document accesses and reach should take into account ancestors link reaches."""
     great_grand_parent = factories.DocumentFactory(link_reach="restricted")
-    grand_parent = factories.DocumentFactory(parent=great_grand_parent, link_reach="authenticated")
+    grand_parent = factories.DocumentFactory(
+        parent=great_grand_parent, link_reach="authenticated"
+    )
     parent = factories.DocumentFactory(parent=grand_parent, link_reach="public")
     document = factories.DocumentFactory(parent=parent, link_reach="restricted")
 
@@ -199,7 +201,11 @@ def test_services_search_indexers_ancestors_users(mock_push):
     assert len(results) == 3
     assert results[str(grand_parent.id)]["users"] == [str(user_gp.sub)]
     assert set(results[str(parent.id)]["users"]) == {str(user_gp.sub), str(user_p.sub)}
-    assert set(results[str(document.id)]["users"]) == {str(user_gp.sub), str(user_p.sub), str(user_d.sub)}
+    assert set(results[str(document.id)]["users"]) == {
+        str(user_gp.sub),
+        str(user_p.sub),
+        str(user_d.sub),
+    }
 
 
 @patch.object(FindDocumentIndexer, "push")

src/backend/core/tests/test_utils.py

@@ -79,24 +79,24 @@ def test_utils_extract_attachments():
 
 def test_utils_get_ancestor_to_descendants_map_single_path():
     """Test ancestor mapping of a single path."""
-    paths = ['000100020005']
+    paths = ["000100020005"]
     result = utils.get_ancestor_to_descendants_map(paths, steplen=4)
 
     assert result == {
-        '0001': {'000100020005'},
-        '00010002': {'000100020005'},
-        '000100020005': {'000100020005'},
+        "0001": {"000100020005"},
+        "00010002": {"000100020005"},
+        "000100020005": {"000100020005"},
     }
 
 
 def test_utils_get_ancestor_to_descendants_map_multiple_paths():
     """Test ancestor mapping of multiple paths with shared prefixes."""
-    paths = ['000100020005', '00010003']
+    paths = ["000100020005", "00010003"]
     result = utils.get_ancestor_to_descendants_map(paths, steplen=4)
 
     assert result == {
-        '0001': {'000100020005', '00010003'},
-        '00010002': {'000100020005'},
-        '000100020005': {'000100020005'},
-        '00010003': {'00010003'},
+        "0001": {"000100020005", "00010003"},
+        "00010002": {"000100020005"},
+        "000100020005": {"000100020005"},
+        "00010003": {"00010003"},
     }

src/backend/core/utils.py

@@ -1,8 +1,8 @@
 """Utils for the core app."""
 
 import base64
-from collections import defaultdict
 import re
+from collections import defaultdict
 
 import pycrdt
 from bs4 import BeautifulSoup