Coverage for backend/django/core/auxiliary/services/scenario_import.py: 81%
145 statements
« prev ^ index » next coverage.py v7.10.7, created at 2026-07-22 05:22 +0000
1import traceback
2from itertools import chain
4from django.db import transaction
5from django.utils import timezone
6from rest_framework.exceptions import NotFound, ValidationError
8from common.models.notifications.payloads import NotificationServiceMessageType
9from common.models.scenario_import import ScenarioCsvImportRequestPayload
10from common.services import messaging
11from core.auxiliary.enums.generalEnums import TaskStatus
12from core.auxiliary.managers.TaskManager import handle_task_running_event
13from core.auxiliary.models.DataColumn import DataColumn
14from core.auxiliary.models.Scenario import Scenario, ScenarioInputModeEnum, ScenarioTabTypeEnum
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 (
23 CsvInspectionError,
24 filter_numeric_headers_by_first_row,
25 parse_numeric_row,
26 stream_csv_rows,
27)
28from core.auxiliary.services.object_storage import s3 as s3_storage
29from core.auxiliary.services.uploads import inspect_upload_session
30from core.auxiliary.services.parameter_sweep import clear_parameter_sweep_definition
31from core.auxiliary.services.mss_input_data import (
32 MssInputCell,
33 MssInputRow,
34 MssInputSource,
35 replace_mss_input_rows,
36)
39PROGRESS_UPDATE_INTERVAL = 1000
42def _send_task_update(task: Task, message_type: NotificationServiceMessageType):
43 messaging.send_flowsheet_notification_message(
44 task.flowsheet_id,
45 TaskSerializer(task).data,
46 message_type,
47 )
50def _get_completed_scenario_upload(
51 upload_session_id: int,
52 user_id: int,
53 scenario_id: int,
54) -> UploadSession:
55 try:
56 upload_session = UploadSession.objects.get(id=upload_session_id)
57 except UploadSession.DoesNotExist as exc:
58 raise ValidationError({"upload_session_id": "Upload session not found."}) from exc
60 if upload_session.created_by_id != user_id: 60 ↛ 61line 60 didn't jump to line 61 because the condition on line 60 was never true
61 raise ValidationError({"upload_session_id": "You do not own this upload session."})
62 if upload_session.status != UploadSessionStatus.COMPLETED: 62 ↛ 63line 62 didn't jump to line 63 because the condition on line 62 was never true
63 raise ValidationError({"upload_session_id": "The upload session must be completed before import."})
64 if upload_session.expires_at is not None and upload_session.expires_at <= timezone.now():
65 upload_session.status = UploadSessionStatus.EXPIRED
66 upload_session.save(update_fields=["status"])
67 raise ValidationError({"upload_session_id": "The upload session expired before import."})
68 if upload_session.purpose != UploadSessionPurpose.SCENARIO_CSV: 68 ↛ 69line 68 didn't jump to line 69 because the condition on line 68 was never true
69 raise ValidationError({"upload_session_id": "The upload session purpose must be scenario_csv."})
70 if upload_session.scenario_id and upload_session.scenario_id != scenario_id: 70 ↛ 71line 70 didn't jump to line 71 because the condition on line 70 was never true
71 raise ValidationError({"upload_session_id": "The upload session belongs to a different scenario."})
73 return upload_session
76def _raise_if_completed_upload_expired(upload_session: UploadSession) -> None:
77 """Prevent delayed imports from consuming objects past their application TTL."""
78 if upload_session.expires_at is not None and upload_session.expires_at <= timezone.now():
79 upload_session.status = UploadSessionStatus.EXPIRED
80 upload_session.save(update_fields=["status"])
81 raise CsvInspectionError("The upload session expired before scenario CSV import could run.")
84def _get_importable_scenario(scenario_id: int) -> Scenario:
85 """Resolve and validate the target scenario before queueing a scenario CSV import."""
86 try:
87 scenario = Scenario.objects.get(id=scenario_id)
88 except Scenario.DoesNotExist as exc:
89 raise NotFound({"scenario_id": "Scenario not found."}) from exc
91 if scenario.state_name not in { 91 ↛ 95line 91 didn't jump to line 95 because the condition on line 91 was never true
92 ScenarioTabTypeEnum.MultiSteadyState,
93 ScenarioTabTypeEnum.Dynamic,
94 }:
95 raise ValidationError(
96 {
97 "scenario_id": (
98 "Only multi steady-state and dynamic scenarios support scenario CSV import."
99 )
100 }
101 )
103 return scenario
106def enqueue_scenario_import(*, user, scenario_id: int, upload_session_id: int) -> Task:
107 """Queue a scenario CSV import unless an active task already owns that upload.
109 Reusing non-terminal tasks prevents duplicate pub/sub deliveries for the same
110 upload session, while failed and cancelled tasks remain retryable.
111 """
112 scenario = _get_importable_scenario(scenario_id)
113 upload_session = _get_completed_scenario_upload(upload_session_id, user.id, scenario_id)
114 stable_flowsheet_id = scenario.flowsheet_state.flowsheet_id
115 if upload_session.flowsheet_state_id != scenario.flowsheet_state_id:
116 raise ValidationError({"upload_session_id": "The upload session belongs to a different flowsheet."})
118 existing_task = (
119 Task.objects.filter(
120 task_type=TaskType.CSV_IMPORT_SCENARIO,
121 flowsheet_id=stable_flowsheet_id,
122 debug__scenario_id=scenario_id,
123 debug__upload_session_id=upload_session_id,
124 status__in=[TaskStatus.Pending, TaskStatus.Running, TaskStatus.Cancelling],
125 )
126 .order_by("-start_time")
127 .first()
128 )
129 if existing_task is not None:
130 return existing_task
132 task = Task.create(
133 user,
134 stable_flowsheet_id,
135 task_type=TaskType.CSV_IMPORT_SCENARIO,
136 status=TaskStatus.Pending,
137 expected_flowsheet_state_id=scenario.flowsheet_state_id,
138 save=True,
139 )
140 task.debug = {
141 "scenario_id": scenario_id,
142 "upload_session_id": upload_session_id,
143 "rows_imported": 0,
144 }
145 task.save(update_fields=["debug"])
147 messaging.send_scenario_csv_import_message(
148 ScenarioCsvImportRequestPayload(
149 task_id=task.id,
150 flowsheet_id=stable_flowsheet_id,
151 scenario_id=scenario_id,
152 upload_session_id=upload_session_id,
153 bucket=upload_session.bucket,
154 object_key=upload_session.object_key,
155 csv_delimiter=upload_session.csv_delimiter or None,
156 requested_by_user_id=user.id,
157 )
158 )
159 return task
162def process_scenario_import(payload: ScenarioCsvImportRequestPayload):
163 """Import scenario CSV data from object storage into scenario data tables.
165 The task is moved from `pending` to `running` by the task manager, then the
166 CSV is streamed from object storage so large uploads do not need to fit in
167 memory. Any import failure is recorded on the task and broadcast to the UI.
168 """
169 task = Task.objects.get(id=payload.task_id)
170 if task.status in {TaskStatus.Completed, TaskStatus.Cancelled, TaskStatus.Running, TaskStatus.Failed}: 170 ↛ 171line 170 didn't jump to line 171 because the condition on line 170 was never true
171 return
173 handle_task_running_event(task.id)
174 try:
175 scenario = _get_importable_scenario(payload.scenario_id)
176 try:
177 upload_session = UploadSession.objects.get(id=payload.upload_session_id)
178 except UploadSession.DoesNotExist as exc:
179 raise CsvInspectionError("The upload session was deleted before the import could run.") from exc
181 if upload_session.status != UploadSessionStatus.COMPLETED: 181 ↛ 182line 181 didn't jump to line 182 because the condition on line 181 was never true
182 raise CsvInspectionError("The upload session must be completed before scenario CSV import can run.")
183 _raise_if_completed_upload_expired(upload_session)
184 if upload_session.purpose != UploadSessionPurpose.SCENARIO_CSV: 184 ↛ 185line 184 didn't jump to line 185 because the condition on line 184 was never true
185 raise CsvInspectionError("The upload session purpose must be scenario_csv for scenario CSV import.")
186 if upload_session.flowsheet_state_id != scenario.flowsheet_state_id: 186 ↛ 187line 186 didn't jump to line 187 because the condition on line 186 was never true
187 raise CsvInspectionError("The upload session belongs to a different flowsheet.")
188 if upload_session.scenario_id and upload_session.scenario_id != scenario.id: 188 ↛ 189line 188 didn't jump to line 189 because the condition on line 188 was never true
189 raise CsvInspectionError("The upload session belongs to a different scenario.")
191 delimiter = payload.csv_delimiter or upload_session.csv_delimiter
192 if not delimiter:
193 delimiter = inspect_upload_session(upload_session).delimiter
195 with transaction.atomic():
196 replacing_parameter_sweep = (
197 scenario.mss_input_mode == ScenarioInputModeEnum.ParameterSweep
198 or hasattr(scenario, "parameterSweepDefinition")
199 )
200 if replacing_parameter_sweep:
201 DataColumn.objects.filter(
202 scenario_id=scenario.id,
203 property_value__isnull=False,
204 ).delete()
205 clear_parameter_sweep_definition(scenario)
206 scenario.mss_input_mode = ScenarioInputModeEnum.Csv
207 scenario.save(update_fields=["mss_input_mode"])
209 body = s3_storage.stream_object(payload.bucket, payload.object_key)
210 headers, reader = stream_csv_rows(body, delimiter)
211 first_row = next(reader, None)
212 if first_row is None: 212 ↛ 213line 212 didn't jump to line 213 because the condition on line 212 was never true
213 raise CsvInspectionError("The uploaded CSV must include at least one data row.")
215 numeric_headers, skipped_headers = filter_numeric_headers_by_first_row(headers, first_row)
216 if not numeric_headers: 216 ↛ 217line 216 didn't jump to line 217 because the condition on line 216 was never true
217 raise CsvInspectionError(
218 "The uploaded CSV does not contain any numeric columns in its first data row."
219 )
221 existing_columns = {
222 column.name: column
223 for column in DataColumn.objects.filter(scenario_id=scenario.id)
224 }
225 missing_headers = [header for header in numeric_headers if header not in existing_columns]
226 if missing_headers: 226 ↛ 243line 226 didn't jump to line 243 because the condition on line 226 was always true
227 DataColumn.objects.bulk_create(
228 [
229 DataColumn(
230 flowsheet_state=scenario.flowsheet_state,
231 scenario_id=scenario.id,
232 name=header,
233 value="",
234 )
235 for header in missing_headers
236 ]
237 )
238 existing_columns = {
239 column.name: column
240 for column in DataColumn.objects.filter(scenario_id=scenario.id)
241 }
243 imported_count = 0
245 def iter_mss_input_rows():
246 nonlocal imported_count
248 for row_index, row in enumerate(chain([first_row], reader)):
249 row_number = row_index + 2
250 parsed_values = parse_numeric_row(row_number, row, numeric_headers)
251 imported_count += 1
252 if imported_count % PROGRESS_UPDATE_INTERVAL == 0: 252 ↛ 253line 252 didn't jump to line 253 because the condition on line 252 was never true
253 task.debug = {
254 **(task.debug or {}),
255 "rows_imported": imported_count,
256 }
257 task.save(update_fields=["debug"])
258 _send_task_update(
259 task,
260 NotificationServiceMessageType.TASK_UPDATED,
261 )
263 yield MssInputRow(
264 index=row_index,
265 cells=[
266 MssInputCell(
267 data_column_id=existing_columns[header].id,
268 value=value,
269 )
270 for header, value in zip(numeric_headers, parsed_values)
271 ],
272 )
274 replace_mss_input_rows(
275 scenario=scenario,
276 rows=iter_mss_input_rows(),
277 source=MssInputSource.csv_import,
278 )
280 task.status = TaskStatus.Completed
281 task.completed_time = timezone.now()
282 task.debug = {
283 **(task.debug or {}),
284 "rows_imported": imported_count,
285 "skipped_non_numeric_headers": skipped_headers,
286 }
287 task.error = None
288 task.save(update_fields=["status", "completed_time", "debug", "error"])
289 except Exception as exc:
290 task.status = TaskStatus.Failed
291 task.completed_time = timezone.now()
292 task.error = {
293 "message": str(exc),
294 "cause": "scenario_csv_import",
295 "traceback": traceback.format_exc(),
296 }
297 task.save(update_fields=["status", "completed_time", "error"])
298 _send_task_update(task, NotificationServiceMessageType.TASK_COMPLETED)
299 return
301 _send_task_update(task, NotificationServiceMessageType.TASK_COMPLETED)
304def mark_scenario_import_delivery_failure(payload: ScenarioCsvImportRequestPayload):
305 """Fail a queued scenario CSV import when Dapr cannot deliver the callback.
307 This is distinct from `process_scenario_import`, which only runs after the
308 callback arrives successfully. Delivery failures need their own task update
309 or the frontend will continue showing the task as pending.
310 """
311 task = Task.objects.get(id=payload.task_id)
312 if task.status in {TaskStatus.Completed, TaskStatus.Cancelled, TaskStatus.Failed}: 312 ↛ 313line 312 didn't jump to line 313 because the condition on line 312 was never true
313 return
315 task.status = TaskStatus.Failed
316 task.completed_time = timezone.now()
317 task.error = {
318 "message": (
319 "The scenario CSV import could not be started. "
320 "Retry the import to queue a new task."
321 ),
322 "cause": "scenario_csv_import_delivery_failed",
323 }
324 task.save(update_fields=["status", "completed_time", "error"])
325 _send_task_update(task, NotificationServiceMessageType.TASK_COMPLETED)