Coverage for backend/core/auxiliary/methods/copy_flowsheet/copy_foreign_keys.py: 89%

1from django.db import transaction 

2from django.db.models import Model, Field 

3from authentication.user.models import User 

4from core.managers import SoftDeleteManager 

5from core.auxiliary.models.Flowsheet import Flowsheet 

6from core.auxiliary.models.PropertyValue import PropertyValue 

7from flowsheetInternals.graphicData.models.groupingModel import Grouping 

8from flowsheetInternals.graphicData.logic.make_group import propagate_streams, propagate_intermediate_streams 

9from flowsheetInternals.unitops.models.SimulationObject import SimulationObject 

10from django.utils import timezone 

11from dataclasses import dataclass 

12from authentication.user.AccessTable import AccessTable 

13from typing import List, Dict, Type 

14from core.validation import flowsheet_ctx 

15from .models_to_copy import custom_module_linked_model_types, linked_model_types 

16from .copy_caching import ModelLookupDict, ModelLookup, get_prefetch_fields 

17from core.managers import include_soft_deleted 

18 

19 

20def _create_foreign_key_lookups(flowsheet: Flowsheet, model_types : list[Type[Model]]) -> ModelLookupDict: 

21 """ 

22 We need to be able to look up the new models by their old primary key, 

23 because when we create the new models, the foreign keys will still point to the old flowsheet's version of the model. 

24 """ 

25 model_lookups: dict[Type[Model], ModelLookup] = {} 

26 

27 # Get all the models linked to the flowsheet, and index them by their primary key. 

28 for Model_Type in model_types: 

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

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

31 prefetch_fields = get_prefetch_fields(Model_Type) 

32 

33 prefetched_objects = include_soft_deleted( 

34 Model_Type.objects 

35 ).filter( 

36 flowsheet=flowsheet 

37 ).prefetch_related( 

38 *prefetch_fields 

39 ) 

40 model_lookups[Model_Type] = ModelLookup(list(prefetched_objects.all())) 

41 

42 return model_lookups 

43 

44def create_foreign_key_lookups(flowsheet:Flowsheet) -> ModelLookupDict: 

45 return _create_foreign_key_lookups(flowsheet, linked_model_types) 

46 

47def create_foreign_key_lookups_for_modules(flowsheet:Flowsheet) -> ModelLookupDict: 

48 """ 

49 We need to be able to look up the new models by their old primary key, 

50 because when we create the new models, the foreign keys will still point to the old flowsheet's version of the model. 

51 This version uses custom_module_linked_model_types for creating modules from templates. 

52 """ 

53 return _create_foreign_key_lookups(flowsheet,custom_module_linked_model_types) 

54 

55def _update_relationships(models: ModelLookup,field:Field,model_lookups: ModelLookupDict) -> None: 

56 for model in models: 

57 # Update foreign keys 

58 # https://docs.djangoproject.com/en/5.2/ref/models/meta/#django.db.models.options.Options.get_fields 

59 Related_Model_Type: Type[Model] = field.related_model 

60 related_model = model_lookups[Related_Model_Type] 

61 

62 if related_model: 62 ↛ 92line 62 didn't jump to line 92 because the condition on line 62 was always true

63 old_object_pk = getattr(model, field.name + "_id") 

64 if old_object_pk is not None: 

65 new_related_object = related_model.get_model(old_object_pk) 

66 if not new_related_object: 66 ↛ 67line 66 didn't jump to line 67 because the condition on line 66 was never true

67 print( 

68 f"ERROR: No related model found for {model}.{field.name} where {Related_Model_Type.__name__} had the pk {old_object_pk}. This should not happen." 

69 ) 

70 raise ValueError( 

71 f"No related model found for {model}.{field.name} to model {Related_Model_Type.__name__} with old pk {old_object_pk}." 

72 ) 

73 if ( 73 ↛ 77line 73 didn't jump to line 77 because the condition on line 73 was never true

74 new_related_object.pk 

75 == old_object_pk 

76 ): 

77 print( 

78 "WARNING: The old object has not updated this field yet!!!" 

79 ) 

80 raise ValueError( 

81 f"The old object {old_object_pk} has not been updated yet in the linked models." 

82 ) 

83 # print(f"Updating {field.name} for model {Model_Type.__name__} from {old_object.pk} to new model {related_model.get_model(old_object.pk).pk}") 

84 setattr( 

85 model, 

86 field.name, 

87 new_related_object, 

88 ) 

89 else: 

90 pass # We don't need to update the field if it is none. 

91 else: 

92 print( 

93 f"ERROR: No related model found for {field.name} in model {model}. This appears to be a bug and should be reported." 

94 ) 

95 raise ValueError( 

96 f"No related model found for {field.name} in model {model}." 

97 ) 

98 

99def update_foreign_keys(model_lookups: ModelLookupDict) -> None: 

100 """ 

101 Update the foreign keys of all the models. 

102 This is done after the primary keys have been updated, 

103 so that we can reference the new models. 

104 """ 

105 

106 # Exclude flowsheet, user and access table from the update 

107 models_to_update = [ 

108 (Model_Type, models) for Model_Type, models in model_lookups.items() 

109 if Model_Type not in [Flowsheet, User, AccessTable] 

110 ] 

111 # Now we need to update the foreign keys, one to one relationships, and many to many relationships of the models that reference other models. 

112 for Model_Type, models in models_to_update: 

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

114 if (field.one_to_one and not field.auto_created): 

115 # Make this reference the new model. 

116 _update_relationships(models, field, model_lookups) 

117 include_soft_deleted(Model_Type.objects).bulk_create(models) 

118 

119 for Model_Type, models in models_to_update: 

120 fields_to_update = [] 

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

122 if (field.many_to_one and not field.auto_created): 

123 fields_to_update.append(field.name) 

124 # Make this reference the new model. 

125 _update_relationships(models, field, model_lookups) 

126 if fields_to_update: 

127 include_soft_deleted(Model_Type.objects).bulk_update(models, fields_to_update) 

128