[pbs-devel] [PATCH proxmox 14/18] api-macro: support flattened parameters

[Wolfgang Bumiller <w.bumiller@proxmox.com>]

Signed-off-by: Wolfgang Bumiller <w.bumiller@proxmox.com>
---
 proxmox-api-macro/src/api/method.rs | 286 +++++++++++++++++++++-------
 proxmox-api-macro/src/api/mod.rs    |  64 ++++++-
 proxmox-api-macro/tests/allof.rs    | 137 ++++++++++++-
 3 files changed, 414 insertions(+), 73 deletions(-)

diff --git a/proxmox-api-macro/src/api/method.rs b/proxmox-api-macro/src/api/method.rs
index ff5d3e0..23501bc 100644
--- a/proxmox-api-macro/src/api/method.rs
+++ b/proxmox-api-macro/src/api/method.rs
@@ -85,7 +85,7 @@ impl TryFrom<JSONObject> for ReturnType {
 ///
 /// See the top level macro documentation for a complete example.
 pub fn handle_method(mut attribs: JSONObject, mut func: syn::ItemFn) -> Result<TokenStream, Error> {
-    let mut input_schema: Schema = match attribs.remove("input") {
+    let input_schema: Schema = match attribs.remove("input") {
         Some(input) => input.into_object("input schema definition")?.try_into()?,
         None => Schema {
             span: Span::call_site(),
@@ -95,6 +95,19 @@ pub fn handle_method(mut attribs: JSONObject, mut func: syn::ItemFn) -> Result<T
         },
     };
 
+    let mut input_schema = if input_schema.as_object().is_some() {
+        input_schema
+    } else {
+        error!(
+            input_schema.span,
+            "method input schema must be an object schema"
+        );
+        let mut schema = Schema::empty_object(input_schema.span);
+        schema.description = input_schema.description;
+        schema.properties = input_schema.properties;
+        schema
+    };
+
     let mut return_type: Option<ReturnType> = attribs
         .remove("returns")
         .map(|ret| ret.into_object("return schema definition")?.try_into())
@@ -158,25 +171,15 @@ pub fn handle_method(mut attribs: JSONObject, mut func: syn::ItemFn) -> Result<T
     // input schema is done, let's give the method body a chance to extract default parameters:
     DefaultParameters(&input_schema).visit_item_fn_mut(&mut func);
 
-    let input_schema = {
-        let mut ts = TokenStream::new();
-        input_schema.to_typed_schema(&mut ts)?;
-        ts
-    };
-
     let vis = &func.vis;
     let func_name = &func.sig.ident;
     let api_method_name = Ident::new(
         &format!("API_METHOD_{}", func_name.to_string().to_uppercase()),
         func.sig.ident.span(),
     );
-    let input_schema_name = Ident::new(
-        &format!(
-            "API_PARAMETER_SCHEMA_{}",
-            func_name.to_string().to_uppercase()
-        ),
-        func.sig.ident.span(),
-    );
+
+    let (input_schema_code, input_schema_parameter) =
+        serialize_input_schema(input_schema, &func.sig.ident, func.sig.span())?;
 
     let mut returns_schema_setter = TokenStream::new();
     if let Some(return_type) = return_type {
@@ -192,13 +195,12 @@ pub fn handle_method(mut attribs: JSONObject, mut func: syn::ItemFn) -> Result<T
     };
 
     Ok(quote_spanned! { func.sig.span() =>
-        pub const #input_schema_name: ::proxmox::api::schema::ObjectSchema =
-            #input_schema;
+        #input_schema_code
 
         #vis const #api_method_name: ::proxmox::api::ApiMethod =
-            ::proxmox::api::ApiMethod::new(
+            ::proxmox::api::ApiMethod::new_full(
                 &#api_handler,
-                &#input_schema_name,
+                #input_schema_parameter,
             )
             #returns_schema_setter
             #access_setter
@@ -538,69 +540,221 @@ fn extract_normal_parameter(
     let name_str = syn::LitStr::new(name.as_str(), span);
     let arg_name = Ident::new(&format!("input_arg_{}", name.as_ident().to_string()), span);
 
+    let default_value = param.entry.schema.find_schema_property("default");
+
     // Optional parameters are expected to be Option<> types in the real function
     // signature, so we can just keep the returned Option from `input_map.remove()`.
-    body.extend(quote_spanned! { span =>
-        let #arg_name = input_map
-            .remove(#name_str)
-            .map(::serde_json::from_value)
-            .transpose()?
-    });
+    match param.entry.flatten {
+        None => {
+            // regular parameter, we just remove it and call `from_value`.
 
-    let default_value = param.entry.schema.find_schema_property("default");
-    if !param.entry.optional {
-        // Non-optional types need to be extracted out of the option though (unless
-        // they have a default):
-        //
-        // Whether the parameter is optional should have been verified by the schema
-        // verifier already, so here we just use anyhow::bail! instead of building a
-        // proper http error!
-        body.extend(quote_spanned! { span =>
-            .ok_or_else(|| ::anyhow::format_err!(
-                "missing non-optional parameter: {}",
-                #name_str,
-            ))?
-        });
-    }
+            body.extend(quote_spanned! { span =>
+                let #arg_name = input_map
+                    .remove(#name_str)
+                    .map(::serde_json::from_value)
+                    .transpose()?
+            });
 
-    let no_option_type = util::is_option_type(param.ty).is_none();
+            if !param.entry.optional {
+                // Non-optional types need to be extracted out of the option though (unless
+                // they have a default):
+                //
+                // Whether the parameter is optional should have been verified by the schema
+                // verifier already, so here we just use anyhow::bail! instead of building a
+                // proper http error!
+                body.extend(quote_spanned! { span =>
+                    .ok_or_else(|| ::anyhow::format_err!(
+                        "missing non-optional parameter: {}",
+                        #name_str,
+                    ))?
+                });
+            }
 
-    if let Some(def) = &default_value {
-        let name_uc = name.as_ident().to_string().to_uppercase();
-        let name = Ident::new(
-            &format!("API_METHOD_{}_PARAM_DEFAULT_{}", func_uc, name_uc),
-            span,
-        );
+            let no_option_type = util::is_option_type(param.ty).is_none();
 
-        // strip possible Option<> from this type:
-        let ty = util::is_option_type(param.ty).unwrap_or(param.ty);
-        default_consts.extend(quote_spanned! { span =>
-            pub const #name: #ty = #def;
-        });
+            if let Some(def) = &default_value {
+                let name_uc = name.as_ident().to_string().to_uppercase();
+                let name = Ident::new(
+                    &format!("API_METHOD_{}_PARAM_DEFAULT_{}", func_uc, name_uc),
+                    span,
+                );
+
+                // strip possible Option<> from this type:
+                let ty = util::is_option_type(param.ty).unwrap_or(param.ty);
+                default_consts.extend(quote_spanned! { span =>
+                    pub const #name: #ty = #def;
+                });
+
+                if param.entry.optional && no_option_type {
+                    // Optional parameter without an Option<T> type requires a default:
+                    body.extend(quote_spanned! { span =>
+                        .unwrap_or(#name)
+                    });
+                }
+            } else if param.entry.optional && no_option_type {
+                // FIXME: we should not be able to reach this without having produced another
+                // error above already anyway?
+                error!(param.ty => "Optional parameter without Option<T> requires a default");
+
+                // we produced an error so just write something that will compile
+                body.extend(quote_spanned! { span =>
+                    .unwrap_or_else(|| unreachable!())
+                });
+            }
 
-        if param.entry.optional && no_option_type {
-            // Optional parameter without an Option<T> type requires a default:
-            body.extend(quote_spanned! { span =>
-                .unwrap_or(#name)
-            });
+            body.extend(quote_spanned! { span => ; });
+        }
+        Some(flatten_span) => {
+            // Flattened parameter, we need ot use our special partial-object deserializer.
+            // Also note that we do not support simply nesting schemas. We need a referenced type.
+            // Otherwise the expanded code here gets ugly and we'd need to make sure we pull out
+            // nested schemas into named variables first... No thanks.
+
+            if default_value.is_some() {
+                error!(
+                    default_value =>
+                    "flattened parameter cannot have a default as it cannot be optional",
+                );
+            }
+
+            if let Some(schema_ref) = param.entry.schema.to_schema_reference() {
+                let ty = param.ty;
+                body.extend(quote_spanned! { span =>
+                    let #arg_name = <#ty as ::serde::Deserialize>::deserialize(
+                        ::proxmox::api::de::ExtractValueDeserializer::try_new(
+                            input_map,
+                            #schema_ref,
+                        )
+                        .ok_or_else(|| ::anyhow::format_err!(
+                            "flattened parameter {:?} has invalid schema",
+                            #name_str,
+                        ))?,
+                    )?;
+                });
+            } else {
+                error!(
+                    flatten_span,
+                    "flattened parameter schema must be a schema reference"
+                );
+                body.extend(quote_spanned! { span => let #arg_name = unreachable!(); });
+            }
         }
-    } else if param.entry.optional && no_option_type {
-        // FIXME: we should not be able to reach this without having produced another
-        // error above already anyway?
-        error!(param.ty => "Optional parameter without Option<T> requires a default");
-
-        // we produced an error so just write something that will compile
-        body.extend(quote_spanned! { span =>
-            .unwrap_or_else(|| unreachable!())
-        });
     }
 
-    body.extend(quote_spanned! { span => ; });
     args.extend(quote_spanned! { span => #arg_name, });
 
     Ok(())
 }
 
+/// Returns a tuple containing the schema code first and the `ParameterSchema` parameter for the
+/// `ApiMethod` second.
+fn serialize_input_schema(
+    mut input_schema: Schema,
+    func_name: &Ident,
+    func_sig_span: Span,
+) -> Result<(TokenStream, TokenStream), Error> {
+    let input_schema_name = Ident::new(
+        &format!(
+            "API_PARAMETER_SCHEMA_{}",
+            func_name.to_string().to_uppercase()
+        ),
+        func_name.span(),
+    );
+
+    let (flattened, has_params) = match &mut input_schema.item {
+        SchemaItem::Object(obj) => {
+            let flattened = obj.drain_filter(|entry| entry.flatten.is_none());
+            (flattened, !obj.is_empty())
+        }
+        _ => (Vec::new(), true),
+    };
+
+    if flattened.is_empty() {
+        let mut ts = TokenStream::new();
+        input_schema.to_typed_schema(&mut ts)?;
+        return Ok((
+            quote_spanned! { func_sig_span =>
+                pub const #input_schema_name: ::proxmox::api::schema::ObjectSchema = #ts;
+            },
+            quote_spanned! { func_sig_span =>
+                ::proxmox::api::router::ParameterSchema::Object(&#input_schema_name)
+            },
+        ));
+    }
+
+    let mut all_of_schemas = TokenStream::new();
+    for entry in flattened {
+        if entry.optional {
+            error!(
+                entry.schema.span,
+                "optional flattened parameters are not supported"
+            );
+        }
+
+        all_of_schemas.extend(quote::quote! {&});
+        entry.schema.to_schema(&mut all_of_schemas)?;
+        all_of_schemas.extend(quote::quote! {,});
+    }
+
+    let description = match input_schema.description.take().ok() {
+        Some(description) => description,
+        None => {
+            error!(input_schema.span, "missing description on api type struct");
+            syn::LitStr::new("<missing description>", input_schema.span)
+        }
+    };
+    // and replace it with a "dummy"
+    input_schema.description = Maybe::Derived(syn::LitStr::new(
+        &format!("<INNER: {}>", description.value()),
+        description.span(),
+    ));
+
+    let (inner_schema, inner_schema_ref) = if has_params {
+        // regular parameters go into the "inner" schema to merge into the AllOfSchema
+        let inner_schema_name = Ident::new(
+            &format!(
+                "API_REGULAR_PARAMETER_SCHEMA_{}",
+                func_name.to_string().to_uppercase()
+            ),
+            func_name.span(),
+        );
+
+        let obj_schema = {
+            let mut ts = TokenStream::new();
+            input_schema.to_schema(&mut ts)?;
+            ts
+        };
+
+        (
+            quote_spanned!(func_sig_span =>
+                const #inner_schema_name: ::proxmox::api::schema::Schema = #obj_schema;
+            ),
+            quote_spanned!(func_sig_span => &#inner_schema_name,),
+        )
+    } else {
+        // otherwise it stays empty
+        (TokenStream::new(), TokenStream::new())
+    };
+
+    Ok((
+        quote_spanned! { func_sig_span =>
+            #inner_schema
+
+            pub const #input_schema_name: ::proxmox::api::schema::AllOfSchema =
+                ::proxmox::api::schema::AllOfSchema::new(
+                    #description,
+                    &[
+                        #inner_schema_ref
+                        #all_of_schemas
+                    ],
+                );
+        },
+        quote_spanned! { func_sig_span =>
+            ::proxmox::api::router::ParameterSchema::AllOf(&#input_schema_name)
+        },
+    ))
+}
+
 struct DefaultParameters<'a>(&'a Schema);
 
 impl<'a> VisitMut for DefaultParameters<'a> {
diff --git a/proxmox-api-macro/src/api/mod.rs b/proxmox-api-macro/src/api/mod.rs
index 4690654..5994f19 100644
--- a/proxmox-api-macro/src/api/mod.rs
+++ b/proxmox-api-macro/src/api/mod.rs
@@ -142,6 +142,17 @@ impl Schema {
         }
     }
 
+    /// Create the token stream for a reference schema (`ExternType` or `ExternSchema`).
+    fn to_schema_reference(&self) -> Option<TokenStream> {
+        match &self.item {
+            SchemaItem::ExternType(path) => {
+                Some(quote_spanned! { path.span() => &#path::API_SCHEMA })
+            }
+            SchemaItem::ExternSchema(path) => Some(quote_spanned! { path.span() => &#path }),
+            _ => None,
+        }
+    }
+
     fn to_typed_schema(&self, ts: &mut TokenStream) -> Result<(), Error> {
         self.item.to_schema(
             ts,
@@ -323,7 +334,8 @@ impl SchemaItem {
             }
             SchemaItem::ExternType(path) => {
                 if !properties.is_empty() {
-                    error!(&properties[0].0 => "additional properties not allowed on external type");
+                    error!(&properties[0].0 =>
+                        "additional properties not allowed on external type");
                 }
                 if let Maybe::Explicit(description) = description {
                     error!(description => "description not allowed on external type");
@@ -385,6 +397,10 @@ pub struct ObjectEntry {
     pub name: FieldName,
     pub optional: bool,
     pub schema: Schema,
+
+    /// This is only valid for methods. Methods should reset this to false after dealing with it,
+    /// as encountering this during schema serialization will always cause an error.
+    pub flatten: Option<Span>,
 }
 
 impl ObjectEntry {
@@ -393,8 +409,14 @@ impl ObjectEntry {
             name,
             optional,
             schema,
+            flatten: None,
         }
     }
+
+    pub fn with_flatten(mut self, flatten: Option<Span>) -> Self {
+        self.flatten = flatten;
+        self
+    }
 }
 
 #[derive(Default)]
@@ -420,6 +442,24 @@ impl SchemaObject {
         &mut self.properties_
     }
 
+    fn drain_filter<F>(&mut self, mut func: F) -> Vec<ObjectEntry>
+    where
+        F: FnMut(&ObjectEntry) -> bool,
+    {
+        let mut out = Vec::new();
+
+        let mut i = 0;
+        while i != self.properties_.len() {
+            if !func(&self.properties_[i]) {
+                out.push(self.properties_.remove(i));
+            } else {
+                i += 1;
+            }
+        }
+
+        out
+    }
+
     fn sort_properties(&mut self) {
         self.properties_.sort_by(|a, b| (a.name).cmp(&b.name));
     }
@@ -445,7 +485,19 @@ impl SchemaObject {
                             .transpose()?
                             .unwrap_or(false);
 
-                        properties.push(ObjectEntry::new(key, optional, schema.try_into()?));
+                        let flatten: Option<Span> = schema
+                            .remove_entry("flatten")
+                            .map(|(field, value)| -> Result<(Span, bool), syn::Error> {
+                                let v: syn::LitBool = value.try_into()?;
+                                Ok((field.span(), v.value))
+                            })
+                            .transpose()?
+                            .and_then(|(span, value)| if value { Some(span) } else { None });
+
+                        properties.push(
+                            ObjectEntry::new(key, optional, schema.try_into()?)
+                                .with_flatten(flatten),
+                        );
 
                         Ok(properties)
                     },
@@ -457,6 +509,14 @@ impl SchemaObject {
 
     fn to_schema_inner(&self, ts: &mut TokenStream) -> Result<(), Error> {
         for element in self.properties_.iter() {
+            if let Some(span) = element.flatten {
+                error!(
+                    span,
+                    "`flatten` attribute is only available on method parameters, \
+                     use #[serde(flatten)] in structs"
+                );
+            }
+
             let key = element.name.as_str();
             let optional = element.optional;
             let mut schema = TokenStream::new();
diff --git a/proxmox-api-macro/tests/allof.rs b/proxmox-api-macro/tests/allof.rs
index 56e86d7..1c1b9a9 100644
--- a/proxmox-api-macro/tests/allof.rs
+++ b/proxmox-api-macro/tests/allof.rs
@@ -1,10 +1,12 @@
 //! Testing the `AllOf` schema on structs and methods.
 
+use anyhow::Error;
+use serde::{Deserialize, Serialize};
+use serde_json::{json, Value};
+
 use proxmox::api::schema;
 use proxmox_api_macro::api;
 
-use serde::{Deserialize, Serialize};
-
 pub const NAME_SCHEMA: schema::Schema = schema::StringSchema::new("Name.").schema();
 pub const VALUE_SCHEMA: schema::Schema = schema::IntegerSchema::new("Value.").schema();
 pub const INDEX_SCHEMA: schema::Schema = schema::IntegerSchema::new("Index.").schema();
@@ -18,7 +20,7 @@ pub const TEXT_SCHEMA: schema::Schema = schema::StringSchema::new("Text.").schem
 )]
 /// Name and value.
 #[derive(Deserialize, Serialize)]
-struct NameValue {
+pub struct NameValue {
     name: String,
     value: u64,
 }
@@ -31,7 +33,7 @@ struct NameValue {
 )]
 /// Index and text.
 #[derive(Deserialize, Serialize)]
-struct IndexText {
+pub struct IndexText {
     index: u64,
     text: String,
 }
@@ -44,7 +46,7 @@ struct IndexText {
 )]
 /// Name, value, index and text.
 #[derive(Deserialize, Serialize)]
-struct Nvit {
+pub struct Nvit {
     #[serde(flatten)]
     nv: NameValue,
 
@@ -116,3 +118,128 @@ fn test_extra() {
 
     assert_eq!(TEST_SCHEMA, WithExtra::API_SCHEMA);
 }
+
+#[api(
+    input: {
+        properties: {
+            nv: { flatten: true, type: NameValue },
+            it: { flatten: true, type: IndexText },
+        },
+    },
+)]
+/// Hello method.
+pub fn hello(it: IndexText, nv: NameValue) -> Result<(NameValue, IndexText), Error> {
+    Ok((nv, it))
+}
+
+#[test]
+fn hello_schema_check() {
+    const TEST_METHOD: ::proxmox::api::ApiMethod = ::proxmox::api::ApiMethod::new_full(
+        &::proxmox::api::ApiHandler::Sync(&api_function_hello),
+        ::proxmox::api::router::ParameterSchema::AllOf(&::proxmox::api::schema::AllOfSchema::new(
+            "Hello method.",
+            &[&IndexText::API_SCHEMA, &NameValue::API_SCHEMA],
+        )),
+    );
+    assert_eq!(TEST_METHOD, API_METHOD_HELLO);
+}
+
+#[api(
+    input: {
+        properties: {
+            nv: { flatten: true, type: NameValue },
+            it: { flatten: true, type: IndexText },
+            extra: { description: "An extra field." },
+        },
+    },
+)]
+/// Extra method.
+pub fn with_extra(
+    it: IndexText,
+    nv: NameValue,
+    extra: String,
+) -> Result<(NameValue, IndexText, String), Error> {
+    Ok((nv, it, extra))
+}
+
+#[test]
+fn with_extra_schema_check() {
+    const INNER_SCHEMA: ::proxmox::api::schema::Schema = ::proxmox::api::schema::ObjectSchema::new(
+        "<INNER: Extra method.>",
+        &[(
+            "extra",
+            false,
+            &::proxmox::api::schema::StringSchema::new("An extra field.").schema(),
+        )],
+    )
+    .schema();
+
+    const TEST_METHOD: ::proxmox::api::ApiMethod = ::proxmox::api::ApiMethod::new_full(
+        &::proxmox::api::ApiHandler::Sync(&api_function_with_extra),
+        ::proxmox::api::router::ParameterSchema::AllOf(&::proxmox::api::schema::AllOfSchema::new(
+            "Extra method.",
+            &[
+                &INNER_SCHEMA,
+                &IndexText::API_SCHEMA,
+                &NameValue::API_SCHEMA,
+            ],
+        )),
+    );
+    assert_eq!(TEST_METHOD, API_METHOD_WITH_EXTRA);
+}
+
+struct RpcEnv;
+impl proxmox::api::RpcEnvironment for RpcEnv {
+    fn result_attrib_mut(&mut self) -> &mut Value {
+        panic!("result_attrib_mut called");
+    }
+
+    fn result_attrib(&self) -> &Value {
+        panic!("result_attrib called");
+    }
+
+    /// The environment type
+    fn env_type(&self) -> proxmox::api::RpcEnvironmentType {
+        panic!("env_type called");
+    }
+
+    /// Set authentication id
+    fn set_auth_id(&mut self, user: Option<String>) {
+        let _ = user;
+        panic!("set_auth_id called");
+    }
+
+    /// Get authentication id
+    fn get_auth_id(&self) -> Option<String> {
+        panic!("get_auth_id called");
+    }
+}
+
+#[test]
+fn test_invocations() {
+    let mut env = RpcEnv;
+    let value = api_function_hello(
+        json!({"name":"Bob", "value":3, "index":4, "text":"Text"}),
+        &API_METHOD_HELLO,
+        &mut env,
+    )
+    .expect("hello function should work");
+
+    assert_eq!(value[0]["name"], "Bob");
+    assert_eq!(value[0]["value"], 3);
+    assert_eq!(value[1]["index"], 4);
+    assert_eq!(value[1]["text"], "Text");
+
+    let value = api_function_with_extra(
+        json!({"name":"Alice", "value":8, "index":2, "text":"Paragraph", "extra":"Some Extra"}),
+        &API_METHOD_WITH_EXTRA,
+        &mut env,
+    )
+    .expect("`with_extra` function should work");
+
+    assert_eq!(value[0]["name"], "Alice");
+    assert_eq!(value[0]["value"], 8);
+    assert_eq!(value[1]["index"], 2);
+    assert_eq!(value[1]["text"], "Paragraph");
+    assert_eq!(value[2], "Some Extra");
+}
-- 
2.20.1