From 362ba90e7dcb65039a4d67626d7fc9a07738ea45 Mon Sep 17 00:00:00 2001
From: Sunli <scott_s829@163.com>
Date: Fri, 19 Jun 2020 12:37:52 +0800
Subject: [PATCH] Add support for extending an entity with computed fields
 (advanced)  #180

---
 async-graphql-derive/src/args.rs         |  5 +++
 async-graphql-derive/src/object.rs       | 56 +++++++++++++++++-------
 async-graphql-derive/src/subscription.rs |  1 +
 3 files changed, 46 insertions(+), 16 deletions(-)

diff --git a/async-graphql-derive/src/args.rs b/async-graphql-derive/src/args.rs
index fb6fa548..3111bb2f 100644
--- a/async-graphql-derive/src/args.rs
+++ b/async-graphql-derive/src/args.rs
@@ -130,6 +130,7 @@ pub struct Argument {
     pub desc: Option<String>,
     pub default: Option<TokenStream>,
     pub validator: TokenStream,
+    pub key: bool, // for entity
 }
 
 impl Argument {
@@ -138,6 +139,7 @@ impl Argument {
         let mut desc = None;
         let mut default = None;
         let mut validator = quote! { None };
+        let mut key = false;
 
         for attr in attrs {
             match attr.parse_meta()? {
@@ -146,6 +148,8 @@ impl Argument {
                         if let NestedMeta::Meta(Meta::Path(p)) = meta {
                             if p.is_ident("default") {
                                 default = Some(quote! { Default::default() });
+                            } else if p.is_ident("key") {
+                                key = true;
                             }
                         } else if let NestedMeta::Meta(Meta::NameValue(nv)) = meta {
                             if nv.path.is_ident("name") {
@@ -185,6 +189,7 @@ impl Argument {
             desc,
             default,
             validator,
+            key,
         })
     }
 }
diff --git a/async-graphql-derive/src/object.rs b/async-graphql-derive/src/object.rs
index 363bdf8d..dc5a2909 100644
--- a/async-graphql-derive/src/object.rs
+++ b/async-graphql-derive/src/object.rs
@@ -113,29 +113,51 @@ pub fn generate(object_args: &args::Object, item_impl: &mut ItemImpl) -> Result<
                 let mut use_keys = Vec::new();
                 let mut keys = Vec::new();
                 let mut keys_str = String::new();
+                let mut requires_getter = Vec::new();
+                let mut one_key = false;
 
-                for (ident, ty, args::Argument { name, .. }) in &args {
+                if args.is_empty() {
+                    return Err(Error::new_spanned(
+                        method,
+                        "Entity need to have at least one key.",
+                    ));
+                } else if args.len() == 1 {
+                    one_key = true;
+                }
+
+                for (ident, ty, args::Argument { name, key, .. }) in &args {
+                    let is_key = one_key || *key;
                     let name = name
                         .clone()
                         .unwrap_or_else(|| ident.ident.to_string().to_camel_case());
 
-                    if !keys_str.is_empty() {
-                        keys_str.push(' ');
-                    }
-                    keys_str.push_str(&name);
+                    if is_key {
+                        if !keys_str.is_empty() {
+                            keys_str.push(' ');
+                        }
+                        keys_str.push_str(&name);
 
-                    key_pat.push(quote! {
-                        Some(#ident)
-                    });
-                    key_getter.push(quote! {
-                        params.get(#name).and_then(|value| {
-                            let value: Option<#ty> = #crate_name::InputValueType::parse(Some(value.clone())).ok();
-                            value
-                        })
-                    });
-                    keys.push(name);
-                    use_keys.push(ident);
+                        key_pat.push(quote! {
+                            Some(#ident)
+                        });
+                        key_getter.push(quote! {
+                            params.get(#name).and_then(|value| {
+                                let value: Option<#ty> = #crate_name::InputValueType::parse(Some(value.clone())).ok();
+                                value
+                            })
+                        });
+                        keys.push(name);
+                        use_keys.push(ident);
+                    } else {
+                        // requires
+                        requires_getter.push(quote! {
+                            let #ident: #ty = #crate_name::InputValueType::parse(params.get(#name).cloned()).
+                                map_err(|err| err.into_error(ctx.position(), <#ty as #crate_name::Type>::qualified_type_name()))?;
+                        });
+                        use_keys.push(ident);
+                    }
                 }
+
                 add_keys.push(quote! { registry.add_keys(&<#entity_type as #crate_name::Type>::type_name(), #keys_str); });
                 create_entity_types.push(
                     quote! { <#entity_type as #crate_name::Type>::create_type_info(registry); },
@@ -163,6 +185,7 @@ pub fn generate(object_args: &args::Object, item_impl: &mut ItemImpl) -> Result<
                     quote! {
                         if typename == &<#entity_type as #crate_name::Type>::type_name() {
                             if let (#(#key_pat),*) = (#(#key_getter),*) {
+                                #(#requires_getter)*
                                 let ctx_obj = ctx.with_selection_set(&ctx.selection_set);
                                 return #crate_name::OutputValueType::resolve(&#do_find, &ctx_obj, ctx.item).await;
                             }
@@ -290,6 +313,7 @@ pub fn generate(object_args: &args::Object, item_impl: &mut ItemImpl) -> Result<
                         desc,
                         default,
                         validator,
+                        ..
                     },
                 ) in args
                 {
diff --git a/async-graphql-derive/src/subscription.rs b/async-graphql-derive/src/subscription.rs
index d73855fb..61852a21 100644
--- a/async-graphql-derive/src/subscription.rs
+++ b/async-graphql-derive/src/subscription.rs
@@ -143,6 +143,7 @@ pub fn generate(object_args: &args::Object, item_impl: &mut ItemImpl) -> Result<
                         desc,
                         default,
                         validator,
+                        ..
                     },
                 ) in args
                 {