1
- from datetime import datetime
2
- from typing import List , Optional , Union , Dict
3
1
import traceback
2
+ from datetime import datetime
3
+ from typing import Dict , List , Optional , Union
4
4
5
5
import aiodocker
6
- import aiopg
7
6
import networkx as nx
8
- from celery .utils .log import get_task_logger
7
+ from aiopg .sa import Engine , SAConnection
8
+ from aiopg .sa .result import RowProxy
9
9
from sqlalchemy import and_ , literal_column
10
10
11
+ from celery .utils .log import get_task_logger
11
12
from simcore_postgres_database .sidecar_models import ( # PENDING,
12
13
FAILED ,
13
14
RUNNING ,
20
21
from simcore_sdk .node_ports import log as node_port_log
21
22
22
23
from . import config , exceptions
24
+ from .db import DBContextManager
23
25
from .executor import Executor
24
26
from .rabbitmq import RabbitMQ
25
27
from .utils import execution_graph , find_entry_point , is_node_ready
26
- from .db import DBContextManager
27
28
28
29
log = get_task_logger (__name__ )
29
30
log .setLevel (config .SIDECAR_LOGLEVEL )
@@ -61,13 +62,12 @@ async def task_required_resources(node_id: str) -> Union[Dict[str, bool], None]:
61
62
62
63
63
64
async def _try_get_task_from_db (
64
- db_connection : aiopg . sa . SAConnection ,
65
+ db_connection : SAConnection ,
65
66
graph : nx .DiGraph ,
66
- job_request_id : int ,
67
+ job_request_id : str ,
67
68
project_id : str ,
68
69
node_id : str ,
69
- ) -> Optional [aiopg .sa .result .RowProxy ]:
70
- task : aiopg .sa .result .RowProxy = None
70
+ ) -> Optional [RowProxy ]:
71
71
# Use SELECT FOR UPDATE TO lock the row
72
72
result = await db_connection .execute (
73
73
query = comp_tasks .select (for_update = True ).where (
@@ -79,7 +79,7 @@ async def _try_get_task_from_db(
79
79
)
80
80
)
81
81
)
82
- task = await result .fetchone ()
82
+ task : RowProxy = await result .fetchone ()
83
83
84
84
if not task :
85
85
log .debug ("No task found" )
@@ -114,9 +114,8 @@ async def _try_get_task_from_db(
114
114
115
115
116
116
async def _get_pipeline_from_db (
117
- db_connection : aiopg .sa .SAConnection , project_id : str ,
118
- ) -> aiopg .sa .result .RowProxy :
119
- pipeline : aiopg .sa .result .RowProxy = None
117
+ db_connection : SAConnection , project_id : str ,
118
+ ) -> RowProxy :
120
119
# get the pipeline
121
120
result = await db_connection .execute (
122
121
comp_pipeline .select ().where (comp_pipeline .c .project_id == project_id )
@@ -126,7 +125,7 @@ async def _get_pipeline_from_db(
126
125
f"Pipeline { result .rowcount } found instead of only one for project_id { project_id } "
127
126
)
128
127
129
- pipeline = await result .first ()
128
+ pipeline : RowProxy = await result .first ()
130
129
if not pipeline :
131
130
raise exceptions .DatabaseError (f"Pipeline { project_id } not found" )
132
131
log .debug ("found pipeline %s" , pipeline )
@@ -135,12 +134,12 @@ async def _get_pipeline_from_db(
135
134
136
135
async def inspect (
137
136
# pylint: disable=too-many-arguments
138
- db_engine : aiopg . sa . Engine ,
137
+ db_engine : Engine ,
139
138
rabbit_mq : RabbitMQ ,
140
- job_request_id : int ,
139
+ job_request_id : str ,
141
140
user_id : str ,
142
141
project_id : str ,
143
- node_id : str ,
142
+ node_id : Optional [ str ] ,
144
143
) -> Optional [List [str ]]:
145
144
log .debug (
146
145
"ENTERING inspect with user %s pipeline:node %s: %s" ,
@@ -149,11 +148,10 @@ async def inspect(
149
148
node_id ,
150
149
)
151
150
152
- pipeline : aiopg .sa .result .RowProxy = None
153
- task : aiopg .sa .result .RowProxy = None
154
- graph : nx .DiGraph = None
151
+ task : Optional [RowProxy ] = None
152
+ graph : Optional [nx .DiGraph ] = None
155
153
async with db_engine .acquire () as connection :
156
- pipeline = await _get_pipeline_from_db (connection , project_id )
154
+ pipeline : RowProxy = await _get_pipeline_from_db (connection , project_id )
157
155
graph = execution_graph (pipeline )
158
156
if not node_id :
159
157
log .debug ("NODE id was zero, this was the entry node id" )
0 commit comments