Merge remote-tracking branch 'origin' into add-sentry-support

Author: jschaff

biothings_annotator/annotator/annotator.py

@@ -9,50 +9,52 @@
 
 import biothings_client
 
-from .exceptions import InvalidCurieError, TRAPIInputError
-from .settings import ANNOTATOR_CLIENTS, SERVICE_PROVIDER_API_HOST
-from .transformer import ResponseTransformer
-from .utils import batched, get_client, get_dotfield_value, group_by_subfield, parse_curie
+from biothings_annotator.annotator.exceptions import InvalidCurieError, TRAPIInputError
+from biothings_annotator.annotator.settings import ANNOTATOR_CLIENTS, SERVICE_PROVIDER_API_HOST
+from biothings_annotator.annotator.transformer import ResponseTransformer, load_atc_cache
+from biothings_annotator.annotator.utils import batched, get_client, get_dotfield_value, group_by_subfield, parse_curie
 
 logger = logging.getLogger(__name__)
 
 
 class Annotator:
-
     def __init__(self):
         self.api_host = os.environ.get("SERVICE_PROVIDER_API_HOST", SERVICE_PROVIDER_API_HOST)
 
-    def query_biothings(
+    async def query_biothings(
         self, node_type: str, query_list: List[str], fields: Optional[Union[str, List[str]]] = None
     ) -> Dict:
         """
         Query biothings client based on node_type for a list of ids
         """
         client = get_client(node_type, self.api_host)
-        if not isinstance(client, biothings_client.BiothingClient):
+        if not isinstance(client, biothings_client.AsyncBiothingClient):
             logger.error("Failed to get the biothings client for %s type. This type is skipped.", node_type)
             return {}
 
         fields = fields or ANNOTATOR_CLIENTS[node_type]["fields"]
         scopes = ANNOTATOR_CLIENTS[node_type]["scopes"]
         logger.info("Querying annotations for %s %ss...", len(query_list), node_type)
-        res = client.querymany(query_list, scopes=scopes, fields=fields)
+        res = await client.querymany(query_list, scopes=scopes, fields=fields)
         logger.info("Done. %s annotation objects returned.", len(res))
         grouped_response = group_by_subfield(collection=res, search_key="query")
         return grouped_response
 
-    def transform(self, res_by_id, node_type):
+    async def transform(self, res_by_id: Dict, node_type: str):
         """
         perform any transformation on the annotation object, but in-place also returned object
         res_by_id is the output of query_biothings, node_type is the same passed to query_biothings
         """
         logger.info("Transforming output annotations for %s %ss...", len(res_by_id), node_type)
-        transformer = ResponseTransformer(res_by_id, node_type)
+        atc_cache = await load_atc_cache(self.api_host)
+        transformer = ResponseTransformer(res_by_id, node_type, self.api_host, atc_cache)
         transformer.transform()
         logger.info("Done.")
         return res_by_id
 
-    def append_extra_annotations(self, node_d: Dict, node_id_subset: Optional[List[str]] = None, batch_n: int = 1000):
+    async def append_extra_annotations(
+        self, node_d: Dict, node_id_subset: Optional[List[str]] = None, batch_n: int = 1000
+    ):
         """
         Append extra annotations to the existing node_d
         """
@@ -61,7 +63,7 @@ def append_extra_annotations(self, node_d: Dict, node_id_subset: Optional[List[s
         cnt = 0
         logger.info("Retrieving extra annotations...")
         for node_id_batch in batched(node_id_list, batch_n):
-            extra_res = extra_api.querymany(node_id_batch, scopes="_id", fields="all")
+            extra_res = await extra_api.querymany(node_id_batch, scopes="_id", fields="all")
             for hit in extra_res:
                 if hit.get("notfound", False):
                     continue
@@ -83,7 +85,7 @@ def append_extra_annotations(self, node_d: Dict, node_id_subset: Optional[List[s
                         cnt += 1
         logger.info("Done. %s extra annotations appended.", cnt)
 
-    def annotate_curie(
+    async def annotate_curie(
         self, curie: str, raw: bool = False, fields: Optional[Union[str, List[str]]] = None, include_extra: bool = True
     ) -> Dict:
         """
@@ -92,18 +94,21 @@ def annotate_curie(
         node_type, _id = parse_curie(curie)
         if not node_type:
             raise InvalidCurieError(curie)
-        res = self.query_biothings(node_type, [_id], fields=fields)
+        res = await self.query_biothings(node_type, [_id], fields=fields)
         if not raw:
-            res = self.transform(res, node_type)
+            res = await self.transform(res, node_type)
             # res = [self.transform(r) for r in res[_id]]
         if res and include_extra:
-            self.append_extra_annotations(res)
-        return {curie: res.get(_id, {})}
+            await self.append_extra_annotations(res)
+
+        curie_annotation = {curie: res.get(_id, {})}
+        return curie_annotation
 
-    def _annotate_node_list_by_type(
+    async def _annotate_node_list_by_type(
         self, node_list_by_type: Dict, raw: bool = False, fields: Optional[Union[str, List[str]]] = None
     ) -> Iterable[tuple]:
-        """This is a helper method re-used in both annotate_curie_list and annotate_trapi methods
+        """
+        This is a helper method re-used in both annotate_curie_list and annotate_trapi methods
         It returns a generator of tuples of (original_node_id, annotation_object) for each node_id,
         passed via node_list_by_type.
         """
@@ -117,11 +122,12 @@ def _annotate_node_list_by_type(
 
             # this is the list of query ids like 1017
             query_list = [parse_curie(_id, return_type=False, return_id=True) for _id in node_list_by_type[node_type]]
+
             # query_id to original id mapping
             node_id_d = dict(zip(query_list, node_list))
-            res_by_id = self.query_biothings(node_type, query_list, fields=fields)
+            res_by_id = await self.query_biothings(node_type, query_list, fields=fields)
             if not raw:
-                res_by_id = self.transform(res_by_id, node_type)
+                res_by_id = await self.transform(res_by_id, node_type)
 
             # map back to original node ids
             # NOTE: we don't want to use `for node_id in res_by_id:` here, since we will mofiify res_by_id in the loop
@@ -131,7 +137,7 @@ def _annotate_node_list_by_type(
                     res_by_id[orig_node_id] = res_by_id.pop(node_id)
                 yield (orig_node_id, res_by_id[orig_node_id])
 
-    def annotate_curie_list(
+    async def annotate_curie_list(
         self,
         curie_list: Union[List[str], Iterable[str]],
         raw: bool = False,
@@ -154,14 +160,15 @@ def annotate_curie_list(
             else:
                 logger.warning("Unsupported Curie prefix: %s. Skipped!", node_id)
 
-        for node_id, res in self._annotate_node_list_by_type(node_list_by_type, raw=raw, fields=fields):
+        async for node_id, res in self._annotate_node_list_by_type(node_list_by_type, raw=raw, fields=fields):
             node_d[node_id] = res
+
         if include_extra:
             # currently, we only need to append extra annotations for chem nodes
             self.append_extra_annotations(node_d, node_id_subset=node_list_by_type.get("chem", []))
         return node_d
 
-    def annotate_trapi(
+    async def annotate_trapi(
         self,
         trapi_input: Dict,
         append: bool = False,
@@ -176,8 +183,8 @@ def annotate_trapi(
         try:
             node_d = get_dotfield_value("message.knowledge_graph.nodes", trapi_input)
             assert isinstance(node_d, dict)
-        except (KeyError, ValueError, AssertionError):
-            raise TRAPIInputError(trapi_input)
+        except (KeyError, ValueError, AssertionError) as access_error:
+            raise TRAPIInputError(trapi_input) from access_error
 
         # if limit is set, we truncate the node_d to that size
         if limit:
@@ -203,7 +210,7 @@ def annotate_trapi(
                 logger.warning("Unsupported Curie prefix: %s. Skipped!", node_id)
 
         _node_d = {}
-        for node_id, res in self._annotate_node_list_by_type(node_list_by_type, raw=raw, fields=fields):
+        async for node_id, res in self._annotate_node_list_by_type(node_list_by_type, raw=raw, fields=fields):
             _node_d[node_id] = res
 
         if include_extra:

biothings_annotator/annotator/transformer.py

@@ -5,10 +5,10 @@
 
 import inspect
 import logging
-import os
+from typing import Dict
 
-from .settings import SERVICE_PROVIDER_API_HOST
-from .utils import get_client
+from biothings_annotator.annotator.settings import SERVICE_PROVIDER_API_HOST
+from biothings_annotator.annotator.utils import get_client
 
 logger = logging.getLogger(__name__)
 
@@ -22,31 +22,33 @@ def append_prefix(id, prefix):
 atc_cache = {}  # The global atc_cache will be load once when Transformer is initialized for the first time
 
 
-def load_atc_cache(api_host: str):
-    """Load WHO atc code-to-name mapping in a dictionary, which will be used in ResponseTransformer._transform_atc_classifications method"""
+async def load_atc_cache(api_host: str) -> Dict:
+    """
+    Load WHO atc code-to-name mapping in a dictionary, which will be used in ResponseTransformer._transform_atc_classifications method
+    """
     global atc_cache
     if not atc_cache:
         logger.info("Loading WHO ATC code-to-name mapping...")
         atc_client = get_client("extra", api_host)
-        atc_li = atc_client.query("_exists_:atc.code", fields="atc.code,atc.name", fetch_all=True)
+        atc_li = await atc_client.query("_exists_:atc.code", fields="atc.code,atc.name", fetch_all=True)
         atc_cache = {}
-        for atc in atc_li:
+        async for atc in atc_li:
             atc_cache[atc["atc"]["code"]] = atc["atc"]["name"]
         logger.info(f"Loaded {len(atc_cache)} WHO ATC code-to-name mappings.")
     return atc_cache
 
 
 class ResponseTransformer:
-    def __init__(self, res_by_id, node_type):
+    def __init__(self, res_by_id: Dict, node_type: str, api_host: str, atc_cache: Dict):
         self.res_by_id = res_by_id
         self.node_type = node_type
-        self.api_host = os.environ.get("SERVICE_PROVIDER_API_HOST", SERVICE_PROVIDER_API_HOST)
+        self.api_host = api_host
 
         self.data_cache = {}  # used to cached required mapping data used for individual transformation
         # typically those data coming from other biothings APIs, we will do a batch
         # query to get them all, and cache them here for later use, to avoid slow
         # one by one queries.
-        self.atc_cache = load_atc_cache(self.api_host)
+        self.atc_cache = atc_cache
 
     def _transform_chembl_drug_indications(self, doc):
         if self.node_type != "chem":
@@ -71,7 +73,9 @@ def _append_mesh_prefix(chembl):
         return doc
 
     def _transform_atc_classifications(self, doc):
-        """add atc_classifications field to chem object based on chembl.atc_classifications and pharmgkb.xrefs.atc fields"""
+        """
+        add atc_classifications field to chem object based on chembl.atc_classifications and pharmgkb.xrefs.atc fields
+        """
         if not self.atc_cache:
             return doc
 

biothings_annotator/annotator/utils.py

@@ -21,13 +21,13 @@ def batched(iterable, n):
 
 import biothings_client
 
-from .exceptions import InvalidCurieError
-from .settings import ANNOTATOR_CLIENTS, BIOLINK_PREFIX_to_BioThings
+from biothings_annotator.annotator.exceptions import InvalidCurieError
+from biothings_annotator.annotator.settings import ANNOTATOR_CLIENTS, BIOLINK_PREFIX_to_BioThings
 
 logger = logging.getLogger(__name__)
 
 
-def get_client(node_type: str, api_host: str) -> Union[biothings_client.BiothingClient, None]:
+def get_client(node_type: str, api_host: str) -> Union[biothings_client.AsyncBiothingClient, None]:
     """
     Attempts to lazy load the biothings-client instance
 
@@ -59,21 +59,20 @@ def get_client(node_type: str, api_host: str) -> Union[biothings_client.Biothing
     client_configuration = client_parameters.get("configuration")
     client_endpoint = client_parameters.get("endpoint")
     client_instance = client_parameters.get("instance")
-
-    if client_instance is not None and isinstance(client_instance, biothings_client.BiothingClient):
+    if client_instance is not None and isinstance(client_instance, biothings_client.AsyncBiothingClient):
         client = client_instance
 
     elif client_configuration is not None and isinstance(client_configuration, dict):
         try:
-            client = biothings_client.get_client(**client_configuration)
+            client = biothings_client.get_async_client(**client_configuration)
         except RuntimeError as runtime_error:
             logger.error("%s [%s]", runtime_error, client_configuration)
             client = None
 
     elif client_endpoint is not None and isinstance(client_endpoint, str):
         client_url = f"{api_host}/{client_endpoint}"
         try:
-            client = biothings_client.get_client(biothing_type=None, instance=True, url=client_url)
+            client = biothings_client.get_async_client(biothing_type=None, instance=True, url=client_url)
         except RuntimeError as runtime_error:
             logger.error("%s [%s]", runtime_error, client_url)
             client = None
@@ -84,15 +83,15 @@ def get_client(node_type: str, api_host: str) -> Union[biothings_client.Biothing
         )
 
     # cache the client
-    if isinstance(client, biothings_client.BiothingClient):
+    if isinstance(client, biothings_client.AsyncBiothingClient):
         ANNOTATOR_CLIENTS[node_type]["client"]["instance"] = client
 
     return client
 
 
 def parse_curie(curie: str, return_type: bool = True, return_id: bool = True):
     """
-    return a both type and if (as a tuple) or either based on the input curie
+    return both type and if (as a tuple) or either based on the input curie
     """
     if ":" not in curie:
         raise InvalidCurieError(curie)
@@ -117,7 +116,7 @@ def group_by_subfield(collection: List[Dict], search_key: str) -> Dict:
     Takes a collection of dictionary entries with a specify subfield key "search_key" and
     extracts the subfield from each entry in the iterable into a dictionary.
 
-    It then bins entries into the dictionary so that identical keys have all results in one
+    It the bins entries into the dictionary so that identical keys have all results in one
     aggregated list across the entire collection of dictionary entries
 
     Example:
@@ -174,4 +173,3 @@ def get_dotfield_value(dotfield: str, d: Dict):
     else:
         first = fields[0]
         return get_dotfield_value(".".join(fields[1:]), d[first])
-        return get_dotfield_value(".".join(fields[1:]), d[first])

biothings_annotator/application/cli/interface.py

@@ -164,9 +164,6 @@ def run(self):
         > Build an HTTP server using the network parameters from the configuration
         and the sanic application
         """
-        if self.inspecting:
-            self._inspector()
-            return
 
         self._precheck()
         try:

biothings_annotator/application/views/annotator.py

@@ -30,7 +30,7 @@ async def head(self, request: Request):
 
         annotator = Annotator()
         try:
-            annotated_node = annotator.annotate_curie(curie, fields=fields, raw=False, include_extra=False)
+            annotated_node = await annotator.annotate_curie(curie, fields=fields, raw=False, include_extra=False)
 
             if "NCBIGene:1017" not in annotated_node:
                 return sanic.json(None, status=500)
@@ -71,7 +71,7 @@ async def get(self, _: Request):
 
         annotator = Annotator()
         try:
-            annotated_node = annotator.annotate_curie(curie, fields=fields, raw=False, include_extra=False)
+            annotated_node = await annotator.annotate_curie(curie, fields=fields, raw=False, include_extra=False)
 
             if "NCBIGene:1017" not in annotated_node:
                 result = {"success": False, "error": "Service unavailable due to a failed data check!"}
@@ -157,7 +157,7 @@ async def get(self, request: Request, curie: str):
         annotator = Annotator()
 
         try:
-            annotated_node = annotator.annotate_curie(curie, fields=fields, raw=raw, include_extra=include_extra)
+            annotated_node = await annotator.annotate_curie(curie, fields=fields, raw=raw, include_extra=include_extra)
             return sanic.json(annotated_node, headers=self.default_headers)
         except InvalidCurieError as curie_err:
             error_context = {
@@ -276,7 +276,7 @@ async def post(self, request: Request):
             return curie_error_response
 
         try:
-            annotated_node = annotator.annotate_curie_list(
+            annotated_node = await annotator.annotate_curie_list(
                 curie_list=curie_list, fields=fields, raw=raw, include_extra=include_extra
             )
             return sanic.json(annotated_node, headers=self.default_headers)
@@ -408,7 +408,7 @@ async def post(self, request: Request):
         annotator = Annotator()
         trapi_body = request.json
         try:
-            annotated_node = annotator.annotate_trapi(
+            annotated_node = await annotator.annotate_trapi(
                 trapi_body, fields=fields, raw=raw, append=append, limit=limit, include_extra=include_extra
             )
             return sanic.json(annotated_node, headers=self.default_headers)