Context records: Optimize record relations - fix equals

1 files changed, 29 insertions, 13 deletions

diff --git a/archaeological_context_records/models.py b/archaeological_context_records/models.py
index 65145be98..454e83e8a 100644
--- a/archaeological_context_records/models.py
+++ b/archaeological_context_records/models.py
@@ -1454,14 +1454,23 @@ class ContextRecordTree(RelationsViews):
         return new_trees
 
     @classmethod
-    def _get_equals(cls, item_id, equal_rel_types):
-        equals = list(RecordRelations.objects.values_list(
+    def _get_equals(cls, item_id, equal_rel_types, exclude=None):
+        if not exclude:
+            exclude = [item_id]
+        q = RecordRelations.objects.values_list(
             "right_record_id", flat=True).filter(
-            left_record_id=item_id, relation_type_id__in=equal_rel_types))
-        equals += list(RecordRelations.objects.values_list(
+            left_record_id=item_id, relation_type_id__in=equal_rel_types)
+        q = q.exclude(right_record_id__in=exclude)
+        equals = list(q)
+        q = RecordRelations.objects.values_list(
             "left_record_id", flat=True).filter(
-            right_record_id=item_id, relation_type_id__in=equal_rel_types))
-        return set(equals)
+            right_record_id=item_id, relation_type_id__in=equal_rel_types)
+        q = q.exclude(left_record_id__in=exclude)
+        equals += list(q)
+        exclude += equals
+        for eq_id in equals:
+            equals += cls._get_equals(eq_id, equal_rel_types, exclude=exclude)
+        return equals
 
     @classmethod
     def _update_equals(cls, item_id, equals):
@@ -1485,24 +1494,32 @@ class ContextRecordTree(RelationsViews):
         equal_rel_types = cls._get_base_equal_relations()
         keys = []
         for child_id, parent_id in relations:
-            equals = set(cls._get_equals(child_id, equal_rel_types))
-            keys += cls._update_equals(child_id, equals)
-            for alt_child in equals:
+            equals_child = set(cls._get_equals(child_id, equal_rel_types))
+            keys += cls._update_equals(child_id, equals_child)
+            for alt_child in equals_child:
                 if alt_child != child_id:
                     cls.objects.get_or_create(
                         key=f"{alt_child}_{parent_id}",
                         cr_id=alt_child, cr_parent_id=parent_id
                     )
                     keys.append((alt_child, parent_id))
-            equals = set(cls._get_equals(parent_id, equal_rel_types))
-            keys += cls._update_equals(parent_id, equals)
-            for alt_parent in equals:
+            equals_parent = set(cls._get_equals(parent_id, equal_rel_types))
+            keys += cls._update_equals(parent_id, equals_parent)
+            for alt_parent in equals_parent:
                 if alt_parent != parent_id:
                     cls.objects.get_or_create(
                         key=f"{child_id}_{alt_parent}",
                         cr_id=child_id, cr_parent_id=alt_parent
                     )
                     keys.append((child_id, alt_parent))
+
+                    for alt_child in equals_child:
+                        if alt_child != child_id:
+                            cls.objects.get_or_create(
+                                key=f"{alt_child}_{alt_parent}",
+                                cr_id=alt_child, cr_parent_id=alt_parent
+                            )
+                            keys.append((alt_child, alt_parent))
         return set(keys)
 
     @classmethod
@@ -1541,7 +1558,6 @@ class ContextRecordTree(RelationsViews):
             equals = set(cls._get_equals(item_id, equal_rel_types))
             all_relations.update(cls._update_equals(item_id, equals))
 
-
         ## delete old relations
         if not already_updated:
             already_updated = [item_id]