refactor: ♻️ simplify and refactor metakg parsing logics

Author: newgene

src/handlers/api.py

@@ -16,12 +16,10 @@
 from controller import SmartAPI
 from controller.exceptions import ControllerError, NotFoundError
 from pipeline import MetaKGQueryPipeline
-from utils.decoder import to_dict
 from utils.downloader import DownloadError, download_async
 from utils.metakg.biolink_helpers import get_expanded_values
 from utils.metakg.cytoscape_formatter import CytoscapeDataFormatter
 from utils.metakg.export import edges2graphml
-from utils.metakg.metakg_errors import MetadataRetrievalError
 from utils.metakg.parser import MetaKGParser
 from utils.metakg.path_finder import MetaKGPathFinder
 from utils.notification import SlackNewAPIMessage, SlackNewTranslatorAPIMessage
@@ -751,19 +749,17 @@ class MetaKGParserHandler(BaseHandler, MetaKGHandlerMixin):
     """
 
     kwargs = {
+        "*": {
+            "api_details": {"type": bool, "default": False},
+            "bte": {"type": bool, "default": False},
+        },
         "GET": {
             "url": {
                 "type": str,
                 "required": True,
                 "max": 1000,
                 "description": "URL of the SmartAPI metadata to parse"
             },
-            "api_details": {"type": bool, "default": False},
-            "bte": {"type": bool, "default": False},
-        },
-        "POST": {
-            "api_details": {"type": bool, "default": False},
-            "bte": {"type": bool, "default": False},
         },
     }
 
@@ -789,95 +785,69 @@ def process_apis(self, apis):
 
     async def get(self, *args, **kwargs):
         url = self.args.url
-        if not url:
-            raise HTTPError(400, reason="A url value is expected for the request, please provide a url.")
-
-        # Set initial args and handle potential errors in query parameters
         parser = MetaKGParser()
 
         try:
-            trapi_data = parser.get_TRAPI_metadatas(data=None, url=url)
-        except MetadataRetrievalError as retrieve_err:
-            raise HTTPError(retrieve_err.status_code, reason=retrieve_err.message)
-        except DownloadError:
-            raise HTTPError(400, reason="There was an error downloading the data from the given input.")
-
-        # Get non-TRAPI metadata
-        try:
-            nontrapi_data = parser.get_non_TRAPI_metadatas(data=None, url=url)
-        except MetadataRetrievalError as retrieve_err:
-            raise HTTPError(retrieve_err.status_code, reason=retrieve_err.message)
+            parsed_metakg = parser.get_metakg(url=url)
         except DownloadError:
-            raise HTTPError(400, reason="There was an error downloading the data from the given input.")
+            self.write_error(400, reason="There was an error downloading the data from the given url.")
+        except (ValueError, TypeError) as err:
+            self.write_error(
+                status_code=400,
+                reason="The data retrived from the given url is not a valid JSON or YAML object.",
+                message=str(err)
+            )
 
         # Apply filtering -- if data found
-        combined_data = trapi_data + nontrapi_data
-        if combined_data:
-            for i, api_dict in enumerate(combined_data):
-                combined_data[i] = self.get_filtered_api(api_dict)
+        if parsed_metakg:
+            for i, api_dict in enumerate(parsed_metakg):
+                parsed_metakg[i] = self.get_filtered_api(api_dict)
 
         # Add url to metadata if api_details is set to 1
         if self.args.api_details:
-            for data_dict in combined_data:
+            for data_dict in parsed_metakg:
                 if "metadata" in data_dict["api"]["smartapi"] and data_dict["api"]["smartapi"]["metadata"] is None:
                     data_dict["api"]["smartapi"]["metadata"] = url
 
         response = {
-            "total": len(combined_data),
-            "hits": combined_data,
+            "total": len(parsed_metakg),
+            "hits": parsed_metakg,
         }
 
         self.finish(response)
 
     async def post(self, *args, **kwargs):
-        raw_body = self.request.body
-        if not raw_body:
-            raise HTTPError(400, reason="Request body cannot be empty.")
-
         content_type = self.request.headers.get("Content-Type", "").lower()
+        if content_type in ["application/json", "application/x-yaml"]:
+            # if content type is set properly, it should have alrady been parsed
+            metadata_from_body = self.args_json or self.args_yaml
+        elif self.request.body:
+            # if request body is provided but no proper content type is set
+            # we will parse it as YAML anyway
+            metadata_from_body = self._parse_yaml()
+        else:
+            metadata_from_body = None
+
+        if metadata_from_body:
+            # Process the parsed metadata
+            parser = MetaKGParser()
+            parsed_metakg = parser.get_metakg(metadata_from_body)
+
+            # Apply filtering to the combined data
+            if parsed_metakg:
+                for i, api_dict in enumerate(parsed_metakg):
+                    parsed_metakg[i] = self.get_filtered_api(api_dict)
+
+            # Send the response back to the client
+            response = {
+                "total": len(parsed_metakg),
+                "hits": parsed_metakg,
+            }
 
-        # Try to parse the request body based on content type
-        try:
-            if content_type == "application/json":
-                data = to_dict(raw_body, ctype="application/json")
-            elif content_type == "application/x-yaml":
-                data = to_dict(raw_body, ctype="application/x-yaml")
-            else:
-                # Default to YAML parsing if the content type is unknown or not specified
-                data = to_dict(raw_body)
-        except ValueError as val_err:
-            if 'mapping values are not allowed here' in str(val_err):
-                raise HTTPError(400, reason="Formatting issue, please consider using --data-binary to maintain YAML format.")
-            else:
-                raise HTTPError(400, reason="Invalid value, please provide a valid YAML object.")
-        except TypeError:
-            raise HTTPError(400, reason="Invalid type, provide valid type metadata.")
-
-        # Ensure the parsed data is a dictionary
-        if not isinstance(data, dict):
-            raise ValueError("Invalid input data type. Please provide a valid JSON/YAML object.")
-
-        # Process the parsed metadata
-        parser = MetaKGParser()
-        try:
-            trapi_data = parser.get_TRAPI_metadatas(data=data)
-            nontrapi_data = parser.get_non_TRAPI_metadatas(data=data)
-        except MetadataRetrievalError as retrieve_err:
-            raise HTTPError(retrieve_err.status_code, reason=retrieve_err.message)
-        except DownloadError:
-            raise HTTPError(400, reason="Error downloading the data from the provided input.")
-
-        combined_data = trapi_data + nontrapi_data
-
-        # Apply filtering to the combined data
-        if combined_data:
-            for i, api_dict in enumerate(combined_data):
-                combined_data[i] = self.get_filtered_api(api_dict)
-
-        # Send the response back to the client
-        response = {
-            "total": len(combined_data),
-            "hits": combined_data,
-        }
-
-        self.finish(response)
+            self.finish(response)
+        else:
+            self.write_error(
+                status_code=400,
+                reason="Request body cannot be empty.",
+                message="Please provide a valid JSON/YAML object in the request body."
+            )

src/pipeline.py

@@ -2,14 +2,10 @@
 from enum import Enum
 from typing import Dict, OrderedDict
 
-from biothings.web.query import (
-    AsyncESQueryBackend,
-    AsyncESQueryPipeline,
-    ESQueryBuilder,
-    ESResultFormatter,
-)
-from controller.base import OpenAPI, Swagger
+from biothings.web.query import AsyncESQueryBackend, AsyncESQueryPipeline, ESQueryBuilder, ESResultFormatter
 from elasticsearch_dsl import Q, Search
+
+from controller.base import OpenAPI, Swagger
 from utils import decoder
 
 
@@ -219,8 +215,8 @@ def apply_extras(self, search, options):
         apply extra filters
         """
         # if not options._source:
