LITE-31232 Init implementation of bulk_relate_cqrs_serialization

Author: maxipavlovic

dj_cqrs/mixins.py

@@ -1,4 +1,4 @@
-#  Copyright © 2023 Ingram Micro Inc. All rights reserved.
+#  Copyright © 2024 Ingram Micro Inc. All rights reserved.
 
 import logging
 
@@ -20,6 +20,7 @@
 from dj_cqrs.managers import MasterManager, ReplicaManager
 from dj_cqrs.metas import MasterMeta, ReplicaMeta
 from dj_cqrs.signals import MasterSignals, post_bulk_create, post_update
+from dj_cqrs.state import cqrs_state
 
 
 logger = logging.getLogger('django-cqrs')
@@ -292,9 +293,16 @@ def _class_serialization(self, using, sync=False):
         if sync:
             instance = self
         else:
+            instance = None
             db = using if using is not None else self._state.db
-            qs = self.__class__._default_manager.using(db)
-            instance = self.relate_cqrs_serialization(qs).get(pk=self.pk)
+
+            bulk_relate_cm = cqrs_state.bulk_relate_cm
+            if bulk_relate_cm:
+                instance = bulk_relate_cm.get_cached_instance(self, db)
+
+            if not instance:
+                qs = self.__class__._default_manager.using(db)
+                instance = self.relate_cqrs_serialization(qs).get(pk=self.pk)
 
         data = self._cqrs_serializer_cls(instance).data
         data['cqrs_revision'] = instance.cqrs_revision

dj_cqrs/signals.py

@@ -1,4 +1,4 @@
-#  Copyright © 2023 Ingram Micro Inc. All rights reserved.
+#  Copyright © 2024 Ingram Micro Inc. All rights reserved.
 
 import logging
 
@@ -9,6 +9,7 @@
 from dj_cqrs.constants import SignalType
 from dj_cqrs.controller import producer
 from dj_cqrs.dataclasses import TransportPayload
+from dj_cqrs.state import cqrs_state
 from dj_cqrs.utils import get_message_expiration_dt
 
 
@@ -64,6 +65,10 @@ def post_save(cls, sender, **kwargs):
 
         using = kwargs['using']
 
+        bulk_relate_cm = cqrs_state.bulk_relate_cm
+        if bulk_relate_cm:
+            bulk_relate_cm.register(instance, using)
+
         sync = kwargs.get('sync', False)
         queue = kwargs.get('queue', None)
 

dj_cqrs/state.py

@@ -0,0 +1,7 @@
+#  Copyright © 2024 Ingram Micro Inc. All rights reserved.
+
+import threading
+
+
+cqrs_state = threading.local()
+cqrs_state.bulk_relate_cm = None

dj_cqrs/utils.py

@@ -1,6 +1,8 @@
-#  Copyright © 2023 Ingram Micro Inc. All rights reserved.
+#  Copyright © 2024 Ingram Micro Inc. All rights reserved.
 
 import logging
+from collections import defaultdict
+from contextlib import ContextDecorator
 from datetime import date, datetime, timedelta
 from uuid import UUID
 
@@ -10,6 +12,7 @@
 
 from dj_cqrs.constants import DB_VENDOR_PG, SUPPORTED_TIMEOUT_DB_VENDORS
 from dj_cqrs.logger import install_last_query_capturer
+from dj_cqrs.state import cqrs_state
 
 
 logger = logging.getLogger('django-cqrs')
@@ -80,3 +83,55 @@ def apply_query_timeouts(model_cls):  # pragma: no cover
         cursor.execute(statement, params=(query_timeout,))
 
     install_last_query_capturer(model_cls)
+
+
+class _BulkRelateCM(ContextDecorator):
+    def __init__(self, cqrs_id=None):
+        self._cqrs_id = cqrs_id
+        self._mapping = defaultdict(lambda: defaultdict(list))
+        self._cache = {}
+
+    def register(self, instance, using):
+        instance_cqrs_id = getattr(instance, 'CQRS_ID', None)
+        if self._cqrs_id and instance_cqrs_id != self._cqrs_id:
+            return
+
+        self._mapping[instance_cqrs_id][using].append(instance.pk)
+
+    def get_cached_instance(self, instance, using):
+        instance_cqrs_id = getattr(instance, 'CQRS_ID', None)
+        if self._cqrs_id and instance_cqrs_id != self._cqrs_id:
+            return
+
+        instance_pk = instance.pk
+        cached_instances = self._cache.get(instance_cqrs_id, {}).get(using, {})
+        if cached_instances:
+            return cached_instances.get(instance_pk)
+
+        cached_pks = self._mapping[instance_cqrs_id][using]
+        if not cached_pks:
+            return
+
+        qs = instance.__class__._default_manager.using(using)
+        instances_cache = {
+            instance.pk: instance
+            for instance in instance.__class__.relate_cqrs_serialization(qs).filter(
+                pk__in=cached_pks,
+            ).order_by().all()
+        }
+        self._cache.update({
+            instance_cqrs_id: {
+                using: instances_cache,
+            },
+        })
+        return instances_cache.get(instance_pk)
+
+    def __enter__(self):
+        cqrs_state.bulk_relate_cm = self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        cqrs_state.bulk_relate_cm = None
+
+
+def bulk_relate_cqrs_serialization(cqrs_id=None):
+    return _BulkRelateCM(cqrs_id=cqrs_id)

tests/test_master/test_signals.py

@@ -1,4 +1,4 @@
-#  Copyright © 2023 Ingram Micro Inc. All rights reserved.
+#  Copyright © 2024 Ingram Micro Inc. All rights reserved.
 
 from datetime import datetime, timezone
 
@@ -8,6 +8,7 @@
 
 from dj_cqrs.constants import SignalType
 from dj_cqrs.signals import post_bulk_create, post_update
+from dj_cqrs.utils import bulk_relate_cqrs_serialization
 from tests.dj_master import models
 from tests.utils import assert_is_sub_dict, assert_publisher_once_called_with_args
 
@@ -127,6 +128,28 @@ def test_manual_post_bulk_create(mocker):
     assert publisher_mock.call_count == 3
 
 
+@pytest.mark.django_db(transaction=True)
+@pytest.mark.parametrize('count', (1, 3, 5))
+def test_bulk_relate_cqrs_serialization(
+    django_assert_num_queries,
+    django_v_trans_q_count_sup,
+    mocker,
+    count,
+):
+    mocker.patch('dj_cqrs.controller.producer.produce')
+
+    opt_query_count = count + 2 + django_v_trans_q_count_sup
+    with django_assert_num_queries(opt_query_count):
+        with bulk_relate_cqrs_serialization():
+            with transaction.atomic(savepoint=False):
+                [models.Author.objects.create(id=i) for i in range(count)]
+
+    not_opt_query_count = count + count * 2 + django_v_trans_q_count_sup
+    with django_assert_num_queries(not_opt_query_count):
+        with transaction.atomic(savepoint=False):
+            [models.Author.objects.create(id=10 + i) for i in range(count)]
+
+
 @pytest.mark.django_db(transaction=True)
 def test_automatic_post_bulk_create(mocker):
     publisher_mock = mocker.patch('dj_cqrs.controller.producer.produce')