Connection based query splitting

lib/graphql/bulk/base_query_splitter.rb

@@ -0,0 +1,43 @@
+module GraphQL
+  module Bulk
+    class BaseQuerySplitter
+      ConnectionQuery = Struct.new(:path_string, :query_string)
+      SplitQuery = Struct.new(:base_query, :connection_queries)
+
+      def initialize(query, schema)
+        @query = query
+        @schema = schema
+      end
+
+      def split_query
+        gql_query_tree = add_type_name_fields
+
+        connection_removal_visitor = Visitors::ConnectionRemovalVisitor.new(gql_query_tree.document)
+        base_query_document = connection_removal_visitor.visit
+
+        connection_queries = connection_removal_visitor.connections.map do |connection|
+          connection_document = Visitors::ConnectionNodeExtractionVisitor.new(gql_query_tree.document, connection)
+          query_document = connection_document.visit
+
+          ConnectionQuery.new(
+            connection.path.map{ |node| node.respond_to?(:name) ? node.name : "" }.join("."),
+            query_document.to_query_string
+          )
+        end
+
+        SplitQuery.new(
+          base_query_document.to_query_string,
+          connection_queries
+        )
+      end
+
+      private
+
+      def add_type_name_fields
+        gql_query_tree = GraphQL::Query.new(@schema, @query)
+        query_with_typename = Visitors::AddTypenameToQueryVisitor.new(gql_query_tree.document).visit
+        GraphQL::Query.new(@schema, query_with_typename.to_query_string)
+      end
+    end
+  end
+end

lib/graphql/bulk/bulk.rb

@@ -0,0 +1,13 @@
+module GraphQL
+  module Bulk
+    require "graphql/bulk/visitors/add_pagination_to_query_visitor"
+    require "graphql/bulk/visitors/add_typename_to_query_visitor"
+    require "graphql/bulk/visitors/connection_node_extraction_visitor"
+    require "graphql/bulk/visitors/connection_removal_visitor"
+    require "graphql/bulk/base_query_splitter"
+    require "graphql/bulk/connection_query_splitter"
+    require "graphql/bulk/debug"
+    require "graphql/bulk/errors"
+    require "graphql/bulk/query_splitter_service"
+  end
+end

lib/graphql/bulk/connection_query_splitter.rb

@@ -0,0 +1,75 @@
+module GraphQL
+  module Bulk
+    class ConnectionQuerySplitter
+      UnrolledConnectionQuery = Struct.new(:rollup_field_name, :unrolled_query)
+      SplitConnectionQuery = Struct.new(:paginated_query, :unrolled_connection_queries)
+
+      def initialize(connection_query, schema)
+        @connection_query = connection_query
+        @schema = schema
+      end
+
+      def split_query
+        gql_query_document = GraphQL::Query.new(@schema, @connection_query.query_string).document
+
+        result = handle_nested_connections(gql_query_document)
+        paginated_query = handle_pagination_splitting(result[:base_query_without_nested_connections])
+
+        SplitConnectionQuery.new(paginated_query, result[:unrolled_connection_queries])
+      end
+
+      private
+
+      def handle_nested_connections(gql_query_document)
+        connection_removal_visitor = Visitors::ConnectionRemovalVisitor.new(gql_query_document, depth: 2)
+        base_query_document = connection_removal_visitor.visit
+
+        unrolled_connection_queries = []
+        connection_removal_visitor.connections.each do |connection|
+          connection_document = Visitors::ConnectionNodeExtractionVisitor.new(gql_query_document, connection)
+          unrolled_connection_queries << unroll_root_connection_query(connection_document.visit)
+        end
+
+        {
+          unrolled_connection_queries: unrolled_connection_queries,
+          base_query_without_nested_connections: base_query_document,
+        }
+      end
+
+      def handle_pagination_splitting(gql_query_document)
+        paginator = Visitors::AddPaginationToQueryVisitor.new(gql_query_document, root_connection_node(gql_query_document))
+        paginator.visit.to_query_string
+      end
+
+      def unroll_root_connection_query(gql_query_document)
+        root_connection_node = root_connection_node(gql_query_document)
+        query_to_inject = extract_unrolled_query(root_connection_node)
+        unrolled_query_string = "query UnrolledQuery($__appPlatformUniqueIdForUnrolling: EncodedId!) {#{root_connection_node.name.singularize}(id: $__appPlatformUniqueIdForUnrolling) { #{query_to_inject} }}"
+
+        UnrolledConnectionQuery.new("#{@connection_query.path_string}.#{root_connection_node.name}", unrolled_query_string)
+      end
+
+      def root_connection_node(gql_query_document)
+        root_connection_finder = Visitors::ConnectionRemovalVisitor.new(gql_query_document, depth: 1)
+        root_connection_finder.visit
+        root_connection_finder.connections.first.node
+      end
+
+      def extract_unrolled_query(root_connection_node)
+        field_node = root_connection_node.selections.first
+
+        if field_node.name == "nodes"
+          # nodes -> <stuff I care about>
+          return field_node.children.first.to_query_string
+        end
+
+        if field_node.name == "edges"
+          # job -> edges -> node -> <stuff I care about>
+          return field_node.children.first.children.first.to_query_string
+        end
+
+        raise Errors::BulkError, "#{root_connection_node.name} node does not have `nodes` or `edges` as it's first child"
+      end
+    end
+  end
+end