-            # by default exclude api.bte or bte field, but can be included by specifying in the fields parameter
-            # options._source = ["-api.bte", "-bte"]
+        #    by default exclude api.bte or bte field, but can be included by specifying in the fields parameter
+        #    options._source = ["-api.bte", "-bte"]
 
         search = super().apply_extras(search, options)
         # apply extra filters from query parameters
@@ -262,6 +258,7 @@ def adjust_index(self, original_index: str, query: str, **options: Dict) -> str:
             query_index = self.indices.get("metakg", None)
         return query_index
 
+
 class MetaKGQueryPipeline(AsyncESQueryPipeline):
     def __init__(self, *args, **kwargs):
         # ns is an instance of BiothingsNamespace

src/utils/metakg/metakg_errors.py

@@ -1,7 +1,8 @@
 import json
 import logging
 from copy import copy
-from utils.metakg.metakg_errors import MetadataRetrievalError
+from typing import Dict, List, Optional, Union
+
 import requests
 
 from .api import API
@@ -13,46 +14,49 @@ class MetaKGParser:
     get_url_timeout = 60
     metakg_errors = None
 
-    def get_non_TRAPI_metadatas(self, data=None, extra_data=None, url=None):
+    def get_metakg(self,
+                   data: Optional[Union[Dict, API]] = None,
+                   extra_data: Optional[Dict] = None,
+                   url: Optional[str] = None) -> List[Dict]:
         """
-        Extract MetaKG edges from a SmartAPI document provided as `data` or fetched from a `url`.
-        Raises an error if no valid input is given, or if parser fails to parse the document.
+        Extract and process metadata from a SmartAPI document or URL.
+        Returns MetaKG edges or propagates errors.
         """
         if not data and not url:
