Repository: Codoc-os/django-opensearch-dsl (language: python)

Issue:
Update Indices on Related Models
I am using django-opensearch-dsl with Wagtail, which relies fairly heavily on relations. I would like to be able to have update run on docs that have related models. It looks as though support for this was dropped in 0.1.0 (2021-12-11). I'm not sure if this was intentional, or if it just hasn't been added into django-opensearch-dsl yet. Unless it was an intentional choice to not include that, it would be a useful feature to bring back. Specifically, in https://github.com/Codoc-os/django-opensearch-dsl/blob/master/django_opensearch_dsl/registries.py support for:

`_get_related_doc()`
`update_related()`
`delete_related()`

For example, if a related page is moved or deleted, the signal should cause the index to update.
Update Indices on Related Models
I am using django-opensearch-dsl with Wagtail, which relies fairly heavily on relations. I would like to be able to have update run on docs that have related models. It looks as though support for this was dropped in 0.1.0 (2021-12-11). I'm not sure if this was intentional, or if it just hasn't been added into django-opensearch-dsl yet. Unless it was an intentional choice to not include that, it would be a useful feature to bring back. Specifically, in https://github.com/Codoc-os/django-opensearch-dsl/blob/master/django_opensearch_dsl/registries.py support for:

`_get_related_doc()`
`update_related()`
`delete_related()`

For example, if a related page is moved or deleted, the signal should cause the index to update.


Discussion / hints:
@qcoumes I can create a fork and work on adding this tomorrow or Monday. Support will also require some tests to be re-added (registries, and signals) and modified (in the signals test file), so I would like to have a brief discussion and make sure we're on the same page when I suggest upstream changes. May I email you and book some time to meet virtually?
@Colin-Seifer Sure, go ahead!
@qcoumes I can create a fork and work on adding this tomorrow or Monday. Support will also require some tests to be re-added (registries, and signals) and modified (in the signals test file), so I would like to have a brief discussion and make sure we're on the same page when I suggest upstream changes. May I email you and book some time to meet virtually?
@Colin-Seifer Sure, go ahead!

Fix patch:
diff --git a/django_opensearch_dsl/documents.py b/django_opensearch_dsl/documents.py
index a34d7fa..6dd398a 100644
--- a/django_opensearch_dsl/documents.py
+++ b/django_opensearch_dsl/documents.py
@@ -49,8 +49,11 @@ class Document(DSLDocument):
 
     _prepared_fields = []
 
-    def __init__(self, **kwargs):
+    def __init__(self, related_instance_to_ignore=None, **kwargs):
         super(Document, self).__init__(**kwargs)
+        # related instances to ignore is required to remove the instance
+        # from related models on deletion.
+        self._related_instance_to_ignore = related_instance_to_ignore
         self._prepared_fields = self.init_prepare()
 
     @classmethod
@@ -132,8 +135,15 @@ def init_prepare(self):
             if not field._path:  # noqa
                 field._path = [name]
 
-            prep_func = getattr(self, "prepare_%s" % name, None)
-            fn = prep_func if prep_func else partial(field.get_value_from_instance)
+            prep_func = getattr(self, "prepare_%s_with_related" % name, None)
+            if prep_func:
+                fn = partial(prep_func, related_to_ignore=self._related_instance_to_ignore)
+            else:
+                prep_func = getattr(self, "prepare_%s" % name, None)
+                if prep_func:
+                    fn = prep_func
+                else:
+                    fn = partial(field.get_value_from_instance, field_value_to_ignore=self._related_instance_to_ignore)
 
             preparers.append((name, field, fn))
 
