Coverage for backend/django/Economics/results/services/comparison/config.py: 100%

1"""Persistence for saved Economics study comparison configurations.""" 

2 

3from __future__ import annotations 

4 

5from django.core.exceptions import PermissionDenied 

6from django.db import transaction 

7 

8from authentication.user.models import User 

9from core.managers import has_flowsheet_write_access 

10from Economics.shared.choices import EconomicsBaselineMode 

11from Economics.studies.models import ( 

12 EconomicsStudy, 

13 EconomicsStudyComparisonConfig, 

14 EconomicsStudyComparisonSelection, 

15) 

16from Economics.studies.services.baseline_access import ( 

17 resolve_study_reference, 

18 stable_flowsheet_id, 

19) 

20 

21from .contracts import ComparisonConfigPayload 

22from .targets import validate_comparison_targets 

23 

24 

25def save_comparison_config( 

26 *, 

27 active_study: EconomicsStudy, 

28 user: User, 

29 comparison_study_ids: list[int], 

30 baseline_study_id: int | None, 

31 baseline_mode: EconomicsBaselineMode = EconomicsBaselineMode.STUDY, 

32) -> EconomicsStudyComparisonConfig: 

33 """Replace the saved comparison configuration for an active study.""" 

34 active_flowsheet_id = stable_flowsheet_id(active_study) 

35 if not has_flowsheet_write_access(user, active_flowsheet_id): 

36 raise PermissionDenied("User does not have write access to this flowsheet.") 

37 selected_baseline_id = _selected_baseline_study_id( 

38 active_study=active_study, 

39 baseline_mode=baseline_mode, 

40 baseline_study_id=baseline_study_id, 

41 ) 

42 targets = validate_comparison_targets( 

43 active_study=active_study, 

44 user=user, 

45 comparison_study_ids=comparison_study_ids, 

46 baseline_mode=baseline_mode, 

47 baseline_study_id=selected_baseline_id, 

48 ) 

49 studies_by_id = {active_study.pk: active_study, **{target.pk: target for target in targets}} 

50 selected_baseline = studies_by_id[selected_baseline_id] 

51 with transaction.atomic(): 

52 EconomicsStudy.objects.select_for_update().only("pk").get( 

53 flowsheet_state=active_study.flowsheet_state, 

54 pk=active_study.pk, 

55 ) 

56 config = ( 

57 EconomicsStudyComparisonConfig.objects.select_for_update() 

58 .filter( 

59 flowsheet_state=active_study.flowsheet_state, 

60 active_study=active_study, 

61 ) 

62 .first() 

63 ) 

64 if config is None: 

65 config = EconomicsStudyComparisonConfig.objects.create( 

66 flowsheet_state=active_study.flowsheet_state, 

67 active_study=active_study, 

68 baseline_mode=baseline_mode, 

69 baseline_flowsheet_id=stable_flowsheet_id(selected_baseline), 

70 baseline_study_lineage_id=selected_baseline.lineage_id, 

71 ) 

72 else: 

73 config.baseline_mode = baseline_mode 

74 config.baseline_flowsheet_id = stable_flowsheet_id(selected_baseline) 

75 config.baseline_study_lineage_id = selected_baseline.lineage_id 

76 config.save( 

77 update_fields=[ 

78 "baseline_mode", 

79 "baseline_flowsheet_id", 

80 "baseline_study_lineage_id", 

81 "updated_at", 

82 ] 

83 ) 

84 

85 EconomicsStudyComparisonSelection.objects.filter( 

86 flowsheet_state=active_study.flowsheet_state, 

87 config=config, 

88 ).delete() 

89 EconomicsStudyComparisonSelection.objects.bulk_create( 

90 [ 

91 EconomicsStudyComparisonSelection( 

92 flowsheet_state=active_study.flowsheet_state, 

93 config=config, 

94 target_study_lineage_id=target.lineage_id, 

95 target_flowsheet_id=stable_flowsheet_id(target), 

96 study_name_snapshot=target.name, 

97 flowsheet_name_snapshot=target.flowsheet_state.flowsheet.name, 

98 sort_order=index, 

99 ) 

100 for index, target in enumerate(targets, start=1) 

101 ] 

102 ) 

103 return config 

104 

105 

106def _selected_baseline_study_id( 

107 *, 

108 active_study: EconomicsStudy, 

109 baseline_mode: EconomicsBaselineMode, 

110 baseline_study_id: int | None, 

111) -> int: 

112 """Normalize saved comparison config so manual mode never depends on another study id.""" 

113 if baseline_mode == EconomicsBaselineMode.MANUAL: 

114 return active_study.pk 

115 return baseline_study_id or active_study.pk 

116 

117 

118def comparison_config(active_study: EconomicsStudy) -> EconomicsStudyComparisonConfig | None: 

119 """Load the saved comparison config, including ordered selections.""" 

120 try: 

121 return ( 

122 EconomicsStudyComparisonConfig.objects.prefetch_related("selections") 

123 .select_related("active_study") 

124 .get( 

125 flowsheet_state=active_study.flowsheet_state, 

126 active_study=active_study, 

127 ) 

128 ) 

129 except EconomicsStudyComparisonConfig.DoesNotExist: 

130 return None 

131 

132 

133def baseline_study_id(*, active_study: EconomicsStudy, config: EconomicsStudyComparisonConfig | None) -> int: 

134 """Return the saved baseline, defaulting to the active study.""" 

135 if config is None: 

136 return active_study.pk 

137 baseline = resolve_study_reference( 

138 owner_state_id=config.flowsheet_state_id, 

139 target_flowsheet_id=config.baseline_flowsheet_id, 

140 target_lineage_id=config.baseline_study_lineage_id, 

141 ) 

142 return baseline.pk if baseline is not None else active_study.pk 

143 

144 

145def baseline_mode(*, config: EconomicsStudyComparisonConfig | None) -> EconomicsBaselineMode: 

146 """Return the saved comparison baseline mode, defaulting to study mode.""" 

147 if config is None: 

148 return EconomicsBaselineMode.STUDY 

149 return EconomicsBaselineMode(config.baseline_mode) 

150 

151 

152def config_payload( 

153 *, 

154 active_study: EconomicsStudy, 

155 config: EconomicsStudyComparisonConfig | None, 

156 selections: list[EconomicsStudyComparisonSelection], 

157) -> ComparisonConfigPayload: 

158 """Serialize the saved config for the comparison API payload.""" 

159 baseline_id = baseline_study_id(active_study=active_study, config=config) 

160 return ComparisonConfigPayload( 

161 id=config.pk if config is not None else None, 

162 active_study=active_study.pk, 

163 baseline_mode=baseline_mode(config=config), 

164 baseline_study_id=baseline_id, 

165 comparison_study_ids=[ 

166 target.pk 

167 for selection in selections 

168 if ( 

169 target := resolve_study_reference( 

170 owner_state_id=selection.flowsheet_state_id, 

171 target_flowsheet_id=selection.target_flowsheet_id, 

172 target_lineage_id=selection.target_study_lineage_id, 

173 ) 

174 ) 

175 is not None 

176 ], 

177 updated_at=config.updated_at if config is not None else None, 

178 )