Manage fix of associated items (after imports or database dump)

1 files changed, 41 insertions, 1 deletions

diff --git a/ishtar_common/models.py b/ishtar_common/models.py
index 53c05d180..75f6adbed 100644
--- a/ishtar_common/models.py
+++ b/ishtar_common/models.py
@@ -1086,7 +1086,46 @@ class FullSearch(models.Model):
         return changed
 
 
-class BaseHistorizedItem(FullSearch, Imported, JsonData):
+class FixAssociated(object):
+    ASSOCIATED = {}
+
+    def fix_associated(self):
+        for key in self.ASSOCIATED:
+            item = getattr(self, key)
+            if not item:
+                continue
+            dct = self.ASSOCIATED[key]
+            for dct_key in dct:
+                subkey, ctype = dct_key
+                expected_values = dct[dct_key]
+                if not isinstance(expected_values, (list, tuple)):
+                    expected_values = [expected_values]
+                if hasattr(ctype, "txt_idx"):
+                    expected_values = [ctype.objects.get(txt_idx=v)
+                                       for v in expected_values]
+                current_vals = getattr(item, subkey)
+                is_many = False
+                if hasattr(current_vals, "all"):
+                    is_many = True
+                    current_vals = current_vals.all()
+                else:
+                    current_vals = [current_vals]
+                is_ok = False
+                for current_val in current_vals:
+                    if current_val in expected_values:
+                        is_ok = True
+                        break
+                if is_ok:
+                    continue
+                # the first value is used
+                new_value = expected_values[0]
+                if is_many:
+                    getattr(item, subkey).add(new_value)
+                else:
+                    setattr(item, subkey, new_value)
+
+
+class BaseHistorizedItem(FullSearch, Imported, JsonData, FixAssociated):
     """
     Historized item with external ID management.
     All historized items are searcheable and have a data json field
@@ -1267,6 +1306,7 @@ class BaseHistorizedItem(FullSearch, Imported, JsonData):
         for dep in self.EXTERNAL_ID_DEPENDENCIES:
             for obj in getattr(self, dep).all():
                 obj.update_external_id(save=True)
+        self.fix_associated()
         return True