diff --git a/django_opensearch_dsl/indices.py b/django_opensearch_dsl/indices.py
new file mode 100644
index 0000000..6426aa9
--- /dev/null
+++ b/django_opensearch_dsl/indices.py
@@ -0,0 +1,26 @@
+from copy import deepcopy
+
+from opensearch_dsl import Index as DSLIndex
+
+from .apps import DODConfig
+from .registries import registry
+
+
+class Index(DSLIndex):
+    """Creates an index and makes a deep copy of default index settings."""
+
+    def __init__(self, *args, **kwargs):
+        super(Index, self).__init__(*args, **kwargs)
+        default_index_settings = deepcopy(DODConfig.default_index_settings())
+        self.settings(**default_index_settings)
+
+    def document(self, document):
+        """Extend to register the document in the global document registry."""
+        document = super(Index, self).document(document)
+        registry.register_document(document)
+        return document
+
+    doc_type = document
+
+    def __str__(self):
+        return self._name
diff --git a/django_opensearch_dsl/registries.py b/django_opensearch_dsl/registries.py
index 2faaeb3..3db1c1d 100644
--- a/django_opensearch_dsl/registries.py
+++ b/django_opensearch_dsl/registries.py
@@ -1,6 +1,7 @@
 from collections import defaultdict
 from copy import deepcopy
 
+from django.core.exceptions import ObjectDoesNotExist
 from django.core.exceptions import ImproperlyConfigured
 from opensearch_dsl import AttrDict
 
@@ -14,11 +15,15 @@ class DocumentRegistry:
     def __init__(self):
         self._indices = defaultdict(set)
         self._models = defaultdict(set)
+        self._related_models = defaultdict(set)
 
     def register(self, index, doc_class):
         """Register the model with the registry."""
         self._models[doc_class.django.model].add(doc_class)
 
+        for related in doc_class.django.related_models:
+            self._related_models[related].add(doc_class.django.model)
+
         for idx, docs in self._indices.items():
             if index._name == idx._name:  # noqa pragma: no cover
                 docs.add(doc_class)
@@ -43,6 +48,7 @@ def register_document(self, document):
                 ),
                 "ignore_signals": getattr(django_meta, "ignore_signals", False),
                 "auto_refresh": getattr(django_meta, "auto_refresh", DODConfig.auto_refresh_enabled()),
+                "related_models": getattr(django_meta, "related_models", []),
             }
         )
         if not django_attr.model:  # pragma: no cover
@@ -79,6 +85,52 @@ def register_document(self, document):
 
         return document
 
+    def _get_related_doc(self, instance):
+        for model in self._related_models.get(instance.__class__, []):
+            for doc in self._models[model]:
+                if instance.__class__ in doc.django.related_models:
+                    yield doc
+
+    def update_related(self, instance, **kwargs):
+        """Update documents related to `instance`.
+
+        Related documents are found using the `get_instances_from_related()`
+        method of each Document classes including the instance's Model is
+        including within their `related_models`.
+        """
+        if not DODConfig.autosync_enabled():
+            return
+
+        for doc in self._get_related_doc(instance):
+            doc_instance = doc()
+            try:
+                related = doc_instance.get_instances_from_related(instance)
+            except ObjectDoesNotExist:
+                related = None
+
+            if related is not None:
+                doc_instance.update(related, **kwargs)
+
+    def delete_related(self, instance, **kwargs):
+        """Remove `instance` from related models.
+
+        `related_instance_to_ignore` ensures that `instance` is only removed
+        from related models. We don't want to update the indexed value for
+        `instance` in this method. Prevents potential orphaned objects.
+        """
+        if not DODConfig.autosync_enabled():
+            return
+
+        for doc in self._get_related_doc(instance):
+            doc_instance = doc(related_instance_to_ignore=instance)
+            try:
+                related = doc_instance.get_instances_from_related(instance)
+            except ObjectDoesNotExist:
+                related = None
+
+            if related is not None:
+                doc_instance.update(related, **kwargs)
+
     def update(self, instance, action="index", **kwargs):
         """Update all the opensearch documents attached to this model.
 
@@ -101,8 +153,11 @@ def delete(self, instance, **kwargs):
         """
         self.update(instance, action="delete", **kwargs)
 
