Coverage for backend/django/flowsheetInternals/unitops/viewsets/PortViewSet.py: 60%
175 statements
« prev ^ index » next coverage.py v7.10.7, created at 2026-07-22 05:22 +0000
1from drf_spectacular.utils import extend_schema, OpenApiParameter, OpenApiTypes
2from core.viewset import ModelViewSet
3from flowsheetInternals.unitops.models.Port import Port
4from flowsheetInternals.unitops.serializers.PortSerializer import PortSerializer
5from rest_framework.response import Response
6from rest_framework import serializers
7import traceback
8from rest_framework.decorators import action
9from django.db import transaction
11from flowsheetInternals.unitops.models.SimulationObject import SimulationObject
12from flowsheetInternals.unitops.models.simulation_object_factory import (
13 SimulationObjectFactory,
14)
15from common.config_types import ConType
16from flowsheetInternals.unitops.models.delete_factory import DeleteFactory
17from flowsheetInternals.unitops.models.FlowsheetEditOperation import (
18 FlowsheetEditOperation,
19)
20from flowsheetInternals.unitops.services.edit_operations.mutation import flowsheet_edit
21from flowsheetInternals.unitops.services.edit_operations.recorder import (
22 tracked_bulk_update,
23)
26class MergeStreamsSerializer(serializers.Serializer):
27 stream1 = serializers.IntegerField(required=True) # material stream
28 stream2 = serializers.IntegerField(required=True) # material stream
31class SplitStreamSerializer(serializers.Serializer):
32 stream = serializers.IntegerField(required=True)
35class AddStreamSerializer(serializers.Serializer):
36 port = serializers.IntegerField(required=True)
39class CreateDNSerializer(serializers.Serializer):
40 stream = serializers.IntegerField(required=True)
43class ConvertToDNSerializer(serializers.Serializer):
44 stream = serializers.IntegerField(required=True)
47class RestoreConnectionsSerializer(serializers.Serializer):
48 connections = serializers.ListField(
49 child=serializers.DictField(child=serializers.JSONField()), required=True
50 )
53class PortViewSet(ModelViewSet):
54 serializer_class = PortSerializer
56 def get_queryset(self):
57 queryset = Port.objects.all()
58 queryset = queryset.filter(unitOp__is_deleted=False)
59 return queryset
61 @extend_schema(
62 parameters=[
63 OpenApiParameter(name="flowsheet", required=True, type=OpenApiTypes.INT),
64 ]
65 )
66 def list(self, request):
67 return super().list(request)
69 def update(self, request, *args, **kwargs):
70 """Update one port as an atomic, reversible topology operation."""
71 port = self.get_object()
72 with flowsheet_edit(
73 flowsheet=port.flowsheet_state.flowsheet,
74 user=request.user,
75 kind=FlowsheetEditOperation.Kind.Connection,
76 label_key="Change connection",
77 ) as mutation:
78 response = super().update(request, *args, **kwargs)
79 return mutation.add_to_response(response)
81 def error_response(self, e):
82 """Return a safe error and roll back any surrounding compound mutation."""
83 if transaction.get_connection().in_atomic_block:
84 transaction.set_rollback(True)
85 traceback.print_exc()
86 return Response(
87 {
88 "status": "error",
89 "message": "The connection change could not be completed.",
90 },
91 status=400,
92 )
94 @extend_schema(request=MergeStreamsSerializer, responses=None)
95 @action(methods=["post"], detail=False, url_path="merge-streams")
96 @transaction.atomic
97 def merge_streams(self, request):
98 """
99 Connects two material streams together, managing the connection/disconnection of UnitOperations.
100 Logic: keep the outlet stream and keep the stream 2 position.
101 For Product/Feed -> Intermediary, Keep stream 1&2, create new outlet stream
102 """
103 try:
104 serializer = MergeStreamsSerializer(data=request.data)
105 serializer.is_valid(raise_exception=True)
106 validated_data = serializer.validated_data
108 ids = [validated_data.get("stream1"), validated_data.get("stream2")]
110 # Fetch both streams in a single query with related data
111 streams = (
112 SimulationObject.objects.filter(id__in=ids)
113 .select_related("flowsheet_state__flowsheet")
114 .prefetch_related("connectedPorts__unitOp", "graphicObject__group")
115 )
117 stream_map = {stream.id: stream for stream in streams}
118 stream1 = stream_map[ids[0]]
119 stream2 = stream_map[ids[1]]
121 with flowsheet_edit(
122 flowsheet=stream1.flowsheet_state.flowsheet,
123 user=request.user,
124 kind=FlowsheetEditOperation.Kind.MergeStreams,
125 label_key="Connect material streams",
126 ) as mutation:
127 stream1.merge_stream(stream2)
129 return Response(
130 mutation.add_to_data({"status": "success"}),
131 status=200,
132 )
133 except Exception as e:
134 return self.error_response(e)
136 @extend_schema(request=SplitStreamSerializer, responses=None)
137 @action(methods=["post"], detail=False, url_path="split-stream")
138 @transaction.atomic
139 def split_stream(self, request):
140 """
141 Splits a stream into two separate streams (one inlet and one outlet - disconnected).
142 """
143 try:
144 serializer = SplitStreamSerializer(data=request.data)
145 serializer.is_valid(raise_exception=True)
146 validated_data = serializer.validated_data
148 stream = SimulationObject.objects.get(id=validated_data.get("stream"))
149 with flowsheet_edit(
150 flowsheet=stream.flowsheet_state.flowsheet,
151 user=request.user,
152 kind=FlowsheetEditOperation.Kind.Connection,
153 label_key="Disconnect material stream",
154 ) as mutation:
155 stream.split_stream()
156 return Response(
157 mutation.add_to_data({"status": "success"}),
158 status=200,
159 )
160 except Exception as e:
161 return self.error_response(e)
163 @extend_schema(request=AddStreamSerializer, responses=None)
164 @action(methods=["post"], detail=False, url_path="add-stream")
165 @transaction.atomic
166 def add_stream(self, request) -> Response:
167 """
168 Creates a new stream for this port (previously with no stream attached).
169 """
170 try:
171 port = Port.objects.get(id=request.data.get("port"))
172 with flowsheet_edit(
173 flowsheet=port.flowsheet_state.flowsheet,
174 user=request.user,
175 kind=FlowsheetEditOperation.Kind.Add,
176 label_key="Add material stream",
177 ) as mutation:
178 SimulationObjectFactory.create_stream_at_port(port)
179 return Response(
180 mutation.add_to_data({"status": "success"}),
181 status=200,
182 )
184 except Exception as e:
185 return self.error_response(e)
187 @extend_schema(request=ConvertToDNSerializer, responses=None)
188 @action(methods=["post"], detail=False, url_path="convert-to-dn")
189 @transaction.atomic
190 def convert_to_dn(self, request) -> Response:
191 """
192 Turns an Intermediary stream into a decision node
193 """
194 try:
195 serializer = CreateDNSerializer(data=request.data)
196 serializer.is_valid(raise_exception=True)
197 validated_data = serializer.validated_data
199 # Get stream and outlets
200 id = validated_data.get("stream", None)
201 stream = SimulationObject.objects.get(id=id)
202 with flowsheet_edit(
203 flowsheet=stream.flowsheet_state.flowsheet,
204 user=request.user,
205 kind=FlowsheetEditOperation.Kind.Connection,
206 label_key="Convert stream to decision node",
207 ) as mutation:
208 stream_port = stream.connectedPorts.first()
209 num_inlets, num_outlets = 1, 1
210 if stream.connectedPorts.count() <= 1: 210 ↛ 215line 210 didn't jump to line 215 because the condition on line 210 was always true
211 if stream_port.direction == ConType.Inlet: 211 ↛ 214line 211 didn't jump to line 214 because the condition on line 211 was always true
212 num_inlets, num_outlets = 0, 1
213 else:
214 num_inlets, num_outlets = 1, 0
215 stream.make_decision_node(
216 num_inlets=num_inlets,
217 num_outlets=num_outlets,
218 )
220 return Response(
221 mutation.add_to_data({"status": "success"}),
222 status=200,
223 )
224 except Exception as e:
225 return self.error_response(e)
227 @transaction.atomic
228 def destroy(self, request, *args, **kwargs):
229 """
230 Override the destroy method to allow for deleting attached streams and updating port indexes.
231 """
232 try:
233 port: Port = self.get_object()
234 stream = port.stream
235 unit_op = port.unitOp
236 with flowsheet_edit(
237 flowsheet=port.flowsheet_state.flowsheet,
238 user=request.user,
239 kind=FlowsheetEditOperation.Kind.Connection,
240 label_key="Delete port",
241 ) as mutation:
242 port.reindex_port_on_delete()
243 unit_op.update_height()
244 if stream and stream.connectedPorts.count() == 0:
245 DeleteFactory.delete_object(stream)
247 return mutation.add_to_response(Response(status=204))
248 except Exception as e:
249 return self.error_response(e)
251 @extend_schema(request=RestoreConnectionsSerializer, responses=None)
252 @action(methods=["post"], detail=False, url_path="restore-connections")
253 def restore_connections(self, request):
254 """
255 Enhanced restoration with detailed error tracking and dependency handling.
256 Restores port connections during undo operations without creating recycle tears.
257 This endpoint bypasses the normal merge_stream logic that can create recycle blocks.
258 """
259 try:
260 with transaction.atomic(): # Ensure all-or-nothing operation
261 serializer = RestoreConnectionsSerializer(data=request.data)
262 serializer.is_valid(raise_exception=True)
263 validated_data = serializer.validated_data
265 connections = validated_data.get("connections", [])
267 # Bulk-fetch DB objects to reduce queries
268 port_ids = {
269 c.get("portId") for c in connections if c.get("portId") is not None
270 }
271 stream_ids = {
272 c.get("streamId")
273 for c in connections
274 if c.get("streamId") is not None
275 }
277 # Fetch all ports and streams in single queries
278 port_map = (
279 Port.objects.filter(id__in=port_ids)
280 .select_related("unitOp__flowsheet_state__flowsheet")
281 .in_bulk()
282 )
283 stream_map = (
284 SimulationObject.objects.filter(id__in=stream_ids)
285 .select_related("flowsheet_state__flowsheet")
286 .in_bulk()
287 )
289 updated_ports = []
290 restored_count = 0
291 failed_connections = []
292 skipped_connections = []
294 for connection in connections:
295 port_id = connection.get("portId")
296 stream_id = connection.get("streamId")
298 if port_id is None:
299 skipped_connections.append(
300 {
301 "portId": port_id,
302 "streamId": stream_id,
303 "reason": "Invalid port ID",
304 }
305 )
306 continue
308 port = port_map.get(port_id)
309 if not port:
310 failed_connections.append(
311 {
312 "portId": port_id,
313 "streamId": stream_id,
314 "error": f"Port {port_id} not found",
315 "retry_recommended": False, # Port truly doesn't exist
316 }
317 )
318 continue
320 # Handle disconnection
321 if stream_id is None:
322 port.stream = None
323 updated_ports.append(port)
324 restored_count += 1
325 continue
327 # Handle connection
328 stream = stream_map.get(stream_id)
329 if not stream:
330 failed_connections.append(
331 {
332 "portId": port_id,
333 "streamId": stream_id,
334 "error": f"Stream {stream_id} not found",
335 "retry_recommended": True,
336 }
337 )
338 continue
340 # Ensure same flowsheet (security check)
341 if port.unitOp.flowsheet_state_id != stream.flowsheet_state_id:
342 failed_connections.append(
343 {
344 "portId": port_id,
345 "streamId": stream_id,
346 "error": "Port and Stream belong to different flowsheets",
347 "retry_recommended": False,
348 }
349 )
350 continue
352 port.stream = stream
353 updated_ports.append(port)
354 restored_count += 1
356 # Bulk update all ports in single query
357 if updated_ports:
358 tracked_bulk_update(Port.objects, updated_ports, ["stream"])
359 updated_ports[0].flowsheet_state.flowsheet.record_content_change()
361 response_data = {
362 "status": "success" if restored_count else "failed",
363 "restored_count": restored_count,
364 "total_connections": len(connections),
365 "failed_connections": failed_connections,
366 "skipped_connections": skipped_connections,
367 "retry_candidates": [
368 f for f in failed_connections if f.get("retry_recommended")
369 ],
370 }
372 return Response(response_data, status=200)
374 except Exception as e:
375 return self.error_response(e)