Merge pull request #14 from VACLab/cohort_creation_background_enhancement

Cohort creation background enhancement

Author: hyi

.gitignore

@@ -3,3 +3,4 @@ config.yaml
 __pycache__/
 dist/
 .coverage
+.ipynb_checkpoints/

biasanalyzer/api.py

@@ -1,3 +1,4 @@
+import time
 from pydantic import ValidationError
 from biasanalyzer.database import OMOPCDMDatabase, BiasDatabase
 from biasanalyzer.cohort import CohortAction
@@ -159,32 +160,41 @@ def display_concept_tree(self, concept_tree: dict, level: int = 0, show_in_text_
             return None
 
 
-    def create_cohort(self, cohort_name, cohort_desc, query_or_yaml_file, created_by):
+    def create_cohort(self, cohort_name: str, cohort_desc: str, query_or_yaml_file: str, created_by: str,
+                      delay: float=0):
         """
-        cohort_name: name of the cohort
-        cohort_desc: description of the cohort
-        query_or_yaml_file: a SQL query or YAML cohort creation file
-        created_by: name of the user that created the cohort
+        API method that allows to create a cohort
+        :param cohort_name: name of the cohort
+        :param cohort_desc: description of the cohort
+        :param query_or_yaml_file: a SQL query or YAML cohort creation file
+        :param created_by: name of the user that created the cohort
+        :param delay: the number of seconds to sleep/delay for simulating long-running task for async testing,
+        default is 0, meaning no delay
+        :return: CohortData object if cohort is created successfully; otherwise, None
         """
+
         c_action = self._set_cohort_action()
         if c_action:
-            created_cohort = c_action.create_cohort(cohort_name, cohort_desc, query_or_yaml_file,
-                                                    created_by)
+            created_cohort = c_action.create_cohort(cohort_name, cohort_desc, query_or_yaml_file, created_by)
             if created_cohort is not None:
+                if delay > 0:
+                    print(f"[DEBUG] Simulating long-running task with {delay} seconds delay...")
+                    time.sleep(delay)
                 print('cohort created successfully')
             return created_cohort
         else:
             print('failed to create a valid cohort action object')
             return None
 
 
-
     def compare_cohorts(self, cohort_id1, cohort_id2):
         c_action = self._set_cohort_action()
         if c_action:
             return c_action.compare_cohorts(cohort_id1, cohort_id2)
         else:
             print('failed to create a valid cohort action object')
+            return None
+
 
     def cleanup(self):
         self.bias_db.close()

biasanalyzer/cohort.py

@@ -79,6 +79,12 @@ def create_cohort(self, cohort_name: str, description: str, query_or_yaml_file:
         Create a new cohort by executing a query on OMOP CDM database
         and storing the result in BiasDatabase. The query can be passed in directly
         or built from a yaml file using a corresponding SQL query template
+        :param cohort_name: the name of the cohort
+        :param description: the description of the cohort
+        :param query_or_yaml_file: the SQL query string or yaml file name for creating a cohort
+        :param created_by: created_by string indicating who created the cohort, it could be 'system',
+        or a username, or whatever metadata to record who created the cohort
+        :return: CohortData object if cohort is created successfully; otherwise, return None
         """
         stages = [
             "Built query",