graphql-fancy_loader 0.1.0

data/lib/graphql/fancy_connection.rb

@@ -0,0 +1,103 @@
+class GraphQL::FancyConnection < GraphQL::Pagination::RelationConnection
+  def initialize(loader, args, key, **super_args)
+    @loader = loader
+    @args = args
+    @key = key
+    @then = nil
+
+    super(nil, **super_args)
+  end
+
+  # @return [Promise<Array<ApplicationRecord>>]
+  def nodes
+    if @then
+      base_nodes.then(@then)
+    else
+      base_nodes
+    end
+  end
+
+  def edges
+    @edges ||= nodes.then do |nodes|
+      nodes.map { |n| @edge_class.new(n, self) }
+    end
+  end
+
+  # @return [Promise<Integer>]
+  def total_count
+    base_nodes.then do |results|
+      if results.first
+        results.first.attributes['total_count']
+      else
+        0
+      end
+    end
+  end
+
+  # @return [Promise<Boolean>]
+  def has_next_page # rubocop:disable Naming/PredicateName
+    base_nodes.then do |results|
+      if results.last
+        results.last.attributes['row_number'] < results.last.attributes['total_count']
+      else
+        false
+      end
+    end
+  end
+
+  # @return [Promise<Boolean>]
+  def has_previous_page # rubocop:disable Naming/PredicateName
+    base_nodes.then do |results|
+      if results.first
+        results.first.attributes['row_number'] > 1
+      else
+        false
+      end
+    end
+  end
+
+  def start_cursor
+    base_nodes.then do |results|
+      cursor_for(results.first)
+    end
+  end
+
+  def end_cursor
+    base_nodes.then do |results|
+      cursor_for(results.last)
+    end
+  end
+
+  def cursor_for(item)
+    item && encode(item.attributes['row_number'].to_s)
+  end
+
+  def then(&block)
+    @then = block
+    self
+  end
+
+  private
+
+  def base_nodes
+    @base_nodes ||= @loader.for(**loader_args).load(@key)
+  end
+
+  def after_offset
+    @after_offset ||= after && decode(after).to_i
+  end
+
+  def before_offset
+    @before_offset ||= before && decode(before).to_i
+  end
+
+  def loader_args
+    @args.merge(
+      before: before_offset,
+      after: after_offset,
+      first: first,
+      last: last,
+      context: context
+    )
+  end
+end

data/lib/graphql/fancy_loader/dsl.rb

@@ -0,0 +1,31 @@
+module GraphQL
+  class FancyLoader
+    module DSL
+      extend ActiveSupport::Concern
+
+      included do
+        class_attribute :model
+        class_attribute :sorts
+        class_attribute :modify_query_lambda
+      end
+
+      class_methods do
+        def from(model)
+          self.model = model
+        end
+
+        def sort(name, transform: nil, on: -> { model.arel_table[name] })
+          self.sorts ||= {}
+          sorts[name] = {
+            transform: transform,
+            column: on
+          }
+        end
+
+        def modify_query(lambda)
+          self.modify_query_lambda = lambda
+        end
+      end
+    end
+  end
+end

data/lib/graphql/fancy_loader/pagination_filter.rb

@@ -0,0 +1,64 @@
+module GraphQL
+  class FancyLoader
+    class PaginationFilter
+      def initialize(query, before: nil, after: nil, first: nil, last: nil)
+        @query = query
+        @before = before
+        @after = after
+        @first = first
+        @last = last
+      end
+
+      def arel
+        [
+          after_filter,
+          before_filter,
+          first_filter,
+          last_filter
+        ].compact.inject(&:and)
+      end
+
+      private
+
+      def row
+        @row ||= @query[:row_number]
+      end
+
+      def count
+        @count ||= @query[:total_count]
+      end
+
+      def after_filter
+        return unless @after
+
+        row.gt(@after)
+      end
+
+      def before_filter
+        return unless @before
+
+        row.lt(@before)
+      end
+
+      def first_filter
+        return unless @first
+
+        if @after
+          row.lteq(@after + @first)
+        else
+          row.lteq(@first)
+        end
+      end
+
+      def last_filter
+        return unless @last
+
+        if @before
+          row.gteq(@before - @last)
+        else
+          row.gt(Arel::Nodes::Subtraction.new(count, @last))
+        end
+      end
+    end
+  end
+end

