⚡️ improve main geo data management

1 files changed, 31 insertions, 34 deletions

diff --git a/ishtar_common/models_common.py b/ishtar_common/models_common.py
index a5b4338f2..daf2d9885 100644
--- a/ishtar_common/models_common.py
+++ b/ishtar_common/models_common.py
@@ -3007,16 +3007,16 @@ def geodata_attached_post_add(model, instance, pk_set):
         m2m_key = f"{child_model._meta.model_name}_id"
         geoitems = {}
         for child_id in query.values_list("id", flat=True):
-            child = None
             for pk in item_pks:
                 q = m2m_model.objects.filter(**{m2m_key: child_id,
                                                 "geovectordata_id": pk})
-                if not q.count():
-                    if not child:
-                        child = model.objects.get(pk=pk)
-                    if pk not in geoitems:
-                        geoitems[pk] = GeoVectorData.objects.get(pk=pk)
-                    child_model.objects.get(pk=child_id).geodata.add(geoitems[pk])
+                if q.count():
+                    continue
+                if pk not in geoitems:
+                    geoitems[pk] = GeoVectorData.objects.get(pk=pk)
+                child = child_model.objects.get(pk=child_id)
+                child.geodata.add(geoitems[pk])
+                child._manage_main_geo()
 
 
 def geodata_attached_remove(model, instance, pk_set=None, clear=False):
@@ -3041,16 +3041,19 @@ def geodata_attached_remove(model, instance, pk_set=None, clear=False):
         m2m_key = f"{child_model._meta.model_name}_id"
         geoitems = {}
         for child_id in query.values_list("id", flat=True):
-            child = None
             for pk in item_pks:
+                q = child_model.objects.filter(pk=child_id, main_geodata_id=pk)
+                if q.count():
+                    q.update(main_geodata_id=None)
                 q = m2m_model.objects.filter(**{m2m_key: child_id,
                                                 "geovectordata_id": pk})
-                if q.count():
-                    if not child:
-                        child = model.objects.get(pk=pk)
-                    if pk not in geoitems:
-                        geoitems[pk] = GeoVectorData.objects.get(pk=pk)
-                    child_model.objects.get(pk=child_id).geodata.remove(geoitems[pk])
+                if not q.count():
+                    continue
+                if pk not in geoitems:
+                    geoitems[pk] = GeoVectorData.objects.get(pk=pk)
+                child = child_model.objects.get(pk=child_id)
+                child.geodata.remove(geoitems[pk])
+                child._manage_main_geo()
 
 
 def geodata_attached_changed(sender, **kwargs):
@@ -3161,15 +3164,7 @@ class GeographicItem(models.Model):
         """
         return []
 
-    def save(
-        self, force_insert=False, force_update=False, using=None, update_fields=None
-    ):
-        super(GeographicItem, self).save(
-            force_insert=force_insert,
-            force_update=force_update,
-            using=using,
-            update_fields=update_fields,
-        )
+    def _manage_main_geo(self):
         if self.main_geodata and not self.geodata.filter(
                 pk=self.main_geodata.pk).count():
             try:
@@ -3177,17 +3172,19 @@ class GeographicItem(models.Model):
                     self.geodata.add(self.main_geodata)
             except (OperationalError, IntegrityError):
                 pass
-        elif not self.main_geodata:
-            try:
-                with transaction.atomic():
-                    if self.geodata.count():
-                        # arbitrary associate the first to geodata
-                        self.main_geodata = self.geodata.order_by("pk").all()[0]
-                        self.skip_history_when_saving = True
-                        self._no_move = True
-                        self.save()
-            except (OperationalError, IntegrityError, IndexError):
-                pass
+        if self.main_geodata or not self.geodata.count():
+            return
+        try:
+            with transaction.atomic():
+                # arbitrary associate the first to geodata
+                data_id = self.geodata.order_by("pk").values_list("pk", flat=True)[0]
+                self.__class__.objects.filter(pk=self.pk).update(main_geodata_id=data_id)
+        except (OperationalError, IntegrityError, IndexError):
+            pass
+
+    def save(self, *args, **kwargs):
+        super().save(*args, **kwargs)
+        self._manage_main_geo()
 
     @property
     def geodata_list(self):