From 69d90572bc0f3bf2909dab6b2c99ea4e68ca7f2b Mon Sep 17 00:00:00 2001
From: "Garen J. Torikian" <gjtorikian@users.noreply.github.com>
Date: Tue, 27 Mar 2018 18:55:36 -0400
Subject: [PATCH] Support `except` and `only` masks

---
 lib/graphql/relay/walker.rb               |  4 ++--
 lib/graphql/relay/walker/client_ext.rb    |  4 ++--
 lib/graphql/relay/walker/query_builder.rb | 25 ++++++++++++++++++-----
 3 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/lib/graphql/relay/walker.rb b/lib/graphql/relay/walker.rb
index 45c0a2d..3cba36f 100644
--- a/lib/graphql/relay/walker.rb
+++ b/lib/graphql/relay/walker.rb
@@ -8,8 +8,8 @@ module GraphQL::Relay
     # schema - The GraphQL::Schema to build a query for.
     #
     # Returns a String query.
-    def self.query_string(schema)
-      QueryBuilder.new(schema).query_string
+    def self.query_string(schema, except: nil, only: nil)
+      QueryBuilder.new(schema, except: except, only: only).query_string
     end
 
     # Start traversing a graph, starting from the given relay node ID.
diff --git a/lib/graphql/relay/walker/client_ext.rb b/lib/graphql/relay/walker/client_ext.rb
index 5164f85..d9105ab 100644
--- a/lib/graphql/relay/walker/client_ext.rb
+++ b/lib/graphql/relay/walker/client_ext.rb
@@ -8,8 +8,8 @@ module GraphQL::Relay::Walker
     # &blk     - A block to call with each Walker::Frame that is visited.
     #
     # Returns nothing.
-    def walk(from_id:, variables: {}, context: {})
-      query_string = GraphQL::Relay::Walker.query_string(schema)
+    def walk(from_id:, except: nil, only: nil, variables: {}, context: {})
+      query_string = GraphQL::Relay::Walker.query_string(schema, except: except, only: only)
       walker_query = parse(query_string)
 
       GraphQL::Relay::Walker.walk(from_id: from_id) do |frame|
diff --git a/lib/graphql/relay/walker/query_builder.rb b/lib/graphql/relay/walker/query_builder.rb
index 8e26f55..cf34b59 100644
--- a/lib/graphql/relay/walker/query_builder.rb
+++ b/lib/graphql/relay/walker/query_builder.rb
@@ -3,7 +3,7 @@ module GraphQL::Relay::Walker
     DEFAULT_ARGUMENTS = { 'first' => 5 }.freeze
     BASE_QUERY = 'query($id: ID!) { node(id: $id) { id } }'.freeze
 
-    attr_reader :schema, :connection_arguments, :ast
+    attr_reader :schema, :connection_arguments, :ast, :except, :only
 
     # Initialize a new QueryBuilder.
     #
@@ -12,8 +12,10 @@ module GraphQL::Relay::Walker
     #                         (optional).
     #
     # Returns nothing.
-    def initialize(schema, connection_arguments: DEFAULT_ARGUMENTS)
+    def initialize(schema, except: nil, only: nil, connection_arguments: DEFAULT_ARGUMENTS)
       @schema = schema
+      @except = except
+      @only   = only
       @connection_arguments = connection_arguments
       @ast = build_query
     end
@@ -37,13 +39,25 @@ module GraphQL::Relay::Walker
         selections = d_ast.definitions.first.selections.first.selections
 
         node_types.each do |type|
-          selections << inline_fragment_ast(type)
+          selections << inline_fragment_ast(type) if include?(type)
         end
 
         selections.compact!
       end
     end
 
+    # Private: Depending on the `except` or `include` filters,
+    # should this item be included a AST of the given type.
+    #
+    # type           - The GraphQL item to identify to make the fragment
+    #
+    # Returns a Boolean.
+    def include?(type)
+      return !@except.call(type, {}) if @except
+      return @only.call(type, {}) if @only
+      true
+    end
+
     # Private: Make a AST of the given type.
     #
     # klass             - The GraphQL::Language::Nodes::AbstractNode subclass
@@ -80,9 +94,10 @@ module GraphQL::Relay::Walker
 
         if with_children
           type.all_fields.each do |field|
-            if node_field?(field)
+            field_type = field.type.unwrap
+            if node_field?(field) && include?(field_type)
               if_ast.selections << node_field_ast(field)
-            elsif connection_field?(field)
+            elsif connection_field?(field) && include?(field_type)
               if_ast.selections << connection_field_ast(field)
             end
           end