async def set_db_errors(error_queue):
async def helper(cur, item: js.Item):
job_id = item.args["job_id"]
protein_id = item.args["protein_id"]
mutation = item.args.get("mutations", "%")
if "_" in mutation:
mutation = mutation.split("_")[-1]
await cur.execute(SET_MUTATION_ERROR_SQL,
(job_id, protein_id, mutation))
async with js.WDBConnection() as conn:
async with conn.cursor() as cur:
try:
if isinstance(error_queue, Queue):
while not error_queue.empty():
item = await error_queue.get()
await helper(cur, item)
logger.debug("set_db_errors for item %s", item)
else:
for item in error_queue:
await helper(cur, item)
logger.debug("set_db_errors for item %s", item)
await conn.commit()
except Exception as e:
logger.error(
"The following error occured while trying to send errors to the database: %s",
e,
)
async def set_job_status(job_id: str) -> int:
"""Mark job with the specified `job_id` as "done"."""
async with js.WDBConnection() as conn:
async with conn.cursor() as cur:
num_rows_affected = await cur.execute(FINALIZE_SUBMISSION_SQL,
(job_id, ))
await conn.commit()
return num_rows_affected
async def get_mutation_info(item: js.Item) -> List[MutationInfo]:
args = item.args
kwargs = {
"protein_id": args["protein_id"],
"mutation": _format_mutation(args["mutations"])
}
if args["job_type"] == "local":
core_mutation_sql = get_core_mutation_local_sql
interface_mutation_sql = get_interface_mutation_local_sql
else:
assert args["job_type"] == "database"
core_mutation_sql = get_core_mutation_database_sql
interface_mutation_sql = get_interface_mutation_database_sql
async with js.WDBConnection() as conn:
async with conn.cursor() as cur:
await cur.execute(core_mutation_sql, kwargs)
core_mutation_values: List[Tuple[int, str, str,
str]] = await cur.fetchall()
await cur.execute(interface_mutation_sql, kwargs)
interface_mutation_values: List[Tuple[int, str, str,
str]] = await cur.fetchall()
mutation_info_list = [
MutationInfo(
domain_or_interface_id=domain_id,
structure_file=structure_file,
chain_id=chain_id,
mutation=mutation,
protein_id=kwargs["protein_id"],
coi=COI.CORE,
) for (domain_id, structure_file, chain_id,
mutation) in core_mutation_values
] + [
MutationInfo(
domain_or_interface_id=interface_id,
structure_file=structure_file,
chain_id=chain_id,
mutation=mutation,
protein_id=kwargs["protein_id"],
coi=COI.INTERFACE,
) for (interface_id, structure_file, chain_id,
mutation) in interface_mutation_values
]
logger.info(
"Obtained the following mutation data for item %s (%s): %s",
item,
kwargs,
mutation_info_list,
)
return mutation_info_list
async def finalize_mutation(item: js.Item):
"""Mark mutations associated with a given `item` as "done" or "error"."""
args = item.args
async with js.WDBConnection() as conn:
async with conn.cursor() as cur:
for mutation in args["mutations"].split(","):
# Local pipelines may have underscore in mutation
if "_" in mutation:
mutation = mutation.split("_")[-1]
await cur.execute(FINALIZE_MUTATION_SQL, {
"protein_id": args["protein_id"],
"mutation": mutation
})
await conn.commit()
system_command = f'bash -c "rm -f "{config.DATA_DIR}/locks/*/*.lock""'
await asyncio.create_subprocess_exec(*shlex.split(system_command))
async def update_mutation_scores(job_type: str,
mutation_scores: List[MutationInfo]) -> None:
if job_type not in ["local", "database"]:
raise ValueError
required_attributes = [
"domain_or_interface_id",
"protein_id",
"mutation",
"protbert_score",
"proteinsolver_score",
"el2_score",
"el2_version",
]
if job_type == "local":
core_mutation_sql = update_core_mutation_local_sql
interface_mutation_sql = update_interface_mutation_local_sql
else:
assert job_type == "database"
core_mutation_sql = update_core_mutation_database_sql
interface_mutation_sql = update_interface_mutation_database_sql
async with js.WDBConnection() as conn:
async with conn.cursor() as cur:
for mutation_score in mutation_scores:
if mutation_score.coi == COI.CORE:
update_mutation_sql = core_mutation_sql
else:
update_mutation_sql = interface_mutation_sql
attributes = {
key: getattr(mutation_score, key)
for key in required_attributes
}
await cur.execute(update_mutation_sql, attributes)
if cur.rowcount != 1:
logger.error(
"Unexpected number of rows modified when updating database: %s "
"(job_type = %s, coi = %s, attributes = %s).",
cur.rowcount,
job_type,
mutation_score.coi.value,
attributes,
)
await conn.commit()