-    def get_indices(self):
-        """Get all indices in the registry."""
+    def get_indices(self, models=None):
+        """Get all indices in the registry or the indices for a list of models."""
+        if models is not None:
+            return set(index for index, docs in self._indices.items() for doc in docs if doc.django.model in models)
+
         return set(self._indices.keys())
 
 
diff --git a/django_opensearch_dsl/signals.py b/django_opensearch_dsl/signals.py
index 6833937..75830ee 100644
--- a/django_opensearch_dsl/signals.py
+++ b/django_opensearch_dsl/signals.py
@@ -49,6 +49,15 @@ def handle_save(self, sender, instance, **kwargs):
         Update the related objects either.
         """
         registry.update(instance)
+        registry.update_related(instance)
+
+    def handle_pre_delete(self, sender, instance, **kwargs):
+        """Handle removing of instance object from related models instance.
+
+        We need to do this before the real delete otherwise the relation
+        doesn't exists anymore and we can't get the related models instance.
+        """
+        registry.delete_related(instance)
 
     def handle_delete(self, sender, instance, **kwargs):
         """Handle delete.
@@ -73,6 +82,7 @@ def setup(self):
 
         # Use to manage related objects update
         models.signals.m2m_changed.connect(self.handle_m2m_changed)
+        models.signals.pre_delete.connect(self.handle_pre_delete)
 
     def teardown(self):
         """Tear down the SignalProcessor."""
@@ -80,6 +90,7 @@ def teardown(self):
         models.signals.post_save.disconnect(self.handle_save)
         models.signals.post_delete.disconnect(self.handle_delete)
         models.signals.m2m_changed.disconnect(self.handle_m2m_changed)
+        models.signals.pre_delete.disconnect(self.handle_pre_delete)
 
 
 # Sent after document indexing is completed