data/lib/graphql/fancy_loader/pundit_middleware.rb

@@ -0,0 +1,13 @@
+module GraphQL
+  class FancyLoader::PunditMiddleware
+    def initialize(key:)
+      @key = key
+    end
+
+    def call(model:, query:, context:)
+      scope = ::Pundit::PolicyFinder.new(model).scope!
+      user = context[@key]
+      scope.new(user, query).resolve
+    end
+  end
+end

data/lib/graphql/fancy_loader/query_generator.rb

@@ -0,0 +1,125 @@
+# @private
+module GraphQL
+  class FancyLoader
+    class QueryGenerator
+      # @param model [ActiveRecord::Model] the model to load from
+      # @param find_by [Symbol, String, Array<Symbol, String>] the key or keys to find by
+      # @param sort [Array<{:column, :transform, :direction => Object}>] The sorts to apply
+      # @param keys [Array] an array of values to find by
+      # @param before [Integer] Filter by rows less than this (one-indexed)
+      # @param after [Integer] Filter by rows greater than this (one-indexed)
+      # @param first [Integer] Filter for first N rows
+      # @param last [Integer] Filter for last N rows
+      # @param where [Hash] a filter to use when querying
+      # @param context [Context] The context of the graphql query. Can be used inside of modify_query.
+      # @param modify_query [Lambda] An escape hatch to FancyLoader to allow modifying
+      #  the base_query before it generates the rest of the query
+      def initialize(
+        model:, find_by:, sort:, keys:,
+        before: nil, after: 0, first: nil, last: nil,
+        where: nil, context: {}, modify_query: nil
+      )
+        @model = model
+        @find_by = find_by
+        @sort = sort
+        @keys = keys
+        @before = before
+        @after = after
+        @first = first
+        @last = last
+        @where = where
+        @context = context
+        @modify_query = modify_query
+      end
+
+      def query
+        # Finally, go *back* to the ActiveRecord model, and do the final select
+        @model.unscoped
+              .select(Arel.star)
+              .from(subquery)
+              .where(pagination_filter(subquery))
+              .order(subquery[:row_number].asc)
+      end
+
+      private
+
+      attr_reader :context
+
+      # The underlying Arel table for the model
+      def table
+        @table ||= @model.arel_table
+      end
+
+      # A window function partition clause to apply the sort within each window
+      #
+      #   PARTITION BY #{find_by} ORDER BY #{orders}
+      def partition
+        @partition ||= begin
+          # Every sort has a column and a direction, apply them
+          orders = @sort.map do |sort|
+            sort[:column].call.public_send(sort[:direction])
+          end
+
+          Arel::Nodes::Window.new.partition(table[@find_by]).order(orders)
+        end
+      end
+
+      # Our actual window function.
+      #
+      #   ROW_NUMBER() OVER (#{partition})
+      def row_number
+        Arel::Nodes::NamedFunction.new('ROW_NUMBER', []).over(partition).as('row_number')
+      end
+
+      # A count window function. Omits sort from the partition to get the total count.
+      #
+      #   COUNT(*) OVER (#{partition})
+      def count
+        count_partition = Arel::Nodes::Window.new.partition(table[@find_by])
+        Arel::Nodes::NamedFunction.new('COUNT', [Arel.star]).over(count_partition).as('total_count')
+      end
+
+      def pagination_filter(query)
+        @pagination_filter ||= GraphQL::FancyLoader::PaginationFilter.new(
+          query,
+          before: @before,
+          after: @after,
+          first: @first,
+          last: @last
+        ).arel
+      end
+
+      # The "base" query. This is the query that would load everything without pagination or sorting,
+      # just auth scoping.
+      def base_query
+        query = @model.where(@find_by => @keys)
+        query = query.where(@where) unless @where.nil?
+        query = middleware(query: query)
+        query.arel
+      end
+
+      def subquery
+        @subquery ||= begin
+          # Apply the sort transforms and add the window function to our projection
+          subquery = @sort.inject(base_query) do |arel, sort|
+            sort[:transform] ? sort[:transform].call(arel, context) : arel
+          end
+
+          subquery = subquery.project(row_number).project(count)
+          subquery = instance_exec(subquery, &@modify_query) unless @modify_query.nil?
+          subquery.as('subquery')
+        end
+      end
+
+      def middleware(query:)
+        return query if GraphQL::FancyLoader.middleware.blank?
+
+        GraphQL::FancyLoader.middleware.each do |klass|
+          query = klass.call(model: @model, query: query, context: context)
+        end
+
+        query
+      end
+    end
+  end
+end

