Coverage for backend/django/core/auxiliary/viewsets/MLColumnMappingsViewSet.py: 88%

1from django.db import IntegrityError, transaction 

2from core.viewset import ModelViewSet 

3from core.auxiliary.models.MLColumnMapping import MLColumnMapping 

4from core.auxiliary.models.Flowsheet import Flowsheet 

5from rest_framework import serializers 

6from core.auxiliary.serializers.MLColumnMappingsSerializer import MLColumnMappingsSerializer 

7from rest_framework.response import Response 

8from drf_spectacular.utils import extend_schema 

9from rest_framework.decorators import action 

10from core.auxiliary.serializers.PropertyInfoSerializer import PropertyInfoSerializer 

11from drf_spectacular.utils import extend_schema 

12from rest_framework.exceptions import NotFound, ValidationError 

13from core.auxiliary.models.MLWizard import create_column_mapping 

14from core.auxiliary.models.MLModel import MLModel 

15from core.auxiliary.services.ml_column_mapping_updates import ( 

16 replace_column_mapping as replace_column_mapping_service, 

17) 

18 

19class ColumnMappingSerializer(serializers.Serializer): 

20 portIndex = serializers.IntegerField() 

21 propertyKey = serializers.CharField() 

22 column = serializers.CharField() 

23 

24class BulkCreateColumnMappingSerializer(serializers.Serializer): 

25 model = serializers.IntegerField() 

26 inlet_mappings = serializers.ListField(child=ColumnMappingSerializer()) 

27 outlet_mappings = serializers.ListField(child=ColumnMappingSerializer()) 

28class PropertyInfoIdsSerializer(serializers.Serializer): 

29 inlets = serializers.ListField(child=PropertyInfoSerializer()) 

30 outlets = serializers.ListField(child=PropertyInfoSerializer()) 

31 

32 

33class MLColumnMappingsViewSet(ModelViewSet): 

34 serializer_class = MLColumnMappingsSerializer 

35 write_actions = {"bulk_create_column_mapping", "replace_column_mapping"} 

36 

37 def get_queryset(self): 

38 return MLColumnMapping.objects.all() 

39 

40 def get_flowsheet_from_request(self, request): 

41 flowsheet_id = request.query_params.get("flowsheet") 

42 if not flowsheet_id: 42 ↛ 43line 42 didn't jump to line 43 because the condition on line 42 was never true

43 raise ValidationError({"flowsheet": "flowsheet query parameter is required."}) 

44 

45 try: 

46 return Flowsheet.objects.get(pk=flowsheet_id) 

47 except Flowsheet.DoesNotExist as exc: 

48 raise NotFound({"flowsheet": "Flowsheet not found."}) from exc 

49 

50 @transaction.atomic 

51 @extend_schema(request=BulkCreateColumnMappingSerializer, responses=None) 

52 @action(detail=False, methods=['post'], url_path='bulk-create-column-mapping') 

53 def bulk_create_column_mapping(self, request): 

54 serializer = BulkCreateColumnMappingSerializer(data=request.data) 

55 serializer.is_valid(raise_exception=True) 

56 validated_data = serializer.validated_data 

57 model = validated_data.get("model") 

58 inlet_mappings = validated_data.get("inlet_mappings") 

59 outlet_mappings = validated_data.get("outlet_mappings") 

60 flowsheet = self.get_flowsheet_from_request(request) 

61 try: 

62 create_column_mapping(flowsheet, model, inlet_mappings, outlet_mappings) 

63 except Exception as e: 

64 print(ValidationError(str(e))) 

65 return Response({'status': 'error'}, status=400) 

66 

67 return Response({'status': 'success'}, status=200) 

68 

69 @extend_schema(request=BulkCreateColumnMappingSerializer, responses=None) 

70 @action(detail=False, methods=['post'], url_path='replace-column-mapping') 

71 def replace_column_mapping(self, request): 

72 serializer = BulkCreateColumnMappingSerializer(data=request.data) 

73 serializer.is_valid(raise_exception=True) 

74 validated_data = serializer.validated_data 

75 model = validated_data.get("model") 

76 inlet_mappings = validated_data.get("inlet_mappings") 

77 outlet_mappings = validated_data.get("outlet_mappings") 

78 flowsheet = self.get_flowsheet_from_request(request) 

79 try: 

80 replace_column_mapping_service( 

81 flowsheet=flowsheet, 

82 model=model, 

83 inlet_mappings=inlet_mappings, 

84 outlet_mappings=outlet_mappings, 

85 ) 

86 except MLModel.DoesNotExist as exc: 

87 raise NotFound({"model": "ML model not found."}) from exc 

88 except IntegrityError as exc: 

89 raise ValidationError({"mappings": "Column mappings could not be replaced."}) from exc 

90 

91 return Response({'status': 'success'}, status=200)