async def _acquire_connection(engine: Engine = Depends( _get_db_engine)) -> SAConnection: logger.debug( "Acquiring pg connection from pool: pool size=%d, acquired=%d, free=%d, reserved=[%d, %d]", engine.size, engine.size - engine.freesize, engine.freesize, engine.minsize, engine.maxsize, ) if engine.freesize <= 1: logger.warning( "Last or no pg connection in pool: pool size=%d, acquired=%d, free=%d, reserved=[%d, %d]", engine.size, engine.size - engine.freesize, engine.freesize, engine.minsize, engine.maxsize, ) async with engine.acquire() as conn: yield conn logger.debug( "Released pg connection: pool size=%d, acquired=%d, free=%d, reserved=[%d, %d]", engine.size, engine.size - engine.freesize, engine.freesize, engine.minsize, engine.maxsize, )
async def _get_repo( engine: Engine = Depends(_get_db_engine), ) -> AsyncGenerator[BaseRepository, None]: logger.debug( "Acquiring pg connection from pool: current=%d, free=%d, reserved=[%d, %d]", engine.size, engine.freesize, engine.minsize, engine.maxsize, ) if engine.freesize <= 1: logger.warning( "Last or no pg connection in pool: current=%d, free=%d, reserved=[%d, %d]", engine.size, engine.freesize, engine.minsize, engine.maxsize, ) async with engine.acquire() as conn: yield repo_type(conn) logger.debug( "Released pg connection: current=%d, free=%d, reserved=[%d, %d]", engine.size, engine.freesize, engine.minsize, engine.maxsize, )
async def save(self, db: Engine): """ Create or update an order. If current order instance isn't created, then inserts new row, and sets it as self.uuid attribute value. Otherwise just updates row in database with order instance data. """ table = db.tables['orders'] data = { 'id_on_exchange': self.id_on_exchange, 'status': self.status, 'pair': self.pair, 'side': self.side, 'price': self.price, 'base': self.base.amount, 'quote': self.quote.amount, 'exchange': self.exchange.name, 'strategy': self.strategy, 'created_at': self.created_at, 'expired_at': self.expired_at, 'executed_at': self.executed_at, } async with db.acquire() as conn: if not self.uuid: self.uuid = await conn.scalar(table.insert().values(**data)) else: await conn.execute(table.update().where( table.c.uuid == self.uuid).values(**data))
async def _set_adjacency_in_pipeline_db(db_engine: Engine, project_id: str, dag_adjacency_list: Dict): # pylint: disable=no-value-for-parameter async with db_engine.acquire() as conn: # READ # get pipeline query = sa.select([comp_pipeline ]).where(comp_pipeline.c.project_id == project_id) result = await conn.execute(query) pipeline = await result.first() # WRITE if pipeline is None: # create pipeline log.debug("No pipeline for project %s, creating one", project_id) query = comp_pipeline.insert().values( project_id=project_id, dag_adjacency_list=dag_adjacency_list, state=UNKNOWN, ) else: # update pipeline log.debug("Found pipeline for project %s, updating it", project_id) query = (comp_pipeline.update().where( comp_pipeline.c.project_id == project_id).values( dag_adjacency_list=dag_adjacency_list, state=UNKNOWN)) await conn.execute(query)
async def inspect( # pylint: disable=too-many-arguments db_engine: Engine, rabbit_mq: RabbitMQ, job_request_id: str, user_id: str, project_id: str, node_id: Optional[str], ) -> Optional[List[str]]: log.debug( "ENTERING inspect with user %s pipeline:node %s: %s", user_id, project_id, node_id, ) task: Optional[RowProxy] = None graph: Optional[nx.DiGraph] = None async with db_engine.acquire() as connection: pipeline: RowProxy = await _get_pipeline_from_db( connection, project_id) graph = execution_graph(pipeline) if not node_id: log.debug("NODE id was zero, this was the entry node id") return find_entry_point(graph) task = await _try_get_task_from_db(connection, graph, job_request_id, project_id, node_id) if not task: log.debug("no task at hand, let's rest...") return # config nodeports node_ports.node_config.USER_ID = user_id node_ports.node_config.NODE_UUID = task.node_id node_ports.node_config.PROJECT_ID = task.project_id # now proceed actually running the task (we do that after the db session has been closed) # try to run the task, return empyt list of next nodes if anything goes wrong run_result = FAILED next_task_nodes = [] try: executor = Executor( db_engine=db_engine, rabbit_mq=rabbit_mq, task=task, user_id=user_id, ) await executor.run() next_task_nodes = list(graph.successors(node_id)) run_result = SUCCESS except asyncio.CancelledError: log.warning("Task has been cancelled") raise finally: await _set_task_status(db_engine, project_id, node_id, run_result) return next_task_nodes
async def raise_if_not_responsive(engine: Engine): async with engine.acquire() as conn: # pylint: disable=protected-access # NOTE: Hacks aiopg.sa.SAConnection interface # to override connection's cursor timeout cursor = await conn._open_cursor() await cursor.execute("SELECT 1 as is_alive", timeout=1)
async def check_credentials(engine: Engine, email: str, password: str) -> bool: async with engine.acquire() as conn: query = users.select().where( sa.and_(users.c.email == email, users.c.status != UserStatus.BANNED)) ret = await conn.execute(query) user = await ret.fetchone() if user is not None: return check_password(password, user["password_hash"]) return False
async def get_all_categories(engine: Engine, ): table: Table = Category.__table__ async with engine.acquire() as conn: async with conn.begin(): result = await conn.execute(table.select()) output = resultproxy_to_dict(result) return output
async def _set_task_status(db_engine: Engine, project_id: str, node_id: str, run_result): async with db_engine.acquire() as connection: 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_id, comp_tasks.c.project_id == project_id, )).values(state=run_result, end=datetime.utcnow()))
async def delete(self, db: Engine): """ Delete current order from database. Throws exception if current order instance doesn't have database id. """ if not self.uuid: raise ValueError(f'Invalid order UUID: {self.uuid}') table = db.tables['orders'] async with db.acquire() as conn: await conn.execute(table.delete().where(table.c.uuid == self.uuid)) self.uuid = None
async def get_entity_by_id(engine: Engine, entity_id: int): table: Table = Entity.__table__ async with engine.acquire() as conn: async with conn.begin(): result = await conn.execute( table.select().where(table.c.id == entity_id)) output = resultproxy_to_dict(result) if len(output) == 0: return None return output[0]
async def _set_tasks_in_tasks_db(db_engine: Engine, project_id: str, tasks: Dict): async with db_engine.acquire() as conn: query = sa.select([ComputationalTask]).\ where(ComputationalTask.__table__.c.project_id == project_id) result = await conn.execute(query) tasks_db = await result.fetchall() # delete tasks that were deleted from the db for task_db in tasks_db: if not task_db.node_id in tasks: query = ComputationalTask.__table__.delete().\ where(and_(ComputationalTask.__table__.c.project_id == project_id, ComputationalTask.__table__.c.node_id == task_db.node_id)) await conn.execute(query) internal_id = 1 for node_id in tasks: task = tasks[node_id] query = sa.select([ComputationalTask]).\ where(and_(ComputationalTask.__table__.c.project_id == project_id, ComputationalTask.__table__.c.node_id == node_id)) result = await conn.execute(query) comp_task = await result.fetchone() if comp_task is None: # add a new one query = ComputationalTask.__table__.insert().\ values(project_id=project_id, node_id=node_id, internal_id=internal_id, image = task["image"], schema = task["schema"], inputs = task["inputs"], outputs = task["outputs"], submit = datetime.datetime.utcnow()) internal_id = internal_id + 1 else: query = ComputationalTask.__table__.update().\ where(and_(ComputationalTask.__table__.c.project_id==project_id, ComputationalTask.__table__.c.node_id==node_id)).\ values(job_id = None, state = 0, image = task["image"], schema = task["schema"], inputs = task["inputs"], outputs = task["outputs"] if "file-picker" in task["image"]["name"] else comp_task.outputs, submit = datetime.datetime.utcnow()) await conn.execute(query)
async def get_entity_by_category_name(engine: Engine, category_title: str): entity_table: Table = Entity.__table__ category_table: Table = Category.__table__ join = sa.join(entity_table, category_table, entity_table.c.category_id == category_table.c.id) query = (sa.select([entity_table], use_labels=False).select_from(join).where( category_table.c.title == category_title)) async with engine.acquire() as conn: async with conn.begin(): result = await conn.execute(query) output = resultproxy_to_dict(result) return output
async def _query_db(uid: str, engine: Engine) -> List[RowProxy]: async with engine.acquire() as conn: query = (sa.select( [ users.c.email, users.c.role, users.c.name, users.c.primary_gid, groups.c.gid, groups.c.name, groups.c.description, groups.c.type, ], use_labels=True, ).select_from( users.join( user_to_groups.join(groups, user_to_groups.c.gid == groups.c.gid), users.c.id == user_to_groups.c.uid, )).where(users.c.id == uid).order_by(sa.asc(groups.c.name))) result = await conn.execute(query) return await result.fetchall()
async def _set_adjacency_in_pipeline_db(db_engine: Engine, project_id: str, dag_adjacency_list: Dict): async with db_engine.acquire() as conn: query = sa.select([ComputationalPipeline]).\ where(ComputationalPipeline.__table__.c.project_id==project_id) result = await conn.execute(query) pipeline = await result.first() if pipeline is None: # let's create one then query = ComputationalPipeline.__table__.insert().\ values(project_id=project_id, dag_adjacency_list=dag_adjacency_list, state=0) log.debug("Pipeline object created") else: # let's modify it log.debug("Pipeline object found") query = ComputationalPipeline.__table__.update().\ where(ComputationalPipeline.__table__.c.project_id == project_id).\ values(state=0, dag_adjacency_list=dag_adjacency_list) await conn.execute(query)
async def _set_tasks_in_tasks_db(db_engine: Engine, project_id: str, tasks: Dict[str, Dict], replace_pipeline=True): # pylint: disable=no-value-for-parameter async def _task_already_exists(conn: SAConnection, project_id: str, node_id: str) -> bool: task_count: int = await conn.scalar( sa.select([sa.func.count()]).where( and_( comp_tasks.c.project_id == project_id, comp_tasks.c.node_id == node_id, ))) assert task_count in ( # nosec 0, 1, ), f"Uniqueness violated: task_count={task_count}" # nosec return task_count != 0 async def _update_task(conn: SAConnection, task: Dict, project_id: str, node_id: str) -> None: # update task's inputs/outputs io_update = {} task_inputs: str = await conn.scalar( sa.select([comp_tasks.c.inputs]).where( and_( comp_tasks.c.project_id == project_id, comp_tasks.c.node_id == node_id, ))) # updates inputs if task_inputs != task["inputs"]: io_update["inputs"] = task["inputs"] # update outputs # NOTE: update ONLY outputs of front-end nodes. The rest are # updated by backend services (e.g. workers, interactive services) if task["outputs"] and task["node_class"] == NodeClass.FRONTEND: io_update["outputs"] = task["outputs"] if io_update: query = (comp_tasks.update().where( and_( comp_tasks.c.project_id == project_id, comp_tasks.c.node_id == node_id, )).values(**io_update)) await conn.execute(query) # MAIN ----------- async with db_engine.acquire() as conn: if replace_pipeline: # get project tasks already stored query = sa.select([comp_tasks ]).where(comp_tasks.c.project_id == project_id) result = await conn.execute(query) tasks_rows = await result.fetchall() # prune database from invalid tasks for task_row in tasks_rows: if not task_row.node_id in tasks: query = comp_tasks.delete().where( and_( comp_tasks.c.project_id == project_id, comp_tasks.c.node_id == task_row.node_id, )) await conn.execute(query) internal_id = 1 for node_id, task in tasks.items(): is_new_task: bool = not await _task_already_exists( conn, project_id, node_id) try: if is_new_task: # create task query = comp_tasks.insert().values( project_id=project_id, node_id=node_id, node_class=task["node_class"], internal_id=internal_id, image=task["image"], schema=task["schema"], inputs=task["inputs"], outputs=task["outputs"] if task["outputs"] else {}, submit=datetime.datetime.utcnow(), ) await conn.execute(query) internal_id = internal_id + 1 except psycopg2.errors.UniqueViolation: # pylint: disable=no-member # avoids race condition is_new_task = False if not is_new_task: if replace_pipeline: # replace task query = (comp_tasks.update().where( and_( comp_tasks.c.project_id == project_id, comp_tasks.c.node_id == node_id, )).values( job_id=None, state=UNKNOWN, node_class=task["node_class"], image=task["image"], schema=task["schema"], inputs=task["inputs"], outputs=task["outputs"] if task["outputs"] else {}, submit=datetime.datetime.utcnow(), )) await conn.execute(query) else: await _update_task(conn, task, project_id, node_id)
async def run_computational_task( # pylint: disable=too-many-arguments db_engine: Engine, rabbit_mq: RabbitMQ, job_request_id: str, user_id: str, project_id: str, node_id: str, retry: int, max_retries: int, sidecar_mode: BootMode, ) -> None: async with db_engine.acquire() as connection: task: Optional[RowProxy] = None graph: Optional[nx.DiGraph] = None try: log.debug( "ENTERING inspect with user %s pipeline:node %s: %s", user_id, project_id, node_id, ) pipeline: RowProxy = await _get_pipeline_from_db( connection, project_id) graph = execution_graph(pipeline) log.debug("NODE id is %s, getting the task from DB...", node_id) task = await _try_get_task_from_db(connection, job_request_id, project_id, node_id) if not task: log.warning( "Worker received task for user %s, project %s, node %s, but the task is already taken! this should not happen. going back to sleep...", user_id, project_id, node_id, ) return except asyncio.CancelledError: log.warning("Task has been cancelled") await _set_task_state(connection, project_id, node_id, StateType.ABORTED) raise run_result = StateType.FAILED try: await rabbit_mq.post_log_message( user_id, project_id, node_id, "[sidecar]Task found: starting...", ) # now proceed actually running the task (we do that after the db session has been closed) # try to run the task, return empyt list of next nodes if anything goes wrong executor = Executor( db_engine=db_engine, rabbit_mq=rabbit_mq, task=task, user_id=user_id, sidecar_mode=sidecar_mode, ) await executor.run() run_result = StateType.SUCCESS except asyncio.CancelledError: log.warning("Task has been cancelled") run_result = StateType.ABORTED raise finally: await rabbit_mq.post_log_message( user_id, project_id, node_id, f"[sidecar]Task completed with result: {run_result.name} [Trial {retry+1}/{max_retries}]", ) if (retry + 1) < max_retries and run_result == StateType.FAILED: # try again! run_result = StateType.PENDING await _set_task_state(connection, project_id, node_id, run_result) if run_result == StateType.FAILED: # set the successive tasks as ABORTED await _set_tasks_state( connection, project_id, nx.bfs_tree(graph, node_id), StateType.ABORTED, offset=1, )
async def db_connection(pg_engine: Engine) -> SAConnection: """fixture for connection""" async with pg_engine.acquire() as conn: yield conn
async def pg_connection(pool: Engine = Depends(pg_pool)) -> SAConnection: async with pool.acquire() as connection: yield connection
async def delete_tables(engine: Engine, tables: list) -> None: """Deletes tables included in list.""" async with engine.acquire() as conn: for table in reversed(tables): delete_table_stmt = DropTable(table.__table__) await conn.execute(delete_table_stmt)
async def create_tables(engine: Engine, tables: list) -> None: """Creates all tables in tables list.""" async with engine.acquire() as conn: for table in tables: create_table_stmt = CreateTable(table.__table__) await conn.execute(create_table_stmt)
async def raise_if_not_responsive(engine: Engine): async with engine.acquire() as conn: await conn.execute("SELECT 1 as is_alive")
async def get_cnx(engine: Engine = Depends(get_engine)): # TODO: problem here is retries?? async with engine.acquire() as conn: yield conn
async def select_shorts(db: Engine): async with db.acquire() as conn: conn: Cursor result = await conn.execute(shorts_table.select()) data = await result.fetchall() return data
async def listen(app: web.Application, db_engine: Engine): listen_query = f"LISTEN {DB_CHANNEL_NAME};" _LISTENING_TASK_BASE_SLEEPING_TIME_S = 1 async with db_engine.acquire() as conn: await conn.execute(listen_query) while True: # NOTE: instead of using await get() we check first if the connection was closed # since aiopg does not reset the await in such a case (if DB was restarted or so) # see aiopg issue: https://github.com/aio-libs/aiopg/pull/559#issuecomment-826813082 if conn.closed: raise ConnectionError("connection with database is closed!") if conn.connection.notifies.empty(): await asyncio.sleep(_LISTENING_TASK_BASE_SLEEPING_TIME_S) continue notification = conn.connection.notifies.get_nowait() log.debug("received update from database: %s", pformat(notification.payload)) # get the data and the info on what changed payload: Dict = json.loads(notification.payload) # FIXME: all this should move to rabbitMQ instead of this task_data = payload.get("data", {}) task_changes = payload.get("changes", []) if not task_data: log.error("task data invalid: %s", pformat(payload)) continue if not task_changes: log.error("no changes but still triggered: %s", pformat(payload)) continue project_uuid = task_data.get("project_id", "undefined") node_uuid = task_data.get("node_id", "undefined") # FIXME: we do not know who triggered these changes. we assume the user had the rights to do so # therefore we'll use the prj_owner user id. This should be fixed when the new sidecar comes in # and comp_tasks/comp_pipeline get deprecated. try: # find the user(s) linked to that project the_project_owner = await _get_project_owner( conn, project_uuid) if any(f in task_changes for f in ["outputs", "run_hash"]): new_outputs = task_data.get("outputs", {}) new_run_hash = task_data.get("run_hash", None) await _update_project_outputs( app, the_project_owner, project_uuid, node_uuid, new_outputs, new_run_hash, ) if "state" in task_changes: new_state = convert_state_from_db(task_data["state"]).value await _update_project_state(app, the_project_owner, project_uuid, node_uuid, new_state) except projects_exceptions.ProjectNotFoundError as exc: log.warning( "Project %s was not found and cannot be updated. Maybe was it deleted?", exc.project_uuid, ) continue except projects_exceptions.ProjectOwnerNotFoundError as exc: log.warning( "Project owner of project %s could not be found, is the project valid?", exc.project_uuid, ) continue except projects_exceptions.NodeNotFoundError as exc: log.warning( "Node %s of project %s not found and cannot be updated. Maybe was it deleted?", exc.node_uuid, exc.project_uuid, ) continue
async def _get_repo( engine: Engine = Depends(_get_db_engine), ) -> AsyncGenerator[BaseRepository, None]: async with engine.acquire() as conn: yield repo_type(conn)