populate criteria missing from the catalogue json

src/catalogue.rs

@@ -1,4 +1,6 @@
 use std::{collections::BTreeMap, sync::Arc, time::Duration};
+use serde::Deserialize;
+use serde::Serialize;
 
 use reqwest::Url;
 use serde_json::{json, Value};
@@ -11,36 +13,63 @@ pub type CriteriaGroup = BTreeMap<String, Criteria>;
 
 pub type CriteriaGroups = BTreeMap<String, CriteriaGroup>;
 
-fn get_element<'a>(count: &'a CriteriaGroups, key1: &'a str, key2: &'a str, key3: &'a str) -> Option<&'a u64> {
-    count.get(key1)
+#[derive(Serialize, Deserialize, Debug, Clone)]
+struct CatalogueCriterion {
+    key: String, 
+    name: String,
+    description: String,
+    count: u64
+}
+
+fn get_element<'a>(
+    counts: &'a CriteriaGroups,
+    key1: &'a str,
+    key2: &'a str,
+    key3: &'a str,
+) -> Option<&'a u64> {
+    counts
+        .get(key1)
         .and_then(|group| group.get(key2))
         .and_then(|criteria| criteria.get(key3))
 }
 
+fn get_stratifier<'a>(
+    counts: &'a CriteriaGroups,
+    key1: &'a str,
+    key2: &'a str
+) -> Option<&'a Criteria> {
+    counts
+        .get(key1)
+        .and_then(|group| group.get(key2))
+}
+
 pub fn spawn_thing(catalogue_url: Url, prism_url: Url) -> Arc<Mutex<Value>> {
     let thing: Arc<Mutex<Value>> = Arc::default();
     let thing1 = thing.clone();
     tokio::spawn(async move {
-        tokio::time::sleep(Duration::from_secs(10)).await;    
+        tokio::time::sleep(Duration::from_secs(10)).await;
         loop {
             match get_extended_json(catalogue_url.clone(), prism_url.clone()).await {
                 Ok(new_value) => {
                     *thing1.lock().await = new_value;
                     info!("Updated Catalogue!");
                     tokio::time::sleep(Duration::from_secs(60 * 60)).await;
-                },
+                }
                 Err(err) => {
                     warn!("Failed to get thing: {err}.\n Retrying in 5s.");
                     tokio::time::sleep(Duration::from_secs(5)).await;
-                },
+                }
             }
         }
     });
 
     thing
 }
 
-pub async fn get_extended_json(catalogue_url: Url, prism_url: Url) -> Result<Value, reqwest::Error> {
+pub async fn get_extended_json(
+    catalogue_url: Url,
+    prism_url: Url,
+) -> Result<Value, reqwest::Error> {
     debug!("Fetching catalogue from {catalogue_url} ...");
 
     let resp = reqwest::Client::new()
@@ -51,7 +80,6 @@ pub async fn get_extended_json(catalogue_url: Url, prism_url: Url) -> Result<Val
 
     let mut json: Value = resp.json().await?;
 
-
     let prism_resp = reqwest::Client::new()
         .post(format!("{}criteria", prism_url))
         .header("Content-Type", "application/json")
@@ -78,10 +106,9 @@ fn recurse(json: &mut Value, counts: &mut CriteriaGroups) {
             for ele in arr {
                 recurse(ele, counts);
             }
-        },
+        }
         Value::Object(obj) => {
-
-            if ! obj.contains_key("childCategories") {
+            if !obj.contains_key("childCategories") {
                 for (_key, child_val) in obj.iter_mut() {
                     recurse(child_val, counts);
                 }
@@ -90,10 +117,15 @@ fn recurse(json: &mut Value, counts: &mut CriteriaGroups) {
                 .expect("Got JSON where a criterion key was not a string. Please check json.").to_owned();
 
                 //TODO consolidate catalogue and MeasureReport group names
-                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 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")
@@ -102,7 +134,7 @@ fn recurse(json: &mut Value, counts: &mut CriteriaGroups) {
                     .unwrap()
                     .iter_mut()
                     .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.").as_str()
                     .expect("Got JSON where a criterion key was not a string. Please check json.").to_owned();
@@ -113,28 +145,50 @@ fn recurse(json: &mut Value, counts: &mut CriteriaGroups) {
                         .as_array_mut()
                         .expect("Got JSON element with childCategory with criteria that are not an array. Please check json.");
 
-                    for criterion in criteria {
-                        let criterion = criterion.as_object_mut()
-                            .expect("Got JSON where a criterion was not an object. Please check json.");
-                        let stratum_key = criterion.get("key")
+                    if !criteria.is_empty() {
+                        for criterion in criteria {
+                            let criterion = criterion.as_object_mut().expect(
+                                "Got JSON where a criterion was not an object. Please check json.",
+                            );
+                            let stratum_key = criterion.get("key")
                             .expect("Got JSON where a criterion did not have a key. Please check json.")
                             .as_str()
                             .expect("Got JSON where a criterion key was not a string. Please check json.");
-
-                        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);
+
+                            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
+                                    );
+                                }
+                            }
+                        }
+                    } else {
+                        let criteria_counts_maybe = get_stratifier(counts, &group_key, &stratifier_key);
+                        if let Some(criteria_counts) = criteria_counts_maybe{
+                            for criterion_count in criteria_counts {
+                                let (key, count) = (criterion_count.0, criterion_count.1);
+                                let catalogue_criterion = CatalogueCriterion {
+                                        key: key.clone(),
+                                        name: key.clone(),
+                                        description: "".into(),
+                                        count: count.clone()
+                                };
+                                criteria.push(json!(catalogue_criterion));
+
                             }
-                        }                        
+                        }
                     }
                 }
             }
-        },
+        }
         _ => {}
     }
 }