Coverage for backend/core/auxiliary/methods/copy_flowsheet/copy_many_to_many.py: 100%

1from django.db.models import Model, Field 

2from authentication.user.models import User 

3from core.auxiliary.models.Flowsheet import Flowsheet 

4from authentication.user.AccessTable import AccessTable 

5from typing import List, Dict, Type 

6from .copy_caching import RelatedModels, ModelLookupDict 

7 

8 

9class ManyToManyLookup: 

10 

11 def __init__(self): 

12 self.model_map: Dict[Field, List[RelatedModels]] = {} 

13 

14 def add(self, field: Field, owner_model: Model, other_model_ids: List[int]) -> None: 

15 if field not in self.model_map: 

16 self.model_map[field] = [] 

17 self.model_map[field].append( 

18 RelatedModels(model=owner_model, related_model_ids=other_model_ids) 

19 ) 

20 

21 

22def create_many_to_many_lookups(model_lookups: ModelLookupDict) -> ManyToManyLookup: 

23 many_to_many_lookup = ManyToManyLookup() 

24 for Model_Type, models in model_lookups.items(): 

25 if Model_Type in [Flowsheet, User, AccessTable]: 

26 continue # We don't want to update these models. 

27 

28 for model in models: 

29 for field in Model_Type._meta.get_fields(): 

30 if field.is_relation and field.many_to_many and not field.auto_created: 

31 # We need to wait to the end to update these, 

32 # as they may reference models that are not yet created. 

33 current_field = getattr(model, field.name) 

34 old_related_objects = list(obj.pk for obj in current_field.all()) 

35 many_to_many_lookup.add(field, model, old_related_objects) 

36 

37 return many_to_many_lookup 

38 

39 

40def update_many_to_many_relationships(many_to_many_lookup: ManyToManyLookup, model_lookups: ModelLookupDict) -> None: 

41 for field, model_list in many_to_many_lookup.model_map.items(): 

42 # Get a reference to the django through model, so that we can bulk create all the new many-to-many relationships. 

43 Field_Through_Type: Type[Model] = field.remote_field.through 

44 # Figure out what the many-to-many field is related to. 

45 Related_Model_Type: Type[Model] = field.related_model 

46 related_model_lookup = model_lookups[Related_Model_Type] 

47 # bulk create everything for this field in one go. 

48 bulk_create_list = [] 

49 for relation_info in model_list: 

50 # For each model that has this many-to-many relationship, 

51 # we need to update the related objects to point to the new models. 

52 model = relation_info.model 

53 # Get the new models based on their old primary keys.  

54 for old_pk in relation_info.related_model_ids: 

55 new_related_object = related_model_lookup.get_model(old_pk) 

56 bulk_create_list.append( 

57 Field_Through_Type(**{ 

58 model._meta.model_name + "_id": model.pk, 

59 new_related_object._meta.model_name + "_id": new_related_object.pk, 

60 }) 

61 ) 

62 

63 Field_Through_Type.objects.bulk_create(bulk_create_list)