lib/graphql/bulk/debug.rb

@@ -0,0 +1,25 @@
+module GraphQL
+  module Bulk
+    class Debug
+      class << self
+        # rubocop:disable Rails/Output
+        def print_string(result, prefix = "")
+          puts "#{prefix}base: #{result[:base].delete("\n")}"
+          puts "#{prefix}nested:"
+          result[:nested].each do |nested|
+            puts "#{prefix}\t ========= Result ======= "
+            puts "#{prefix}\t path: #{nested[:path]}"
+            puts "#{prefix}\t paginated: #{nested[:paginated].delete("\n")}"
+            puts "#{prefix}\t unrolled:"
+            nested[:unrolled].each do |u|
+              puts "#{prefix}\t\t rollup_field: #{u[:rollup_field_name]}"
+              print_string u[:queries], "#{prefix}\t\t "
+            end
+            puts "#{prefix}\t ========================"
+          end
+          # rubocop:enable Rails/Output
+        end
+      end
+    end
+  end
+end

lib/graphql/bulk/errors.rb

@@ -0,0 +1,21 @@
+module GraphQL
+  module Bulk
+    module Errors
+      class BulkError < StandardError; end
+
+      class FragmentError < BulkError; end
+
+      class QueryInvalidError < BulkError
+        attr_reader :query
+
+        def initialize(query)
+          @query = query
+          errors = query.static_errors
+          message = errors.map(&:message).join(", ")
+
+          super(message)
+        end
+      end
+    end
+  end
+end

lib/graphql/bulk/query_splitter_service.rb

@@ -0,0 +1,58 @@
+module GraphQL
+    module Bulk
+      class QuerySplitterService
+        class << self
+        def split(schema, query, debug: false)
+            validate(schema, query)
+
+            result = split_all_bulk_queries(schema, query)
+
+            Debug.print_string(result) if debug
+
+            result
+          end
+
+          def validate(schema, query)
+            query = GraphQL::Query.new(schema, query)
+
+            raise Errors::QueryInvalidError, query unless query.valid?
+            raise Errors::FragmentError, "Bulk operations on a query with fragments is not supported" if query.fragments.any?
+          end
+
+          private
+
+          def split_all_bulk_queries(schema, query)
+            queries = {}
+
+            base_splitter = BaseQuerySplitter.new(query, schema)
+            base_results = base_splitter.split_query
+
+            queries[:base] = base_results[:base_query]
+            queries[:nested] = []
+
+            base_results.connection_queries.each do |cq|
+              connection_splitter = ConnectionQuerySplitter.new(cq, schema)
+              connection_splitter_results = connection_splitter.split_query
+
+              paginated = connection_splitter_results[:paginated_query]
+
+              unrolled = connection_splitter_results.unrolled_connection_queries.map do |unrolled_query|
+                {
+                  rollup_field_name: unrolled_query.rollup_field_name,
+                  queries: split_all_bulk_queries(schema, unrolled_query.unrolled_query),
+                }
+              end
+
+              queries[:nested] << {
+                path: cq.path_string,
+                paginated: paginated,
+                unrolled: unrolled,
+              }
+            end
+
+            queries
+          end
+        end
+      end
+    end
+  end

