SCHED-704: Add test sub query plan creation

scheduler/graphql_mid/schema.py

@@ -1,9 +1,9 @@
 # Copyright (c) 2016-2024 Association of Universities for Research in Astronomy, Inc. (AURA)
 # For license information see LICENSE or https://opensource.org/licenses/BSD-3-Clause
-
+import asyncio
 import os
 from datetime import datetime
-from typing import List
+from typing import List, AsyncGenerator
 
 import strawberry # noqa
 from astropy.time import Time
@@ -24,72 +24,18 @@
 from .inputs import CreateNewScheduleInput, UseFilesSourceInput
 from .scalars import SOrigin
 
-
-# TODO: This variables need a Redis cache to work with different mutations correctly.
-# TODO: This should NOT be 3, but the actual number of nights.
-sources = Sources()
-event_queue = EventQueue(frozenset([i for i in range(3)]), ALL_SITES)
 REDIS_URL = os.environ.get("REDISCLOUD_URL")
 redis = aioredis.from_url(REDIS_URL) if REDIS_URL else None
 
+task_queue = asyncio.Queue()
 # TODO: All times need to be in UTC. This is done here but converted from the Optimizer plans, where it should be done.
 
 
-@strawberry.type
-class Mutation:
-    """
-    @strawberry.mutation
-    def change_mode():
-        pass
-    """
-
-    @strawberry.mutation
-    async def load_sources_files(self, files_input: UseFilesSourceInput) -> SourceFileHandlerResponse:
-        service = Services[files_input.service]
-
-        match service:
-            case Services.RESOURCE:
-                calendar = await files_input.calendar.read()
-                gmos_fpu = await files_input.gmos_fpus.read()
-                gmos_gratings = await files_input.gmos_gratings.read()
-                faults = await files_input.faults.read()
-                eng_tasks = await files_input.eng_tasks.read()
-                weather_closures = await files_input.weather_closures.read()
-
-                loaded = sources.use_file(files_input.sites,
-                                          service,
-                                          calendar,
-                                          gmos_fpu,
-                                          gmos_gratings,
-                                          faults,
-                                          eng_tasks,
-                                          weather_closures)
-                if loaded:
-                    return SourceFileHandlerResponse(service=files_input.service,
-                                                     loaded=loaded,
-                                                     msg=f'Files were loaded for service: {service}')
-                else:
-                    return SourceFileHandlerResponse(service=files_input.service,
-                                                     loaded=loaded,
-                                                     msg='Files failed to load!')
-            case Services.ENV:
-                return SourceFileHandlerResponse(service=files_input.service,
-                                                 loaded=False,
-                                                 msg='Handler not implemented yet!')
-
-    @strawberry.mutation
-    def change_origin(self, new_origin: SOrigin, mode: SchedulerModes) -> ChangeOriginSuccess:
-
-        old = str(sources.origin)
-        new = str(new_origin)
-        if new == 'OCS' and mode is SchedulerModes.SIMULATION:
-            raise ValueError('Simulation mode can only work with GPP origin source.')
-        elif new == 'GPP' and mode is SchedulerModes.VALIDATION:
-            raise ValueError('Validation mode can only work with OCS origin source.')
-        if old == str(new_origin):
-            return ChangeOriginSuccess(from_origin=old, to_origin=old)
-        sources.set_origin(new_origin)
-        return ChangeOriginSuccess(from_origin=old, to_origin=str(new_origin))
+def synchronous_task():
+    import time
+    print('task')
+    time.sleep(65)
+    return 'new plan'
 
 
 @strawberry.type
@@ -144,3 +90,19 @@ async def schedule(self,
         except RuntimeError as e:
             raise RuntimeError(f'Schedule query error: {e}')
         return NewNightPlans(night_plans=s_timelines, plans_summary=plan_summary)
+
+    @strawberry.field
+    async def test_sub_query(self) -> str:
+        task = asyncio.to_thread(synchronous_task)
+        await task_queue.put(task)
+        return 'Plan is on the queue!'
+
+
+@strawberry.type
+class Subscription:
+    @strawberry.subscription
+    async def queue_schedule(self) -> AsyncGenerator[str, None]:
+        while True:
+            item = await task_queue.get()  # Wait for item from the queue
+            result = await item
+            yield result  # Yield item to the subscription

scheduler/graphql_mid/server.py

@@ -4,7 +4,7 @@
 import strawberry # noqa
 from strawberry.asgi import GraphQL # noqa
 
-from .schema import Query, Mutation
+from .schema import Query, Subscription
 
-schema = strawberry.Schema(query=Query, mutation=Mutation)
+schema = strawberry.Schema(query=Query, subscription=Subscription)
 graphql_server = GraphQL(schema)