1use rustlavel_core::Json;
31use std::collections::BTreeMap;
32
33#[derive(Debug, Clone, PartialEq)]
38pub struct Query(Json);
39
40impl Query {
41 pub fn match_all() -> Query {
43 Query(clause("match_all", empty()))
44 }
45
46 pub fn match_none() -> Query {
49 Query(clause("match_none", empty()))
50 }
51
52 pub fn matching(field: impl Into<String>, text: impl Into<String>) -> Match {
58 Match {
59 kind: "match",
60 field: field.into(),
61 text: text.into(),
62 options: BTreeMap::new(),
63 }
64 }
65
66 pub fn match_phrase(field: impl Into<String>, text: impl Into<String>) -> Match {
68 Match {
69 kind: "match_phrase",
70 field: field.into(),
71 text: text.into(),
72 options: BTreeMap::new(),
73 }
74 }
75
76 pub fn multi_match<S: Into<String>>(
81 text: impl Into<String>,
82 fields: impl IntoIterator<Item = S>,
83 ) -> Query {
84 let fields: Vec<Json> = fields.into_iter().map(|f| Json::String(f.into())).collect();
85 Query(clause(
86 "multi_match",
87 Json::object([("query", Json::String(text.into())), ("fields", Json::Array(fields))]),
88 ))
89 }
90
91 pub fn term(field: impl Into<String>, value: impl Into<Json>) -> Query {
99 Query(clause(
100 "term",
101 Json::object([(field.into(), Json::object([("value", value.into())]))]),
102 ))
103 }
104
105 pub fn terms<V: Into<Json>>(
107 field: impl Into<String>,
108 values: impl IntoIterator<Item = V>,
109 ) -> Query {
110 let values: Vec<Json> = values.into_iter().map(Into::into).collect();
111 Query(clause("terms", Json::object([(field.into(), Json::Array(values))])))
112 }
113
114 pub fn exists(field: impl Into<String>) -> Query {
120 Query(clause("exists", Json::object([("field", Json::String(field.into()))])))
121 }
122
123 pub fn prefix(field: impl Into<String>, value: impl Into<String>) -> Query {
125 Query(clause(
126 "prefix",
127 Json::object([(field.into(), Json::object([("value", Json::String(value.into()))]))]),
128 ))
129 }
130
131 pub fn wildcard(field: impl Into<String>, pattern: impl Into<String>) -> Query {
134 Query(clause(
135 "wildcard",
136 Json::object([(field.into(), Json::object([("value", Json::String(pattern.into()))]))]),
137 ))
138 }
139
140 pub fn ids<S: Into<String>>(ids: impl IntoIterator<Item = S>) -> Query {
142 let values: Vec<Json> = ids.into_iter().map(|id| Json::String(id.into())).collect();
143 Query(clause("ids", Json::object([("values", Json::Array(values))])))
144 }
145
146 pub fn range(field: impl Into<String>) -> Range {
148 Range { field: field.into(), bounds: BTreeMap::new() }
149 }
150
151 pub fn bool() -> BoolQuery {
153 BoolQuery::default()
154 }
155
156 pub fn all_of<Q: Into<Query>>(clauses: impl IntoIterator<Item = Q>) -> BoolQuery {
158 let mut combined = BoolQuery::default();
159 for clause in clauses {
160 combined = combined.must(clause);
161 }
162 combined
163 }
164
165 pub fn any_of<Q: Into<Query>>(clauses: impl IntoIterator<Item = Q>) -> BoolQuery {
172 let mut combined = BoolQuery::default();
173 for clause in clauses {
174 combined = combined.should(clause);
175 }
176 combined.minimum_should_match(1)
177 }
178
179 pub fn raw(json: Json) -> Query {
181 Query(json)
182 }
183
184 pub fn json(&self) -> &Json {
185 &self.0
186 }
187
188 pub fn into_json(self) -> Json {
189 self.0
190 }
191}
192
193#[derive(Debug, Clone, PartialEq)]
195pub struct Match {
196 kind: &'static str,
197 field: String,
198 text: String,
199 options: BTreeMap<String, Json>,
200}
201
202impl Match {
203 pub fn and(self) -> Match {
210 self.option("operator", Json::String("and".into()))
211 }
212
213 pub fn fuzzy(self) -> Match {
216 self.fuzziness("AUTO")
217 }
218
219 pub fn fuzziness(self, fuzziness: impl Into<String>) -> Match {
221 self.option("fuzziness", Json::String(fuzziness.into()))
222 }
223
224 pub fn boost(self, boost: f64) -> Match {
226 self.option("boost", Json::Number(boost))
227 }
228
229 pub fn minimum_should_match(self, value: impl Into<Json>) -> Match {
231 self.option("minimum_should_match", value.into())
232 }
233
234 pub fn option(mut self, name: impl Into<String>, value: Json) -> Match {
236 self.options.insert(name.into(), value);
237 self
238 }
239}
240
241impl From<Match> for Query {
242 fn from(value: Match) -> Query {
243 let mut body = BTreeMap::new();
247 body.insert("query".to_string(), Json::String(value.text));
248 body.extend(value.options);
249
250 Query(clause(value.kind, Json::object([(value.field, Json::Object(body))])))
251 }
252}
253
254#[derive(Debug, Clone, PartialEq)]
256pub struct Range {
257 field: String,
258 bounds: BTreeMap<String, Json>,
259}
260
261impl Range {
262 pub fn gte(self, value: impl Into<Json>) -> Range {
264 self.bound("gte", value)
265 }
266
267 pub fn gt(self, value: impl Into<Json>) -> Range {
269 self.bound("gt", value)
270 }
271
272 pub fn lte(self, value: impl Into<Json>) -> Range {
274 self.bound("lte", value)
275 }
276
277 pub fn lt(self, value: impl Into<Json>) -> Range {
279 self.bound("lt", value)
280 }
281
282 pub fn format(self, format: impl Into<String>) -> Range {
285 self.bound("format", Json::String(format.into()))
286 }
287
288 pub fn time_zone(self, zone: impl Into<String>) -> Range {
293 self.bound("time_zone", Json::String(zone.into()))
294 }
295
296 fn bound(mut self, name: &str, value: impl Into<Json>) -> Range {
297 self.bounds.insert(name.to_string(), value.into());
298 self
299 }
300}
301
302impl From<Range> for Query {
303 fn from(value: Range) -> Query {
304 Query(clause("range", Json::object([(value.field, Json::Object(value.bounds))])))
305 }
306}
307
308#[derive(Debug, Clone, Default, PartialEq)]
319pub struct BoolQuery {
320 must: Vec<Query>,
321 should: Vec<Query>,
322 must_not: Vec<Query>,
323 filter: Vec<Query>,
324 minimum_should_match: Option<Json>,
325}
326
327impl BoolQuery {
328 pub fn must(mut self, query: impl Into<Query>) -> BoolQuery {
330 self.must.push(query.into());
331 self
332 }
333
334 pub fn should(mut self, query: impl Into<Query>) -> BoolQuery {
337 self.should.push(query.into());
338 self
339 }
340
341 pub fn must_not(mut self, query: impl Into<Query>) -> BoolQuery {
343 self.must_not.push(query.into());
344 self
345 }
346
347 pub fn filter(mut self, query: impl Into<Query>) -> BoolQuery {
352 self.filter.push(query.into());
353 self
354 }
355
356 pub fn minimum_should_match(mut self, value: impl Into<Json>) -> BoolQuery {
359 self.minimum_should_match = Some(value.into());
360 self
361 }
362}
363
364impl From<BoolQuery> for Query {
365 fn from(value: BoolQuery) -> Query {
366 let mut body = BTreeMap::new();
367
368 for (name, clauses) in [
372 ("must", value.must),
373 ("should", value.should),
374 ("must_not", value.must_not),
375 ("filter", value.filter),
376 ] {
377 if !clauses.is_empty() {
378 body.insert(
379 name.to_string(),
380 Json::Array(clauses.into_iter().map(Query::into_json).collect()),
381 );
382 }
383 }
384
385 if let Some(minimum) = value.minimum_should_match {
386 body.insert("minimum_should_match".to_string(), minimum);
387 }
388
389 Query(clause("bool", Json::Object(body)))
390 }
391}
392
393#[derive(Debug, Clone, Copy, PartialEq, Eq)]
395pub enum Order {
396 Asc,
397 Desc,
398}
399
400impl Order {
401 fn as_str(self) -> &'static str {
402 match self {
403 Order::Asc => "asc",
404 Order::Desc => "desc",
405 }
406 }
407}
408
409#[derive(Debug, Clone, PartialEq)]
415pub struct Aggregation {
416 kind: String,
417 options: BTreeMap<String, Json>,
418 subs: BTreeMap<String, Aggregation>,
419}
420
421impl Aggregation {
422 pub fn terms(field: impl Into<String>) -> Aggregation {
429 Aggregation::on("terms", field)
430 }
431
432 pub fn stats(field: impl Into<String>) -> Aggregation {
434 Aggregation::on("stats", field)
435 }
436
437 pub fn min(field: impl Into<String>) -> Aggregation {
438 Aggregation::on("min", field)
439 }
440
441 pub fn max(field: impl Into<String>) -> Aggregation {
442 Aggregation::on("max", field)
443 }
444
445 pub fn avg(field: impl Into<String>) -> Aggregation {
446 Aggregation::on("avg", field)
447 }
448
449 pub fn sum(field: impl Into<String>) -> Aggregation {
450 Aggregation::on("sum", field)
451 }
452
453 pub fn cardinality(field: impl Into<String>) -> Aggregation {
456 Aggregation::on("cardinality", field)
457 }
458
459 pub fn value_count(field: impl Into<String>) -> Aggregation {
461 Aggregation::on("value_count", field)
462 }
463
464 pub fn date_histogram(field: impl Into<String>, interval: impl Into<String>) -> Aggregation {
470 Aggregation::on("date_histogram", field)
471 .option("calendar_interval", Json::String(interval.into()))
472 }
473
474 pub fn raw(kind: impl Into<String>, body: Json) -> Aggregation {
476 let options = match body {
477 Json::Object(map) => map,
478 other => BTreeMap::from([("value".to_string(), other)]),
479 };
480 Aggregation { kind: kind.into(), options, subs: BTreeMap::new() }
481 }
482
483 pub fn size(self, size: u32) -> Aggregation {
490 self.option("size", Json::from(size))
491 }
492
493 pub fn min_doc_count(self, count: u64) -> Aggregation {
495 self.option("min_doc_count", Json::from(count))
496 }
497
498 pub fn missing(self, value: impl Into<Json>) -> Aggregation {
501 self.option("missing", value.into())
502 }
503
504 pub fn option(mut self, name: impl Into<String>, value: Json) -> Aggregation {
506 self.options.insert(name.into(), value);
507 self
508 }
509
510 pub fn aggregate(mut self, name: impl Into<String>, sub: impl Into<Aggregation>) -> Aggregation {
512 self.subs.insert(name.into(), sub.into());
513 self
514 }
515
516 pub fn body(&self) -> Json {
518 let mut body = BTreeMap::new();
519 body.insert(self.kind.clone(), Json::Object(self.options.clone()));
520
521 if !self.subs.is_empty() {
522 let subs =
523 self.subs.iter().map(|(name, agg)| (name.clone(), agg.body())).collect();
524 body.insert("aggs".to_string(), Json::Object(subs));
525 }
526
527 Json::Object(body)
528 }
529
530 fn on(kind: &str, field: impl Into<String>) -> Aggregation {
531 Aggregation {
532 kind: kind.to_string(),
533 options: BTreeMap::from([("field".to_string(), Json::String(field.into()))]),
534 subs: BTreeMap::new(),
535 }
536 }
537}
538
539#[derive(Debug, Clone, Default, PartialEq)]
541enum Source {
542 #[default]
543 All,
544 None,
545 Some {
546 includes: Vec<String>,
547 excludes: Vec<String>,
548 },
549}
550
551#[derive(Debug, Clone, Default, PartialEq)]
553pub struct Search {
554 query: Option<Query>,
555 from: Option<u64>,
556 size: Option<u64>,
557 sort: Vec<Json>,
558 source: Source,
559 aggregations: BTreeMap<String, Aggregation>,
560 track_total_hits: Option<Json>,
561 min_score: Option<f64>,
562}
563
564impl Search {
565 pub fn new() -> Search {
567 Search::default()
568 }
569
570 pub fn query_only(query: impl Into<Query>) -> Search {
572 Search::new().query(query)
573 }
574
575 pub fn query(mut self, query: impl Into<Query>) -> Search {
576 self.query = Some(query.into());
577 self
578 }
579
580 pub fn from(mut self, from: u64) -> Search {
588 self.from = Some(from);
589 self
590 }
591
592 pub fn size(mut self, size: u64) -> Search {
597 self.size = Some(size);
598 self
599 }
600
601 pub fn sort_by(mut self, field: impl Into<String>, order: Order) -> Search {
602 self.sort.push(Json::object([(
603 field.into(),
604 Json::object([("order", Json::String(order.as_str().into()))]),
605 )]));
606 self
607 }
608
609 pub fn sort_asc(self, field: impl Into<String>) -> Search {
610 self.sort_by(field, Order::Asc)
611 }
612
613 pub fn sort_desc(self, field: impl Into<String>) -> Search {
614 self.sort_by(field, Order::Desc)
615 }
616
617 pub fn sort_by_score(self) -> Search {
623 self.sort_by("_score", Order::Desc)
624 }
625
626 pub fn source<S: Into<String>>(mut self, fields: impl IntoIterator<Item = S>) -> Search {
628 let includes: Vec<String> = fields.into_iter().map(Into::into).collect();
629 self.source = match self.source {
630 Source::Some { excludes, .. } => Source::Some { includes, excludes },
631 _ => Source::Some { includes, excludes: Vec::new() },
632 };
633 self
634 }
635
636 pub fn source_excluding<S: Into<String>>(
638 mut self,
639 fields: impl IntoIterator<Item = S>,
640 ) -> Search {
641 let excludes: Vec<String> = fields.into_iter().map(Into::into).collect();
642 self.source = match self.source {
643 Source::Some { includes, .. } => Source::Some { includes, excludes },
644 _ => Source::Some { includes: Vec::new(), excludes },
645 };
646 self
647 }
648
649 pub fn without_source(mut self) -> Search {
651 self.source = Source::None;
652 self
653 }
654
655 pub fn aggregate(mut self, name: impl Into<String>, aggregation: impl Into<Aggregation>) -> Search {
656 self.aggregations.insert(name.into(), aggregation.into());
657 self
658 }
659
660 pub fn track_total_hits(mut self, exact: bool) -> Search {
670 self.track_total_hits = Some(Json::Bool(exact));
671 self
672 }
673
674 pub fn track_total_hits_up_to(mut self, limit: u64) -> Search {
679 self.track_total_hits = Some(Json::from(limit));
680 self
681 }
682
683 pub fn min_score(mut self, score: f64) -> Search {
685 self.min_score = Some(score);
686 self
687 }
688
689 pub fn body(&self) -> Json {
691 let mut body = BTreeMap::new();
692
693 body.insert("query".to_string(), self.query_json());
694
695 if let Some(from) = self.from {
696 body.insert("from".to_string(), Json::from(from));
697 }
698 if let Some(size) = self.size {
699 body.insert("size".to_string(), Json::from(size));
700 }
701 if !self.sort.is_empty() {
702 body.insert("sort".to_string(), Json::Array(self.sort.clone()));
703 }
704 match &self.source {
705 Source::All => {}
706 Source::None => {
707 body.insert("_source".to_string(), Json::Bool(false));
708 }
709 Source::Some { includes, excludes } if excludes.is_empty() => {
710 body.insert("_source".to_string(), strings(includes));
711 }
712 Source::Some { includes, excludes } => {
713 body.insert(
714 "_source".to_string(),
715 Json::object([
716 ("includes", strings(includes)),
717 ("excludes", strings(excludes)),
718 ]),
719 );
720 }
721 }
722 if !self.aggregations.is_empty() {
723 let aggs = self
724 .aggregations
725 .iter()
726 .map(|(name, agg)| (name.clone(), agg.body()))
727 .collect();
728 body.insert("aggs".to_string(), Json::Object(aggs));
729 }
730 if let Some(track) = &self.track_total_hits {
731 body.insert("track_total_hits".to_string(), track.clone());
732 }
733 if let Some(score) = self.min_score {
734 body.insert("min_score".to_string(), Json::Number(score));
735 }
736
737 Json::Object(body)
738 }
739
740 pub fn count_body(&self) -> Json {
745 Json::object([("query", self.query_json())])
746 }
747
748 fn query_json(&self) -> Json {
749 self.query.as_ref().map(|q| q.0.clone()).unwrap_or_else(|| clause("match_all", empty()))
750 }
751}
752
753fn clause(name: &str, body: Json) -> Json {
754 Json::object([(name, body)])
755}
756
757fn empty() -> Json {
758 Json::Object(BTreeMap::new())
759}
760
761fn strings(values: &[String]) -> Json {
762 Json::Array(values.iter().map(|v| Json::String(v.clone())).collect())
763}
764
765#[cfg(test)]
766mod tests {
767 use super::*;
768
769 fn body_of(search: &Search) -> String {
770 search.body().to_string()
771 }
772
773 #[test]
774 fn an_empty_search_is_match_all_rather_than_an_empty_body() {
775 assert_eq!(body_of(&Search::new()), r#"{"query":{"match_all":{}}}"#);
776 }
777
778 #[test]
779 fn a_match_always_uses_the_expanded_form_so_options_have_somewhere_to_go() {
780 let plain: Query = Query::matching("title", "rust").into();
781 assert_eq!(plain.json().to_string(), r#"{"match":{"title":{"query":"rust"}}}"#);
782
783 let strict: Query = Query::matching("title", "rust web").and().fuzzy().into();
784 assert_eq!(
785 strict.json().to_string(),
786 r#"{"match":{"title":{"fuzziness":"AUTO","operator":"and","query":"rust web"}}}"#
787 );
788 }
789
790 #[test]
791 fn a_term_query_is_the_expanded_form_too() {
792 let query = Query::term("status", "published");
793 assert_eq!(query.json().to_string(), r#"{"term":{"status":{"value":"published"}}}"#);
794 assert_eq!(
795 Query::term("archived", true).json().to_string(),
796 r#"{"term":{"archived":{"value":true}}}"#
797 );
798 }
799
800 #[test]
801 fn terms_and_ids_take_a_list() {
802 assert_eq!(
803 Query::terms("tags", ["rust", "web"]).json().to_string(),
804 r#"{"terms":{"tags":["rust","web"]}}"#
805 );
806 assert_eq!(
807 Query::ids(["1", "2"]).json().to_string(),
808 r#"{"ids":{"values":["1","2"]}}"#
809 );
810 }
811
812 #[test]
813 fn a_range_collects_its_bounds_under_the_field() {
814 let query: Query = Query::range("age").gte(18).lt(65).into();
815 assert_eq!(query.json().to_string(), r#"{"range":{"age":{"gte":18,"lt":65}}}"#);
816 }
817
818 #[test]
819 fn a_bool_leaves_out_the_slots_nobody_filled() {
820 let query: Query = Query::bool()
821 .must(Query::matching("title", "rust"))
822 .filter(Query::term("status", "published"))
823 .into();
824
825 let json = query.json().to_string();
826 assert!(json.contains(r#""must":[{"match""#), "got {json}");
827 assert!(json.contains(r#""filter":[{"term""#), "got {json}");
828 assert!(!json.contains("must_not"), "an empty slot reads as a bug in a log: {json}");
829 assert!(!json.contains("should"), "got {json}");
830 }
831
832 #[test]
833 fn any_of_makes_should_a_requirement_because_on_its_own_it_is_not_one() {
834 let bare: Query = Query::bool()
838 .must(Query::match_all())
839 .should(Query::term("tags", "rust"))
840 .into();
841 assert!(!bare.json().to_string().contains("minimum_should_match"));
842
843 let required: Query =
844 Query::any_of([Query::term("tags", "rust"), Query::term("tags", "web")])
845 .must(Query::match_all())
846 .into();
847 assert!(
848 required.json().to_string().contains(r#""minimum_should_match":1"#),
849 "got {required:?}"
850 );
851 }
852
853 #[test]
854 fn all_of_puts_everything_in_must() {
855 let query: Query =
856 Query::all_of([Query::term("a", 1), Query::term("b", 2)]).into();
857 let json = query.json().to_string();
858 assert!(json.contains(r#""must":[{"term":{"a""#), "got {json}");
859 assert!(!json.contains("should"), "got {json}");
860 }
861
862 #[test]
863 fn a_nested_bool_is_just_a_query_in_a_slot() {
864 let query: Query = Query::bool()
865 .filter(Query::any_of([Query::term("tags", "rust"), Query::term("tags", "web")]))
866 .into();
867
868 assert_eq!(
869 query.json().to_string(),
870 r#"{"bool":{"filter":[{"bool":{"minimum_should_match":1,"should":[{"term":{"tags":{"value":"rust"}}},{"term":{"tags":{"value":"web"}}}]}}]}}"#
871 );
872 }
873
874 #[test]
875 fn paging_and_sorting_land_beside_the_query_and_not_inside_it() {
876 let search = Search::new().from(20).size(10).sort_desc("published_at").sort_asc("title");
879
880 assert_eq!(
881 body_of(&search),
882 r#"{"from":20,"query":{"match_all":{}},"size":10,"sort":[{"published_at":{"order":"desc"}},{"title":{"order":"asc"}}]}"#
883 );
884 }
885
886 #[test]
887 fn source_filtering_uses_the_shape_that_matches_what_was_asked() {
888 assert!(body_of(&Search::new().source(["title", "author"]))
889 .contains(r#""_source":["title","author"]"#));
890 assert!(body_of(&Search::new().without_source()).contains(r#""_source":false"#));
891 assert!(
892 body_of(&Search::new().source(["a"]).source_excluding(["b"]))
893 .contains(r#""_source":{"excludes":["b"],"includes":["a"]}"#)
894 );
895 }
896
897 #[test]
898 fn aggregations_nest_under_aggs_and_can_hold_each_other() {
899 let search = Search::new().size(0).aggregate(
900 "by_tag",
901 Aggregation::terms("tags").size(5).aggregate("price", Aggregation::stats("price")),
902 );
903
904 assert_eq!(
905 body_of(&search),
906 r#"{"aggs":{"by_tag":{"aggs":{"price":{"stats":{"field":"price"}}},"terms":{"field":"tags","size":5}}},"query":{"match_all":{}},"size":0}"#
907 );
908 }
909
910 #[test]
911 fn a_date_histogram_asks_for_a_calendar_interval_rather_than_a_fixed_one() {
912 let agg = Aggregation::date_histogram("created_at", "day");
914 assert_eq!(
915 agg.body().to_string(),
916 r#"{"date_histogram":{"calendar_interval":"day","field":"created_at"}}"#
917 );
918 }
919
920 #[test]
921 fn tracking_the_total_is_opt_in_because_it_costs_something() {
922 assert!(!body_of(&Search::new()).contains("track_total_hits"));
923 assert!(body_of(&Search::new().track_total_hits(true)).contains(r#""track_total_hits":true"#));
924 assert!(
925 body_of(&Search::new().track_total_hits_up_to(1000))
926 .contains(r#""track_total_hits":1000"#)
927 );
928 }
929
930 #[test]
931 fn a_count_body_drops_everything_count_would_reject() {
932 let search = Search::new()
933 .query(Query::matching("title", "rust"))
934 .size(10)
935 .from(5)
936 .sort_desc("x")
937 .aggregate("by_tag", Aggregation::terms("tags"));
938
939 assert_eq!(
940 search.count_body().to_string(),
941 r#"{"query":{"match":{"title":{"query":"rust"}}}}"#
942 );
943 }
944
945 #[test]
946 fn a_raw_clause_goes_through_untouched() {
947 let raw = Json::parse(r#"{"span_first":{"match":{"span_term":{"user":"kimchy"}},"end":3}}"#)
948 .unwrap();
949 let query: Query = Query::bool().must(Query::raw(raw.clone())).into();
950
951 assert_eq!(query.json().get("bool.must.0"), Some(&raw));
952 }
953}