data/lib/graphql/fancy_loader/rank_query_generator.rb

@@ -0,0 +1,40 @@
+# @private
+module GraphQL
+  class FancyLoader
+    class RankQueryGenerator
+      # @param column [Symbol] The table column to rank by
+      # @param partition_by [Symbol] The find_by key for the table
+      # @param table [Arel::Table]
+      # @param name_suffix [String] The suffix the be used for the column name
+      def initialize(column:, partition_by:, table:, name_suffix: '_rank')
+        @column = column
+        @partition_by = partition_by
+        @table = table
+        @name_suffix = name_suffix
+      end
+
+      # Our actual window function.
+      #
+      #   ROW_NUMBER() OVER (#{partition})
+      def arel
+        Arel::Nodes::NamedFunction.new('ROW_NUMBER', []).over(partition).as(name)
+      end
+
+      private
+
+      def name
+        return @column if @name_suffix.blank?
+
+        "#{@column}#{@name_suffix}"
+      end
+
+      def partition
+        @partition ||= Arel::Nodes::Window.new.partition(@table[@partition_by]).order(order)
+      end
+
+      def order
+        @table[@column].asc
+      end
+    end
+  end
+end

data/lib/graphql/fancy_loader/type_generator.rb

@@ -0,0 +1,42 @@
+##
+# Generate parameter types for sorting
+module GraphQL
+  class FancyLoader
+    class TypeGenerator
+      def initialize(loader, name: loader.model.name)
+        @loader = loader
+        @name = name
+      end
+
+      def sorts_enum
+        @sorts_enum ||= begin
+          sorts = @loader.sorts
+          name = "#{@name}SortEnum"
+
+          Class.new(GraphQL::Schema::Enum) do
+            graphql_name name
+            sorts.each_key do |sort_name|
+              value(sort_name.to_s.underscore.upcase, value: sort_name)
+            end
+          end
+        end
+      end
+
+      def sorts_option
+        @sorts_option ||= begin
+          enum = sorts_enum
+          name = "#{@name}SortOption"
+          Class.new(GraphQL::Schema::InputObject) do
+            graphql_name name
+            argument :on, enum, required: true
+            argument :direction, GraphQL::SortDirection, required: true
+          end
+        end
+      end
+
+      def sorts_list
+        @sorts_list ||= GraphQL::Schema::List.new(sorts_option)
+      end
+    end
+  end
+end

data/lib/graphql/fancy_loader/version.rb

@@ -0,0 +1,5 @@
+module GraphQL
+  class FancyLoader
+    VERSION = '0.1.0'.freeze
+  end
+end

data/lib/graphql/fancy_loader.rb

