rs_es/operations/search/
count.rs

1/*
2 * Copyright 2015-2019 Ben Ashford
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *     http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17//! Implementations of the Count API
18
19use reqwest::StatusCode;
20
21use serde::{Deserialize, Serialize};
22
23use crate::{
24    error::EsError,
25    json::ShouldSkip,
26    operations::{
27        common::{OptionVal, Options},
28        format_indexes_and_types, ShardCountResult,
29    },
30    query::Query,
31    Client, EsResponse,
32};
33
34/// Representing a count operation
35#[derive(Debug)]
36pub struct CountURIOperation<'a, 'b> {
37    client: &'a mut Client,
38    indexes: &'b [&'b str],
39    doc_types: &'b [&'b str],
40    options: Options<'b>,
41}
42
43impl<'a, 'b> CountURIOperation<'a, 'b> {
44    pub fn new(client: &'a mut Client) -> CountURIOperation<'a, 'b> {
45        CountURIOperation {
46            client,
47            indexes: &[],
48            doc_types: &[],
49            options: Options::default(),
50        }
51    }
52
53    pub fn with_indexes(&'b mut self, indexes: &'b [&'b str]) -> &'b mut Self {
54        self.indexes = indexes;
55        self
56    }
57
58    pub fn with_types(&'b mut self, doc_types: &'b [&'b str]) -> &'b mut Self {
59        self.doc_types = doc_types;
60        self
61    }
62
63    pub fn with_query<S: Into<String>>(&'b mut self, qs: S) -> &'b mut Self {
64        self.options.push("q", qs.into());
65        self
66    }
67
68    add_option!(with_df, "df");
69    add_option!(with_analyzer, "analyzer");
70    add_option!(with_default_operator, "default_operator");
71    add_option!(with_lenient, "lenient");
72    add_option!(with_analyze_wildcard, "analyze_wildcard");
73    add_option!(with_terminate_after, "terminate_after");
74
75    pub fn send(&'b mut self) -> Result<CountResult, EsError> {
76        let url = format!(
77            "/{}/_count{}",
78            format_indexes_and_types(&self.indexes, &self.doc_types),
79            self.options
80        );
81        log::info!("Counting with: {}", url);
82        let response = self.client.get_op(&url)?;
83        match response.status_code() {
84            StatusCode::OK => Ok(response.read_response()?),
85            status_code => Err(EsError::EsError(format!(
86                "Unexpected status: {}",
87                status_code
88            ))),
89        }
90    }
91}
92
93#[derive(Debug, Default, Serialize)]
94struct CountQueryOperationBody<'b> {
95    /// The query
96    #[serde(skip_serializing_if = "ShouldSkip::should_skip")]
97    query: Option<&'b Query>,
98}
99
100#[derive(Debug)]
101pub struct CountQueryOperation<'a, 'b> {
102    /// The HTTP client
103    client: &'a mut Client,
104
105    /// The indexes to which this query applies
106    indexes: &'b [&'b str],
107
108    /// The types to which the query applies
109    doc_types: &'b [&'b str],
110
111    /// Optionals
112    options: Options<'b>,
113
114    /// The query body
115    body: CountQueryOperationBody<'b>,
116}
117
118impl<'a, 'b> CountQueryOperation<'a, 'b> {
119    pub fn new(client: &'a mut Client) -> Self {
120        CountQueryOperation {
121            client,
122            indexes: &[],
123            doc_types: &[],
124            options: Options::new(),
125            body: Default::default(),
126        }
127    }
128
129    pub fn with_indexes(&'b mut self, indexes: &'b [&'b str]) -> &'b mut Self {
130        self.indexes = indexes;
131        self
132    }
133
134    pub fn with_types(&'b mut self, doc_types: &'b [&'b str]) -> &'b mut Self {
135        self.doc_types = doc_types;
136        self
137    }
138
139    pub fn with_query(&'b mut self, query: &'b Query) -> &'b mut Self {
140        self.body.query = Some(query);
141        self
142    }
143
144    add_option!(with_df, "df");
145    add_option!(with_analyzer, "analyzer");
146    add_option!(with_default_operator, "default_operator");
147    add_option!(with_lenient, "lenient");
148    add_option!(with_analyze_wildcard, "analyze_wildcard");
149    add_option!(with_terminate_after, "terminate_after");
150
151    /// Performs the count with the specified query and options
152    pub fn send(&'b mut self) -> Result<CountResult, EsError> {
153        let url = format!(
154            "/{}/_count{}",
155            format_indexes_and_types(&self.indexes, &self.doc_types),
156            self.options
157        );
158        let response = self.client.post_body_op(&url, &self.body)?;
159        match response.status_code() {
160            StatusCode::OK => Ok(response.read_response()?),
161            status_code => Err(EsError::EsError(format!(
162                "Unexpected status: {}",
163                status_code
164            ))),
165        }
166    }
167}
168
169impl Client {
170    /// Count via the query parameter
171    ///
172    /// See: https://www.elastic.co/guide/en/elasticsearch/reference/1.x/search-uri-request.html
173    pub fn count_uri(&mut self) -> CountURIOperation {
174        CountURIOperation::new(self)
175    }
176
177    /// Count via the query DSL
178    ///
179    /// See: https://www.elastic.co/guide/en/elasticsearch/reference/1.x/search-request-body.html
180    pub fn count_query(&mut self) -> CountQueryOperation {
181        CountQueryOperation::new(self)
182    }
183}
184
185#[derive(Debug, Deserialize)]
186pub struct CountResult {
187    pub count: u64,
188
189    #[serde(rename = "_shards")]
190    pub shards: ShardCountResult,
191}
192
193#[cfg(test)]
194mod tests {
195    use crate::tests::setup_test_data;
196    use crate::tests::{clean_db, make_client};
197
198    use crate::query::Query;
199
200    use super::CountResult;
201
202    #[test]
203    fn test_count_uri() {
204        let index_name = "test_count_uri";
205        let mut client = make_client();
206
207        clean_db(&mut client, index_name);
208        setup_test_data(&mut client, index_name);
209
210        let all_results: CountResult = client
211            .count_uri()
212            .with_indexes(&[index_name])
213            .send()
214            .unwrap();
215        assert_eq!(3, all_results.count);
216
217        let doc_1: CountResult = client
218            .count_uri()
219            .with_indexes(&[index_name])
220            .with_query("str_field:1ABC")
221            .send()
222            .unwrap();
223        assert_eq!(1, doc_1.count);
224
225        let not_found_doc: CountResult = client
226            .count_uri()
227            .with_indexes(&[index_name])
228            .with_query("str_field:lolol")
229            .send()
230            .unwrap();
231        assert_eq!(0, not_found_doc.count);
232    }
233
234    #[test]
235    fn test_count_query() {
236        let index_name = "test_count_query";
237        let mut client = make_client();
238
239        clean_db(&mut client, index_name);
240        setup_test_data(&mut client, index_name);
241
242        let all_results: CountResult = client
243            .count_query()
244            .with_indexes(&[index_name])
245            .with_query(&Query::build_match_all().build())
246            .send()
247            .unwrap();
248        assert_eq!(3, all_results.count);
249
250        let doc_1: CountResult = client
251            .count_query()
252            .with_indexes(&[index_name])
253            .with_query(
254                &Query::build_range("int_field")
255                    .with_gte(2)
256                    .with_lte(3)
257                    .build(),
258            )
259            .send()
260            .unwrap();
261        assert_eq!(2, doc_1.count);
262
263        let not_found_doc: CountResult = client
264            .count_query()
265            .with_indexes(&[index_name])
266            .with_query(&Query::build_range("int_field").with_gte(99).build())
267            .send()
268            .unwrap();
269        assert_eq!(0, not_found_doc.count);
270    }
271}
rs_es/operations/search/count.rs

rs_es/operations/search/
count.rs
