Skip to main content

mongodb/
atlas_search.rs

1//! Helpers for building Atlas Search aggregation pipelines.  Use one of the constructor functions
2//! and chain optional value setters, and then convert to a pipeline stage [`Document`] via
3//! [`into_stage`](SearchOperator::into_stage).
4//!
5//! ```no_run
6//! # async fn wrapper() -> mongodb::error::Result<()> {
7//! # use mongodb::{Collection, bson::{Document, doc}};
8//! # let collection: Collection<Document> = todo!();
9//! use mongodb::atlas_search;
10//! let cursor = collection.aggregate(vec![
11//!     atlas_search::autocomplete("title", "pre")
12//!         .fuzzy(doc! { "maxEdits": 1, "prefixLength": 1, "maxExpansions": 256 })
13//!         .into_stage(),
14//!     doc! {
15//!         "$limit": 10,
16//!    },
17//!    doc! {
18//!        "$project": {
19//!            "_id": 0,
20//!            "title": 1,
21//!         }
22//!    },
23//! ]).await?;
24//! # Ok(())
25//! # }
26mod gen;
27
28pub use gen::*;
29
30use std::marker::PhantomData;
31
32use crate::bson::{doc, Bson, DateTime, Document};
33use mongodb_internal_macros::{export_doc, options_doc};
34
35/// A helper to build the aggregation stage for Atlas Search.
36pub struct SearchOperator<T> {
37    pub(crate) name: &'static str,
38    pub(crate) spec: Document,
39    _t: PhantomData<T>,
40}
41
42impl<T> SearchOperator<T> {
43    fn new(name: &'static str, spec: Document) -> Self {
44        Self {
45            name,
46            spec,
47            _t: PhantomData,
48        }
49    }
50
51    /// Finalize this search operator as a `$search` aggregation stage document.
52    pub fn into_stage(self) -> Document {
53        search(self).into_stage()
54    }
55
56    /// Finalize this search operator as a `$searchMeta` aggregation stage document.
57    pub fn into_stage_meta(self) -> Document {
58        search_meta(self).into_stage()
59    }
60
61    /// Erase the type of this builder.  Not typically needed, but can be useful to include builders
62    /// of different types in a single `Vec`:
63    /// ```no_run
64    /// # async fn wrapper() -> mongodb::error::Result<()> {
65    /// # use mongodb::{Collection, bson::{Document, doc}};
66    /// # let collection: Collection<Document> = todo!();
67    /// use mongodb::atlas_search;
68    /// let cursor = collection.aggregate(vec![
69    ///     atlas_search::compound()
70    ///         .must(vec![
71    ///             atlas_search::text("description", "varieties").unit(),
72    ///             atlas_search::compound()
73    ///                 .should(atlas_search::text("description", "Fuji"))
74    ///                 .unit(),
75    ///         ])
76    ///         .into_stage(),
77    /// ]).await?;
78    /// # }
79    /// ```
80    pub fn unit(self) -> SearchOperator<()> {
81        SearchOperator {
82            name: self.name,
83            spec: self.spec,
84            _t: PhantomData,
85        }
86    }
87}
88
89/// Finalize a search operator as a pending `$search` aggregation stage, allowing
90/// options to be set.
91/// ```no_run
92/// # async fn wrapper() -> mongodb::error::Result<()> {
93/// # use mongodb::{Collection, bson::{Document, doc}};
94/// # let collection: Collection<Document> = todo!();
95/// use mongodb::atlas_search::{autocomplete, search};
96/// let cursor = collection.aggregate(vec![
97///     search(
98///         autocomplete("title", "pre")
99///             .fuzzy(doc! { "maxEdits": 1, "prefixLength": 1, "maxExpansions": 256 })
100///     )
101///     .index("movies")
102///     .into_stage(),
103///     doc! {
104///         "$limit": 10,
105///     },
106///     doc! {
107///         "$project": {
108///             "_id": 0,
109///             "title": 1,
110///          }
111///     },
112/// ]).await?;
113/// # Ok(())
114/// # }
115/// ```
116#[options_doc(atlas_search, "into_stage")]
117pub fn search<T>(op: SearchOperator<T>) -> AtlasSearch {
118    AtlasSearch {
119        stage: doc! { op.name: op.spec },
120    }
121}
122
123/// A pending `$search` aggregation stage.  Construct with [`search`].
124pub struct AtlasSearch {
125    stage: Document,
126}
127
128#[export_doc(atlas_search)]
129impl AtlasSearch {
130    /// Parallelize search across segments on dedicated search nodes.
131    pub fn concurrent(mut self, value: bool) -> Self {
132        self.stage.insert("concurrent", value);
133        self
134    }
135
136    /// Document that specifies the count options for retrieving a count of the results.
137    pub fn count(mut self, value: Document) -> Self {
138        self.stage.insert("count", value);
139        self
140    }
141
142    /// Document that specifies the highlighting options for displaying search terms in their
143    /// original context.
144    pub fn highlight(mut self, value: Document) -> Self {
145        self.stage.insert("highlight", value);
146        self
147    }
148
149    /// Name of the Atlas Search index to use.
150    pub fn index(mut self, value: impl Into<String>) -> Self {
151        self.stage.insert("index", value.into());
152        self
153    }
154
155    /// Flag that specifies whether to perform a full document lookup on the backend database or
156    /// return only stored source fields directly from Atlas Search.
157    pub fn return_stored_source(mut self, value: bool) -> Self {
158        self.stage.insert("returnStoredSource", value);
159        self
160    }
161
162    /// Reference point for retrieving results.
163    pub fn search_after(mut self, value: impl Into<String>) -> Self {
164        self.stage.insert("searchAfter", value.into());
165        self
166    }
167
168    /// Reference point for retrieving results.
169    pub fn search_before(mut self, value: impl Into<String>) -> Self {
170        self.stage.insert("searchBefore", value.into());
171        self
172    }
173
174    /// Flag that specifies whether to retrieve a detailed breakdown of the score for the documents
175    /// in the results.
176    pub fn score_details(mut self, value: bool) -> Self {
177        self.stage.insert("scoreDetails", value);
178        self
179    }
180
181    /// Document that specifies the fields to sort the Atlas Search results by in ascending or
182    /// descending order.
183    pub fn sort(mut self, value: Document) -> Self {
184        self.stage.insert("sort", value);
185        self
186    }
187
188    /// Convert to an aggregation stage document.
189    pub fn into_stage(self) -> Document {
190        doc! { "$search": self.stage }
191    }
192}
193
194/// Finalize a search operator as a pending `$searchMeta` aggregation stage, allowing
195/// options to be set.
196/// ```no_run
197/// # async fn wrapper() -> mongodb::error::Result<()> {
198/// # use mongodb::{Collection, bson::{DateTime, Document, doc}};
199/// # let collection: Collection<Document> = todo!();
200/// # let start: DateTime = todo!();
201/// # let end: DateTime = todo!();
202/// use mongodb::atlas_search::{facet, range, search_meta};
203/// let cursor = collection.aggregate(vec![
204///     search_meta(
205///         facet(doc! {
206///             "directorsFacet": facet::string("directors").num_buckets(7),
207///             "yearFacet": facet::number("year", [2000, 2005, 2010, 2015]),
208///         })
209///        .operator(range("released").gte(start).lte(end))
210///     )
211///     .index("movies")
212///     .into_stage(),
213///     doc! {
214///         "$limit": 10,
215///     },
216/// ]).await?;
217/// # Ok(())
218/// # }
219/// ```
220#[options_doc(atlas_search_meta, "into_stage")]
221pub fn search_meta<T>(op: SearchOperator<T>) -> AtlasSearchMeta {
222    AtlasSearchMeta {
223        stage: doc! { op.name: op.spec },
224    }
225}
226
227/// A pending `$searchMeta` aggregation stage.  Construct with [`search_meta`].
228pub struct AtlasSearchMeta {
229    stage: Document,
230}
231
232#[export_doc(atlas_search_meta)]
233impl AtlasSearchMeta {
234    /// Document that specifies the count options for retrieving a count of the results.
235    pub fn count(mut self, value: Document) -> Self {
236        self.stage.insert("count", value);
237        self
238    }
239
240    /// Name of the Atlas Search index to use.
241    pub fn index(mut self, value: impl Into<String>) -> Self {
242        self.stage.insert("index", value.into());
243        self
244    }
245
246    /// Convert to an aggregation stage document.
247    pub fn into_stage(self) -> Document {
248        doc! { "$searchMeta": self.stage }
249    }
250}
251
252impl<T> IntoIterator for SearchOperator<T> {
253    type Item = SearchOperator<T>;
254
255    type IntoIter = std::iter::Once<SearchOperator<T>>;
256
257    fn into_iter(self) -> Self::IntoIter {
258        std::iter::once(self)
259    }
260}
261
262/// Order in which to search for tokens.
263#[derive(Debug, Clone, PartialEq)]
264#[non_exhaustive]
265pub enum TokenOrder {
266    /// Indicates tokens in the query can appear in any order in the documents.
267    Any,
268    /// Indicates tokens in the query must appear adjacent to each other or in the order specified
269    /// in the query in the documents.
270    Sequential,
271    /// Fallback for future compatibility.
272    Other(String),
273}
274
275impl TokenOrder {
276    fn name(&self) -> &str {
277        match self {
278            Self::Any => "any",
279            Self::Sequential => "sequential",
280            Self::Other(s) => s.as_str(),
281        }
282    }
283}
284
285/// Criteria to use to match the terms in the query.
286#[derive(Debug, Clone, PartialEq)]
287#[non_exhaustive]
288pub enum MatchCriteria {
289    /// Return documents that contain any of the terms from the query field.
290    Any,
291    /// Only return documents that contain all of the terms from the query field.
292    All,
293    /// Fallback for future compatibility.
294    Other(String),
295}
296
297impl MatchCriteria {
298    fn name(&self) -> &str {
299        match self {
300            Self::Any => "any",
301            Self::All => "all",
302            Self::Other(s) => s.as_str(),
303        }
304    }
305}
306
307mod private {
308    use crate::bson::{doc, Bson};
309
310    /// An Atlas Search operator parameter that can accept multiple types.
311    pub trait Parameter {
312        fn to_bson(self) -> Bson;
313    }
314
315    impl<T: Into<Bson>> Parameter for T {
316        fn to_bson(self) -> Bson {
317            self.into()
318        }
319    }
320
321    impl<T> Parameter for super::SearchOperator<T> {
322        fn to_bson(self) -> Bson {
323            Bson::Document(doc! { self.name: self.spec })
324        }
325    }
326}
327
328/// An Atlas Search operator parameter that can be either a string or array of strings.
329pub trait StringOrArray: private::Parameter {}
330impl StringOrArray for &str {}
331impl StringOrArray for String {}
332#[cfg(feature = "bson-3")]
333impl<const N: usize> StringOrArray for [&str; N] {}
334impl StringOrArray for &[&str] {}
335impl StringOrArray for &[String] {}
336
337/// An Atlas Search operator parameter that is itself a search operator.
338pub trait SearchOperatorParam: private::Parameter {}
339impl<T> SearchOperatorParam for SearchOperator<T> {}
340impl SearchOperatorParam for Document {}
341
342/// Facet definitions.  These can be used when constructing a facet definition doc:
343/// ```
344/// # use mongodb::bson::doc;
345/// use mongodb::atlas_search::facet;
346/// let search = facet(doc! {
347///   "directorsFacet": facet::string("directors").num_buckets(7),
348///   "yearFacet": facet::number("year", [2000, 2005, 2010, 2015]),
349/// });
350/// ```
351pub mod facet {
352    use crate::bson::{doc, Bson, Document};
353    use std::marker::PhantomData;
354
355    /// A facet definition; see the [facet docs](https://www.mongodb.com/docs/atlas/atlas-search/facet/) for more details.
356    pub struct Facet<T> {
357        inner: Document,
358        _t: PhantomData<T>,
359    }
360
361    impl<T> From<Facet<T>> for Bson {
362        fn from(value: Facet<T>) -> Self {
363            Bson::Document(value.inner)
364        }
365    }
366
367    /// A string facet.  Construct with [`facet::string`](string).
368    pub struct String;
369    /// String facets allow you to narrow down Atlas Search results based on the most frequent
370    /// string values in the specified string field.
371    pub fn string(path: impl AsRef<str>) -> Facet<String> {
372        Facet {
373            inner: doc! {
374                "type": "string",
375                "path": path.as_ref(),
376            },
377            _t: PhantomData,
378        }
379    }
380    impl Facet<String> {
381        /// Maximum number of facet categories to return in the results. Value must be less than or
382        /// equal to 1000.
383        pub fn num_buckets(mut self, num: i32) -> Self {
384            self.inner.insert("numBuckets", num);
385            self
386        }
387    }
388
389    /// A number facet.  Construct with [`facet::number`](number).
390    pub struct Number;
391    /// Numeric facets allow you to determine the frequency of numeric values in your search results
392    /// by breaking the results into separate ranges of numbers.
393    pub fn number(
394        path: impl AsRef<str>,
395        boundaries: impl IntoIterator<Item = impl Into<Bson>>,
396    ) -> Facet<Number> {
397        Facet {
398            inner: doc! {
399                "type": "number",
400                "path": path.as_ref(),
401                "boundaries": boundaries.into_iter().map(Into::into).collect::<Vec<_>>(),
402            },
403            _t: PhantomData,
404        }
405    }
406    impl Facet<Number> {
407        /// Name of an additional bucket that counts documents returned from the operator that do
408        /// not fall within the specified boundaries.
409        pub fn default_bucket(mut self, bucket: impl AsRef<str>) -> Self {
410            self.inner.insert("default", bucket.as_ref());
411            self
412        }
413    }
414
415    /// A date facet.  Construct with [`facet::date`](date).
416    pub struct Date;
417    /// Date facets allow you to narrow down search results based on a date.
418    pub fn date(
419        path: impl AsRef<str>,
420        boundaries: impl IntoIterator<Item = crate::bson::DateTime>,
421    ) -> Facet<Date> {
422        Facet {
423            inner: doc! {
424                "type": "date",
425                "path": path.as_ref(),
426                "boundaries": boundaries.into_iter().collect::<Vec<_>>(),
427            },
428            _t: PhantomData,
429        }
430    }
431    impl Facet<Date> {
432        /// Name of an additional bucket that counts documents returned from the operator that do
433        /// not fall within the specified boundaries.
434        pub fn default_bucket(mut self, bucket: impl AsRef<str>) -> Self {
435            self.inner.insert("default", bucket.as_ref());
436            self
437        }
438    }
439}
440
441/// Relation of the query shape geometry to the indexed field geometry.
442#[derive(Debug, Clone, PartialEq)]
443#[non_exhaustive]
444pub enum Relation {
445    /// Indicates that the indexed geometry contains the query geometry.
446    Contains,
447    /// Indicates that both the query and indexed geometries have nothing in common.
448    Disjoint,
449    /// Indicates that both the query and indexed geometries intersect.
450    Intersects,
451    /// Indicates that the indexed geometry is within the query geometry. You can't use within with
452    /// LineString or Point.
453    Within,
454    /// Fallback for future compatibility.
455    Other(String),
456}
457
458impl Relation {
459    fn name(&self) -> &str {
460        match self {
461            Self::Contains => "contains",
462            Self::Disjoint => "disjoint",
463            Self::Intersects => "intersects",
464            Self::Within => "within",
465            Self::Other(s) => s,
466        }
467    }
468}
469
470/// An Atlas Search operator parameter that can be either a document or array of documents.
471pub trait DocumentOrArray: private::Parameter {}
472impl DocumentOrArray for Document {}
473#[cfg(feature = "bson-3")]
474impl<const N: usize> DocumentOrArray for [Document; N] {}
475impl DocumentOrArray for &[Document] {}
476
477macro_rules! numeric {
478    ($trait:ty) => {
479        impl $trait for i32 {}
480        impl $trait for i64 {}
481        impl $trait for u32 {}
482        impl $trait for f32 {}
483        impl $trait for f64 {}
484    };
485}
486
487/// An Atlas Search operator parameter that can be a date, number, or GeoJSON point.
488pub trait NearOrigin: private::Parameter {}
489impl NearOrigin for DateTime {}
490impl NearOrigin for Document {}
491numeric! { NearOrigin }
492
493/// An Atlas Search operator parameter that can be any BSON numeric type.
494pub trait BsonNumber: private::Parameter {}
495numeric! { BsonNumber }
496
497/// An Atlas Search operator parameter that can be compared using [`range`].
498pub trait RangeValue: private::Parameter {}
499numeric! { RangeValue }
500impl RangeValue for DateTime {}
501impl RangeValue for &str {}
502impl RangeValue for &String {}
503impl RangeValue for String {}
504impl RangeValue for crate::bson::oid::ObjectId {}