elastic_query_builder/aggregation/
terms_aggregation.rs

1use crate::aggregation::AggregationTrait;
2use crate::merge;
3use serde::ser::{SerializeMap, SerializeStruct};
4use serde::{Serialize, Serializer};
5use serde_json::{json, Value};
6
7#[derive(Default)]
8pub struct TermsAggregation {
9    name: String,
10    // aggregation: Option<T>,
11    value: TermsValue,
12    aggregation: Value,
13}
14
15#[derive(Default, Clone)]
16struct TermsOrder {
17    order_field: String,
18    order_value: String,
19}
20
21#[derive(Default)]
22struct TermsValue {
23    script: String,
24    field: String,
25    size: i64,
26    order: Option<TermsOrder>,
27    include: Vec<String>,
28    exclude: Vec<String>,
29}
30
31impl TermsAggregation {
32    pub fn new(name: &str) -> Self {
33        let mut term = TermsAggregation {
34            name: name.to_string(),
35            ..Default::default()
36        };
37        term.value.size = 10;
38        term
39    }
40    pub fn set_field(mut self, field: &str) -> Self {
41        self.value.field = field.to_string();
42        self
43    }
44    pub fn set_size(mut self, size: i64) -> Self {
45        self.value.size = size;
46        self
47    }
48    pub fn set_script(mut self, script: &str) -> Self {
49        self.value.script = script.to_string();
50        self
51    }
52    pub fn set_order(mut self, order_field: &str, order_value: &str) -> Self {
53        self.value.order = Some(TermsOrder {
54            order_field: order_field.to_string(),
55            order_value: order_value.to_string(),
56        });
57        self
58    }
59    pub fn set_include(mut self, include: &[&str]) -> Self {
60        self.value.include = include.iter().map(|q| q.to_string()).collect();
61        self
62    }
63    pub fn set_exclude(mut self, exclude: &[&str]) -> Self {
64        self.value.exclude = exclude.iter().map(|q| q.to_string()).collect();
65        self
66    }
67    pub fn set_aggregation<T>(mut self, aggregation: T) -> Self
68        where
69            T: AggregationTrait,
70    {
71        self.aggregation = aggregation.build();
72        self
73    }
74    pub fn append_aggregation<T>(mut self, query: T) -> Self
75        where
76            T: AggregationTrait,
77    {
78        let mut values = self.aggregation.clone();
79        merge(&mut values, &query.build());
80        self.aggregation = json!(values);
81        return self;
82    }
83}
84
85impl Serialize for TermsValue {
86    fn serialize<S>(&self, serializer: S) -> Result<<S as Serializer>::Ok, <S as Serializer>::Error>
87        where
88            S: Serializer,
89    {
90        let mut state = serializer.serialize_struct("TermsBuilder", 0)?;
91
92        if self.script.is_empty() {
93            state.serialize_field("field", &self.field)?;
94            state.serialize_field("size", &self.size)?;
95        } else {
96            // TODO:
97        }
98
99        if self.order.is_some() {
100            state.serialize_field("order", &self.order.as_ref().unwrap())?;
101        }
102        if !self.include.is_empty() {
103            state.serialize_field("include", &self.include)?;
104        }
105        if !self.exclude.is_empty() {
106            state.serialize_field("exclude", &self.exclude)?;
107        }
108        state.end()
109    }
110}
111
112impl Serialize for TermsOrder {
113    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
114        where
115            S: Serializer,
116    {
117        let mut map = serializer.serialize_map(Some(1))?;
118        map.serialize_entry(&self.order_field, &self.order_value)?;
119        map.end()
120    }
121}
122
123impl Serialize for TermsAggregation {
124    fn serialize<S>(&self, serializer: S) -> Result<<S as Serializer>::Ok, <S as Serializer>::Error>
125        where
126            S: Serializer,
127    {
128        let mut state = serializer.serialize_struct("BoolQuery", 0)?;
129        state.serialize_field("terms", &self.value)?;
130        if !(self.aggregation.is_null() || self.aggregation.to_string().is_empty()) {
131            state.serialize_field("aggs", &self.aggregation)?;
132        }
133        // }
134        state.end()
135    }
136}
137
138// impl Serialize for TermsOrder {
139//     fn serialize<S>(&self, serializer: S) -> Result<<S as Serializer>::Ok, <S as Serializer>::Error>
140//         where
141//             S: Serializer,
142//     {
143//         let mut state = serializer.serialize_struct("BoolQuery", 0)?;
144//         state.serialize_field(&self.order_field.to_string(), &self.order_value.to_string());
145//         state.end()
146//     }
147// }
148
149impl AggregationTrait for TermsAggregation {
150    fn name(&self) -> &str {
151        self.name.as_str()
152    }
153
154    fn build(&self) -> Value {
155        let name = self.name.to_string();
156        json!({ name: self })
157    }
158
159    fn query_name(&self) -> String {
160        "terms".to_string()
161    }
162}
163
164#[cfg(test)]
165mod tests {
166    use super::*;
167    use crate::aggregation::AggregationTrait;
168
169    #[test]
170    fn test_terms_aggregation() {
171        let agg = TermsAggregation::new("hoge")
172            .set_field("field_name")
173            .set_include(&["field_a", "field_b"])
174            .set_order("field_a","asc")
175            .set_aggregation(TermsAggregation::new("agg"));
176
177        let json = agg.build();
178        println!("{}", json);
179        // assert_eq!(json["test"]["terms"]["field"], "test");
180        // assert_eq!(json["test"]["terms"]["size"], 1);
181    }
182}