es_query.rs

extern crate serialize;

use std::collections::TreeMap;
use serialize::json::ToJson;
use serialize::json;

use super::case_folds::CaseFolds;
use super::{ Category, HotelsRequest };

#[deriving(Show)]
pub struct EsQuery {
  mhcs:     Option<Vec<String>>,
  tlcs:     Option<Vec<String>>,
  category: Option<Category>,
  size: i32
}

impl EsQuery {
  pub fn new(req: HotelsRequest) -> EsQuery {
    EsQuery { mhcs: req.mhcs, tlcs: req.tlcs, category: req.category, size: 5000 }
  }

  fn mhcs(&self) -> Option<Vec<String>> {
    match self.mhcs {
      Some(ref mhcs) => Some(mhcs.iter().map(|h| h.to_lowercase().to_string()).collect()),
      None => None
    }
  }

  fn tlcs(&self) -> Option<Vec<String>> {
    match self.tlcs {
      Some(ref tlcs) => Some(tlcs.iter().map(|h| h.to_lowercase().to_string()).collect()),
      None       => None
    }
  }
}

impl ToJson for EsQuery {
  fn to_json(&self) -> json::Json {

    fn j(name: &str, o: json::Json) -> json::Json {
      let mut t = TreeMap::new();
      t.insert(name.to_string(), o);
      t.to_json()
    };

    fn terms(name: &str, terms: &Vec<String>) -> json::Json {
      j("terms",
        j(name, terms.to_json())
      )
    };

    fn js(leafs: Vec<(&str, json::Json)>) -> json::Json {
      let mut d = TreeMap::new();
      for &(ref name, ref value) in leafs.iter() {
        d.insert(name.to_string(), value.to_json());
      }
      d.to_json()
    }

    fn ja(name: &str, leafs: Vec<json::Json>) -> json::Json {
      j(name, leafs.to_json())
    }

    let mut filters : Vec<json::Json> = vec!();
    if let Some(ref mhcs) = self.mhcs() {
      filters.push(terms("mhc", mhcs));
    }

    if let Some(ref tlcs) = self.tlcs() {
      filters.push(terms("tlc", tlcs));
    }

    let mut ands = vec!(j("or", j("filters", filters.to_json())));

    if let Some(ref category) = self.category {
      let mut d = TreeMap::new();

      if let Some(ref min) = category.min {
        d.insert("from".to_string(), min.to_json());
      }

      if let Some(ref max) = category.max {
        d.insert("to".to_string(), max.to_json());
      }

      ands.push(
        j("range",
          j("category", d.to_json())
        )
      );
    }

    js(vec!(
        ("size", self.size.to_json()),
        ("filter",
          j("and",
            ja("filters",ands)
          )
        )
      )
    )
  }
}

#[cfg(test)]
mod tests {
  use super::EsQuery;
  use super::super::HotelsRequest;
  use serialize::json::ToJson;

  #[test]
  fn mhcs() {
    let hotels   = vec!("PAR00DMD".to_string());

    let expected = r#"{"filter":{"and":{"filters":[{"or":{"filters":[{"terms":{"mhc":["par00dmd"]}}]}}]}},"size":5000}"#;
    let hr = HotelsRequest::empty().mhcs(hotels);
    let es_query = EsQuery::new(hr);
    assert_eq!(es_query.to_json().to_string().as_slice(), expected);
  }

  #[test]
  fn tlc() {
    let tlcs   = vec!("MUC".to_string());

    let expected = r#"{"filter":{"and":{"filters":[{"or":{"filters":[{"terms":{"tlc":["muc"]}}]}}]}},"size":5000}"#;
    let hr = HotelsRequest::empty().tlcs(tlcs);
    let es_query = EsQuery::new(hr);
    assert_eq!(es_query.to_json().to_string().as_slice(), expected);
  }
}