add async loading option for reports

pori_python/ipr/connection.py

@@ -4,8 +4,9 @@
 import os
 import zlib
 from typing import Dict, List
-
+import time
 from .constants import DEFAULT_URL
+from .util import logger
 
 IMAGE_MAX = 20  # cannot upload more than 20 images at a time
 
@@ -62,8 +63,46 @@ def post(self, uri: str, data: Dict = {}, **kwargs) -> Dict:
             **kwargs,
         )
 
-    def upload_report(self, content: Dict) -> Dict:
-        return self.post('reports', content)
+    def get(self, uri: str, data: Dict = {}, **kwargs) -> Dict:
+        """Convenience method for making get requests"""
+        return self.request(
+            uri,
+            method='GET',
+            data=zlib.compress(json.dumps(data, allow_nan=False).encode('utf-8')),
+            **kwargs,
+        )
+
+    def upload_report(self, content: Dict, mins_to_wait: int = 5, async_upload: bool = False) -> Dict:
+        if async_upload:
+            initial_result = self.post('reports-async', content)
+            report_id = initial_result["ident"]
+            for i in range(5):
+                logger.info('checking report loading status in 5 seconds')
+                time.sleep(5)
+                current_status = self.get(f'reports-async/{report_id}')
+                if current_status['state'] not in ['active', 'ready']:
+                    raise Exception(f'async report upload in unexpected state: {current_status}')
+                if current_status['state'] == 'ready':
+                    return current_status
+            for i in range(5):
+                logger.info('checking report loading status in 30 seconds')
+                time.sleep(30)
+                current_status = self.get(f'reports-async/{report_id}')
+                if current_status['state'] not in ['active', 'ready']:
+                    raise Exception(f'async report upload in unexpected state: {current_status}')
+                if current_status['state'] == 'ready':
+                    return current_status
+            for i in range(mins_to_wait):
+                logger.info('checking report loading status in 1 minute')
+                time.sleep(60)
+                current_status = self.get(f'reports-async/{report_id}')
+                if current_status['state'] not in ['active', 'ready']:
+                    raise Exception(f'async report upload in unexpected state: {current_status}')
+                if current_status['state'] == 'ready':
+                    return current_status
+            raise Exception(f'async report upload taking longer than expected: {current_status}')
+        else:
+            return self.post('reports', content)
 
     def set_analyst_comments(self, report_id: str, data: Dict) -> Dict:
         """

pori_python/ipr/main.py

@@ -216,6 +216,8 @@ def ipr_report(
     generate_comments: bool = True,
     match_germline: bool = False,
     custom_kb_match_filter=None,
+    async_upload: bool = False,
+    mins_to_wait: int = 5,
 ) -> Dict:
     """Run the matching and create the report JSON for upload to IPR.
 
@@ -234,6 +236,8 @@ def ipr_report(
         generate_comments: create the analyst comments section for upload with the report
         match_germline: match only germline statements to germline events and non-germline statements to non-germline events.
         custom_kb_match_filter: function(List[kbMatch]) -> List[kbMatch]
+        async_upload: use report_async endpoint to upload reports
+        mins_to_wait: if using report_async, number of minutes to wait for success before exception raised
 
     Returns:
         ipr_conn.upload_report return dictionary
@@ -444,7 +448,7 @@ def ipr_report(
     if ipr_upload:
         try:
             logger.info(f'Uploading to IPR {ipr_conn.url}')
-            ipr_result = ipr_conn.upload_report(output)
+            ipr_result = ipr_conn.upload_report(output, async_upload, mins_to_wait)
             logger.info(ipr_result)
             output.update(ipr_result)
         except Exception as err: