src/autobigs/engine/analysis/bigsdb.py

@@ -9,7 +9,7 @@ import shutil
 import tempfile
 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.structures.alignment import PairwiseAlignment
@@ -43,12 +43,11 @@ class BIGSdbMLSTProfiler(AbstractAsyncContextManager):
 
 class RemoteBIGSdbMLSTProfiler(BIGSdbMLSTProfiler):
 
-    def __init__(self, database_api: str, database_name: str, scheme_id: int, retry_requests: int = 5):
-        self._retry_limit = retry_requests
+    def __init__(self, database_api: str, database_name: str, scheme_id: int):
         self._database_name = database_name
         self._scheme_id = 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):
         return self
@@ -58,19 +57,11 @@ class RemoteBIGSdbMLSTProfiler(BIGSdbMLSTProfiler):
         uri_path = "sequence"
         if isinstance(query_sequence_strings, str) or isinstance(query_sequence_strings, NamedString):
             query_sequence_strings = [query_sequence_strings]
-        
         for sequence_string in query_sequence_strings:
-            attempts = 0
-            success = False
-            last_error = None
-            while not success and attempts < self._retry_limit:
-                attempts += 1
-                request = self._http_client.post(uri_path, json={
+            async with self._http_client.post(uri_path, json={
                 "sequence": sequence_string if isinstance(sequence_string, str) else sequence_string.sequence,
                 "partial_matches": True
-                })
-                try:
-                    async with request as response:
+            }) as response:
                 sequence_response: dict = await response.json()
 
                 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)
                 else:
                     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]:
         uri_path = "designations"
@@ -133,13 +113,6 @@ class RemoteBIGSdbMLSTProfiler(BIGSdbMLSTProfiler):
         request_json = {
             "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:
             response_json: dict = await response.json()
             allele_set: Set[Allele] = set()
@@ -156,19 +129,6 @@ class RemoteBIGSdbMLSTProfiler(BIGSdbMLSTProfiler):
             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)
             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]:
         alleles = self.determine_mlst_allele_variants(query_sequence_strings)