DIG-2136: a few more optimizations

htsget_server/beacon_operations.py

@@ -257,27 +257,28 @@ def search(raw_req):
         headers = {
             "X-Service-Token": create_service_token()
         }
-        resp = requests.get(url=f"{os.getenv("DRS_URL")}/ga4gh/drs/v1/objects", headers=headers)
-        drs_obj_dict = {}
+        resp = requests.post(url=f"{os.getenv("DRS_URL")}/ga4gh/drs/v1/experiments", headers=headers, json={})
+
+        exp_lookup = {}
         if resp.status_code == 200:
-            drs_objects = resp.json()
-            for obj in drs_objects:
-                drs_obj_dict[obj["id"]] = obj
+            for exp in resp.json():
+                for v in exp["variants"]:
+                    if v not in exp_lookup:
+                        exp_lookup[v] = []
+                    exp_lookup[v].append(exp)
+        else:
+            raise Exception(f"couldn't get experiments: {resp.text} {resp.status_code}")
 
         hits_to_search = []
         for i in range(len(potential_hits)):
             drs_obj_id = potential_hits[i]['drs_object_id']
-
-            drs_obj = drs_obj_dict[drs_obj_id]
-            if "program" in drs_obj:
-                if drs_obj["program"] not in results:
-                    results[drs_obj["program"]] = []
-                for c in drs_obj["contents"]:
-                    if c["id"] not in ["analysis", "index"]:
-                        # this is a ExperimentContentObject
-                        res = {"submitter_sample_id": c["name"], "variant_count": potential_hits[i]["variantcount"]}
-                        results[drs_obj["program"]].append(res)
-                        hits_to_search.append(potential_hits[i])
+            experiments = exp_lookup[drs_obj_id]
+            for experiment in experiments:
+                if experiment["program"] not in results:
+                    results[experiment["program"]] = []
+                res = {"submitter_sample_id": experiment["experiment_id"], "variant_count": potential_hits[i]["variantcount"]}
+                results[experiment["program"]].append(res)
+                hits_to_search.append(potential_hits[i])
 
         search_json = {
             "potential_hits": hits_to_search,
@@ -418,11 +419,12 @@ def full_beacon_search(search_json, headers=None):
 
     # if the request granularity was "record", check to see that the user is actually authorized to see any programs:
     response['beaconHandovers'] = []
+    headers = {
+        "X-Service-Token": create_service_token()
+    }
+
     for drs_obj_id in variants_by_file.keys():
         # look for experiments and programs for all drs objects, even if user is not authorized
-        headers = {
-            "X-Service-Token": create_service_token()
-        }
         resp = requests.get(url=f"{os.getenv("DRS_URL")}/ga4gh/drs/v1/objects/{drs_obj_id}", headers=headers)
         if resp.status_code == 200:
             drs_obj = resp.json()

htsget_server/htsget_operations.py

@@ -246,8 +246,9 @@ def get_experiment(id_=None):
 
 async def get_multiple_experiments():
     req = await connexion.request.json()
-    return _get_experiments(req["experiments"]), 200
-
+    if "experiments" in req:
+        return _get_experiments(req["experiments"]), 200
+    return _get_experiments(None), 200
 
 def get_program_experiments(program=None):
     headers = {
@@ -310,15 +311,21 @@ def get_pysam_obj(object_id, headers=None):
 
 def _get_experiments(experiments):
     result = []
-    experiments_by_program = {}
-    for experiment in experiments:
-        res, status_code = _get_experiment(experiment)
-        if status_code == 200:
+    headers = {
+        "X-Service-Token": create_service_token()
+    }
+    if experiments is None:
+        resp = requests.post(f"{os.getenv("DRS_URL")}/ga4gh/drs/v1/experiments", headers=headers, json={})
+    else:
+        resp = requests.post(f"{os.getenv("DRS_URL")}/ga4gh/drs/v1/experiments", headers=headers, json={"submitter_sample_ids": experiments})
+    if resp.status_code == 200:
+        experiments_by_program = {}
+        for res in resp.json():
             if res["program"] not in experiments_by_program:
                 experiments_by_program[res["program"]] = []
             experiments_by_program[res["program"]].append(res)
-    for program in experiments_by_program.keys():
-        result.extend(experiments_by_program[program])
+        for program in experiments_by_program.keys():
+            result.extend(experiments_by_program[program])
     return result
 
 
@@ -335,24 +342,12 @@ def _get_experiment(id_=None):
     headers = {
         "X-Service-Token": create_service_token()
     }
-    resp = requests.get(url=f"{os.getenv("DRS_URL")}/ga4gh/drs/v1/objects/{id_}", headers=headers)
+    resp = requests.post(f"{os.getenv("DRS_URL")}/ga4gh/drs/v1/experiments", headers=headers, json={"submitter_sample_ids": [id_]})
+
     if resp.status_code == 200:
-        experiment_drs_obj = resp.json()
-        if experiment_drs_obj is not None and "contents" in experiment_drs_obj and experiment_drs_obj["description"] in ["wgs", "wts"]:
-            if experiment_drs_obj["description"] == "wgs":
-                result["genomes"].append(experiment_drs_obj["id"])
-            elif experiment_drs_obj["description"] == "wts":
-                result["transcriptomes"].append(experiment_drs_obj["id"])
-            result["program"] = experiment_drs_obj["program"]
-            for contents_obj in experiment_drs_obj["contents"]:
-                drs_obj = _describe_drs_object(contents_obj["id"])
-                if drs_obj is not None and "type" in drs_obj:
-                    if drs_obj["type"] == "variant":
-                        result["variants"].append(drs_obj["name"])
-                    elif drs_obj["type"] == "read":
-                        result["reads"].append(drs_obj["name"])
-            return result, 200
-        return f"{id_} is not an Experiment", 404
+        if len(resp.json()) == 0:
+            return f"{id_} is not an Experiment", 404
+        return resp.json(), 200
     return resp.text, resp.status_code
 
 

htsget_server/variants.py

@@ -73,6 +73,9 @@ def parse_vcf_file(drs_object_id, reference_name=None, start=None, end=None, hea
     if 'contig' in headers:
         variants_by_file['contig'] = headers.pop('contig')
     experiment_dict = {}
+    req_headers = {
+        "X-Service-Token": create_service_token()
+    }
     for r in records:
         experiments = []
         for vcf_sample in r.samples:
@@ -81,10 +84,7 @@ def parse_vcf_file(drs_object_id, reference_name=None, start=None, end=None, hea
                 experiment_id = analysis_obj['experiments'][vcf_sample]
                 if experiment_id not in experiment_dict:
                     experiment_id = analysis_obj['experiments'][vcf_sample]
-                    headers = {
-                        "X-Service-Token": create_service_token()
-                    }
-                    response = requests.get(url=f"{os.getenv("DRS_URL")}/ga4gh/drs/v1/objects/{experiment_id}", headers=headers)
+                    response = requests.get(url=f"{os.getenv("DRS_URL")}/ga4gh/drs/v1/objects/{experiment_id}", headers=req_headers)
                     experiment_obj = None
                     if response.status_code == 200:
                         experiment_obj = response.json()