From 9f8097cb14b7ae4385749b70948a50379d48b0af Mon Sep 17 00:00:00 2001 From: Enola Knezevic Date: Wed, 3 Apr 2024 17:05:04 +0200 Subject: [PATCH] getting counts from prism --- src/catalogue.rs | 81 ++++++++++++++++++++++++++++++------------------ src/config.rs | 8 +++-- src/main.rs | 2 +- 3 files changed, 58 insertions(+), 33 deletions(-) diff --git a/src/catalogue.rs b/src/catalogue.rs index 01f64ab..a222cd7 100644 --- a/src/catalogue.rs +++ b/src/catalogue.rs @@ -1,31 +1,24 @@ -use std::time::Duration; +use std::{collections::BTreeMap, time::Duration}; use reqwest::Url; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; use tracing::{debug, info}; -#[derive(Serialize, Deserialize)] -#[serde(rename_all = "UPPERCASE")] -enum ChildCategoryType { - Equals, - SomethingElse -} +pub type Criteria = BTreeMap; -#[derive(Serialize, Deserialize)] -struct ChildCategoryCriterion { - key: String, - count: Option -} +pub type CriteriaGroup = BTreeMap; -#[derive(Serialize, Deserialize)] -struct ChildCategory { - key: String, - type_: ChildCategoryType, - criteria: Vec +pub type CriteriaGroups = BTreeMap; + +fn get_element<'a>(count: &'a CriteriaGroups, key1: &'a str, key2: &'a str, key3: &'a str) -> Option<&'a u64> { + count.get(key1) + .and_then(|group| group.get(key2)) + .and_then(|criteria| criteria.get(key3)) } -pub async fn get_extended_json(catalogue_url: Url) -> Value { + +pub async fn get_extended_json(catalogue_url: Url, prism_url: Url) -> Value { debug!("Fetching catalogue from {catalogue_url} ..."); let resp = reqwest::Client::new() @@ -38,11 +31,25 @@ pub async fn get_extended_json(catalogue_url: Url) -> Value { let mut json: Value = resp.json().await .expect("Unable to parse catalogue from upstream; please check URL specified in config."); - // TODO: Query prism for counts here. - recurse(&mut json); + let prism_resp = reqwest::Client::new() + .post(format!("{}criteria", prism_url)) + .header("Content-Type", "application/json") + .body("{\"sites\": []}") + .timeout(Duration::from_secs(300)) + .send() + .await + .expect("Unable to fetch response from Prism; please check it's running."); + + let mut counts: CriteriaGroups = prism_resp.json().await + .expect("Unable to parse response from Prism into CriteriaGroups"); - // println!("{}", serde_json::to_string_pretty(&json).unwrap()); + + //dbg!(&counts); + + recurse(&mut json, &mut counts); //TODO remove from counts once copied into catalogue to make it O(n log n) + + //println!("{}", serde_json::to_string_pretty(&json).unwrap()); info!("Catalogue built successfully."); @@ -52,7 +59,7 @@ pub async fn get_extended_json(catalogue_url: Url) -> Value { /// Key order: group key (e.g. patient) /// \-- stratifier key (e.g. admin_gender) /// \-- stratum key (e.g. male, other) -fn recurse(json: &mut Value) { +fn recurse(json: &mut Value, counts: &mut CriteriaGroups) { match json { Value::Null => (), Value::Bool(_) => (), @@ -60,16 +67,22 @@ fn recurse(json: &mut Value) { Value::String(_) => (), Value::Array(arr) => { for ele in arr { - recurse(ele); + recurse(ele, counts); } }, Value::Object(obj) => { if ! obj.contains_key("childCategories") { for (_key, child_val) in obj.iter_mut() { - recurse(child_val); + recurse(child_val, counts); } } else { - let group_key = obj.get("key").expect("Got JSON element with childCategories but without (group) key. Please check json."); + let group_key = obj.get("key").expect("Got JSON element with childCategories but without (group) key. Please check json.").as_str() + .expect("Got JSON where a criterion key was not a string. Please check json.").to_owned(); + + let group_key = if group_key == "patient" {"patients"} + else if group_key == "tumor_classification" {"diagnosis"} + else if group_key == "biosamples" {"specimen"} + else {&group_key}; let children_cats = obj .get_mut("childCategories") @@ -80,7 +93,9 @@ fn recurse(json: &mut Value) { .filter(|item| item.get("type").unwrap_or(&Value::Null) == "EQUALS"); for child_cat in children_cats { - let stratifier_key = child_cat.get("key").expect("Got JSON element with childCategory that does not contain a (stratifier) key. Please check json."); + let stratifier_key = child_cat.get("key").expect("Got JSON element with childCategory that does not contain a (stratifier) key. Please check json.").as_str() + .expect("Got JSON where a criterion key was not a string. Please check json.").to_owned(); + let criteria = child_cat .get_mut("criteria") .expect("Got JSON element with childCategory that does not contain a criteria array. Please check json.") @@ -95,10 +110,16 @@ fn recurse(json: &mut Value) { .as_str() .expect("Got JSON where a criterion key was not a string. Please check json."); - // fetch from Prism output - let count_from_prism = 10; - - criterion.insert("count".into(), json!(count_from_prism)); + let count_from_prism = get_element(counts, &group_key, &stratifier_key, stratum_key); + + match count_from_prism { + Some(count) => { + criterion.insert("count".into(), json!(count)); + }, + None => { + debug!("No count from Prism for {}, {}, {}", group_key, stratifier_key, stratum_key); + } + } } } } diff --git a/src/config.rs b/src/config.rs index cdf1dbf..c14e3db 100644 --- a/src/config.rs +++ b/src/config.rs @@ -29,12 +29,16 @@ pub struct Config { pub project: Option, /// The socket address this server will bind to - #[clap(long, env, default_value = "0.0.0.0:8080")] + #[clap(long, env, default_value = "0.0.0.0:8055")] pub bind_addr: SocketAddr, /// URL to catalogue.json file #[clap(long, env)] - pub catalogue_url: Url + pub catalogue_url: Url, + + /// URL to prism + #[clap(long, env, default_value= "http://localhost:8066")] + pub prism_url: Url } fn parse_cors(v: &str) -> Result { diff --git a/src/main.rs b/src/main.rs index 5b72eea..3fad023 100644 --- a/src/main.rs +++ b/src/main.rs @@ -47,7 +47,7 @@ async fn main() { info!("{:#?}", Lazy::force(&CONFIG)); - let extended_json = catalogue::get_extended_json(CONFIG.catalogue_url.clone()).await; + let extended_json = catalogue::get_extended_json(CONFIG.catalogue_url.clone(), CONFIG.prism_url.clone()).await; let state = SharedState { extended_json }; // TODO: Add check for reachability of beam-proxy