elastic_query_builder/aggregation/
top_hits_aggregation.rs

1use crate::aggregation::AggregationTrait;
2use crate::merge;
3use serde::ser::SerializeStruct;
4use serde::{Serialize, Serializer};
5use serde_json::{json, Value};
6use std::collections::HashMap;
7
8#[derive(Default)]
9pub struct TopHitsAggregation {
10    name: String,
11    value: TopHitsValue,
12    aggregation: Value,
13}
14
15#[derive(Default)]
16struct TopHitsValue {
17    script: String,
18    size: i64,
19    sort: Vec<(String, String)>,
20}
21
22impl TopHitsAggregation {
23    pub fn new(name: &str) -> Self {
24        let mut term = TopHitsAggregation {
25            name: name.to_string(),
26            ..Default::default()
27        };
28        term.value.size = 10;
29        term
30    }
31    pub fn add_sort(mut self, field: &str, order: &str) -> Self {
32        self.value.sort.push((field.to_string(), order.to_string()));
33        self
34    }
35    pub fn set_size(mut self, size: i64) -> Self {
36        self.value.size = size;
37        self
38    }
39    pub fn set_script(mut self, script: &str) -> Self {
40        self.value.script = script.to_string();
41        self
42    }
43    pub fn set_aggregation<T>(mut self, aggregation: T) -> Self
44    where
45        T: AggregationTrait,
46    {
47        self.aggregation = aggregation.build();
48        self
49    }
50    pub fn append_aggregation<T>(mut self, query: T) -> Self
51    where
52        T: AggregationTrait,
53    {
54        let mut values = self.aggregation.clone();
55        merge(&mut values, &query.build());
56        self.aggregation = json!(values);
57        return self;
58    }
59}
60
61impl Serialize for TopHitsValue {
62    fn serialize<S>(&self, serializer: S) -> Result<<S as Serializer>::Ok, <S as Serializer>::Error>
63    where
64        S: Serializer,
65    {
66        let mut state = serializer.serialize_struct("TopHitsBuilder", 0)?;
67
68        if self.script.is_empty() {
69            state.serialize_field("size", &self.size)?;
70            if !self.sort.is_empty() {
71                let mut sort_map = HashMap::new();
72                for (field, order) in self.sort.iter() {
73                    sort_map.insert(field.to_string(), order.to_string());
74                }
75                state.serialize_field("sort", &sort_map)?;
76            }
77        } else {
78            let mut source = HashMap::new();
79            source.insert("includes", self.script.clone());
80            state.serialize_field("_source", &source)?;
81            state.serialize_field("size", &self.size)?;
82            // TODO:
83        }
84        state.end()
85    }
86}
87
88impl Serialize for TopHitsAggregation {
89    fn serialize<S>(&self, serializer: S) -> Result<<S as Serializer>::Ok, <S as Serializer>::Error>
90    where
91        S: Serializer,
92    {
93        let mut state = serializer.serialize_struct("TopHitsAggregation", 0)?;
94        state.serialize_field("top_hits", &self.value)?;
95        if !(self.aggregation.is_null() || self.aggregation.to_string().is_empty()) {
96            state.serialize_field("aggs", &self.aggregation)?;
97        }
98        // }
99        state.end()
100    }
101}
102
103impl AggregationTrait for TopHitsAggregation {
104    fn name(&self) -> &str {
105        self.name.as_str()
106    }
107
108    fn build(&self) -> Value {
109        let name = self.name.to_string();
110        json!({ name: self })
111    }
112
113    fn query_name(&self) -> String {
114        "top_hits".to_string()
115    }
116}
117
118#[cfg(test)]
119mod tests {
120    use super::*;
121    use crate::aggregation::AggregationTrait;
122
123    #[test]
124    fn test_terms_aggregation() {
125        let agg = TopHitsAggregation::new("hoge");
126
127        let json = agg.build();
128        println!("{}", json);
129        // assert_eq!(json["test"]["terms"]["field"], "test");
130        // assert_eq!(json["test"]["terms"]["size"], 1);
131    }
132}