much cleaner queries

crates/lib/elasticsearch_schema/src/searches/search_model_weights.rs

@@ -2,10 +2,95 @@ use elasticsearch::{Elasticsearch, SearchParts};
 use serde_json::{json, Value};
 
 use enums::by_table::model_weights::weights_types::WeightsType;
-use errors::AnyhowResult;
+use errors::{anyhow, AnyhowResult};
 
 use crate::documents::model_weight_document::{MODEL_WEIGHT_INDEX, ModelWeightDocument};
 
+pub struct SearchModelWeightsQuery<'a> {
+  pub search_term: &'a str,
+  pub maybe_language_subtag: Option<&'a str>,
+  pub maybe_weights_type: Option<WeightsType>,
+}
+
+impl <'a>SearchModelWeightsQuery<'a> {
+
+  fn base_query(&self) -> Value {
+    json!({
+      "query": {
+        "bool": {
+          "must": [
+            {
+              "bool": {
+                "should": [
+                  {
+                    "fuzzy": {
+                      "title": {
+                        "value": self.search_term,
+                        "fuzziness": 2
+                      }
+                    }
+                  },
+                  {
+                    "match": {
+                      "title": {
+                        "query": self.search_term,
+                        "boost": 1
+                      }
+                    }
+                  },
+                  {
+                    "multi_match": {
+                      "query": self.search_term,
+                      "type": "bool_prefix",
+                      "fields": [
+                        "title",
+                        "title._2gram",
+                        "title._3gram"
+                      ],
+                      "boost": 50
+                    }
+                  }
+                ]
+              }
+            }
+          ]
+        }
+      }
+    })
+  }
+
+  pub fn query(&self) -> AnyhowResult<Value> {
+    let mut search_json = self.base_query();
+
+    let must_clause = search_json.pointer_mut("/query/bool/must")
+        .map(|pointer| pointer.as_array_mut())
+        .flatten()
+        .ok_or_else(|| anyhow!("could not get pointer to must clause"))?;
+
+    if let Some(maybe_language_subtag) = &self.maybe_language_subtag {
+      must_clause.push(json!(
+        {
+          "match": {
+            "maybe_ietf_primary_language_subtag": maybe_language_subtag
+          }
+        }
+      ));
+    }
+
+    if let Some(weights_type) = &self.maybe_weights_type {
+      must_clause.push(json!(
+        {
+          "match": {
+            "weights_type": weights_type.to_string(),
+          }
+        }
+      ));
+    }
+
+    Ok(search_json)
+  }
+}
+
 pub async fn search_model_weights(
   client: &Elasticsearch,
   search_term: &str,
@@ -13,12 +98,14 @@ pub async fn search_model_weights(
   maybe_weights_type: Option<WeightsType>,
 ) -> AnyhowResult<Vec<ModelWeightDocument>> {
 
-  let search_json = match (maybe_language_subtag, maybe_weights_type) {
-    (Some(language), None) => query_model_weights_with_required_language(search_term, language),
-    (None, Some(weights_type)) => query_model_weights_with_model_weights_type(search_term, weights_type),
-    _ => query_model_weights(search_term),
+  let query = SearchModelWeightsQuery {
+    search_term,
+    maybe_language_subtag,
+    maybe_weights_type,
   };
 
+  let search_json = query.query()?;
+
   let search_response = client
       .search(SearchParts::Index(&[MODEL_WEIGHT_INDEX]))
       .body(search_json)
@@ -102,6 +189,52 @@ fn query_model_weights(search_term: &str) -> Value {
   })
 }
 
+fn base_search(search_term: &str) -> Value {
+  json!({
+    "query": {
+      "bool": {
+        "must": [
+          {
+            "bool": {
+              "should": [
+                {
+                  "fuzzy": {
+                    "title": {
+                      "value": search_term,
+                      "fuzziness": 2
+                    }
+                  }
+                },
+                {
+                  "match": {
+                    "title": {
+                      "query": search_term,
+                      "boost": 1
+                    }
+                  }
+                },
+                {
+                  "multi_match": {
+                    "query": search_term,
+                    "type": "bool_prefix",
+                    "fields": [
+                      "title",
+                      "title._2gram",
+                      "title._3gram"
+                    ],
+                    "boost": 50
+                  }
+                }
+              ]
+            }
+          }
+        ]
+      }
+    }
+  })
+}
+
+
 fn query_model_weights_with_required_language(search_term: &str, language_tag: &str) -> Value {
   json!({
     "query": {
@@ -205,12 +338,31 @@ fn query_model_weights_with_model_weights_type(search_term: &str, weight_type: W
 #[cfg(test)]
 mod tests {
   use regex::Regex;
-  use crate::searches::search_model_weights::query_model_weights;
+  use enums::by_table::model_weights::weights_types::WeightsType;
+  use errors::AnyhowResult;
+  use crate::searches::search_model_weights::{query_model_weights, SearchModelWeightsQuery};
+
+  fn query_to_json_string(query: SearchModelWeightsQuery<'_>) -> AnyhowResult<String> {
+    let json = query.query()?;
+    let json = serde_json::to_string(&json)?;
+    Ok(json)
+  }
+
+  fn compact_json(json: &str) -> String {
+    let regex = Regex::new("\\s+").expect("regex should parse");
+    let json = regex.replace_all(&json, "");
+    json.to_string()
+  }
 
   #[test]
   fn default_search_only_keyword() {
-    let json = query_model_weights("FOO_BAR_BAZ");
-    let json = serde_json::to_string(&json).unwrap();
+    let query = SearchModelWeightsQuery {
+      search_term: "FOO_BAR_BAZ",
+      maybe_language_subtag: None,
+      maybe_weights_type: None,
+    };
+
+    let json = query_to_json_string(query).unwrap();
 
     // NB: Keys in emitted JSON are sorted.
     let expected_json = r#"
@@ -258,9 +410,78 @@ mod tests {
       }
     "#;
 
-    let regex = Regex::new("\\s+").expect("regex should parse");
+    let expected_json = compact_json(&expected_json);
 
-    let expected_json = regex.replace_all(&expected_json, "");
+    assert_eq!(&json, &expected_json);
+  }
+
+  #[test]
+  fn search_with_language_and_weights() {
+    let query = SearchModelWeightsQuery {
+      search_term: "FOO_BAR_BAZ",
+      maybe_language_subtag: Some("en"),
+      maybe_weights_type: Some(WeightsType::SoVitsSvc),
+    };
+
+    let json = query_to_json_string(query).unwrap();
+
+    // NB: Keys in emitted JSON are sorted.
+    let expected_json = r#"
+      {
+        "query": {
+          "bool": {
+            "must": [
+              {
+                "bool": {
+                  "should": [
+                    {
+                      "fuzzy": {
+                        "title": {
+                          "fuzziness": 2,
+                          "value": "FOO_BAR_BAZ"
+                        }
+                      }
+                    },
+                    {
+                      "match": {
+                        "title": {
+                          "boost": 1,
+                          "query": "FOO_BAR_BAZ"
+                        }
+                      }
+                    },
+                    {
+                      "multi_match": {
+                        "boost": 50,
+                        "fields": [
+                          "title",
+                          "title._2gram",
+                          "title._3gram"
+                        ],
+                        "query": "FOO_BAR_BAZ",
+                        "type": "bool_prefix"
+                      }
+                    }
+                  ]
+                }
+              },
+              {
+                "match": {
+                  "maybe_ietf_primary_language_subtag": "en"
+                }
+              },
+              {
+                "match": {
+                  "weights_type": "so_vits_svc"
+                }
+              }
+            ]
+          }
+        }
+      }
+    "#;
+
+    let expected_json = compact_json(&expected_json);
 
     assert_eq!(&json, &expected_json);
   }