Geodata redesign: container, warehouse migration

3 files changed, 117 insertions, 12 deletions

diff --git a/archaeological_operations/models.py b/archaeological_operations/models.py
index 263aa5f55..90a4043b6 100644
--- a/archaeological_operations/models.py
+++ b/archaeological_operations/models.py
@@ -198,11 +198,12 @@ class GeographicTownItem(GeoItem):
         current_town_geo = None
         if q_towns_nb == 1:
             current_town_geo = q_towns.all()[0]
-            if not q_geodata_town.filter(pk=current_town_geo.pk).count():
-                for geo in q_geodata_town.exclude(source_id=current_town_geo.pk).all():
-                    self.geodata.remove(geo)
-                    if self.main_geodata == geo:
-                        self.main_geodata = None
+            for geo in q_geodata_town.exclude(source_id=current_town_geo.pk).all():
+                self.geodata.remove(geo)
+                if self.main_geodata == geo:
+                    self.main_geodata = None
+                changed = True
+            if not q_geodata_town.filter(source_id=current_town_geo.pk).count():
                 self.geodata.add(current_town_geo.main_geodata)
                 changed = True
 
diff --git a/archaeological_warehouse/models.py b/archaeological_warehouse/models.py
index 8494b5bdc..996925d0a 100644
--- a/archaeological_warehouse/models.py
+++ b/archaeological_warehouse/models.py
@@ -52,6 +52,7 @@ from ishtar_common.models_common import (
     QuickAction,
     MainItem,
     Merge,
+    GeoVectorData,
 )
 from ishtar_common.model_merging import merge_model_objects
 from ishtar_common.utils import (
@@ -411,6 +412,48 @@ class Warehouse(
     def __str__(self):
         return self.name
 
+    def post_save_geo(self, save=True):
+        # manage geodata towns
+        if getattr(self, "_post_save_geo_ok", False):
+            # prevent infinite loop - should not happen, but...
+            return
+        self._post_save_geo_ok = True
+        if not self.precise_town:
+            q_geotown = GeoVectorData.objects.filter(pk=None)
+        else:
+            q_geotown = GeoVectorData.objects.filter(
+                source_content_type__model="town",
+                source_content_type__app_label="ishtar_common",
+                source_id=self.precise_town.pk,
+                multi_polygon__isnull=False)
+        q_geodata_town = self.geodata.filter(
+            source_content_type__model="town",
+            source_content_type__app_label="ishtar_common",
+        )
+        changed = False
+        if not q_geotown.count():
+            # no simple town - clean
+            for geo in q_geodata_town.all():
+                self.geodata.remove(geo)
+                if self.main_geodata == geo:
+                    self.main_geodata = None
+                changed = True
+        else:
+            current_geo_town = q_geotown.all()[0]
+            for geo in q_geodata_town.exclude(pk=current_geo_town.pk).all():
+                self.geodata.remove(geo)
+                if self.main_geodata == geo:
+                    self.main_geodata = None
+            changed = True
+            if not q_geodata_town.filter(pk=current_geo_town.pk).count():
+                self.geodata.add(current_geo_town)
+                changed = True
+
+        if changed and save:
+            self.skip_history_when_saving = True
+            self._no_move = True
+            self.save()
+
     def get_container_type_by_place(self, place: int):
         """
         Container type by place based on the default organisation of the
@@ -1131,6 +1174,49 @@ class Container(
     def __str__(self):
         return self.cached_label or ""
 
+    def post_save_geo(self, save=True):
+        # manage geodata towns
+        if getattr(self, "_post_save_geo_ok", False):
+            # prevent infinite loop - should not happen, but...
+            return
+        self._post_save_geo_ok = True
+
+        if not self.location.precise_town:
+            town_id = None
+            q_geotown = GeoVectorData.objects.filter(pk=None)
+        else:
+            town_id = self.location.precise_town.pk
+            q_geotown = GeoVectorData.objects.filter(
+                source_content_type__model="town",
+                source_content_type__app_label="ishtar_common",
+                source_id=self.location.precise_town.pk,
+                multi_polygon__isnull=False)
+        q_geodata_town = self.geodata.filter(
+            source_content_type__model="town",
+            source_content_type__app_label="ishtar_common",
+        )
+        changed = False
+        if not q_geotown.count():
+            for geo in q_geodata_town.all():
+                self.geodata.remove(geo)
+                if self.main_geodata == geo:
+                    self.main_geodata = None
+                changed = True
+        else:
+            for geo in q_geodata_town.exclude(source_id=town_id).all():
+                self.geodata.remove(geo)
+                if self.main_geodata == geo:
+                    self.main_geodata = None
+            changed = True
+            if not q_geodata_town.filter(source_id=town_id).count():
+                self.geodata.add(q_geotown.all()[0])
+                changed = True
+
+        if changed and save:
+            self.skip_history_when_saving = True
+            self._no_move = True
+            self.save()
+
     @property
     def start_division_number(self):
         depth = 1
diff --git a/ishtar_common/management/commands/migrate_to_geo_v4.py b/ishtar_common/management/commands/migrate_to_geo_v4.py
index 2c27eebe9..2a316c2a4 100644
--- a/ishtar_common/management/commands/migrate_to_geo_v4.py
+++ b/ishtar_common/management/commands/migrate_to_geo_v4.py
@@ -15,6 +15,8 @@ from ishtar_common.utils import ugettext_lazy as _, get_log_time, get_percent, g
 from ishtar_common import models_common, models
 from archaeological_operations.models import Operation, ArchaeologicalSite
 from archaeological_context_records.models import ContextRecord
+from archaeological_finds.models import BaseFind
+from archaeological_warehouse.models import Warehouse, Container
 
 
 log_path = os.sep.join([settings.ROOT_PATH, "logs"])
@@ -183,7 +185,7 @@ def migrate(quiet=False, log=True):
             "basefind",
             "mobilier d'origine",
             "du mobilier d'origine",
-            ContextRecord,
+            BaseFind,
         ),
     ]
     for app, model_slug, model_name, model_full_name, model in model_list:
@@ -199,7 +201,7 @@ def migrate(quiet=False, log=True):
             defaults={"label": f"Centre {model_full_name}"},
         )
         ref_time = datetime.datetime.now()
-        for idx, obj in enumerate(q.all()[:100]):
+        for idx, obj in enumerate(q.all()):
             if not quiet:
                 write_output(model_name, idx, nb, ref_time)
             obj._no_move = True
@@ -268,11 +270,27 @@ def migrate(quiet=False, log=True):
                     changed.append(
                         ["geovectordata", data.name, data.pk, f"Point {model_name}"]
                     )
-    if not quiet and nb:
-        sys.stdout.write(
-            f"\r[{get_log_time()}] {model_name.capitalize()} migrated" + " " * 20 + "\n"
-        )
-        sys.stdout.flush()
+        if not quiet and nb:
+            sys.stdout.write(
+                f"\r[{get_log_time()}] {model_name.capitalize()} migrated"
+                + " " * 20
+                + "\n"
+            )
+            sys.stdout.flush()
+    model_list = [Warehouse, Container]
+    for model in model_list:
+        ref_time = datetime.datetime.now()
+        q = model.objects.exclude(main_geodata__isnull=False)
+        nb = q.count()
+        for idx, obj in enumerate(q.all()):
+            if not quiet:
+                write_output(model.__name__, idx, nb, ref_time)
+            obj.save()
+        if not quiet and nb:
+            sys.stdout.write(
+                f"\r[{get_log_time()}] {model.__name__.capitalize()} migrated" + " " * 20 + "\n"
+            )
+            sys.stdout.flush()
 
     if log and changed:
         filename = f"geo_migration-created-{get_log_time().replace(':', '')}.csv"