lib/graphql/bulk/visitors/add_pagination_to_query_visitor.rb

@@ -0,0 +1,63 @@
+module GraphQL
+  module Bulk
+    module Visitors
+      class AddPaginationToQueryVisitor < GraphQL::Language::Visitor
+        def initialize(document, connection_node)
+          super(document)
+          @connection_node = connection_node
+        end
+
+        def on_operation_definition(node, parent)
+          modified_node = node.merge_variable(
+            name: "__appPlatformCursor",
+            type: GraphQL::Language::Nodes::TypeName.new(name: "String!"),
+          )
+          super(modified_node, parent)
+        end
+
+        def on_field(node, parent)
+          if node == @connection_node
+            old_arguments = node.arguments
+            new_arguments = [
+              GraphQL::Language::Nodes::Argument.new(
+                name: "after",
+                value: GraphQL::Language::Nodes::VariableIdentifier.new(name: "__appPlatformCursor"),
+              ),
+              GraphQL::Language::Nodes::Argument.new(
+                name: "first",
+                value: 50,
+              ),
+            ]
+
+            old_selections = node.selections
+            new_selections = [
+              GraphQL::Language::Nodes::Field.new(
+                name: "pageInfo",
+                field_alias: "__appPlatformPageInfo",
+                selections: [
+                  GraphQL::Language::Nodes::Field.new(
+                    name: "hasNextPage"
+                  ),
+                  GraphQL::Language::Nodes::Field.new(
+                    name: "endCursor"
+                  ),
+                ]
+              ),
+            ]
+
+            modified_node = node.merge(
+              {
+                arguments: old_arguments + new_arguments,
+                selections: old_selections + new_selections,
+              }
+            )
+
+            return super(modified_node, parent)
+          end
+
+          super
+        end
+      end
+    end
+  end
+end

lib/graphql/bulk/visitors/add_typename_to_query_visitor.rb

@@ -0,0 +1,25 @@
+module GraphQL
+  module Bulk
+    module Visitors
+      class AddTypenameToQueryVisitor < GraphQL::Language::Visitor
+        def on_field(node, parent)
+          unless node.selections.empty?
+            has_typename = false
+            node.selections.each do |selection|
+              has_typename = true if selection.name == "__typename"
+            end
+
+            return super if has_typename
+
+            modified_node = node.merge_selection(
+              name: "__typename"
+            )
+            return super(modified_node, parent)
+          end
+
+          super
+        end
+      end
+    end
+  end
+end

lib/graphql/bulk/visitors/connection_node_extraction_visitor.rb

@@ -0,0 +1,23 @@
+module GraphQL
+  module Bulk
+    module Visitors
+      class ConnectionNodeExtractionVisitor < GraphQL::Language::Visitor
+        def initialize(document, connection_node)
+          super(document)
+          @connection_node = connection_node
+        end
+
+        def on_field(node, parent)
+          # Stop traversing once we find the node. No need to go any further
+          return if node == @connection_node.node
+
+          # Keep this node if it's in the path
+          return super if @connection_node.path.include?(node)
+
+          # Destroy all other nodes
+          super(DELETE_NODE, parent)
+        end
+      end
+    end
+  end
+end