Add support for Federation nested keys.

CHANGELOG.md

@@ -4,6 +4,12 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [Unreleased]
+
+### Added
+
+- Add support for Federation [nested keys](https://www.apollographql.com/docs/federation/entities/#defining-a-compound-primary-key)
+
 ## [2.5.4] - 2021-02-15
 
 ### Fixed

derive/src/input_object.rs

@@ -52,6 +52,7 @@ pub fn generate(object_args: &args::InputObject) -> GeneratorResult<TokenStream>
     let mut fields = Vec::new();
     let mut schema_fields = Vec::new();
     let mut flatten_fields = Vec::new();
+    let mut federation_fields = Vec::new();
 
     for field in &s.fields {
         let ident = field.ident.as_ref().unwrap();
@@ -70,6 +71,8 @@ pub fn generate(object_args: &args::InputObject) -> GeneratorResult<TokenStream>
             continue;
         }
 
+        federation_fields.push((ty, name.clone()));
+
         if field.flatten {
             flatten_fields.push((ident, ty));
 
@@ -172,6 +175,23 @@ pub fn generate(object_args: &args::InputObject) -> GeneratorResult<TokenStream>
 
     let visible = visible_fn(&object_args.visible);
 
+    let get_federation_fields = {
+        let fields = federation_fields.into_iter().map(|(ty, name)| {
+            quote! {
+                if let ::std::option::Option::Some(fields) = <#ty as #crate_name::InputType>::federation_fields() {
+                    res.push(::std::format!("{} {}", #name, fields));
+                } else {
+                    res.push(::std::string::ToString::to_string(#name));
+                }
+            }
+        });
+        quote! {
+            let mut res = ::std::vec::Vec::new();
+            #(#fields)*
+            ::std::option::Option::Some(::std::format!("{{ {} }}", res.join(" ")))
+        }
+    };
+
     let expanded = if object_args.concretes.is_empty() {
         quote! {
             #[allow(clippy::all, clippy::pedantic)]
@@ -210,6 +230,10 @@ pub fn generate(object_args: &args::InputObject) -> GeneratorResult<TokenStream>
                     #(#put_fields)*
                     #crate_name::Value::Object(map)
                 }
+
+                fn federation_fields() -> ::std::option::Option<::std::string::String> {
+                    #get_federation_fields
+                }
             }
 
             impl #crate_name::InputObjectType for #ident {}
@@ -247,6 +271,10 @@ pub fn generate(object_args: &args::InputObject) -> GeneratorResult<TokenStream>
                     #(#put_fields)*
                     #crate_name::Value::Object(map)
                 }
+
+                fn __internal_federation_fields() -> ::std::option::Option<::std::string::String> {
+                    #get_federation_fields
+                }
             }
         });
 
@@ -276,6 +304,10 @@ pub fn generate(object_args: &args::InputObject) -> GeneratorResult<TokenStream>
                     fn to_value(&self) -> #crate_name::Value {
                         self.__internal_to_value()
                     }
+
+                    fn federation_fields() -> ::std::option::Option<::std::string::String> {
+                        Self::__internal_federation_fields()
+                    }
                 }
 
                 impl #crate_name::InputObjectType for #concrete_type {}

derive/src/object.rs

@@ -131,8 +131,7 @@ pub fn generate(
                 let mut key_pat = Vec::new();
                 let mut key_getter = Vec::new();
                 let mut use_keys = Vec::new();
-                let mut keys = Vec::new();
-                let mut keys_str = String::new();
+                let mut get_federation_key = Vec::new();
                 let mut requires_getter = Vec::new();
                 let all_key = args.iter().all(|(_, _, arg)| !arg.key);
 
@@ -153,10 +152,13 @@ pub fn generate(
                     });
 
                     if is_key {
-                        if !keys_str.is_empty() {
-                            keys_str.push(' ');
-                        }
-                        keys_str.push_str(&name);
+                        get_federation_key.push(quote! {
+                            if let Some(fields) = <#ty as #crate_name::InputType>::federation_fields() {
+                                key_str.push(format!("{} {}", #name, fields));                                
+                            } else {
+                                key_str.push(#name.to_string());
+                            }
+                        });
 
                         key_pat.push(quote! {
                             ::std::option::Option::Some(#ident)
@@ -167,7 +169,6 @@ pub fn generate(
                                 value
                             })
                         });
-                        keys.push(name);
                         use_keys.push(ident);
                     } else {
                         // requires
@@ -179,7 +180,13 @@ pub fn generate(
                     }
                 }
 
-                add_keys.push(quote! { registry.add_keys(&<#entity_type as #crate_name::Type>::type_name(), #keys_str); });
+                add_keys.push(quote! {
+                    {
+                        let mut key_str = Vec::new();
+                        #(#get_federation_key)*
+                        registry.add_keys(&<#entity_type as #crate_name::Type>::type_name(), &key_str.join(" "));
+                    }
+                });
                 create_entity_types.push(
                     quote! { <#entity_type as #crate_name::Type>::create_type_info(registry); },
                 );

src/base.rs

@@ -45,6 +45,12 @@ pub trait InputType: Type + Send + Sync + Sized {
 
     /// Convert to a `Value` for introspection.
     fn to_value(&self) -> Value;
+
+    /// Get the federation fields, only for InputObject.
+    #[doc(hidden)]
+    fn federation_fields() -> Option<String> {
+        None
+    }
 }
 
 /// Represents a GraphQL output value.

tests/federation.rs

@@ -2,73 +2,132 @@
 
 use async_graphql::*;
 
-struct User {
-    id: ID,
-}
-
-#[Object(extends)]
-impl User {
-    #[graphql(external)]
-    async fn id(&self) -> &ID {
-        &self.id
+#[async_std::test]
+pub async fn test_nested_key() {
+    #[derive(InputObject)]
+    struct MyInputA {
+        a: i32,
+        b: i32,
+        c: MyInputB,
     }
 
-    async fn reviews(&self) -> Vec<Review> {
-        todo!()
-    }
-}
-
-struct Review;
-
-#[Object]
-impl Review {
-    async fn body(&self) -> String {
-        todo!()
+    #[derive(InputObject)]
+    struct MyInputB {
+        v: i32,
     }
 
-    #[graphql(provides = "username")]
-    async fn author(&self) -> User {
-        todo!()
+    assert_eq!(MyInputB::federation_fields().as_deref(), Some("{ v }"));
+    assert_eq!(
+        MyInputA::federation_fields().as_deref(),
+        Some("{ a b c { v } }")
+    );
+
+    struct QueryRoot;
+
+    #[derive(SimpleObject)]
+    struct MyObj {
+        a: i32,
+        b: i32,
+        c: i32,
     }
 
-    async fn product(&self) -> Product {
-        todo!()
-    }
-}
-
-struct Product {
-    upc: String,
-}
-
-#[Object(extends)]
-impl Product {
-    #[graphql(external)]
-    async fn upc(&self) -> &str {
-        &self.upc
+    #[Object]
+    impl QueryRoot {
+        #[graphql(entity)]
+        async fn find_obj(&self, input: MyInputA) -> MyObj {
+            MyObj {
+                a: input.a,
+                b: input.b,
+                c: input.c.v,
+            }
+        }
     }
 
-    async fn reviews(&self) -> Vec<Review> {
-        todo!()
-    }
-}
-
-struct QueryRoot;
-
-#[Object]
-impl QueryRoot {
-    #[graphql(entity)]
-    async fn find_user_by_id(&self, id: ID) -> User {
-        User { id }
-    }
-
-    #[graphql(entity)]
-    async fn find_product_by_upc(&self, upc: String) -> Product {
-        Product { upc }
-    }
+    let schema = Schema::new(QueryRoot, EmptyMutation, EmptySubscription);
+    let query = r#"{
+            _entities(representations: [{__typename: "MyObj", input: {a: 1, b: 2, c: { v: 3 }}}]) {
+                __typename
+                ... on MyObj {
+                    a b c
+                }
+            }
+        }"#;
+    assert_eq!(
+        schema.execute(query).await.into_result().unwrap().data,
+        value!({
+            "_entities": [
+                {"__typename": "MyObj", "a": 1, "b": 2, "c": 3},
+            ]
+        })
+    );
 }
 
 #[async_std::test]
 pub async fn test_federation() {
+    struct User {
+        id: ID,
+    }
+
+    #[Object(extends)]
+    impl User {
+        #[graphql(external)]
+        async fn id(&self) -> &ID {
+            &self.id
+        }
+
+        async fn reviews(&self) -> Vec<Review> {
+            todo!()
+        }
+    }
+
+    struct Review;
+
+    #[Object]
+    impl Review {
+        async fn body(&self) -> String {
+            todo!()
+        }
+
+        async fn author(&self) -> User {
+            todo!()
+        }
+
+        async fn product(&self) -> Product {
+            todo!()
+        }
+    }
+
+    struct Product {
+        upc: String,
+    }
+
+    #[Object(extends)]
+    impl Product {
+        #[graphql(external)]
+        async fn upc(&self) -> &str {
+            &self.upc
+        }
+
+        async fn reviews(&self) -> Vec<Review> {
+            todo!()
+        }
+    }
+
+    struct QueryRoot;
+
+    #[Object]
+    impl QueryRoot {
+        #[graphql(entity)]
+        async fn find_user_by_id(&self, id: ID) -> User {
+            User { id }
+        }
+
+        #[graphql(entity)]
+        async fn find_product_by_upc(&self, upc: String) -> Product {
+            Product { upc }
+        }
+    }
+
     let schema = Schema::new(QueryRoot, EmptyMutation, EmptySubscription);
     let query = r#"{
             _entities(representations: [{__typename: "Product", upc: "B00005N5PF"}]) {