Merge pull request #109 from postatum/97736564_refactor_indexing

Implement relationships bulk indexation

nefertari/elasticsearch.py

@@ -2,6 +2,7 @@
 import json
 import logging
 from functools import partial
+from collections import defaultdict
 
 import elasticsearch
 from elasticsearch import helpers
@@ -631,3 +632,25 @@ def index_relations(cls, db_obj, request=None, **kwargs):
             if getattr(model_cls, '_index_enabled', False) and documents:
                 cls(model_cls.__name__).index(
                     to_dicts(documents), request=request)
+
+    @classmethod
+    def bulk_index_relations(cls, items, request=None, **kwargs):
+        """ Index objects related to :items: in bulk.
+
+        Related items are first grouped in map
+        {model_name: {item1, item2, ...}} and then indexed.
+
+        :param items: Sequence of DB objects related objects if which
+            should be indexed.
+        :param request: Pyramid Request instance.
+        """
+        index_map = defaultdict(set)
+        for item in items:
+            relations = item.get_related_documents(**kwargs)
+            for model_cls, related_items in relations:
+                indexable = getattr(model_cls, '_index_enabled', False)
+                if indexable and related_items:
+                    index_map[model_cls.__name__].update(related_items)
+
+        for model_name, instances in index_map.items():
+            cls(model_name).index(to_dicts(instances), request=request)

tests/test_elasticsearch.py

@@ -854,3 +854,24 @@ class Foo(object):
         mock_settings.index_name = 'foo'
         es.ES.index_relations(db_obj)
         assert not mock_ind.called
+
+    @patch('nefertari.elasticsearch.ES.settings')
+    @patch('nefertari.elasticsearch.ES.index')
+    def test_bulk_index_relations(self, mock_index, mock_settings):
+        mock_settings.index_name = 'foo'
+
+        class Foo(int):
+            _index_enabled = True
+
+        doc1 = Foo(1)
+        doc2 = Foo(2)
+
+        db_object1 = Mock()
+        db_object1.get_related_documents.return_value = [
+            (Foo, [doc1])]
+        db_object2 = Mock()
+        db_object2.get_related_documents.return_value = [
+            (Foo, [doc2])]
+
+        es.ES.bulk_index_relations([db_object1, db_object2])
+        mock_index.assert_called_once_with(sorted([doc1, doc2]), request=None)