-            raise MetadataRetrievalError(400, "Either data or url value is expected for this request, please provide data or a url.")
+            raise ValueError("Either data or url value is expected for this request, please provide data or a url.")
 
+        # if both data and url are provided, prefer data
         if data:
-            parser = API(smartapi_doc=data)
+            _api = data if isinstance(data, API) else API(data)
         elif url:
-            parser = API(url=url)
+            _api = API(url=url)
+
+        if _api.is_trapi:
+            return self.get_TRAPI_metadatas(data=_api, extra_data=extra_data)
         else:
-            raise MetadataRetrievalError(404, "No metadata available from provided data or url.")
+            return self.get_non_TRAPI_metadatas(data=_api, extra_data=extra_data)
 
-        mkg = self.extract_metakgedges(parser.metadata["operations"], extra_data=extra_data)
+    def get_non_TRAPI_metadatas(self, data: Union[Dict, API], extra_data: Optional[Dict] = None) -> List[Dict]:
+        """
+        Extract MetaKG edges from a SmartAPI document provided as `data` or fetched from a `url`.
+        Raises an error if no valid input is given, or if parser fails to parse the document.
+        """
+        _api = data if isinstance(data, API) else API(data)
+        mkg = self.extract_metakgedges(_api.metadata["operations"], extra_data=extra_data)
         no_nodes = len({x["subject"] for x in mkg} | {x["object"] for x in mkg})
         no_edges = len({x["predicate"] for x in mkg})
         logger.info("Done [%s nodes, %s edges]", no_nodes, no_edges)
         return mkg
 
-    def get_TRAPI_metadatas(self, data=None, extra_data=None, url=None):
+    def get_TRAPI_metadatas(self, data: Union[Dict, API], extra_data: Optional[Dict] = None) -> List[Dict]:
         """
         Extract and process TRAPI metadata from a SmartAPI document or URL.
         Returns MetaKG edges or propagates errors.
         """
-        if not data and not url:
-            raise MetadataRetrievalError(400, "Either data or url value is expected for this request, please provide data or a url.")
-
-        try:
-            if data:
-                metadata_list = self.get_TRAPI_with_metakg_endpoint(data=data)
-            else:
-                metadata_list = self.get_TRAPI_with_metakg_endpoint(url=url)
-        except MetadataRetrievalError:
-            raise MetadataRetrievalError(404, "No metadata available from provided data or url.")
-
+        ops = []
+        metadata_list = self.get_TRAPI_with_metakg_endpoint(data)
         count_metadata_list = len(metadata_list)
         self.metakg_errors = {}
-        ops = []
 
         for i, metadata in enumerate(metadata_list):
             ops.extend(self.get_ops_from_metakg_endpoint(metadata, f"[{i + 1}/{count_metadata_list}]"))
@@ -63,27 +67,22 @@ def get_TRAPI_metadatas(self, data=None, extra_data=None, url=None):
 
         return self.extract_metakgedges(ops, extra_data=extra_data)
 
-    def get_TRAPI_with_metakg_endpoint(self, data=None, url=None):
+    def get_TRAPI_with_metakg_endpoint(self, data: Union[Dict, API]):
         """
         Retrieve TRAPI metadata from a SmartAPI document or URL.
         Returns metadata if TRAPI endpoints are found, else an empty list.
         """
-        if not data and not url:
-            raise MetadataRetrievalError(400, "Either data or url value is expected for this request, please provide data or a url.")
+        metadatas = []
+        _api = data if isinstance(data, API) else API(data)
 
-        # Initialize API with either data or URL
-        parser = API(smartapi_doc=data) if data else API(url=url)
-
-        # Download the metadata
-        metadata = parser.metadata
+        metadata = _api.metadata
         _paths = metadata.get("paths", {})
         _team = metadata.get("x-translator", {}).get("team")
 
         # Check for required TRAPI paths
         if "/meta_knowledge_graph" in _paths and "/query" in _paths and _team:
-            return [metadata]
-        else:
-            return []
+            metadatas.append(metadata)
+        return metadatas
 
     def construct_query_url(self, server_url):
         if server_url.endswith("/"):