From 6d3644401fe4d5f6ac0539ced8e28e08b0cd74a5 Mon Sep 17 00:00:00 2001
From: Gordon Krieger <gordon.krieger@gmail.com>
Date: Fri, 13 Dec 2024 18:07:58 +0000
Subject: [PATCH] beacon query handling improvements

---
 bento_beacon/endpoints/individuals.py | 9 +++++----
 bento_beacon/endpoints/variants.py    | 8 ++++----
 bento_beacon/utils/beacon_request.py  | 6 +++---
 bento_beacon/utils/censorship.py      | 6 +++---
 4 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/bento_beacon/endpoints/individuals.py b/bento_beacon/endpoints/individuals.py
index 87806f4..87707e5 100644
--- a/bento_beacon/endpoints/individuals.py
+++ b/bento_beacon/endpoints/individuals.py
@@ -31,10 +31,11 @@
 
 @route_with_optional_project_id("/individuals", methods=["GET", "POST"])
 async def get_individuals(project_id=None):
-    variants_query = g.beacon_query_parameters["variants_query"]
-    phenopacket_filters = g.beacon_query_parameters["phenopacket_filters"]
-    experiment_filters = g.beacon_query_parameters["experiment_filters"]
-    config_filters = g.beacon_query_parameters["config_filters"]
+    variants_query = g.beacon_query["variants_query"]
+    phenopacket_filters = g.beacon_query["phenopacket_filters"]
+    experiment_filters = g.beacon_query["experiment_filters"]
+    config_filters = g.beacon_query["config_filters"]
+    dataset_ids = g.beacon_query["dataset_ids"]
 
     no_query = not (variants_query or phenopacket_filters or experiment_filters or config_filters)
     search_sample_ids = variants_query or experiment_filters
diff --git a/bento_beacon/endpoints/variants.py b/bento_beacon/endpoints/variants.py
index 168c9bc..e35bd63 100644
--- a/bento_beacon/endpoints/variants.py
+++ b/bento_beacon/endpoints/variants.py
@@ -12,10 +12,10 @@
 @variants.route("/g_variants", methods=["GET", "POST"])
 @authz_middleware.deco_public_endpoint  # TODO: for now. eventually, return more depending on permissions
 async def get_variants():
-    variants_query = g.beacon_query_parameters["variants_query"]
-    phenopacket_filters = g.beacon_query_parameters["phenopacket_filters"]
-    experiment_filters = g.beacon_query_parameters["experiment_filters"]
-    config_filters = g.beacon_query_parameters["config_filters"]
+    variants_query = g.beacon_query["variants_query"]
+    phenopacket_filters = g.beacon_query["phenopacket_filters"]
+    experiment_filters = g.beacon_query["experiment_filters"]
+    config_filters = g.beacon_query["config_filters"]
     has_filters = phenopacket_filters or experiment_filters or config_filters
 
     # if no query, return total count of variants
diff --git a/bento_beacon/utils/beacon_request.py b/bento_beacon/utils/beacon_request.py
index 6333737..b84ba5c 100644
--- a/bento_beacon/utils/beacon_request.py
+++ b/bento_beacon/utils/beacon_request.py
@@ -27,7 +27,7 @@ def parse_query_params(request_data):
     phenopacket_filters = list(filter(lambda f: f["id"].startswith("phenopacket."), filters))
     experiment_filters = list(filter(lambda f: f["id"].startswith("experiment."), filters))
     config_filters = [f for f in filters if f not in phenopacket_filters and f not in experiment_filters]
-    datasetIds = request_data.get("datasets", {}).get("datasetIds") or []
+    dataset_ids = request_data.get("datasets", {}).get("datasetIds") or []
 
     # strip filter prefixes and convert remaining ids to bento format
     phenopacket_filters = list(
@@ -55,7 +55,7 @@ def parse_query_params(request_data):
         "phenopacket_filters": phenopacket_filters,
         "experiment_filters": experiment_filters,
         "config_filters": config_filters,
-        "datasetIds": datasetIds,
+        "dataset_ids": dataset_ids,
     }
 
 
@@ -146,7 +146,7 @@ async def save_request_data():
     g.request_data = request_data
 
     # parsed query components
-    g.beacon_query_parameters = parse_query_params(request_data)
+    g.beacon_query = parse_query_params(request_data)
 
 
 def validate_request():
diff --git a/bento_beacon/utils/censorship.py b/bento_beacon/utils/censorship.py
index 7ca2af3..30136e9 100644
--- a/bento_beacon/utils/censorship.py
+++ b/bento_beacon/utils/censorship.py
@@ -16,7 +16,7 @@ async def censorship_settings_lookup() -> tuple[int, int]:
     # we're "before request" but view args are available from request object
     view_args = request.view_args if request.view_args else {}
     project_id = view_args.get("project_id")
-    dataset_id = view_args.get("dataset_id")
+    dataset_id = request.args.get("datasetIds")
 
     max_filters, count_threshold = await katsu_censorship_settings(project_id=project_id, dataset_id=dataset_id)
     if max_filters is None or count_threshold is None:
@@ -82,11 +82,11 @@ async def censored_chart_data(data) -> list[dict[str, int]]:
 
 
 def query_has_phenopacket_filter() -> bool:
-    return bool(g.beacon_query_parameters["phenopacket_filters"])
+    return bool(g.beacon_query["phenopacket_filters"])
 
 
 def query_has_experiment_filter() -> bool:
-    return bool(g.beacon_query_parameters["experiment_filters"])
+    return bool(g.beacon_query["experiment_filters"])
 
 
 # some anonymous queries are not permitted