Coverage for backend/django/core/auxiliary/services/pinch_import.py: 80%

1import traceback 

2 

3from django.db import transaction 

4from django.utils import timezone 

5from rest_framework.exceptions import NotFound, ValidationError 

6 

7from PinchAnalysis.models.InputModels import PinchUtility 

8from PinchAnalysis.models.StreamDataProject import StreamDataProject 

9from common.models.notifications.payloads import NotificationServiceMessageType 

10from common.models.pinch_import import PinchUtilityCsvImportRequestPayload 

11from common.services import messaging 

12from core.auxiliary.enums import pinchEnums 

13from core.auxiliary.enums.generalEnums import TaskStatus 

14from core.auxiliary.managers.TaskManager import handle_task_running_event 

15from core.auxiliary.models.Task import Task, TaskType 

16from core.auxiliary.models.UploadSession import ( 

17 UploadSession, 

18 UploadSessionPurpose, 

19 UploadSessionStatus, 

20) 

21from core.auxiliary.serializers import TaskSerializer 

22from core.auxiliary.services.csv_inspect import CsvInspectionError, parse_float_cell, stream_csv_rows 

23from core.auxiliary.services.object_storage import s3 as s3_storage 

24from core.auxiliary.services.uploads import inspect_upload_session 

25 

26 

27ROW_BATCH_SIZE = 1000 

28PROGRESS_UPDATE_INTERVAL = 1000 

29REQUIRED_COLUMNS = ("name", "t_supply", "t_target") 

30OPTIONAL_COLUMNS = ("heat_flow", "dt_cont", "htc", "price", "type") 

31SUPPORTED_COLUMNS = REQUIRED_COLUMNS + OPTIONAL_COLUMNS 

32VALID_UTILITY_TYPES = {choice for choice, _ in pinchEnums.StreamType.choices if choice} 

33 

34 

35def _send_task_update(task: Task, message_type: NotificationServiceMessageType): 

36 messaging.send_flowsheet_notification_message( 

37 task.flowsheet_id, 

38 TaskSerializer(task).data, 

39 message_type, 

40 ) 

41 

42 

43def _get_pinch_project(project_id: int) -> StreamDataProject: 

44 try: 

45 return StreamDataProject.objects.select_related( 

46 "flowsheet_state__flowsheet", "Inputs" 

47 ).get(id=project_id) 

48 except StreamDataProject.DoesNotExist as exc: 

49 raise NotFound({"project_id": "Pinch project not found."}) from exc 

50 

51 

52def _get_completed_pinch_upload(upload_session_id: int, user_id: int) -> UploadSession: 

53 try: 

54 upload_session = UploadSession.objects.get(id=upload_session_id) 

55 except UploadSession.DoesNotExist as exc: 

56 raise ValidationError({"upload_session_id": "Upload session not found."}) from exc 

57 

58 if upload_session.created_by_id != user_id: 58 ↛ 59line 58 didn't jump to line 59 because the condition on line 58 was never true

59 raise ValidationError({"upload_session_id": "You do not own this upload session."}) 

60 if upload_session.status != UploadSessionStatus.COMPLETED: 60 ↛ 61line 60 didn't jump to line 61 because the condition on line 60 was never true

61 raise ValidationError({"upload_session_id": "The upload session must be completed before import."}) 

62 if upload_session.expires_at is not None and upload_session.expires_at <= timezone.now(): 

63 upload_session.status = UploadSessionStatus.EXPIRED 

64 upload_session.save(update_fields=["status"]) 

65 raise ValidationError({"upload_session_id": "The upload session expired before import."}) 

66 if upload_session.purpose != UploadSessionPurpose.PINCH_UTILITY_CSV: 66 ↛ 67line 66 didn't jump to line 67 because the condition on line 66 was never true

67 raise ValidationError({"upload_session_id": "The upload session purpose must be pinch_utility_csv."}) 

68 return upload_session 

69 

70 

71def _raise_if_completed_upload_expired(upload_session: UploadSession) -> None: 

72 """Prevent delayed imports from consuming objects past their application TTL.""" 

73 if upload_session.expires_at is not None and upload_session.expires_at <= timezone.now(): 

74 upload_session.status = UploadSessionStatus.EXPIRED 

75 upload_session.save(update_fields=["status"]) 

76 raise CsvInspectionError("The upload session expired before Pinch utility import could run.") 

77 

78 

79def enqueue_pinch_utility_import(*, user, project_id: int, upload_session_id: int) -> Task: 

80 """Queue a Pinch utility CSV import unless the same upload already has an active task.""" 

81 project = _get_pinch_project(project_id) 

82 upload_session = _get_completed_pinch_upload(upload_session_id, user.id) 

83 

84 stable_flowsheet_id = project.flowsheet_state.flowsheet_id 

85 if upload_session.flowsheet_state_id != project.flowsheet_state_id: 

86 raise ValidationError({"upload_session_id": "The upload session belongs to a different flowsheet."}) 

87 

