async def test_listen_comp_tasks_task( mock_project_subsystem: Dict, comp_task_listening_task: asyncio.Task, client, upd_value: Dict[str, Any], exp_calls: List[str], task_class: NodeClass, ): db_engine: aiopg.sa.Engine = client.app[APP_DB_ENGINE_KEY] async with db_engine.acquire() as conn: # let's put some stuff in there now result = await conn.execute( comp_tasks.insert() .values(outputs=json.dumps({}), node_class=task_class) .returning(literal_column("*")) ) row: RowProxy = await result.fetchone() task = dict(row) # let's update some values await conn.execute( comp_tasks.update() .values(**upd_value) .where(comp_tasks.c.task_id == task["task_id"]) ) for key, mock_fct in mock_project_subsystem.items(): if key in exp_calls: await _wait_for_call(mock_fct) else: mock_fct.assert_not_called()
async def write_ports_configuration( self, json_configuration: str, project_id: str, node_uuid: str ): message = ( f"Writing port configuration to database for " f"project={project_id} node={node_uuid}: {json_configuration}" ) log.debug(message) node_configuration = json.loads(json_configuration) async with DBContextManager(self._db_engine) as engine: async with engine.acquire() as connection: # update the necessary parts await connection.execute( # FIXME: E1120:No value for argument 'dml' in method call # pylint: disable=E1120 comp_tasks.update() .where( and_( comp_tasks.c.node_id == node_uuid, comp_tasks.c.project_id == project_id, ) ) .values( schema=node_configuration["schema"], inputs=node_configuration["inputs"], outputs=node_configuration["outputs"], run_hash=node_configuration.get("run_hash"), ) )
async def test_listen_comp_tasks_task( mock_project_subsystem: Dict, comp_task_listening_task: asyncio.Task, client, update_values: Dict[str, Any], expected_calls: List[str], task_class: NodeClass, ): db_engine: aiopg.sa.Engine = client.app[APP_DB_ENGINE_KEY] async with db_engine.acquire() as conn: # let's put some stuff in there now result = await conn.execute(comp_tasks.insert().values( outputs=json.dumps({}), node_class=task_class).returning(literal_column("*"))) row: RowProxy = await result.fetchone() task = dict(row) # let's update some values await conn.execute(comp_tasks.update().values(**update_values).where( comp_tasks.c.task_id == task["task_id"])) # tests whether listener gets hooked calls executed for call_name, mocked_call in mock_project_subsystem.items(): if call_name in expected_calls: async for attempt in _async_retry_if_fails(): with attempt: mocked_call.assert_awaited() else: mocked_call.assert_not_called()
async def _set_task_state( aiopg_engine: Iterator[aiopg.sa.engine.Engine], node_id: str, state: StateType # type: ignore ): async with aiopg_engine.acquire() as conn: # type: ignore await conn.execute(comp_tasks.update().where( comp_tasks.c.node_id == node_id).values(state=state))
async def set_comp_task_outputs( aiopg_engine: aiopg.sa.engine.Engine, node_id: NodeID, outputs_schema: Dict[str, Any], outputs: Dict[str, Any] # type: ignore ): async with aiopg_engine.acquire() as conn: # type: ignore await conn.execute( # pylint: disable=no-value-for-parameter comp_tasks.update().where(comp_tasks.c.node_id == f"{node_id}" ).values(outputs=outputs, schema={ "outputs": outputs_schema, "inputs": {} }))
async def write_ports_configuration(self, json_configuration: str, node_uuid: str): log.debug("Writing ports configuration to database") node_configuration = json.loads(json_configuration) async with DBContextManager(self._db_engine) as engine: async with engine.acquire() as connection: # update the necessary parts await connection.execute( # FIXME: E1120:No value for argument 'dml' in method call # pylint: disable=E1120 comp_tasks.update().where( and_( comp_tasks.c.node_id == node_uuid, comp_tasks.c.project_id == config.PROJECT_ID, )).values( schema=node_configuration["schema"], inputs=node_configuration["inputs"], outputs=node_configuration["outputs"], ))
async def test_listen_comp_tasks_task( mock_project_subsystem: Dict, comp_task_listening_task: None, client, update_values: Dict[str, Any], expected_calls: List[str], task_class: NodeClass, ): db_engine: aiopg.sa.Engine = client.app[APP_DB_ENGINE_KEY] async with db_engine.acquire() as conn: # let's put some stuff in there now result = await conn.execute( comp_tasks.insert() .values(outputs=json.dumps({}), node_class=task_class) .returning(literal_column("*")) ) row: RowProxy = await result.fetchone() task = dict(row) # let's update some values await conn.execute( comp_tasks.update() .values(**update_values) .where(comp_tasks.c.task_id == task["task_id"]) ) # tests whether listener gets hooked calls executed for call_name, mocked_call in mock_project_subsystem.items(): if call_name in expected_calls: async for attempt in AsyncRetrying( wait=wait_fixed(1), stop=stop_after_delay(10), retry=retry_if_exception_type(AssertionError), before_sleep=before_sleep_log(logger, logging.INFO), reraise=True, ): with attempt: mocked_call.assert_awaited() else: mocked_call.assert_not_called()
def update_configuration( postgres_engine: sa.engine.Engine, project_id: str, node_uuid: str, new_configuration: Dict, ) -> None: log.debug( "Update configuration of pipeline %s, node %s", project_id, node_uuid, ) with postgres_engine.connect() as conn: conn.execute(comp_tasks.update() # pylint: disable=no-value-for-parameter .where((comp_tasks.c.project_id == project_id) & (comp_tasks.c.node_id == node_uuid)).values( schema=new_configuration["schema"], inputs=new_configuration["inputs"], outputs=new_configuration["outputs"], )) log.debug("Updated configuration")
async def _update_comp_task_with(conn: SAConnection, task: Dict, **kwargs): await conn.execute(comp_tasks.update().values(**kwargs).where( comp_tasks.c.task_id == task["task_id"]))