diff --git a/docs/document.md b/docs/document.md
index 45f6691..715893d 100644
--- a/docs/document.md
+++ b/docs/document.md
@@ -10,6 +10,8 @@ The `Django` subclass contains parameters related to Django's side of the docume
   into this list. See [Document Field Reference](fields.md) for how to manually define fields.
 * `queryset_pagination` (*optional*) - Size of the chunk when indexing,
   override [`OPENSEARCH_DSL_QUERYSET_PAGINATION`](settings.md#opensearch_dsl_queryset_pagination.md).
+* `related_models` (*optional*) - List of related Django models. Specifies a relation between models that allows for
+  index updating based on these defined relationships.
 
 ```python
 class Country(models.Model):
@@ -29,6 +31,7 @@ class CountryDocument(Document):
             'area',
             'population',
         ]
+        related_models = []
 
     id = fields.LongField()
     continent = fields.ObjectField(properties={
diff --git a/docs/fields.md b/docs/fields.md
index 8e3703c..951ffa9 100644
--- a/docs/fields.md
+++ b/docs/fields.md
@@ -102,6 +102,29 @@ class CarDocument(Document):
         return " ".join(instance.foos)
 ```
 
+## Using `prepare_[field]_with_related`
+
+Allows you to do extra prepping before a field with related models should be saved to Opensearch. This ensures that
+the index is updated appropriately. You can add a `prepare_[field]_with_related(self, instance)` method to a 
+`Document` (where `foo` is the name of the field), and that method will be called when the field needs to be saved.
+
+```python
+# documents.py
+
+class CarDocument(Document):
+    # ... #
+
+    class Django:
+        fields = ["name", "price"]
+        model = Car
+        related_models = [Manufacturer]
+
+    foo = TextField()
+
+    def prepare_foo_with_related(self, instance):
+        return " ".join(instance.foos)
+```
+
 ## Handle relationship with `NestedField` / `ObjectField`
 
 To represent relationships, you can use `NestedField` for :


Test patch:
diff --git a/tests/tests/fixtures.py b/tests/tests/fixtures.py
new file mode 100644
index 0000000..a84aed0
--- /dev/null
+++ b/tests/tests/fixtures.py
@@ -0,0 +1,50 @@
+from unittest.mock import Mock
+
+from django.db import models
+
+from django_opensearch_dsl.documents import Document
+
+
+class WithFixturesMixin(object):
+    """Used for generating mock documents based on the models defined here."""
+
+    class ModelA(models.Model):
+        class Meta:
+            app_label = "foo"
+
+    class ModelB(models.Model):
+        class Meta:
+            app_label = "foo"
+
+    class ModelC(models.Model):
+        class Meta:
+            app_label = "bar"
+
+    class ModelD(models.Model):
+        class Meta:
+            app_label = "foo"
+
+    class ModelE(models.Model):
+        class Meta:
+            app_label = "foo"
+
+    def _generate_doc_mock(self, _model, index=None, mock_qs=None, _ignore_signals=False, _related_models=None):
+        _index = index
+
+        class Doc(Document):
+            class Django:
+                model = _model
+                related_models = _related_models if _related_models is not None else []
+                ignore_signals = _ignore_signals
+
+        if _index:
+            _index.document(Doc)
+            self.registry.register_document(Doc)
+
+        Doc.update = Mock()
+        if mock_qs:
+            Doc.get_queryset = Mock(return_value=mock_qs)
+        if _related_models:
+            Doc.get_instances_from_related = Mock()
+
+        return Doc
diff --git a/tests/tests/test_documents.py b/tests/tests/test_documents.py
index 5373c9b..0c0b01f 100644
--- a/tests/tests/test_documents.py
+++ b/tests/tests/test_documents.py
@@ -61,6 +61,7 @@ class Meta:
     class Django:
         fields = ["name", "price"]
         model = Car
+        related_models = [Manufacturer]
 
     class Index:
         name = "car_index"
@@ -100,6 +101,10 @@ def test_fields_populated(self):
         mapping = CarDocument._doc_type.mapping
         self.assertEqual(set(mapping.properties.properties.to_dict().keys()), {"color", "name", "price", "type"})
 
+    def test_related_models_added(self):
+        related_models = CarDocument.django.related_models
+        self.assertEqual([Manufacturer], related_models)
+
     def test_duplicate_field_names_not_allowed(self):
         with self.assertRaises(RedeclaredFieldError):
 
diff --git a/tests/tests/test_indices.py b/tests/tests/test_indices.py
new file mode 100644
index 0000000..162f7dd
--- /dev/null
+++ b/tests/tests/test_indices.py
@@ -0,0 +1,45 @@
+from unittest import TestCase
+from unittest.mock import patch
+
+from django.conf import settings
+from django.test import override_settings
+
+from django_opensearch_dsl.indices import Index
+from django_opensearch_dsl.registries import DocumentRegistry
+
+from .fixtures import WithFixturesMixin
+
+
+class IndexTestCase(WithFixturesMixin, TestCase):
+    def setUp(self):
+        self.registry = DocumentRegistry()
+
+    def test_documents_add_to_register(self):
+        registry = self.registry
+        with patch("django_opensearch_dsl.indices.registry", new=registry):
+            index = Index("test")
+            doc_a1 = self._generate_doc_mock(self.ModelA)
+            index.document(doc_a1)
+            indices = list(registry.get_indices())
+            self.assertEqual(len(indices), 1)
+            self.assertEqual(indices[0], index)
+
+    def test__str__(self):
+        index = Index("test")
+        self.assertEqual(index.__str__(), "test")
+
+    @override_settings(
+        OPENSEARCH_DSL_INDEX_SETTINGS={
+            "number_of_replicas": 0,
+            "number_of_shards": 2,
+        }
+    )
+    def test__init__(self):
+        index = Index("test")
+        self.assertEqual(
+            index._settings,
+            {
+                "number_of_replicas": 0,
+                "number_of_shards": 2,
+            },
+        )
diff --git a/tests/tests/test_registries.py b/tests/tests/test_registries.py
new file mode 100644
index 0000000..9eb7516
--- /dev/null
+++ b/tests/tests/test_registries.py
@@ -0,0 +1,179 @@
+from unittest import mock
+from unittest.mock import Mock
+from unittest import TestCase
+
+from django.core.exceptions import ObjectDoesNotExist
+from django.conf import settings
+from django.test import override_settings
+
+from django_opensearch_dsl.registries import DocumentRegistry
+from django_opensearch_dsl.indices import Index
+
+from .fixtures import WithFixturesMixin
+
+
+class DocumentRegistryTestCase(WithFixturesMixin, TestCase):
+    def setUp(self):
+        self.registry = DocumentRegistry()
+        self.index_1 = Index(name="index_1")
+        self.index_2 = Index(name="index_2")
+
+        self.doc_a1 = self._generate_doc_mock(self.ModelA, self.index_1)
+        self.doc_a2 = self._generate_doc_mock(self.ModelA, self.index_1)
+        self.doc_b1 = self._generate_doc_mock(self.ModelB, self.index_2)
+        self.doc_c1 = self._generate_doc_mock(self.ModelC, self.index_1)
+        self.doc_d1 = self._generate_doc_mock(self.ModelD, self.index_1, _related_models=[self.ModelE])
+        self.doc_e1 = self._generate_doc_mock(self.ModelE, self.index_1)
+
+    def test_empty_registry(self):
+        registry = DocumentRegistry()
+        self.assertEqual(registry._indices, {})
+        self.assertEqual(registry._models, {})
+
+    def test_register(self):
+        self.assertEqual(self.registry._models[self.ModelA], set([self.doc_a1, self.doc_a2]))
+        self.assertEqual(self.registry._models[self.ModelB], set([self.doc_b1]))
+
+        self.assertEqual(
+            self.registry._indices[self.index_1], set([self.doc_a1, self.doc_a2, self.doc_c1, self.doc_d1, self.doc_e1])
+        )
+        self.assertEqual(self.registry._indices[self.index_2], set([self.doc_b1]))
+
+    def test_register_with_related_models(self):
+        self.assertEqual(self.registry._related_models[self.ModelE], set([self.ModelD]))
+
+    def test_get_related_doc(self):
+        instance = self.ModelE()
+        related_set = set()
+        for doc in self.registry._get_related_doc(instance):
+            related_set.add(doc)
+        self.assertEqual(related_set, set([self.doc_d1]))
+
+    def test_get_indices(self):
+        self.assertEqual(self.registry.get_indices(), set([self.index_1, self.index_2]))
+
+    def test_get_indices_by_model(self):
+        self.assertEqual(self.registry.get_indices([self.ModelA]), set([self.index_1]))
+
+    def test_get_indices_by_unregister_model(self):
+        ModelC = Mock()
+        self.assertFalse(self.registry.get_indices([ModelC]))
+
+    def test_update_instance(self):
+        doc_a3 = self._generate_doc_mock(self.ModelA, self.index_1, _ignore_signals=True)
+
+        instance = self.ModelA()
+        self.registry.update(instance)
+
+        self.assertFalse(doc_a3.update.called)
+        self.assertFalse(self.doc_b1.update.called)
+        self.doc_a1.update.assert_called_once_with(instance, "index")
+        self.doc_a2.update.assert_called_once_with(instance, "index")
+
+    def test_update_related_instances(self):
+        doc_d1 = self._generate_doc_mock(self.ModelD, self.index_1, _related_models=[self.ModelE, self.ModelB])
+        doc_d2 = self._generate_doc_mock(self.ModelD, self.index_1, _related_models=[self.ModelE])
+
+        instance_e = self.ModelE()
+        instance_b = self.ModelB()
+        related_instance = self.ModelD()
+
+        doc_d2.get_instances_from_related.return_value = related_instance
+        doc_d1.get_instances_from_related.return_value = related_instance
+        self.registry.update_related(instance_e)
+
+        doc_d1.get_instances_from_related.assert_called_once_with(instance_e)
+        doc_d1.update.assert_called_once_with(related_instance)
+        doc_d2.get_instances_from_related.assert_called_once_with(instance_e)
+        doc_d2.update.assert_called_once_with(related_instance)
+
+        doc_d1.get_instances_from_related.reset_mock()
+        doc_d1.update.reset_mock()
+        doc_d2.get_instances_from_related.reset_mock()
+        doc_d2.update.reset_mock()
+
+        self.registry.update_related(instance_b)
+        doc_d1.get_instances_from_related.assert_called_once_with(instance_b)
+        doc_d1.update.assert_called_once_with(related_instance)
+        doc_d2.get_instances_from_related.assert_not_called()
+        doc_d2.update.assert_not_called()
+
+    def test_update_related_instances_not_defined(self):
+        doc_d1 = self._generate_doc_mock(_model=self.ModelD, index=self.index_1, _related_models=[self.ModelE])
+
+        instance = self.ModelE()
+
+        doc_d1.get_instances_from_related.side_effect = ObjectDoesNotExist(Mock(return_value=None))
+        self.registry.update_related(instance)
+
+        doc_d1.get_instances_from_related.assert_called_once_with(instance)
+        doc_d1.update.assert_not_called()
+
+    def test_delete_instance(self):
+        doc_a3 = self._generate_doc_mock(self.ModelA, self.index_1, _ignore_signals=True)
+
+        instance = self.ModelA()
+        self.registry.delete(instance)
+
+        self.assertFalse(doc_a3.update.called)
+        self.assertFalse(self.doc_b1.update.called)
+        self.doc_a1.update.assert_called_once_with(instance, "delete")
+        self.doc_a2.update.assert_called_once_with(instance, "delete")
+
+    def test_delete_related_instance(self):
+        """
+        Test that update is called when related_instance is defined.
+
+        doc_d1 is a related doc for instance. If we call delete_related on
+        instance, then update should be called with related_instance.
+        """
+        doc_d1 = self._generate_doc_mock(self.ModelD, self.index_1, _related_models=[self.ModelE])
+
+        instance = self.ModelE()
+        related_instance = self.ModelD()
+
+        doc_d1.get_instances_from_related.return_value = related_instance
+        self.registry.delete_related(instance)
+
+        doc_d1.get_instances_from_related.assert_called_once_with(instance)
+        doc_d1.update.assert_called_once_with(related_instance)
+
+    def test_delete_related_instance_not_defined(self):
+        """
+        Test that update is not called if related not defined.
+
+        If we try to call delete_related on instance when
+        get_instances_from_related returns None on doc_d1, then update
+        should not be called on doc_d1.
+        """
+        doc_d1 = self._generate_doc_mock(_model=self.ModelD, index=self.index_1, _related_models=[self.ModelE])
+
+        instance = self.ModelE()
+
+        doc_d1.get_instances_from_related.side_effect = ObjectDoesNotExist(Mock(return_value=None))
+        self.registry.delete_related(instance)
+
+        doc_d1.get_instances_from_related.assert_called_once_with(instance)
+        doc_d1.update.assert_not_called()
+
+    @mock.patch("django_opensearch_dsl.apps.DODConfig.autosync_enabled")
+    def test_delete_related_autosync_disabled(self, autosync_mock):
+        """If autosync is disabled, delete_related should just return."""
+
+        autosync_mock.return_value = False
+
+        doc_d1 = self._generate_doc_mock(self.ModelD, self.index_1, _related_models=[self.ModelE])
+
+        instance = self.ModelE()
+
+        doc_d1.get_instances_from_related.return_value = None
+        self.registry.delete_related(instance)
+
+        doc_d1.get_instances_from_related.assert_not_called()
+        doc_d1.update.assert_not_called()
+
+    @override_settings(OPENSEARCH_DSL_AUTOSYNC=False)
+    def test_autosync(self):
+        instance = self.ModelA()
+        self.registry.update(instance)
+        self.assertFalse(self.doc_a1.update.called)