88 existing_task = ( 

89 Task.objects.filter( 

90 task_type=TaskType.CSV_IMPORT_PINCH_UTILITIES, 

91 flowsheet_id=stable_flowsheet_id, 

92 debug__project_id=project_id, 

93 debug__upload_session_id=upload_session_id, 

94 ) 

95 .exclude(status__in=[TaskStatus.Failed, TaskStatus.Cancelled]) 

96 .order_by("-start_time") 

97 .first() 

98 ) 

99 if existing_task is not None: 

100 return existing_task 

101 

102 task = Task.create( 

103 user, 

104 stable_flowsheet_id, 

105 task_type=TaskType.CSV_IMPORT_PINCH_UTILITIES, 

106 status=TaskStatus.Pending, 

107 expected_flowsheet_state_id=project.flowsheet_state_id, 

108 save=True, 

109 ) 

110 task.debug = { 

111 "project_id": project_id, 

112 "upload_session_id": upload_session_id, 

113 "utilities_imported": 0, 

114 } 

115 task.save(update_fields=["debug"]) 

116 

117 messaging.send_pinch_utility_csv_import_message( 

118 PinchUtilityCsvImportRequestPayload( 

119 task_id=task.id, 

120 flowsheet_id=stable_flowsheet_id, 

121 project_id=project_id, 

122 upload_session_id=upload_session_id, 

123 bucket=upload_session.bucket, 

124 object_key=upload_session.object_key, 

125 csv_delimiter=upload_session.csv_delimiter or None, 

126 requested_by_user_id=user.id, 

127 ) 

128 ) 

129 return task 

130 

131 

132def _row_has_any_supported_value(row: dict[str, str | None]) -> bool: 

133 return any((row.get(column) or "").strip() for column in SUPPORTED_COLUMNS) 

134 

135 

136def _parse_float( 

137 row_number: int, 

138 row: dict[str, str], 

139 column_name: str, 

140 *, 

141 required: bool, 

142) -> float | None: 

143 parsed = parse_float_cell(row_number, row, column_name, required=required) 

144 if parsed is None and required: 144 ↛ 145line 144 didn't jump to line 145 because the condition on line 144 was never true

145 raise CsvInspectionError(f"Row {row_number}, column '{column_name}' is required.") 

146 return parsed 

147 

148 

149def _parse_type(row_number: int, row: dict[str, str]) -> str: 

150 raw_value = (row.get("type") or "").strip() 

151 if not raw_value: 

152 return pinchEnums.StreamType.Both 

153 if raw_value not in VALID_UTILITY_TYPES: 153 ↛ 154line 153 didn't jump to line 154 because the condition on line 153 was never true

154 raise CsvInspectionError( 

155 f"Row {row_number}, column 'type' must be one of: {', '.join(sorted(VALID_UTILITY_TYPES))}." 

156 ) 

157 return raw_value 

158 

159 

160def _build_utility(row_number: int, row: dict[str, str], project: StreamDataProject) -> PinchUtility | None: 

161 if not _row_has_any_supported_value(row): 161 ↛ 162line 161 didn't jump to line 162 because the condition on line 161 was never true

162 return None 

163 

164 name = (row.get("name") or "").strip() 

165 if not name: 165 ↛ 166line 165 didn't jump to line 166 because the condition on line 165 was never true

166 raise CsvInspectionError(f"Row {row_number}, column 'name' is required.") 

167 

168 return PinchUtility( 

169 flowsheet_state=project.flowsheet_state, 

170 input_owner=project.Inputs, 

171 name=name, 

172 t_supply=_parse_float(row_number, row, "t_supply", required=True), 

173 t_target=_parse_float(row_number, row, "t_target", required=True), 

174 heat_flow=_parse_float(row_number, row, "heat_flow", required=False), 

175 dt_cont=_parse_float(row_number, row, "dt_cont", required=False), 

176 htc=_parse_float(row_number, row, "htc", required=False), 

177 price=_parse_float(row_number, row, "price", required=False), 

178 type=_parse_type(row_number, row), 

179 ) 

180 

181 

182def _flush_utilities( 

183 *, 

184 task: Task, 

185 utility_buffer: list[PinchUtility], 

186 imported_count: int, 

187) -> int: 

188 if not utility_buffer: 188 ↛ 189line 188 didn't jump to line 189 because the condition on line 188 was never true

189 return imported_count 

190 

191 PinchUtility.objects.bulk_create(utility_buffer) 

192 imported_count += len(utility_buffer) 

193 if imported_count % PROGRESS_UPDATE_INTERVAL == 0: 193 ↛ 194line 193 didn't jump to line 194 because the condition on line 193 was never true

194 task.debug = {**(task.debug or {}), "utilities_imported": imported_count} 

195 task.save(update_fields=["debug"]) 

196 _send_task_update(task, NotificationServiceMessageType.TASK_UPDATED) 

197 

198 utility_buffer.clear() 

199 return imported_count 

200 

201 

202def process_pinch_utility_import(payload: PinchUtilityCsvImportRequestPayload): 

