Merge pull request #501 from gemini-hlsw/GSCHED-731-union-types

Added union types to handle subscription errors

scheduler/graphql_mid/schema.py

@@ -15,7 +15,7 @@
 from scheduler.db.planmanager import PlanManager
 from scheduler.services.logger_factory import create_logger
 
-from .types import (SPlans, NewNightPlans, SNightTimelines)
+from .types import (SPlans, SNightTimelines, NewNightPlans, NightPlansError, NightPlansResponse)
 from .inputs import CreateNewScheduleInput
 
 
@@ -134,7 +134,7 @@ async def test_sub_query(self, schedule_id: str, new_schedule_input: CreateNewSc
 @strawberry.type
 class Subscription:
     @strawberry.subscription
-    async def queue_schedule(self, schedule_id: str) -> AsyncGenerator[NewNightPlans, None]:
+    async def queue_schedule(self, schedule_id: str) -> AsyncGenerator[NightPlansResponse, None]:
         if schedule_id not in active_subscriptions:
             queue = asyncio.Queue()
             active_subscriptions[schedule_id] = queue
@@ -143,9 +143,15 @@ async def queue_schedule(self, schedule_id: str) -> AsyncGenerator[NewNightPlans
 
         try:
             while True:
-                item = await queue.get()  # Wait for item from the queue
-                result = await item
-                yield result  # Yield item to the subscription
+                try:
+                    print(f'Queueing {schedule_id}')
+                    item = await queue.get()  # Wait for item from the queue
+                    print(f'Run ID: {schedule_id}')
+                    result = await item
+                    yield result  # Yield item to the subscription
+                except Exception as e:
+                    _logger.error(f'Queue error: {e}')
+                    yield NightPlansError(error=f'Queue error: {e}')
         finally:
             if schedule_id in active_subscriptions:
                 del active_subscriptions[schedule_id]

scheduler/graphql_mid/types.py

@@ -7,6 +7,7 @@
 from zoneinfo import ZoneInfo
 
 import strawberry  # noqa
+from typing import Annotated, Union
 import astropy.units as u
 from astropy.coordinates import Angle
 from strawberry.scalars import JSON  # noqa
@@ -192,6 +193,11 @@ class NewNightPlans:
     night_plans: SNightTimelines
     plans_summary: JSON
 
+@strawberry.type
+class NightPlansError:
+    error: str
+
+NightPlansResponse = Annotated[Union[NewNightPlans, NightPlansError], strawberry.union("NightPlansResponse")]
 
 @strawberry.type
 class NewScheduleSuccess: