Merge pull request #813 from tamasfe/feat-flatten-context

feat: Pass context to resolvers with flatten attribute

derive/src/complex_object.rs

@@ -125,6 +125,9 @@ pub fn generate(
             let cfg_attrs = get_cfg_attrs(&method.attrs);
 
             if method_args.flatten {
+                // Only used to inject the context placeholder if required.
+                extract_input_args(&crate_name, method)?;
+
                 let ty = match &method.sig.output {
                     ReturnType::Type(_, ty) => OutputType::parse(ty)?,
                     ReturnType::Default => {
@@ -149,7 +152,7 @@ pub fn generate(
 
                 resolvers.push(quote! {
                     #(#cfg_attrs)*
-                    if let ::std::option::Option::Some(value) = #crate_name::ContainerType::resolve_field(&self.#ident().await, ctx).await? {
+                    if let ::std::option::Option::Some(value) = #crate_name::ContainerType::resolve_field(&self.#ident(ctx).await, ctx).await? {
                         return ::std::result::Result::Ok(std::option::Option::Some(value));
                     }
                 });

derive/src/object.rs

@@ -261,6 +261,9 @@ pub fn generate(
                 let cfg_attrs = get_cfg_attrs(&method.attrs);
 
                 if method_args.flatten {
+                    // Only used to inject the context placeholder if required.
+                    extract_input_args(&crate_name, method)?;
+
                     let ty = match &method.sig.output {
                         ReturnType::Type(_, ty) => OutputType::parse(ty)?,
                         ReturnType::Default => {
@@ -285,7 +288,7 @@ pub fn generate(
 
                     resolvers.push(quote! {
                         #(#cfg_attrs)*
-                        if let ::std::option::Option::Some(value) = #crate_name::ContainerType::resolve_field(&self.#ident().await, ctx).await? {
+                        if let ::std::option::Option::Some(value) = #crate_name::ContainerType::resolve_field(&self.#ident(ctx).await, ctx).await? {
                             return ::std::result::Result::Ok(std::option::Option::Some(value));
                         }
                     });

tests/complex_object.rs

@@ -365,6 +365,70 @@ async fn test_flatten() {
     );
 }
 
+#[tokio::test]
+async fn test_flatten_with_context() {
+    #[derive(SimpleObject)]
+    struct A {
+        a: i32,
+        b: i32,
+    }
+
+    #[derive(SimpleObject)]
+    #[graphql(complex)]
+    struct B {
+        #[graphql(skip)]
+        a: A,
+        c: i32,
+    }
+
+    #[ComplexObject]
+    impl B {
+        #[graphql(flatten)]
+        async fn a(&self, _ctx: &Context<'_>) -> &A {
+            &self.a
+        }
+    }
+
+    struct Query;
+
+    #[Object]
+    impl Query {
+        async fn obj(&self) -> B {
+            B {
+                a: A { a: 100, b: 200 },
+                c: 300,
+            }
+        }
+    }
+
+    let schema = Schema::new(Query, EmptyMutation, EmptySubscription);
+    let query = "{ __type(name: \"B\") { fields { name } } }";
+    assert_eq!(
+        schema.execute(query).await.data,
+        value!({
+            "__type": {
+                "fields": [
+                    {"name": "c"},
+                    {"name": "a"},
+                    {"name": "b"}
+                ]
+            }
+        })
+    );
+
+    let query = "{ obj { a b c } }";
+    assert_eq!(
+        schema.execute(query).await.data,
+        value!({
+            "obj": {
+                "a": 100,
+                "b": 200,
+                "c": 300,
+            }
+        })
+    );
+}
+
 #[tokio::test]
 async fn test_flatten_with_result() {
     #[derive(SimpleObject)]

tests/object.rs

@@ -58,3 +58,62 @@ async fn test_flatten() {
         })
     );
 }
+
+#[tokio::test]
+async fn test_flatten_with_context() {
+    #[derive(SimpleObject)]
+    struct A {
+        a: i32,
+        b: i32,
+    }
+
+    struct B;
+
+    #[Object]
+    impl B {
+        #[graphql(flatten)]
+        async fn a(&self, _ctx: &Context<'_>) -> A {
+            A { a: 100, b: 200 }
+        }
+
+        async fn c(&self) -> i32 {
+            300
+        }
+    }
+
+    struct Query;
+
+    #[Object]
+    impl Query {
+        async fn obj(&self) -> B {
+            B
+        }
+    }
+
+    let schema = Schema::new(Query, EmptyMutation, EmptySubscription);
+    let query = "{ __type(name: \"B\") { fields { name } } }";
+    assert_eq!(
+        schema.execute(query).await.data,
+        value!({
+            "__type": {
+                "fields": [
+                    {"name": "a"},
+                    {"name": "b"},
+                    {"name": "c"}
+                ]
+            }
+        })
+    );
+
+    let query = "{ obj { a b c } }";
+    assert_eq!(
+        schema.execute(query).await.data,
+        value!({
+            "obj": {
+                "a": 100,
+                "b": 200,
+                "c": 300,
+            }
+        })
+    );
+}