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 {}