Coverage for backend/django/core/auxiliary/views/UploadMSSData.py: 87%

1from rest_framework.response import Response 

2from django.db import transaction 

3from core.auxiliary.models.DataCell import DataCell 

4from core.auxiliary.models.DataRow import DataRow 

5from drf_spectacular.utils import extend_schema 

6from core.auxiliary.models.DataColumn import DataColumn 

7from core.auxiliary.models.Scenario import Scenario, ScenarioInputModeEnum 

8from core.auxiliary.services.parameter_sweep import clear_mss_input_data 

9from rest_framework.decorators import api_view 

10from rest_framework import serializers, status 

11from core.validation import api_view_validate 

12from core.managers import get_flowsheet_access 

13 

14class UploadDataSerializer(serializers.Serializer): 

15 # THe data format is e.g: 

16 # { 

17 # "data": { 

18 # "heater_enthalpy": [1, 2, 3, 4, 5], 

19 # "heater_temperature": [1, 2, 3, 4, 5] 

20 # } 

21 # "flowsheet": 1 

22 # } 

23 flowsheet = serializers.IntegerField() 

24 scenario=serializers.IntegerField() 

25 data = serializers.DictField( # column name 

26 child=serializers.ListField( # List of values 

27 child=serializers.FloatField() # Value 

28 ) 

29 ) 

30 

31 

32@api_view_validate 

33@extend_schema(request=UploadDataSerializer, responses=None) 

34@api_view(['POST']) 

35def upload_data(request) -> Response: 

36 flowsheet_id = request.GET.get("flowsheet") 

37 access_state = get_flowsheet_access(request.user, flowsheet_id) 

38 if access_state.has_read_access and not access_state.has_write_access: 

39 return Response( 

40 {"error": "This flowsheet is shared with read-only access."}, 

41 status=status.HTTP_403_FORBIDDEN, 

42 ) 

43 

44 try: 

45 serializer = UploadDataSerializer(data=request.data) 

46 serializer.is_valid(raise_exception=True) 

47 validated_data = serializer.validated_data 

48 data = validated_data.get('data') 

49 flowsheet_id = validated_data.get('flowsheet') 

50 scenario_id = validated_data.get('scenario') 

51 except Exception as e: 

52 return Response(status=400, data=f"Invalid csv data: {e}") 

53 

54 if not data: 54 ↛ 55line 54 didn't jump to line 55 because the condition on line 54 was never true

55 return Response(status=400, data="CSV data must include at least one column.") 

56 

57 num_rows = len(list(next(iter(data.values())))) 

58 if any(len(values) != num_rows for values in data.values()): 

59 return Response(status=400, data="CSV columns must all have the same number of rows.") 

60 

61 try: 

62 scenario = Scenario.objects.get(id=scenario_id, flowsheet_id=flowsheet_id) 

63 except Scenario.DoesNotExist: 

64 return Response(status=404, data="Scenario not found.") 

65 

66 with transaction.atomic(): 

67 replacing_parameter_sweep = ( 

68 scenario.mss_input_mode == ScenarioInputModeEnum.ParameterSweep 

69 or hasattr(scenario, "parameterSweepDefinition") 

70 ) 

71 if replacing_parameter_sweep: 71 ↛ 72line 71 didn't jump to line 72 because the condition on line 71 was never true

72 clear_mss_input_data(scenario) 

73 

74 scenario.mss_input_mode = ScenarioInputModeEnum.Csv 

75 scenario.save(update_fields=["mss_input_mode"]) 

76 

77 # Step 1: Create any missing data column under the given optimization 

78 data_columns = [] 

79 for key in data: 

80 data_columns.append(DataColumn(name=key, scenario_id=scenario_id,flowsheet_id=flowsheet_id)) 

81 

82 DataColumn.objects.bulk_create(data_columns, ignore_conflicts=True) 

83 

84 # Step 2: Get existing data rows by optimization 

85 existing_rows = list(DataRow.objects.filter(scenario_id=scenario_id).order_by("index")) 

86 existing_indices = {r.index for r in existing_rows} 

87 

88 # Step 3: Create missing data rows 

89 new_rows = [ 

90 DataRow(index=i, flowsheet_id=flowsheet_id, scenario_id=scenario_id) 

91 for i in range(num_rows) 

92 if i not in existing_indices 

93 ] 

94 if new_rows: 94 ↛ 98line 94 didn't jump to line 98 because the condition on line 94 was always true

95 DataRow.objects.bulk_create(new_rows) 

96 

97 # Refresh data row list 

98 data_rows = list(DataRow.objects.filter(scenario_id=scenario_id).order_by("index")) 

99 data_row_map = {r.index: r for r in data_rows} 

100 

101 # Step 4: Get updated data columns 

102 data_columns = DataColumn.objects.filter(scenario_id=scenario_id).prefetch_related("dataCells") 

103 column_map = {column.name: column for column in data_columns} 

104 

105 # Step 5: Build mapping of existing values 

106 existing_values = { 

107 column.name: {sv.data_row.index: sv for sv in column.dataCells.all()} 

108 for column in data_columns 

109 } 

110 

111 # Step 6: Insert or update DataCells 

112 for i in range(num_rows): 

113 data_row = data_row_map[i] 

114 for column_name, values in data.items(): 

115 value = values[i] 

116 column = column_map[column_name] 

117 existing = existing_values.get(column_name, {}).get(i) 

118 

119 if existing: 119 ↛ 120line 119 didn't jump to line 120 because the condition on line 119 was never true

120 existing.value = value 

121 existing.save() 

122 else: 

123 DataCell.objects.create(value=value, data_column=column, data_row=data_row, flowsheet_id=flowsheet_id) 

124 

125 return Response(status=200, data="Data uploaded successfully")