Coverage for backend/django/core/auxiliary/models/UploadSession.py: 68%

1from django.conf import settings 

2from django.db import models 

3 

4from core.managers import IdentityAccessControlManager 

5from flowsheetInternals.unitops.models import SimulationObject 

6 

7 

8class UploadSessionPurpose(models.TextChoices): 

9 ML_TRAINING_CSV = "ml_training_csv" 

10 SCENARIO_CSV = "scenario_csv" 

11 PINCH_UTILITY_CSV = "pinch_utility_csv" 

12 

13 

14class UploadSessionStatus(models.TextChoices): 

15 INITIATED = "initiated" 

16 UPLOADING = "uploading" 

17 COMPLETED = "completed" 

18 ABORTED = "aborted" 

19 FAILED = "failed" 

20 EXPIRED = "expired" 

21 

22 

23class UploadSession(models.Model): 

24 created_by = models.ForeignKey( 

25 settings.AUTH_USER_MODEL, 

26 on_delete=models.CASCADE, 

27 related_name="upload_sessions", 

28 ) 

29 flowsheet = models.ForeignKey( 

30 "Flowsheet", 

31 on_delete=models.CASCADE, 

32 related_name="upload_sessions", 

33 ) 

34 flowsheet_state = models.ForeignKey( 

35 "FlowsheetState", 

36 on_delete=models.SET_NULL, 

37 related_name="upload_sessions", 

38 null=True, 

39 blank=True, 

40 ) 

41 scenario = models.ForeignKey( 

42 "Scenario", 

43 on_delete=models.SET_NULL, 

44 related_name="upload_sessions", 

45 null=True, 

46 blank=True, 

47 ) 

48 simulationObject = models.ForeignKey( 

49 SimulationObject, 

50 on_delete=models.SET_NULL, 

51 related_name="upload_sessions", 

52 null=True, 

53 blank=True, 

54 ) 

55 purpose = models.CharField(max_length=32, choices=UploadSessionPurpose.choices) 

56 bucket = models.CharField(max_length=255) 

57 object_key = models.CharField(max_length=1024, unique=True) 

58 s3_upload_id = models.CharField(max_length=512) 

59 original_filename = models.CharField(max_length=255) 

60 content_type = models.CharField(max_length=255, blank=True, default="") 

61 size_bytes = models.BigIntegerField() 

62 status = models.CharField( 

63 max_length=16, 

64 choices=UploadSessionStatus.choices, 

65 default=UploadSessionStatus.INITIATED, 

66 ) 

67 csv_headers = models.JSONField(default=list, blank=True) 

68 csv_delimiter = models.CharField(max_length=1, blank=True, default="") 

69 csv_warnings = models.JSONField(default=list, blank=True) 

70 csv_inspected_at = models.DateTimeField(null=True, blank=True) 

71 created_at = models.DateTimeField(auto_now_add=True) 

72 completed_at = models.DateTimeField(null=True, blank=True) 

73 expires_at = models.DateTimeField(null=True, blank=True) 

74 

75 objects = IdentityAccessControlManager() 

76 

77 def clean(self): 

78 """Require operational provenance to resolve within one stable flowsheet.""" 

79 

80 super().clean() 

81 if self.flowsheet_state_id is None: 

82 return 

83 if self.flowsheet_state.flowsheet_id != self.flowsheet_id: 

84 from django.core.exceptions import ValidationError 

85 

86 raise ValidationError( 

87 {"flowsheet_state": "Upload state must belong to the upload flowsheet."} 

88 ) 

89 for field_name in ("scenario", "simulationObject"): 

90 related = getattr(self, field_name, None) 

91 if related is not None and related.flowsheet_state_id != self.flowsheet_state_id: 

92 from django.core.exceptions import ValidationError 

93 

94 raise ValidationError( 

95 {field_name: "Upload source must belong to the upload state."} 

96 ) 

97 

98 class Meta: 

99 indexes = [ 

100 models.Index(fields=["purpose", "status"]), 

101 models.Index(fields=["expires_at"]), 

102 ]