Faster next page check for GraphQL-ruby

smart_relation_connection.rb

# frozen_string_literal: true
module GraphQL
  module Relay
    # A connection implementation to expose SQL collection objects.
    # It works for:
    # - `ActiveRecord::Relation`
    # - `Sequel::Dataset`
    class SmartRelationConnection < BaseConnection
      def cursor_from_node(item)
        item_index = paged_nodes.index(item)
        if item_index.nil?
          raise("Can't generate cursor, item not found in connection: #{item}")
        else
          offset = item_index + 1 + ((paged_nodes_offset || 0) - (relation_offset(sliced_nodes) || 0))

          if after
            offset += offset_from_cursor(after)
          end

          encode(offset.to_s)
        end
      end

      def has_next_page
        if first
          @has_next_page
        else
          false
        end
      end

      def first
        return @first if defined? @first

        @first = get_limited_arg(:first)
        @first = max_page_size if !first || (@first && max_page_size && @first > max_page_size)
        @first
      end

      private

      # apply first limit results
      # @return [Array]
      def paged_nodes
        return @paged_nodes if defined? @paged_nodes

        items = sliced_nodes.limit(first + 1)
        items_array = items.to_a
        if items_array.count > first
          @has_next_page = true
          items_array = items_array[0..-2]
        else
          @has_next_page = false
        end

        # Store this here so we can convert the relation to an Array
        # (this avoids an extra DB call on Sequel)
        @paged_nodes_offset = relation_offset(items)

        @paged_nodes = items_array
      end

      def paged_nodes_offset
        paged_nodes && @paged_nodes_offset
      end

      def relation_offset(relation)
        if relation.respond_to?(:offset_value)
          relation.offset_value
        else
          relation.opts[:offset]
        end
      end

      # Apply cursors to edges
      def sliced_nodes
        return @sliced_nodes if defined? @sliced_nodes

        @sliced_nodes = nodes

        if after
          offset = (relation_offset(@sliced_nodes) || 0) + offset_from_cursor(after)
          @sliced_nodes = @sliced_nodes.offset(offset)
        end

        @sliced_nodes
      end

      def offset_from_cursor(cursor)
        decode(cursor).to_i
      end
    end

    if defined?(ActiveRecord::Relation)
      BaseConnection.register_connection_implementation(ActiveRecord::Relation, SmartRelationConnection)
    end
    if defined?(Sequel::Dataset)
      BaseConnection.register_connection_implementation(Sequel::Dataset, SmartRelationConnection)
    end
  end
end