Support Directives on Subscriptions

This supports the same directives that are supported on Object on Subscriptions. My main use-case here is to support WunderGraph EDFS, which requires directives on Subscription fields: https://cosmo-docs.wundergraph.com/router/event-driven-federated-subscriptions-edfs
async-graphql · Mar 19, 2024 · f2ebc98 · f2ebc98
1 parent 0cd495e
commit f2ebc98
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 8 deletions.
diff --git a/derive/src/args.rs b/derive/src/args.rs
@@ -688,6 +688,8 @@ pub struct Subscription {
     pub visible: Option<Visible>,
     #[darling(default)]
     pub guard: Option<Expr>,
+    #[darling(default, multiple, rename = "directive")]
+    pub directives: Vec<Expr>,
 }
 
 #[derive(FromMeta, Default)]
@@ -713,6 +715,8 @@ pub struct SubscriptionField {
     pub guard: Option<Expr>,
     pub visible: Option<Visible>,
     pub complexity: Option<Expr>,
+    #[darling(default, multiple, rename = "directive")]
+    pub directives: Vec<Expr>,
 }
 
 #[derive(FromField)]

diff --git a/derive/src/subscription.rs b/derive/src/subscription.rs
@@ -6,12 +6,13 @@ use syn::{
 };
 
 use crate::{
-    args::{self, RenameRuleExt, RenameTarget, SubscriptionField},
+    args::{self, RenameRuleExt, RenameTarget, SubscriptionField, TypeDirectiveLocation},
     output_type::OutputType,
     utils::{
-        extract_input_args, gen_deprecation, generate_default, generate_guards, get_cfg_attrs,
-        get_crate_name, get_rustdoc, get_type_path_and_name, parse_complexity_expr,
-        parse_graphql_attrs, remove_graphql_attrs, visible_fn, GeneratorResult,
+        extract_input_args, gen_deprecation, gen_directive_calls, generate_default,
+        generate_guards, get_cfg_attrs, get_crate_name, get_rustdoc, get_type_path_and_name,
+        parse_complexity_expr, parse_graphql_attrs, remove_graphql_attrs, visible_fn,
+        GeneratorResult,
     },
 };
 
@@ -24,6 +25,8 @@ pub fn generate(
     let generics = &item_impl.generics;
     let where_clause = &item_impl.generics.where_clause;
     let extends = subscription_args.extends;
+    let directives =
+        gen_directive_calls(&subscription_args.directives, TypeDirectiveLocation::Object);
 
     let gql_typename = if !subscription_args.name_type {
         let name = subscription_args
@@ -244,6 +247,9 @@ pub fn generate(
                 quote! { ::std::option::Option::None }
             };
 
+            let directives =
+                gen_directive_calls(&field.directives, TypeDirectiveLocation::FieldDefinition);
+
             schema_fields.push(quote! {
                 #(#cfg_attrs)*
                 fields.insert(::std::borrow::ToOwned::to_owned(#field_name), #crate_name::registry::MetaField {
@@ -266,7 +272,7 @@ pub fn generate(
                     inaccessible: false,
                     tags: ::std::default::Default::default(),
                     compute_complexity: #complexity,
-                    directive_invocations: ::std::default::Default::default(),
+                    directive_invocations: ::std::vec![ #(#directives),* ]
                 });
             });
 
@@ -420,7 +426,7 @@ pub fn generate(
                     tags: ::std::default::Default::default(),
                     is_subscription: true,
                     rust_typename: ::std::option::Option::Some(::std::any::type_name::<Self>()),
-                    directive_invocations: ::std::default::Default::default(),
+                    directive_invocations: ::std::vec![ #(#directives),* ]
                 })
             }
 

diff --git a/tests/schemas/test_fed2_compose.schema.graphql b/tests/schemas/test_fed2_compose.schema.graphql
@@ -12,6 +12,11 @@ type SimpleValue @testDirective(scope: "simple object type", input: 1, opt: 3) {
 }
 
 
+type Subscription @testDirective(scope: "object type", input: 3) {
+	value: String! @testDirective(scope: "object field", input: 4) @noArgsDirective
+	anotherValue: SimpleValue!
+}
+
 directive @include(if: Boolean!) on FIELD | FRAGMENT_SPREAD | INLINE_FRAGMENT
 directive @noArgsDirective on FIELD_DEFINITION
 directive @skip(if: Boolean!) on FIELD | FRAGMENT_SPREAD | INLINE_FRAGMENT

diff --git a/tests/type_directive.rs b/tests/type_directive.rs
@@ -1,5 +1,6 @@
-use async_graphql::{EmptyMutation, EmptySubscription, SDLExportOptions, Schema};
+use async_graphql::{EmptyMutation, EmptySubscription, SDLExportOptions, Schema, Subscription};
 use async_graphql_derive::{Object, SimpleObject, TypeDirective};
+use futures_util::{stream, Stream};
 
 #[test]
 pub fn test_type_directive() {
@@ -48,10 +49,36 @@ pub fn test_type_directive() {
         }
     }
 
-    let schema = Schema::build(Query, EmptyMutation, EmptySubscription).finish();
+    struct Subscription;
+
+    #[Subscription(
+        directive = testDirective::apply("object type".to_string(), 3, None),
+    )]
+    impl Subscription {
+        #[graphql(
+        directive = testDirective::apply("object field".to_string(), 4, None),
+        directive = noArgsDirective::apply())
+        ]
+        async fn value(&self) -> impl Stream<Item = &'static str> {
+            stream::iter(vec!["abc"])
+        }
+
+        async fn another_value(&self) -> impl Stream<Item = SimpleValue> {
+            stream::iter(vec![SimpleValue {
+                some_data: "data".to_string(),
+            }])
+        }
+    }
+
+    let schema = Schema::build(Query, EmptyMutation, Subscription)
+        .enable_subscription_in_federation()
+        .finish();
 
     let sdl = schema.sdl_with_options(SDLExportOptions::new().federation().compose_directive());
 
+    println!("----");
+    println!("{}", sdl);
+
     let expected = include_str!("schemas/test_fed2_compose.schema.graphql");
     assert_eq!(expected, &sdl)
 }