Geodata: cascade add

1 files changed, 60 insertions, 1 deletions

diff --git a/ishtar_common/models_common.py b/ishtar_common/models_common.py
index 19ee2c613..d32f835bd 100644
--- a/ishtar_common/models_common.py
+++ b/ishtar_common/models_common.py
@@ -2581,6 +2581,47 @@ class GeoVectorData(Imported, OwnPerms):
 post_save.connect(post_save_geodata, sender=GeoVectorData)
 
 
+def geodata_attached_changed(sender, **kwargs):
+    # manage main geoitem and cascade association
+    instance = kwargs.get("instance", None)
+    model = kwargs.get("model", None)
+    pk_set = kwargs.get("pk_set", None)
+    if not instance or not model or not pk_set:
+        return
+    item_pks = list(model.objects.filter(pk__in=pk_set).values_list("pk", flat=True))
+    if not item_pks:
+        return
+
+    if not hasattr(instance, "_geodata"):  # use a cache to manage during geodata attach
+        instance._geodata = []
+    if not instance.main_geodata_id:
+        instance.main_geodata_id = item_pks[0]
+        instance.skip_history_when_saving = True
+        instance._no_move = True
+        if not hasattr(instance, "_geodata"):
+            instance._geodata = []
+        instance._geodata += [pk for pk in item_pks if pk not in instance._geodata]
+        instance.save()
+
+    # for all sub item verify that the geo items are present
+    for query in instance.geodata_child_item_queries():
+        child_model = query.model
+        m2m_model = child_model.geodata.through
+        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 not pk in geoitems:
+                        geoitems[pk] = GeoVectorData.objects.get(pk=pk)
+                    child_model.objects.get(pk=child_id).geodata.add(geoitems[pk])
+
+
 class GeographicItem(models.Model):
     main_geodata = models.ForeignKey(
         GeoVectorData,
@@ -2596,6 +2637,16 @@ class GeographicItem(models.Model):
     class Meta:
         abstract = True
 
+    def geodata_child_item_queries(self):
+        """
+        :return: list of queries associated geographically with this item. When
+        geographic data is add to this item all sub items get the geographic data.
+        For instance an operation return the list of context records associated, so
+        when you add the syrvey limit, it is associated to all context records of
+        the operation.
+        """
+        return []
+
     def save(
         self, force_insert=False, force_update=False, using=None, update_fields=None
     ):
@@ -2605,8 +2656,12 @@ class GeographicItem(models.Model):
             using=using,
             update_fields=update_fields,
         )
-        if self.main_geodata and not self.geodata.filter(pk=self.main_geodata.pk):
+        if not hasattr(self, "_geodata"):  # use a cache to manage during geodata attach
+            self._geodata = []
+        if self.main_geodata and not self.geodata.filter(pk=self.main_geodata.pk) and\
+                self.main_geodata.pk not in self._geodata:
             self.geodata.add(self.main_geodata)
+            self._geodata.append(self.main_geodata.pk)
         elif not self.main_geodata and self.geodata.count():
             # arbitrary associate the first to geodata
             self.main_geodata = self.geodata.order_by("pk").all()[0]
@@ -2707,6 +2762,9 @@ class Town(GeographicItem, Imported, models.Model):
     def __str__(self):
         return self.cached_label or ""
 
+    def geodata_child_item_queries(self):
+        return [self.sites, self.operations]
+
     @property
     def label_with_areas(self):
         label = [self.name]
@@ -2808,6 +2866,7 @@ def post_save_town(sender, **kwargs):
 
 
 post_save.connect(post_save_town, sender=Town)
+m2m_changed.connect(geodata_attached_changed, sender=Town.geodata.through)
 
 
 def town_child_changed(sender, **kwargs):