@@ -0,0 +1,103 @@
+# +FancyLoader+ allows for easily batching complex custom sorts and pagination. It does so through
+# heavy use of complex Arel and Postgres window functions, but has performance attributes that make
+# it worthwhile.
+#
+# To use +FancyLoader+, you'll make a subclass to define your sorts and source model. You can then
+# create a field which uses your subclass to load data.
+
+require 'graphql/batch'
+require 'active_support'
+require 'active_support/concern'
+require 'active_support/configurable'
+require 'active_support/core_ext/class/attribute'
+
+require 'graphql/sort_direction'
+require 'graphql/fancy_connection'
+# FancyLoader
+require 'graphql/fancy_loader/dsl'
+require 'graphql/fancy_loader/pagination_filter'
+require 'graphql/fancy_loader/query_generator'
+require 'graphql/fancy_loader/rank_query_generator'
+require 'graphql/fancy_loader/type_generator'
+# Middleware
+require 'graphql/fancy_loader/pundit_middleware'
+
+module GraphQL
+  class FancyLoader < GraphQL::Batch::Loader
+    include ActiveSupport::Configurable
+    include GraphQL::FancyLoader::DSL
+
+    config_accessor :middleware
+
+    # Get an autogenerated GraphQL type for an order input
+    def self.sort_argument
+      @sort_argument ||= GraphQL::FancyLoader::TypeGenerator.new(self).sorts_list
+    end
+
+    # Get a FancyConnection wrapping this Loader
+    def self.connection_for(args, key)
+      GraphQL::FancyConnection.new(self, args.except(:context), key, **args.slice(:context))
+    end
+
+    # Initialize a FancyLoader. This takes all the keys which are used to batch, which is a *lot* of
+    # them. Thanks to the design of GraphQL, however, the frequently-called fields also tend to have
+    # the same parameters each time. This means that we can get away with this less-than-ideal
+    # batching and still have significant performance gains.
+    #
+    # The pagination parameters have some odd interactions to be aware of! They are *intersected*, so
+    # if you pass before and after, you're specifying after < row < before. That's pretty logical,
+    # but first+last are weirder, because when combined they will return the *middle*, due to that
+    # intersection-driven logic. That is, given a set of 10 rows, first=6 & last=6 will return rows
+    # 4, 5, and 6 because they are the only ones in both sets. This isn't a particularly useful
+    # behavior, but the Relay spec is pretty clear that you shouldn't expect good results if you pass
+    # both first and last to the same field.
+    #
+    # @param find_by [Symbol, String] the key to find by
+    # @param before [Integer] Filter by rows less than this
+    # @param after [Integer] Filter by rows greater than this
+    # @param first [Integer] Filter for first N rows
+    # @param last [Integer] Filter for last N rows
+    # @param sort [Array<{:on, :direction => Symbol}>] The sorts to apply while loading
+    def initialize(find_by:, sort:, before: nil, after: 0, first: nil, last: nil, where: nil, context: {})
+      @find_by = find_by
+      @sort = sort.map(&:to_h)
+      @before = before
+      @after = after
+      @first = first
+      @last = last
+      @where = where
+      @context = context
+    end
+
+    # Perform the loading. Uses {GraphQL::FancyLoader::QueryGenerator} to build a query, then groups
+    # the results by the @find_by column, then fulfills all the Promises.
+    def perform(keys)
+      query = QueryGenerator.new(
+        model: model,
+        find_by: @find_by,
+        before: @before,
+        after: @after,
+        first: @first,
+        last: @last,
+        sort: sort,
+        keys: keys,
+        where: @where,
+        context: @context,
+        modify_query: modify_query_lambda
+      ).query
+
+      results = query.to_a.group_by { |rec| rec[@find_by] }
+      keys.each do |key|
+        fulfill(key, results[key] || [])
+      end
+    end
+
+    private
+
+    def sort
+      @sort.map do |sort|
+        sorts[sort[:on]].merge(direction: sort[:direction])
+      end
+    end
+  end
+end

data/lib/graphql/sort_direction.rb

@@ -0,0 +1,4 @@
+class GraphQL::SortDirection < GraphQL::Schema::Enum
+  value :ASCENDING, value: :asc
+  value :DESCENDING, value: :desc
+end