have both asyncpg and aiopg

Author: sanderegg

services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/_scheduler_base.py

@@ -33,6 +33,7 @@
 from servicelib.logging_utils import log_catch, log_context
 from servicelib.rabbitmq import RabbitMQClient, RabbitMQRPCClient
 from servicelib.redis import RedisClientSDK
+from sqlalchemy.ext.asyncio import AsyncEngine
 
 from ...constants import UNDEFINED_STR_METADATA
 from ...core.errors import (
@@ -116,7 +117,7 @@ class SortedTasks:
 
 
 async def _triage_changed_tasks(
-    changed_tasks: list[tuple[_Previous, _Current]]
+    changed_tasks: list[tuple[_Previous, _Current]],
 ) -> SortedTasks:
     started_tasks = [
         current
@@ -159,6 +160,7 @@ async def _triage_changed_tasks(
 @dataclass
 class BaseCompScheduler(ABC):
     db_engine: Engine
+    asyncpg_db_engine: AsyncEngine
     rabbitmq_client: RabbitMQClient
     rabbitmq_rpc_client: RabbitMQRPCClient
     settings: ComputationalBackendSettings

services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/_scheduler_dask.py

@@ -282,7 +282,7 @@ async def _process_task_result(
                 if isinstance(result, TaskOutputData):
                     # success!
                     await parse_output_data(
-                        self.db_engine,
+                        self.asyncpg_db_engine,
                         task.job_id,
                         result,
                     )
@@ -307,7 +307,7 @@ async def _process_task_result(
                             simcore_platform_status = SimcorePlatformStatus.BAD
                     # we need to remove any invalid files in the storage
                     await clean_task_output_and_log_files_if_invalid(
-                        self.db_engine, user_id, project_id, node_id
+                        self.asyncpg_db_engine, user_id, project_id, node_id
                     )
             except TaskSchedulingError as err:
                 task_final_state = RunningState.FAILED

services/director-v2/src/simcore_service_director_v2/modules/comp_scheduler/_scheduler_factory.py

@@ -6,7 +6,7 @@
 
 from ...core.settings import AppSettings
 from ..dask_clients_pool import DaskClientsPool
-from ..db import get_db_engine
+from ..db import get_asyncpg_engine, get_db_engine
 from ..rabbitmq import get_rabbitmq_client, get_rabbitmq_rpc_client
 from ..redis import get_redis_client_manager
 from ._scheduler_base import BaseCompScheduler
@@ -27,5 +27,6 @@ def create_scheduler(app: FastAPI) -> BaseCompScheduler:
             rabbitmq_rpc_client=get_rabbitmq_rpc_client(app),
             redis_client=get_redis_client_manager(app).client(RedisDatabase.LOCKS),
             db_engine=get_db_engine(app),
+            asyncpg_db_engine=get_asyncpg_engine(app),
             service_runtime_heartbeat_interval=app_settings.SERVICE_TRACKING_HEARTBEAT,
         )