Coverage for backend/django/core/auxiliary/services/uploads.py: 83%
281 statements
« prev ^ index » next coverage.py v7.10.7, created at 2026-07-22 05:22 +0000
1import logging
2import os
3import re
4import uuid
5from datetime import timedelta
6from typing import Any
8from botocore.exceptions import ClientError
9from django.conf import settings
10from django.db import transaction
11from django.utils import timezone
12from opentelemetry import trace
13from rest_framework.exceptions import NotFound, PermissionDenied, ValidationError
15from core.auxiliary.models.MLModel import MLModel
16from core.auxiliary.models.Flowsheet import Flowsheet
17from core.auxiliary.models.Scenario import Scenario, ScenarioTabTypeEnum
18from core.auxiliary.models.UploadSession import (
19 UploadSession,
20 UploadSessionPurpose,
21 UploadSessionStatus,
22)
23from core.auxiliary.services.csv_inspect import (
24 DEFAULT_INSPECTION_BYTES,
25 CsvInspectionResult,
26 inspect_csv_sample,
27)
28from core.auxiliary.services.csv_lifecycle import (
29 SHORT_CSV_LIFECYCLE_TTL,
30 completed_csv_lifecycle_ttl,
31 ensure_csv_upload_lifecycle_configuration,
32 expires_at_from_ttl,
33)
34from core.auxiliary.services.ml_column_mapping_updates import (
35 clear_ml_model_column_mappings,
36 stage_ml_model_update_snapshot,
37)
38from core.auxiliary.services.object_storage import s3 as s3_storage
39from core.auxiliary.services.ml_artifacts import (
40 delete_ml_artifact_if_unreferenced,
41)
42from core.auxiliary.services.upload_types import (
43 CompletedMultipartPart,
44 CompleteUploadResult,
45 MultipartInitiationResult,
46 MultipartPartUrlsResult,
47 MultipartUploadStatusResult,
48 UploadCleanupSummary,
49)
50from flowsheetInternals.unitops.models import SimulationObject
53MULTIPART_PART_SIZE_BYTES = 8 * 1024 * 1024
54PRESIGNED_PART_URL_EXPIRY_SECONDS = 60 * 60
55MAX_NORMALIZED_FILENAME_LENGTH = 128
58logger = logging.getLogger(__name__)
59tracer = trace.get_tracer(settings.OPEN_TELEMETRY_TRACER_NAME)
62def _set_span_attributes(span, attributes: dict[str, object | None]) -> None:
63 """Attach non-empty CSV upload diagnostic attributes to the active span."""
64 for key, value in attributes.items():
65 if value is not None:
66 span.set_attribute(key, value)
69def _csv_expires_at(*, ttl: timedelta, reference_time=None):
70 """Return the application-level CSV expiry timestamp for an upload object."""
71 return expires_at_from_ttl(ttl, reference_time=reference_time)
74def _can_resume_upload_session(upload_session: UploadSession, *, now=None) -> bool:
75 """Return whether an upload session is still eligible for browser resume calls."""
76 current_time = now or timezone.now()
77 return (
78 upload_session.status in {UploadSessionStatus.INITIATED, UploadSessionStatus.UPLOADING}
79 and (upload_session.expires_at is None or upload_session.expires_at > current_time)
80 )
83def _get_in_progress_upload_session(*, upload_session_id: int, user_id: int) -> UploadSession:
84 """Load, reconcile, and validate an upload session before accepting more browser parts."""
85 upload_session = _synchronise_upload_session_state(
86 _get_owned_upload_session(upload_session_id, user_id)
87 )
88 if upload_session.status == UploadSessionStatus.EXPIRED:
89 raise ValidationError(
90 {
91 "upload_session_id": (
92 "This upload session expired and cannot be resumed. Choose the CSV again to start over."
93 )
94 }
95 )
96 return upload_session
99def _validate_flowsheet_linkage(
100 *,
101 flowsheet_id: int,
102 scenario: Scenario | None = None,
103 simulation_object: SimulationObject | None = None,
104) -> None:
105 """Ensure the provided flowsheet id matches the linked scenario or simulation object."""
106 linked_flowsheet_id = None
107 if scenario is not None:
108 linked_flowsheet_id = scenario.flowsheet_state.flowsheet_id
109 elif simulation_object is not None:
110 linked_flowsheet_id = simulation_object.flowsheet_state.flowsheet_id
112 if linked_flowsheet_id is not None and linked_flowsheet_id != flowsheet_id: 112 ↛ 113line 112 didn't jump to line 113 because the condition on line 112 was never true
113 raise ValidationError({"flowsheet_id": "flowsheet_id does not match the referenced scenario or simulation object."})
116def _normalise_filename(filename: str) -> str:
117 """Sanitize a client-supplied filename so it is safe to embed into an object key."""
118 base_name = filename.rsplit("/", 1)[-1].rsplit("\\", 1)[-1]
119 cleaned = re.sub(r"[^A-Za-z0-9._-]", "_", base_name).strip("._")
120 if not cleaned: 120 ↛ 121line 120 didn't jump to line 121 because the condition on line 120 was never true
121 return "upload.csv"
123 stem, suffix = os.path.splitext(cleaned)
124 suffix = suffix[:16]
125 max_stem_length = max(1, MAX_NORMALIZED_FILENAME_LENGTH - len(suffix))
126 truncated_stem = stem[:max_stem_length].rstrip("._-") or "upload"
127 return f"{truncated_stem}{suffix}"
130def _build_object_key(purpose: str, flowsheet_id: int, original_filename: str) -> str:
131 """Build a unique object key for a CSV upload scoped to a flowsheet and purpose."""
132 safe_name = _normalise_filename(original_filename)
133 return f"csv-uploads/{purpose}/flowsheet-{flowsheet_id}/{uuid.uuid4().hex}-{safe_name}"
136def _get_simulation_object(simulation_object_id: int) -> SimulationObject:
137 try:
138 return SimulationObject.objects.get(id=simulation_object_id)
139 except SimulationObject.DoesNotExist as exc:
140 raise NotFound({"simulationObject_id": "SimulationObject not found."}) from exc
143def _get_scenario(scenario_id: int) -> Scenario:
144 try:
145 return Scenario.objects.get(id=scenario_id)
146 except Scenario.DoesNotExist as exc:
147 raise NotFound({"scenario_id": "Scenario not found."}) from exc
150def _get_owned_upload_session(
151 upload_session_id: int,
152 user_id: int,
153 *,
154 for_update: bool = False,
155) -> UploadSession:
156 """Load an upload session and enforce ownership checks for the calling user."""
157 queryset = UploadSession.objects.select_related("scenario", "simulationObject")
158 if for_update:
159 # The related targets are nullable, so PostgreSQL cannot lock the
160 # outer-joined rows. Only the upload session itself coordinates
161 # attachment and lifecycle transitions.
162 queryset = queryset.select_for_update(of=("self",))
163 try:
164 upload_session = queryset.get(id=upload_session_id)
165 except UploadSession.DoesNotExist as exc:
166 raise NotFound({"upload_session_id": "Upload session not found."}) from exc
168 if upload_session.created_by_id != user_id: 168 ↛ 169line 168 didn't jump to line 169 because the condition on line 168 was never true
169 raise PermissionDenied("You do not own this upload session.")
171 return upload_session
174@tracer.start_as_current_span("csv_upload.initiate")
175def initiate_upload_session(
176 *,
177 user,
178 flowsheet_id: int,
179 purpose: str,
180 original_filename: str,
181 content_type: str,
182 size_bytes: int,
183 scenario_id: int | None = None,
184 simulation_object_id: int | None = None,
185) -> MultipartInitiationResult:
186 """Create an upload session and the matching multipart upload in object storage."""
187 span = trace.get_current_span()
188 _set_span_attributes(
189 span,
190 {
191 "csv_upload.purpose": purpose,
192 "csv_upload.flowsheet_id": flowsheet_id,
193 "csv_upload.scenario_id": scenario_id,
194 "csv_upload.simulation_object_id": simulation_object_id,
195 "csv_upload.size_bytes": size_bytes,
196 "csv_upload.content_type": content_type or "text/csv",
197 },
198 )
199 if size_bytes <= 0: 199 ↛ 200line 199 didn't jump to line 200 because the condition on line 199 was never true
200 raise ValidationError({"size_bytes": "size_bytes must be greater than zero."})
202 scenario = None
203 simulation_object = None
204 if purpose == UploadSessionPurpose.ML_TRAINING_CSV:
205 if simulation_object_id is None: 205 ↛ 206line 205 didn't jump to line 206 because the condition on line 205 was never true
206 raise ValidationError({"simulationObject_id": "simulationObject_id is required for ML uploads."})
207 simulation_object = _get_simulation_object(simulation_object_id)
208 elif purpose == UploadSessionPurpose.SCENARIO_CSV:
209 if scenario_id is None: 209 ↛ 210line 209 didn't jump to line 210 because the condition on line 209 was never true
210 raise ValidationError({"scenario_id": "scenario_id is required for scenario CSV uploads."})
211 scenario = _get_scenario(scenario_id)
212 if scenario.state_name not in { 212 ↛ 216line 212 didn't jump to line 216 because the condition on line 212 was never true
213 ScenarioTabTypeEnum.MultiSteadyState,
214 ScenarioTabTypeEnum.Dynamic,
215 }:
216 raise ValidationError(
217 {
218 "scenario_id": (
219 "Only multi steady-state and dynamic scenarios support scenario CSV uploads."
220 )
221 }
222 )
223 elif purpose == UploadSessionPurpose.PINCH_UTILITY_CSV: 223 ↛ 231line 223 didn't jump to line 231 because the condition on line 223 was always true
224 if simulation_object_id is not None or scenario_id is not None: 224 ↛ 225line 224 didn't jump to line 225 because the condition on line 224 was never true
225 raise ValidationError(
226 {
227 "purpose": "Pinch utility uploads do not accept scenario_id or simulationObject_id linkage."
228 }
229 )
230 else:
231 raise ValidationError({"purpose": "Unsupported upload purpose."})
233 _validate_flowsheet_linkage(
234 flowsheet_id=flowsheet_id,
235 scenario=scenario,
236 simulation_object=simulation_object,
237 )
238 if scenario is not None:
239 flowsheet_state = scenario.flowsheet_state
240 elif simulation_object is not None:
241 flowsheet_state = simulation_object.flowsheet_state
242 else:
243 flowsheet_state = Flowsheet.objects.only("current_state_id").get(
244 pk=flowsheet_id
245 ).current_state
247 bucket = s3_storage.get_bucket_name()
248 ensure_csv_upload_lifecycle_configuration(bucket)
249 object_key = _build_object_key(purpose, flowsheet_id, original_filename)
250 upload_id = s3_storage.create_multipart_upload(bucket, object_key, content_type)
251 initial_expires_at = _csv_expires_at(
252 ttl=SHORT_CSV_LIFECYCLE_TTL,
253 )
255 upload_session = UploadSession.objects.create(
256 created_by=user,
257 flowsheet_id=flowsheet_id,
258 flowsheet_state=flowsheet_state,
259 purpose=purpose,
260 scenario=scenario,
261 simulationObject=simulation_object,
262 bucket=bucket,
263 object_key=object_key,
264 s3_upload_id=upload_id,
265 original_filename=_normalise_filename(original_filename),
266 content_type=content_type or "text/csv",
267 size_bytes=size_bytes,
268 status=UploadSessionStatus.INITIATED,
269 expires_at=initial_expires_at,
270 )
271 span.set_attribute("csv_upload.upload_session_id", upload_session.id)
273 return MultipartInitiationResult(
274 upload_session_id=upload_session.id,
275 bucket=upload_session.bucket,
276 object_key=upload_session.object_key,
277 s3_upload_id=upload_session.s3_upload_id,
278 part_size_bytes=MULTIPART_PART_SIZE_BYTES,
279 expires_in_seconds=PRESIGNED_PART_URL_EXPIRY_SECONDS,
280 )
283@tracer.start_as_current_span("csv_upload.complete")
284def complete_upload(*, upload_session_id: int, parts: list[dict[str, object]], user_id: int) -> CompleteUploadResult:
285 """Finalize a multipart upload session and persist any expiry metadata."""
286 span = trace.get_current_span()
287 _set_span_attributes(
288 span,
289 {
290 "csv_upload.upload_session_id": upload_session_id,
291 "csv_upload.part_count": len(parts),
292 },
293 )
294 upload_session = _get_in_progress_upload_session(
295 upload_session_id=upload_session_id,
296 user_id=user_id,
297 )
298 _set_span_attributes(
299 span,
300 {
301 "csv_upload.purpose": upload_session.purpose,
302 "csv_upload.status": upload_session.status,
303 "csv_upload.size_bytes": upload_session.size_bytes,
304 },
305 )
306 if upload_session.status == UploadSessionStatus.COMPLETED: 306 ↛ 307line 306 didn't jump to line 307 because the condition on line 306 was never true
307 span.set_attribute("csv_upload.already_completed", True)
308 return CompleteUploadResult(
309 upload_session_id=upload_session.id,
310 bucket=upload_session.bucket,
311 object_key=upload_session.object_key,
312 size_bytes=upload_session.size_bytes,
313 etag=None,
314 location=None,
315 )
317 if upload_session.status not in {UploadSessionStatus.INITIATED, UploadSessionStatus.UPLOADING}: 317 ↛ 318line 317 didn't jump to line 318 because the condition on line 317 was never true
318 raise ValidationError({"upload_session_id": "Only in-progress uploads can be completed."})
320 validated_parts = _validate_completed_parts(parts)
321 completion = s3_storage.complete_multipart_upload(
322 upload_session.bucket,
323 upload_session.object_key,
324 upload_session.s3_upload_id,
325 validated_parts,
326 )
328 upload_session.status = UploadSessionStatus.COMPLETED
329 upload_session.completed_at = timezone.now()
330 upload_session.expires_at = _csv_expires_at(
331 ttl=completed_csv_lifecycle_ttl(upload_session.purpose),
332 reference_time=upload_session.completed_at,
333 )
334 upload_session.save(update_fields=["status", "completed_at", "expires_at"])
335 span.set_attribute("csv_upload.expires_at", upload_session.expires_at.isoformat())
337 return CompleteUploadResult(
338 upload_session_id=upload_session.id,
339 bucket=completion.bucket,
340 object_key=completion.object_key,
341 size_bytes=upload_session.size_bytes,
342 etag=completion.etag,
343 location=completion.location,
344 )
347def get_part_urls(*, upload_session_id: int, part_numbers: list[int], user_id: int) -> MultipartPartUrlsResult:
348 """Generate presigned URLs for the requested multipart part numbers."""
349 upload_session = _get_in_progress_upload_session(
350 upload_session_id=upload_session_id,
351 user_id=user_id,
352 )
353 if upload_session.status not in {UploadSessionStatus.INITIATED, UploadSessionStatus.UPLOADING}: 353 ↛ 354line 353 didn't jump to line 354 because the condition on line 353 was never true
354 raise ValidationError({"upload_session_id": "Part URLs can only be requested while an upload is in progress."})
356 if upload_session.status == UploadSessionStatus.INITIATED: 356 ↛ 360line 356 didn't jump to line 360 because the condition on line 356 was always true
357 upload_session.status = UploadSessionStatus.UPLOADING
358 upload_session.save(update_fields=["status"])
360 return MultipartPartUrlsResult(
361 urls={
362 str(part_number): s3_storage.presign_upload_part(
363 upload_session.bucket,
364 upload_session.object_key,
365 upload_session.s3_upload_id,
366 part_number,
367 PRESIGNED_PART_URL_EXPIRY_SECONDS,
368 )
369 for part_number in part_numbers
370 }
371 )
374def get_upload_session(*, upload_session_id: int, user_id: int) -> UploadSession:
375 """Return an owned upload session after reconciling any stale object-storage state."""
376 upload_session = _get_owned_upload_session(upload_session_id, user_id)
377 return _synchronise_upload_session_state(upload_session)
380def get_upload_status(*, upload_session_id: int, user_id: int) -> MultipartUploadStatusResult:
381 """Return the current upload-session status and any uploaded multipart parts."""
382 now = timezone.now()
383 upload_session = _synchronise_upload_session_state(
384 _get_owned_upload_session(upload_session_id, user_id)
385 )
386 if upload_session.status in {
387 UploadSessionStatus.COMPLETED,
388 UploadSessionStatus.ABORTED,
389 UploadSessionStatus.FAILED,
390 UploadSessionStatus.EXPIRED,
391 }:
392 return MultipartUploadStatusResult(
393 upload_session_id=upload_session.id,
394 status=upload_session.status,
395 can_resume=_can_resume_upload_session(upload_session, now=now),
396 expires_at=upload_session.expires_at,
397 uploaded_parts=[],
398 part_size_bytes=MULTIPART_PART_SIZE_BYTES,
399 bucket=upload_session.bucket,
400 object_key=upload_session.object_key,
401 s3_upload_id=upload_session.s3_upload_id,
402 )
404 uploaded_parts = s3_storage.list_uploaded_parts(
405 upload_session.bucket,
406 upload_session.object_key,
407 upload_session.s3_upload_id,
408 )
409 return MultipartUploadStatusResult(
410 upload_session_id=upload_session.id,
411 status=upload_session.status,
412 can_resume=_can_resume_upload_session(upload_session, now=now),
413 expires_at=upload_session.expires_at,
414 uploaded_parts=uploaded_parts,
415 part_size_bytes=MULTIPART_PART_SIZE_BYTES,
416 bucket=upload_session.bucket,
417 object_key=upload_session.object_key,
418 s3_upload_id=upload_session.s3_upload_id,
419 )
422def _validate_completed_parts(parts: list[dict[str, object]]) -> list[CompletedMultipartPart]:
423 """Reject malformed multipart completion payloads before sending them to S3."""
424 if not parts:
425 raise ValidationError({"parts": "At least one uploaded part is required."})
427 seen_part_numbers: set[int] = set()
428 validated_parts: list[CompletedMultipartPart] = []
429 for raw_part in parts:
430 part = CompletedMultipartPart.model_validate(raw_part)
431 part_number = part.part_number
432 etag = part.etag.strip()
433 if part_number in seen_part_numbers:
434 raise ValidationError({"parts": "Duplicate part numbers are not allowed in multipart completion."})
435 if not etag: 435 ↛ 436line 435 didn't jump to line 436 because the condition on line 435 was never true
436 raise ValidationError({"parts": f"Part {part_number} is missing an etag."})
438 seen_part_numbers.add(part_number)
439 validated_parts.append(CompletedMultipartPart(part_number=part_number, etag=etag))
441 return validated_parts
444def _abort_multipart_upload_best_effort(upload_session: UploadSession) -> None:
445 try:
446 s3_storage.abort_multipart_upload(
447 upload_session.bucket,
448 upload_session.object_key,
449 upload_session.s3_upload_id,
450 )
451 except ClientError:
452 logger.warning(
453 "Failed to abort multipart upload during session reconciliation.",
454 extra={"upload_session_id": upload_session.id},
455 exc_info=True,
456 )
459def _synchronise_upload_session_state(upload_session: UploadSession) -> UploadSession:
460 """Update persisted session state from age-based and object-storage-derived facts."""
461 now = timezone.now()
463 if upload_session.status in {UploadSessionStatus.INITIATED, UploadSessionStatus.UPLOADING}:
464 # The bucket lifecycle eventually removes abandoned parts, but the API must
465 # stop resumptions as soon as the one-hour window has elapsed.
466 if upload_session.expires_at is not None and upload_session.expires_at <= now:
467 _abort_multipart_upload_best_effort(upload_session)
468 upload_session.status = UploadSessionStatus.EXPIRED
469 upload_session.save(update_fields=["status"])
470 return upload_session
472 stale_cutoff = now - timedelta(hours=settings.SEAWEED_STALE_MULTIPART_UPLOAD_HOURS)
473 if upload_session.created_at < stale_cutoff: 473 ↛ 474line 473 didn't jump to line 474 because the condition on line 473 was never true
474 _abort_multipart_upload_best_effort(upload_session)
475 upload_session.status = UploadSessionStatus.ABORTED
476 upload_session.save(update_fields=["status"])
477 return upload_session
479 if (
480 upload_session.status == UploadSessionStatus.COMPLETED
481 and upload_session.expires_at is not None
482 and upload_session.expires_at <= now
483 and not s3_storage.object_exists(upload_session.bucket, upload_session.object_key)
484 ):
485 upload_session.status = UploadSessionStatus.EXPIRED
486 upload_session.save(update_fields=["status"])
488 return upload_session
491def abort_upload(*, upload_session_id: int, user_id: int) -> UploadSession:
492 """Abort an in-progress multipart upload session."""
493 upload_session = _synchronise_upload_session_state(
494 _get_owned_upload_session(upload_session_id, user_id)
495 )
496 if upload_session.status == UploadSessionStatus.ABORTED:
497 return upload_session
499 if upload_session.status not in {UploadSessionStatus.INITIATED, UploadSessionStatus.UPLOADING}:
500 raise ValidationError({"upload_session_id": "Only in-progress uploads can be aborted."})
502 s3_storage.abort_multipart_upload(
503 upload_session.bucket,
504 upload_session.object_key,
505 upload_session.s3_upload_id,
506 )
507 upload_session.status = UploadSessionStatus.ABORTED
508 upload_session.save(update_fields=["status"])
509 return upload_session
512def inspect_upload_session(upload_session: UploadSession) -> CsvInspectionResult:
513 """Inspect an uploaded CSV object and cache its header metadata on the session."""
514 if upload_session.status != UploadSessionStatus.COMPLETED: 514 ↛ 515line 514 didn't jump to line 515 because the condition on line 514 was never true
515 raise ValidationError({"upload_session_id": "The upload session must be completed before CSV inspection can run."})
517 sample_bytes = s3_storage.get_range(
518 upload_session.bucket,
519 upload_session.object_key,
520 f"0-{DEFAULT_INSPECTION_BYTES - 1}",
521 )
522 inspection = inspect_csv_sample(sample_bytes)
524 if upload_session.csv_headers and upload_session.csv_delimiter:
525 return CsvInspectionResult(
526 headers=list(upload_session.csv_headers),
527 delimiter=upload_session.csv_delimiter,
528 preview_rows=inspection.preview_rows,
529 warnings=list(upload_session.csv_warnings or inspection.warnings),
530 )
532 upload_session.csv_headers = inspection.headers
533 upload_session.csv_delimiter = inspection.delimiter
534 upload_session.csv_warnings = inspection.warnings
535 upload_session.csv_inspected_at = timezone.now()
536 upload_session.save(
537 update_fields=[
538 "csv_headers",
539 "csv_delimiter",
540 "csv_warnings",
541 "csv_inspected_at",
542 ]
543 )
544 return inspection
547@transaction.atomic
548def attach_upload_to_ml_model(*, ml_model: MLModel, upload_session_id: int, user_id: int) -> MLModel:
549 """Create or reuse an ML model bound to a completed CSV upload session."""
550 upload_session = _get_owned_upload_session(
551 upload_session_id,
552 user_id,
553 for_update=True,
554 )
555 if upload_session.status != UploadSessionStatus.COMPLETED: 555 ↛ 556line 555 didn't jump to line 556 because the condition on line 555 was never true
556 raise ValidationError(
557 {
558 "upload_session_id": (
559 "The upload session must be completed before it can be "
560 "attached to an ML model."
561 )
562 }
563 )
564 if upload_session.purpose != UploadSessionPurpose.ML_TRAINING_CSV: 564 ↛ 565line 564 didn't jump to line 565 because the condition on line 564 was never true
565 raise ValidationError(
566 {
567 "upload_session_id": (
568 "The upload session purpose must be ml_training_csv."
569 )
570 }
571 )
573 ml_model = MLModel.objects.select_for_update().get(id=ml_model.id)
574 simulation_object = _get_simulation_object(ml_model.simulationObject_id)
575 if ( 575 ↛ 579line 575 didn't jump to line 579 because the condition on line 575 was never true
576 upload_session.simulationObject_id
577 and upload_session.simulationObject_id != simulation_object.id
578 ):
579 raise ValidationError(
580 {
581 "upload_session_id": (
582 "This upload session belongs to a different simulation object."
583 )
584 }
585 )
586 if upload_session.flowsheet_state_id != simulation_object.flowsheet_state_id: 586 ↛ 587line 586 didn't jump to line 587 because the condition on line 586 was never true
587 raise ValidationError(
588 {
589 "upload_session_id": (
590 "This upload session belongs to a different flowsheet."
591 )
592 }
593 )
595 inspection = inspect_upload_session(upload_session)
596 if not inspection.preview_rows:
597 raise ValidationError(
598 {
599 "upload_session_id": (
600 "The uploaded CSV must include at least one data row."
601 )
602 }
603 )
605 snapshot_staged = stage_ml_model_update_snapshot(ml_model)
606 clear_ml_model_column_mappings(ml_model)
608 ml_model.csv_file_name = upload_session.original_filename
609 ml_model.csv_bucket = upload_session.bucket
610 ml_model.csv_object_key = upload_session.object_key
611 ml_model.csv_headers = inspection.headers
612 ml_model.csv_delimiter = inspection.delimiter
613 ml_model.csv_upload_session = upload_session
614 ml_model.is_resetting = ml_model.is_resetting
615 ml_model.is_updating = ml_model.is_updating
616 ml_model.completed_steps = []
617 ml_model.surrogate_model = {}
618 ml_model.charts = []
619 ml_model.metrics = []
620 ml_model.result_state = MLModel.ResultState.PENDING
621 ml_model.test_results_bucket = ""
622 ml_model.test_results_key = ""
624 update_fields = [
625 "csv_file_name",
626 "csv_bucket",
627 "csv_object_key",
628 "csv_headers",
629 "csv_delimiter",
630 "csv_upload_session",
631 "is_resetting",
632 "is_updating",
633 "completed_steps",
634 "surrogate_model",
635 "charts",
636 "metrics",
637 "result_state",
638 "test_results_bucket",
639 "test_results_key",
640 ]
641 if snapshot_staged:
642 update_fields.append("mapping_update_snapshot")
644 ml_model.save(update_fields=update_fields)
646 return ml_model
649def cleanup_upload_sessions() -> UploadCleanupSummary:
650 """Abort stale multipart uploads and delete expired unowned ML objects."""
651 now = timezone.now()
652 stale_cutoff = now - timedelta(hours=settings.SEAWEED_STALE_MULTIPART_UPLOAD_HOURS)
653 aborted_count = 0
654 expired_ml_count = 0
656 stale_uploads = UploadSession.objects.filter(
657 status__in=[UploadSessionStatus.INITIATED, UploadSessionStatus.UPLOADING],
658 created_at__lt=stale_cutoff,
659 )
660 for upload_session in stale_uploads:
661 s3_storage.abort_multipart_upload(
662 upload_session.bucket,
663 upload_session.object_key,
664 upload_session.s3_upload_id,
665 )
666 upload_session.status = UploadSessionStatus.ABORTED
667 upload_session.save(update_fields=["status"])
668 aborted_count += 1
670 expired_ml_session_ids = list(
671 UploadSession.objects.filter(
672 purpose=UploadSessionPurpose.ML_TRAINING_CSV,
673 status=UploadSessionStatus.COMPLETED,
674 expires_at__lte=now,
675 ).values_list("pk", flat=True)
676 )
677 for upload_session_id in expired_ml_session_ids:
678 with transaction.atomic():
679 upload_session = UploadSession.objects.select_for_update().get(
680 pk=upload_session_id
681 )
682 if ( 682 ↛ 687line 682 didn't jump to line 687 because the condition on line 682 was never true
683 upload_session.status != UploadSessionStatus.COMPLETED
684 or upload_session.expires_at is None
685 or upload_session.expires_at > now
686 ):
687 continue
688 deleted = delete_ml_artifact_if_unreferenced(
689 bucket=upload_session.bucket,
690 key=upload_session.object_key,
691 )
692 if not deleted:
693 continue
694 upload_session.status = UploadSessionStatus.EXPIRED
695 upload_session.save(update_fields=["status"])
696 expired_ml_count += 1
698 return UploadCleanupSummary(
699 aborted_stale_uploads=aborted_count,
700 expired_ml_uploads=expired_ml_count,
701 )