Compare commits

..

No commits in common. "34bf02c75a53593a0b271e7f8217dff46b74e46d" and "af7edf0942cba82aaa342776014b1f03e8ae5d5b" have entirely different histories.

View File

@ -9,7 +9,7 @@ import shutil
import tempfile import tempfile
from typing import Any, AsyncGenerator, AsyncIterable, Coroutine, Iterable, Mapping, Sequence, Set, Union from typing import Any, AsyncGenerator, AsyncIterable, Coroutine, Iterable, Mapping, Sequence, Set, Union
from aiohttp import ClientOSError, ClientSession, ClientTimeout, ServerDisconnectedError from aiohttp import ClientSession, ClientTimeout
from autobigs.engine.reading import read_fasta from autobigs.engine.reading import read_fasta
from autobigs.engine.structures.alignment import PairwiseAlignment from autobigs.engine.structures.alignment import PairwiseAlignment
@ -43,12 +43,11 @@ class BIGSdbMLSTProfiler(AbstractAsyncContextManager):
class RemoteBIGSdbMLSTProfiler(BIGSdbMLSTProfiler): class RemoteBIGSdbMLSTProfiler(BIGSdbMLSTProfiler):
def __init__(self, database_api: str, database_name: str, scheme_id: int, retry_requests: int = 5): def __init__(self, database_api: str, database_name: str, scheme_id: int):
self._retry_limit = retry_requests
self._database_name = database_name self._database_name = database_name
self._scheme_id = scheme_id self._scheme_id = scheme_id
self._base_url = f"{database_api}/db/{self._database_name}/schemes/{self._scheme_id}/" self._base_url = f"{database_api}/db/{self._database_name}/schemes/{self._scheme_id}/"
self._http_client = ClientSession(self._base_url, timeout=ClientTimeout(300)) self._http_client = ClientSession(self._base_url, timeout=ClientTimeout(60))
async def __aenter__(self): async def __aenter__(self):
return self return self
@ -58,19 +57,11 @@ class RemoteBIGSdbMLSTProfiler(BIGSdbMLSTProfiler):
uri_path = "sequence" uri_path = "sequence"
if isinstance(query_sequence_strings, str) or isinstance(query_sequence_strings, NamedString): if isinstance(query_sequence_strings, str) or isinstance(query_sequence_strings, NamedString):
query_sequence_strings = [query_sequence_strings] query_sequence_strings = [query_sequence_strings]
for sequence_string in query_sequence_strings: for sequence_string in query_sequence_strings:
attempts = 0 async with self._http_client.post(uri_path, json={
success = False
last_error = None
while not success and attempts < self._retry_limit:
attempts += 1
request = self._http_client.post(uri_path, json={
"sequence": sequence_string if isinstance(sequence_string, str) else sequence_string.sequence, "sequence": sequence_string if isinstance(sequence_string, str) else sequence_string.sequence,
"partial_matches": True "partial_matches": True
}) }) as response:
try:
async with request as response:
sequence_response: dict = await response.json() sequence_response: dict = await response.json()
if "exact_matches" in sequence_response: if "exact_matches" in sequence_response:
@ -100,17 +91,6 @@ class RemoteBIGSdbMLSTProfiler(BIGSdbMLSTProfiler):
yield result_allele if isinstance(sequence_string, str) else (sequence_string.name, result_allele) yield result_allele if isinstance(sequence_string, str) else (sequence_string.name, result_allele)
else: else:
raise NoBIGSdbMatchesException(self._database_name, self._scheme_id, sequence_string.name if isinstance(sequence_string, NamedString) else None) raise NoBIGSdbMatchesException(self._database_name, self._scheme_id, sequence_string.name if isinstance(sequence_string, NamedString) else None)
except (ConnectionError, ServerDisconnectedError, ClientOSError) as e: # Errors we will retry
last_error = e
success = False
await asyncio.sleep(5) # In case the connection issue is due to rate issues
else:
success = True
if not success and last_error is not None:
try:
raise last_error
except (ConnectionError, ServerDisconnectedError, ClientOSError) as e: # Non-fatal errors
yield Allele("error", "error", None)
async def determine_mlst_st(self, alleles: Union[AsyncIterable[Union[Allele, tuple[str, Allele]]], Iterable[Union[Allele, tuple[str, Allele]]]]) -> Union[MLSTProfile, NamedMLSTProfile]: async def determine_mlst_st(self, alleles: Union[AsyncIterable[Union[Allele, tuple[str, Allele]]], Iterable[Union[Allele, tuple[str, Allele]]]]) -> Union[MLSTProfile, NamedMLSTProfile]:
uri_path = "designations" uri_path = "designations"
@ -133,13 +113,6 @@ class RemoteBIGSdbMLSTProfiler(BIGSdbMLSTProfiler):
request_json = { request_json = {
"designations": allele_request_dict "designations": allele_request_dict
} }
attempts = 0
success = False
last_error = None
while attempts < self._retry_limit and not success:
attempts += 1
try:
async with self._http_client.post(uri_path, json=request_json) as response: async with self._http_client.post(uri_path, json=request_json) as response:
response_json: dict = await response.json() response_json: dict = await response.json()
allele_set: Set[Allele] = set() allele_set: Set[Allele] = set()
@ -156,19 +129,6 @@ class RemoteBIGSdbMLSTProfiler(BIGSdbMLSTProfiler):
if len(names_list) > 0: if len(names_list) > 0:
result_mlst_profile = NamedMLSTProfile(str(tuple(names_list)) if len(set(names_list)) > 1 else names_list[0], result_mlst_profile) result_mlst_profile = NamedMLSTProfile(str(tuple(names_list)) if len(set(names_list)) > 1 else names_list[0], result_mlst_profile)
return result_mlst_profile return result_mlst_profile
except (ConnectionError, ServerDisconnectedError, ClientOSError) as e:
last_error = e
success = False
await asyncio.sleep(5)
else:
success = True
try:
if last_error is not None:
raise last_error
except (ConnectionError, ServerDisconnectedError, ClientOSError) as e:
result_mlst_profile = NamedMLSTProfile((str(tuple(names_list)) if len(set(names_list)) > 1 else names_list[0]) + ":Error", None)
raise ValueError("Last error was not recorded.")
async def profile_string(self, query_sequence_strings: Iterable[Union[NamedString, str]]) -> Union[NamedMLSTProfile, MLSTProfile]: async def profile_string(self, query_sequence_strings: Iterable[Union[NamedString, str]]) -> Union[NamedMLSTProfile, MLSTProfile]:
alleles = self.determine_mlst_allele_variants(query_sequence_strings) alleles = self.determine_mlst_allele_variants(query_sequence_strings)