From aa1e72e68e12d323127c0f71fac3d28d4ee6eb5f Mon Sep 17 00:00:00 2001
From: Sunli <scott_s829@163.com>
Date: Sat, 20 Mar 2021 17:22:18 +0800
Subject: [PATCH] Add `SchemaBuilder::enable_subscription_in_federation`
 method.  #449

---
 CHANGELOG.md               | 4 ++++
 src/registry/export_sdl.rs | 6 ++++++
 src/registry/mod.rs        | 1 +
 src/schema.rs              | 9 +++++++++
 4 files changed, 20 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index f886c4a2..9241308e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,10 @@ 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).
 
+## [2.6.2] - 2021-03-20
+
+- Add `SchemaBuilder::enable_subscription_in_federation` method.  [#449](https://github.com/async-graphql/async-graphql/issues/449)
+
 ## [2.6.1] - 2021-03-19
 
 Fix tracing extension doesn't work with async code. [#448](https://github.com/async-graphql/async-graphql/issues/448)
diff --git a/src/registry/export_sdl.rs b/src/registry/export_sdl.rs
index e42af07e..06316632 100644
--- a/src/registry/export_sdl.rs
+++ b/src/registry/export_sdl.rs
@@ -116,6 +116,12 @@ impl Registry {
                     return;
                 }
 
+                if let Some(subscription_type) = &self.subscription_type {
+                    if name == subscription_type && !self.federation_subscription {
+                        return;
+                    }
+                }
+
                 if description.is_some() {
                     writeln!(sdl, "\"\"\"\n{}\n\"\"\"", description.unwrap()).ok();
                 }
diff --git a/src/registry/mod.rs b/src/registry/mod.rs
index 6e3d595e..9213e39f 100644
--- a/src/registry/mod.rs
+++ b/src/registry/mod.rs
@@ -341,6 +341,7 @@ pub struct Registry {
     pub subscription_type: Option<String>,
     pub disable_introspection: bool,
     pub enable_federation: bool,
+    pub federation_subscription: bool,
 }
 
 impl Registry {
diff --git a/src/schema.rs b/src/schema.rs
index 2a3034b7..549fc1a1 100644
--- a/src/schema.rs
+++ b/src/schema.rs
@@ -105,6 +105,14 @@ impl<Query, Mutation, Subscription> SchemaBuilder<Query, Mutation, Subscription>
         self
     }
 
+    /// Make the Federation SDL include subscriptions.
+    ///
+    /// Note: Not included by default, in order to be compatible with Apollo Server.
+    pub fn enable_subscription_in_federation(mut self) -> Self {
+        self.registry.federation_subscription = true;
+        self
+    }
+
     /// Override the name of the specified type.
     pub fn override_description<T: Type>(mut self, desc: &'static str) -> Self {
         self.registry.set_description::<T>(desc);
@@ -251,6 +259,7 @@ where
             },
             disable_introspection: false,
             enable_federation: false,
+            federation_subscription: false,
         };
 
         registry.add_directive(MetaDirective {