203 """Import Pinch utility rows from object storage into the target project. 

204 

205 The payload must reference an existing pending task, a completed upload session, 

206 and a Pinch project owned by the same flowsheet. The CSV is streamed from object 

207 storage, validated row-by-row, and inserted in batches while task progress 

208 notifications are emitted. Invalid CSV structure or row values raise 

209 `CsvInspectionError` or `ValueError`, which are captured on the task as a failure. 

210 """ 

211 task = Task.objects.get(id=payload.task_id) 

212 if task.status in {TaskStatus.Completed, TaskStatus.Cancelled, TaskStatus.Running, TaskStatus.Failed}: 212 ↛ 213line 212 didn't jump to line 213 because the condition on line 212 was never true

213 return 

214 

215 handle_task_running_event(task.id) 

216 try: 

217 project = _get_pinch_project(payload.project_id) 

218 try: 

219 upload_session = UploadSession.objects.get(id=payload.upload_session_id) 

220 except UploadSession.DoesNotExist as exc: 

221 raise CsvInspectionError("The upload session was deleted before the import could run.") from exc 

222 

223 if upload_session.status != UploadSessionStatus.COMPLETED: 223 ↛ 224line 223 didn't jump to line 224 because the condition on line 223 was never true

224 raise CsvInspectionError("The upload session must be completed before Pinch utility import can run.") 

225 _raise_if_completed_upload_expired(upload_session) 

226 if upload_session.purpose != UploadSessionPurpose.PINCH_UTILITY_CSV: 226 ↛ 227line 226 didn't jump to line 227 because the condition on line 226 was never true

227 raise CsvInspectionError("The upload session purpose must be pinch_utility_csv for Pinch import.") 

228 if upload_session.flowsheet_state_id != project.flowsheet_state_id: 228 ↛ 229line 228 didn't jump to line 229 because the condition on line 228 was never true

229 raise CsvInspectionError("The upload session belongs to a different flowsheet.") 

230 

231 delimiter = payload.csv_delimiter or upload_session.csv_delimiter 

232 if not delimiter: 

233 delimiter = inspect_upload_session(upload_session).delimiter 

234 

235 with transaction.atomic(): 

236 PinchUtility.objects.filter(input_owner=project.Inputs).delete() 

237 

238 body = s3_storage.stream_object(payload.bucket, payload.object_key) 

239 _, reader = stream_csv_rows(body, delimiter) 

240 

241 utility_buffer: list[PinchUtility] = [] 

242 imported_count = 0 

243 

244 for row_index, row in enumerate(reader): 

245 row_number = row_index + 2 

246 utility = _build_utility(row_number, row, project) 

247 if utility is None: 247 ↛ 248line 247 didn't jump to line 248 because the condition on line 247 was never true

248 continue 

249 utility_buffer.append(utility) 

250 

251 if len(utility_buffer) >= ROW_BATCH_SIZE: 251 ↛ 252line 251 didn't jump to line 252 because the condition on line 251 was never true

252 imported_count = _flush_utilities( 

253 task=task, 

254 utility_buffer=utility_buffer, 

255 imported_count=imported_count, 

256 ) 

257 

258 imported_count = _flush_utilities( 

259 task=task, 

260 utility_buffer=utility_buffer, 

261 imported_count=imported_count, 

262 ) 

263 

264 if imported_count == 0: 264 ↛ 265line 264 didn't jump to line 265 because the condition on line 264 was never true

265 raise CsvInspectionError("The uploaded CSV must include at least one utility row.") 

266 

267 task.status = TaskStatus.Completed 

268 task.completed_time = timezone.now() 

269 task.debug = {**(task.debug or {}), "utilities_imported": imported_count} 

270 task.error = None 

271 task.save(update_fields=["status", "completed_time", "debug", "error"]) 

272 except Exception as exc: 

273 task.status = TaskStatus.Failed 

274 task.completed_time = timezone.now() 

275 task.error = { 

276 "message": str(exc), 

277 "cause": "pinch_utility_csv_import", 

278 "traceback": traceback.format_exc(), 

279 } 

280 task.save(update_fields=["status", "completed_time", "error"]) 

281 _send_task_update(task, NotificationServiceMessageType.TASK_COMPLETED) 

282 return 

283 

284 _send_task_update(task, NotificationServiceMessageType.TASK_COMPLETED) 

285 

286 

287def mark_pinch_utility_import_delivery_failure(payload: PinchUtilityCsvImportRequestPayload): 

288 """Fail a queued Pinch utility CSV import when Dapr cannot deliver the callback.""" 

289 task = Task.objects.get(id=payload.task_id) 

290 if task.status in {TaskStatus.Completed, TaskStatus.Cancelled, TaskStatus.Failed}: 290 ↛ 291line 290 didn't jump to line 291 because the condition on line 290 was never true

291 return 

292 

293 task.status = TaskStatus.Failed 

294 task.completed_time = timezone.now() 

295 task.error = { 

296 "message": ( 

297 "The utility CSV import could not be started. " 

298 "Retry the import to queue a new task." 

299 ), 

300 "cause": "pinch_utility_csv_import_delivery_failed", 

301 } 

302 task.save(update_fields=["status", "completed_time", "error"]) 

303 _send_task_update(task, NotificationServiceMessageType.TASK_COMPLETED)