RUST-2157 Complete set of Atlas Search helpers (#1462)

Author: abr-egn

Cargo.lock

@@ -1,3 +1,5 @@
+use std::path::Path;
+
 use convert_case::{Case, Casing};
 use proc_macro2::TokenStream;
 use quote::format_ident;
@@ -39,31 +41,43 @@ impl Operator {
 
     fn gen_helper(&self) -> TokenStream {
         let name_text = &self.name;
-        let name_ident = format_ident!("{}", name_text.to_case(Case::Pascal));
-        let constr_ident = format_ident!("{}", name_text.to_case(Case::Snake));
+        let ident_base = match name_text.as_str() {
+            "in" => "searchIn",
+            _ => name_text,
+        };
+        let name_ident = format_ident!("{}", ident_base.to_case(Case::Pascal));
+        let constr_ident = format_ident!("{}", ident_base.to_case(Case::Snake));
 
         let mut required_args = TokenStream::new();
+        let mut required_arg_names = TokenStream::new();
         let mut init_doc = TokenStream::new();
         let mut setters = TokenStream::new();
 
         for arg in &self.arguments {
-            let ident = format_ident!("{}", arg.name.to_case(Case::Snake));
-            let rust_type = arg.rust_type();
+            let ident = format_ident!(
+                "{}",
+                match arg.name.as_str() {
+                    // `box` is a reserved word
+                    "box" => "geo_box".to_owned(),
+                    _ => arg.name.to_case(Case::Snake),
+                }
+            );
+            let rust_type = arg.rust_type(&self.name);
             let type_ = rust_type.tokens();
             let arg_name = &arg.name;
             let init_expr = rust_type.bson_expr(&ident);
 
             if arg.optional.unwrap_or(false) {
-                let tvars = rust_type.variables();
                 setters.push(parse_quote! {
                     #[allow(missing_docs)]
-                    pub fn #ident<#tvars>(mut self, #ident: #type_) -> Self {
-                        self.stage.insert(#arg_name, #init_expr);
+                    pub fn #ident(mut self, #ident: #type_) -> Self {
+                        self.spec.insert(#arg_name, #init_expr);
                         self
                     }
                 });
             } else {
                 required_args.push(parse_quote! { #ident : #type_, });
+                required_arg_names.push(parse_quote! { #ident, });
                 init_doc.push(parse_quote! { #arg_name : #init_expr, });
             }
         }
@@ -73,21 +87,27 @@ impl Operator {
             "For more details, see the [{name_text} operator reference]({}).",
             self.link
         );
+        let struct_doc = format!(
+            "`{name_text}` Atlas Search operator.  Construct with \
+             [`{constr_ident}`]({constr_ident}())."
+        );
         parse_quote! {
-            #[allow(missing_docs)]
+            #[doc = #struct_doc]
             pub struct #name_ident;
 
-            impl AtlasSearch<#name_ident> {
-                #[doc = #desc]
-                #[doc = ""]
-                #[doc = #link]
-                pub fn #constr_ident(#required_args) -> Self {
-                    AtlasSearch {
-                        name: #name_text,
-                        stage: doc! { #init_doc },
-                        _t: PhantomData,
-                    }
-                }
+            #[doc = #desc]
+            #[doc = ""]
+            #[doc = #link]
+            #[options_doc(#constr_ident, "into_stage")]
+            pub fn #constr_ident(#required_args) -> SearchOperator<#name_ident> {
+                SearchOperator::new(
+                    #name_text,
+                    doc! { #init_doc },
+                )
+            }
+
+            #[export_doc(#constr_ident)]
+            impl SearchOperator<#name_ident> {
                 #setters
             }
         }
@@ -107,80 +127,114 @@ struct Argument {
 #[derive(Debug, Deserialize)]
 #[serde(rename_all = "camelCase")]
 enum ArgumentType {
-    String,
-    Object,
-    SearchScore,
-    SearchPath,
-    SearchOperator,
+    Any,
     Array,
+    BinData,
+    Bool,
+    Date,
+    Geometry,
     Int,
+    Null,
+    Number,
+    Object,
+    ObjectId,
+    SearchOperator,
+    SearchPath,
+    SearchScore,
+    String,
 }
 
-static QUERY: &str = "query";
-static TOKEN_ORDER: &str = "tokenOrder";
-static MATCH_CRITERIA: &str = "matchCriteria";
-
 impl Argument {
-    fn rust_type(&self) -> ArgumentRustType {
-        if self.name == QUERY {
-            return ArgumentRustType::StringOrArray;
-        }
-        if self.name == TOKEN_ORDER {
-            return ArgumentRustType::TokenOrder;
-        }
-        if self.name == MATCH_CRITERIA {
-            return ArgumentRustType::MatchCriteria;
+    fn rust_type(&self, operator: &str) -> ArgumentRustType {
+        match (operator, self.name.as_str()) {
+            ("autocomplete" | "text", "query") => return ArgumentRustType::StringOrArray,
+            ("autocomplete", "tokenOrder") => return ArgumentRustType::TokenOrder,
+            ("text", "matchCriteria") => return ArgumentRustType::MatchCriteria,
+            ("equals", "value") => return ArgumentRustType::IntoBson,
+            ("geoShape", "relation") => return ArgumentRustType::Relation,
+            ("range", "gt" | "gte" | "lt" | "lte") => return ArgumentRustType::RangeValue,
+            ("near", "origin") => return ArgumentRustType::NearOrigin,
+            _ => (),
         }
+        use ArgumentType::*;
         match self.type_.as_slice() {
-            [ArgumentType::String] => ArgumentRustType::String,
-            [ArgumentType::Object] => ArgumentRustType::Document,
-            [ArgumentType::SearchScore] => ArgumentRustType::Document,
-            [ArgumentType::SearchPath] => ArgumentRustType::StringOrArray,
-            [ArgumentType::SearchOperator, ArgumentType::Array] => ArgumentRustType::Operator,
-            [ArgumentType::Int] => ArgumentRustType::I32,
+            [String] => ArgumentRustType::String,
+            [Object] => ArgumentRustType::Document,
+            [SearchScore] => ArgumentRustType::Document,
+            [SearchPath] => ArgumentRustType::StringOrArray,
+            [SearchOperator] => ArgumentRustType::SearchOperator,
+            [SearchOperator, Array] => ArgumentRustType::SeachOperatorIter,
+            [Int] => ArgumentRustType::I32,
+            [Geometry] => ArgumentRustType::Document,
+            [Any, Array] => ArgumentRustType::IntoBson,
+            [Object, Array] => ArgumentRustType::DocumentOrArray,
+            [Number] => ArgumentRustType::BsonNumber,
+            [String, Array] => ArgumentRustType::StringOrArray,
+            [Bool] => ArgumentRustType::Bool,
             _ => panic!("Unexpected argument types: {:?}", self.type_),
         }
     }
 }
 
 enum ArgumentRustType {
-    String,
+    Bool,
+    BsonNumber,
     Document,
+    DocumentOrArray,
+    I32,
+    IntoBson,
+    MatchCriteria,
+    NearOrigin,
+    RangeValue,
+    Relation,
+    SearchOperator,
+    SeachOperatorIter,
+    String,
     StringOrArray,
     TokenOrder,
-    MatchCriteria,
-    Operator,
-    I32,
 }
 
 impl ArgumentRustType {
     fn tokens(&self) -> syn::Type {
         match self {
-            Self::String => parse_quote! { impl AsRef<str> },
+            Self::Bool => parse_quote! { bool },
+            Self::BsonNumber => parse_quote! { impl BsonNumber },
             Self::Document => parse_quote! { Document },
+            Self::DocumentOrArray => parse_quote! { impl DocumentOrArray },
+            Self::I32 => parse_quote! { i32 },
+            Self::IntoBson => parse_quote! { impl Into<Bson> },
+            Self::MatchCriteria => parse_quote! { MatchCriteria },
+            Self::NearOrigin => parse_quote! { impl NearOrigin },
+            Self::RangeValue => parse_quote! { impl RangeValue },
+            Self::Relation => parse_quote! { Relation },
+            Self::SearchOperator => parse_quote! { impl SearchOperatorParam },
+            Self::SeachOperatorIter => {
+                parse_quote! { impl IntoIterator<Item = impl SearchOperatorParam> }
+            }
+            Self::String => parse_quote! { impl AsRef<str> },
             Self::StringOrArray => parse_quote! { impl StringOrArray },
             Self::TokenOrder => parse_quote! { TokenOrder },
-            Self::MatchCriteria => parse_quote! { MatchCriteria },
-            Self::Operator => parse_quote! { impl IntoIterator<Item = AtlasSearch<T>> },
-            Self::I32 => parse_quote! { i32 },
-        }
-    }
-
-    fn variables(&self) -> TokenStream {
-        match self {
-            Self::Operator => parse_quote! { T },
-            _ => parse_quote! {},
         }
     }
 
     fn bson_expr(&self, ident: &syn::Ident) -> syn::Expr {
         match self {
+            Self::Document | Self::I32 | Self::Bool => parse_quote! { #ident },
+            Self::IntoBson => parse_quote! { #ident.into() },
+            Self::SeachOperatorIter => {
+                parse_quote! { #ident.into_iter().map(|o| o.to_bson()).collect::<Vec<_>>() }
+            }
             Self::String => parse_quote! { #ident.as_ref() },
-            Self::StringOrArray => parse_quote! { #ident.to_bson() },
-            Self::TokenOrder | Self::MatchCriteria => parse_quote! { #ident.name() },
-            Self::Document | Self::I32 => parse_quote! { #ident },
-            Self::Operator => {
-                parse_quote! { #ident.into_iter().map(Document::from).collect::<Vec<_>>() }
+            Self::StringOrArray
+            | Self::DocumentOrArray
+            | Self::SearchOperator
+            | Self::NearOrigin
+            | Self::RangeValue
+            | Self::BsonNumber => {
+                parse_quote! { #ident.to_bson() }
+            }
+            Self::TokenOrder | Self::MatchCriteria | Self::Relation => {
+                parse_quote! { #ident.name() }
             }
         }
     }
@@ -200,11 +254,14 @@ impl TokenStreamExt for TokenStream {
 
 fn main() {
     let mut operators = TokenStream::new();
-    for path in [
-        "yaml/search/autocomplete.yaml",
-        "yaml/search/text.yaml",
-        "yaml/search/compound.yaml",
-    ] {
+    let mut paths = Path::new("yaml/search")
+        .read_dir()
+        .unwrap()
+        .map(|e| e.unwrap().path())
+        .filter(|p| p.extension().is_some_and(|e| e == "yaml"))
+        .collect::<Vec<_>>();
+    paths.sort();
+    for path in paths {
         let contents = std::fs::read_to_string(path).unwrap();
         let parsed = serde_yaml::from_str::<Operator>(&contents)
             .unwrap()
@@ -215,6 +272,7 @@ fn main() {
     let file = parse_quote! {
         //! This file was autogenerated.  Do not manually edit.
         use super::*;
+        use mongodb_internal_macros::{export_doc, options_doc};
 
         #operators
     };

etc/gen_atlas_search/src/main.rs

@@ -1 +1,2 @@
 /target/
+Cargo.lock

macros/.gitignore

@@ -30,6 +30,9 @@ pub fn deeplink(
     crate::rustdoc::deeplink(attr, item)
 }
 
+/// Generate setters for the given options struct.
+/// Arguments:
+///   * the fully-qualified path of the struct type
 #[import_tokens_attr]
 #[with_custom_parsing(crate::option::OptionSettersArgs)]
 #[proc_macro_attribute]
@@ -40,6 +43,10 @@ pub fn option_setters(
     crate::option::option_setters(attr, item, __custom_tokens)
 }
 
+/// Export the setters in this `impl` block so they can be used in `options_doc`.
+/// Arguments:
+///   * an identifier for the exported list
+///   * an optional `extra = [fn_name[,..]]` list of additional setters to include
 #[proc_macro_attribute]
 pub fn export_doc(
     attr: proc_macro::TokenStream,
@@ -48,6 +55,10 @@ pub fn export_doc(
     crate::rustdoc::export_doc(attr, item)
 }
 
+/// Include options documentation generated by `export_doc` in the rustdoc for this method:
+/// Arguments:
+///   * the doc identifier given to `export_doc`
+///   * an optional `sync` keyword that alters the documentation to be appropriate for a sync action
 #[import_tokens_attr]
 #[with_custom_parsing(crate::rustdoc::OptionsDocArgs)]
 #[proc_macro_attribute]

macros/src/lib.rs

@@ -153,7 +153,10 @@ pub(crate) fn options_doc(
     });
     let preamble = format!(
         "These methods can be chained before `{}` to set options:",
-        if args.is_async() { ".await" } else { "run" }
+        args.term
+            .as_ref()
+            .map(|(_, b)| b.as_str())
+            .unwrap_or(".await")
     );
     impl_fn.attrs.push(parse_quote! {
         #[doc = #preamble]
@@ -169,34 +172,29 @@ pub(crate) fn options_doc(
 
 pub(crate) struct OptionsDocArgs {
     foreign_path: syn::Path,
-    sync: Option<(Token![,], Ident)>,
-}
-
-impl OptionsDocArgs {
-    fn is_async(&self) -> bool {
-        self.sync.is_none()
-    }
+    term: Option<(Token![,], String)>,
 }
 
 impl Parse for OptionsDocArgs {
     fn parse(input: ParseStream) -> syn::Result<Self> {
         let foreign_path = input.parse()?;
-        let sync = if input.is_empty() {
+        let term = if input.is_empty() {
             None
         } else {
-            Some((input.parse()?, parse_name(input, "sync")?))
+            let (comma, lit) = (input.parse()?, input.parse::<syn::LitStr>()?);
+            Some((comma, lit.value()))
         };
 
-        Ok(Self { foreign_path, sync })
+        Ok(Self { foreign_path, term })
     }
 }
 
 impl ToTokens for OptionsDocArgs {
     fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
         tokens.extend(self.foreign_path.to_token_stream());
-        if let Some((comma, ident)) = &self.sync {
+        if let Some((comma, lit)) = &self.term {
             tokens.extend(comma.to_token_stream());
-            tokens.extend(ident.to_token_stream());
+            tokens.extend(lit.to_token_stream());
         }
     }
 }