Optionally use the new keyset implementation

This MR uses the new keyset implementation in our GraphQL API when the
order columns can be easily determined.

config/feature_flags/development/new_graphql_keyset_pagination.yml

+---
+name: new_graphql_keyset_pagination
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/56751
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/323730
+milestone: '13.10'
+type: development
+group: group::optimize
+default_enabled: false

lib/gitlab/graphql/pagination/keyset/generic_keyset_pagination.rb

@@ -10,6 +10,8 @@ module Gitlab
           extend ActiveSupport::Concern
 
           def ordered_items
+            raise ArgumentError, 'Relation must have a primary key' unless items.primary_key.present?
+
             return super unless Gitlab::Pagination::Keyset::Order.keyset_aware?(items)
 
             items
@@ -40,6 +42,17 @@ module Gitlab
             sliced = slice_nodes(sliced, after, :after) if after.present?
             sliced
           end
+
+          def items
+            original_items = super
+            return original_items if Gitlab::Pagination::Keyset::Order.keyset_aware?(original_items) || Feature.disabled?(:new_graphql_keyset_pagination)
+
+            strong_memoize(:generic_keyset_pagination_items) do
+              rebuilt_items_with_keyset_order, success = Gitlab::Pagination::Keyset::SimpleOrderBuilder.build(original_items)
+
+              success ? rebuilt_items_with_keyset_order : original_items
+            end
+          end
         end
       end
     end

lib/gitlab/pagination/keyset/order.rb

@@ -170,6 +170,8 @@ module Gitlab
           self.class.build(column_definitions.map(&:reverse))
         end
 
+        alias_method :to_sql, :to_s
+
         private
 
         # Adds extra columns to the SELECT clause

lib/gitlab/pagination/keyset/simple_order_builder.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Pagination
+    module Keyset
+      # This class transforms the `order()` values from an Activerecord scope into a
+      # Gitlab::Pagination::Keyset::Order instance so the query later can be  used in
+      # keyset pagination.
+      #
+      # Return values:
+      # [transformed_scope, true] # true indicates that the new scope was successfully built
+      # [orginal_scope, false] # false indicates that the order values are not supported in this class
+      class SimpleOrderBuilder
+        def self.build(scope)
+          new(scope: scope).build
+        end
+
+        def initialize(scope:)
+          @scope = scope
+          @order_values = scope.order_values
+          @model_class = scope.model
+          @arel_table = @model_class.arel_table
+          @primary_key = @model_class.primary_key
+        end
+
+        def build
+          order = if order_values.empty?
+                    primary_key_descending_order
+                  elsif ordered_by_primary_key?
+                    primary_key_order
+                  elsif ordered_by_other_column?
+                    column_with_tie_breaker_order
+                  elsif ordered_by_other_column_with_tie_breaker?
+                    tie_breaker_attribute = order_values.second
+
+                    tie_breaker_column_order = Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
+                      attribute_name: model_class.primary_key,
+                      order_expression: tie_breaker_attribute
+                    )
+
+                    column_with_tie_breaker_order(tie_breaker_column_order)
+                  end
+
+          order ? [scope.reorder!(order), true] : [scope, false] # [scope, success]
+        end
+
+        private
+
+        attr_reader :scope, :order_values, :model_class, :arel_table, :primary_key
+
+        def primary_key_descending_order
+          Gitlab::Pagination::Keyset::Order.build([
+            Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
+              attribute_name: model_class.primary_key,
+              order_expression: arel_table[primary_key].desc
+            )
+          ])
+        end
+
+        def primary_key_order
+          Gitlab::Pagination::Keyset::Order.build([
+            Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
+              attribute_name: model_class.primary_key,
+              order_expression: order_values.first
+            )
+          ])
+        end
+
+        def column_with_tie_breaker_order(tie_breaker_column_order = default_tie_breaker_column_order)
+          order_expression = order_values.first
+          attribute_name = order_expression.expr.name
+
+          column_nullable = model_class.columns.find { |column| column.name == attribute_name }.null
+
+          nullable = if column_nullable && order_expression.is_a?(Arel::Nodes::Ascending)
+                       :nulls_last
+                     elsif column_nullable && order_expression.is_a?(Arel::Nodes::Descending)
+                       :nulls_first
+                     else
+                       :not_nullable
+                     end
+
+          Gitlab::Pagination::Keyset::Order.build([
+            Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
+              attribute_name: attribute_name,
+              order_expression: order_expression,
+              nullable: nullable,
+              distinct: false
+            ),
+            tie_breaker_column_order
+          ])
+        end
+
+        def ordered_by_primary_key?
+          return unless order_values.one?
+
+          attribute = order_values.first.try(:expr)
+
+          return unless attribute
+
+          arel_table[primary_key].to_s == attribute.to_s
+        end
+
+        def ordered_by_other_column?
+          return unless order_values.one?
+
+          attribute = order_values.first.try(:expr)
+
+          return unless attribute
+          return unless attribute.try(:name)
+
+          model_class.column_names.include?(attribute.name.to_s)
+        end
+
+        def ordered_by_other_column_with_tie_breaker?
+          return unless order_values.size == 2
+
+          attribute = order_values.first.try(:expr)
+          tie_breaker_attribute = order_values.second.try(:expr)
+
+          return unless attribute
+          return unless tie_breaker_attribute
+
+          model_class.column_names.include?(attribute.name.to_s) &&
+            arel_table[primary_key].to_s == tie_breaker_attribute.to_s
+        end
+
+        def default_tie_breaker_column_order
+          Gitlab::Pagination::Keyset::ColumnOrderDefinition.new(
+            attribute_name: model_class.primary_key,
+            order_expression: arel_table[primary_key].desc
+          )
+        end
+      end
+    end
+  end
+end

spec/lib/gitlab/graphql/pagination/keyset/connection_spec.rb

@@ -357,9 +357,10 @@ RSpec.describe Gitlab::Graphql::Pagination::Keyset::Connection do
 
       it 'is added to end' do
         sliced = subject.sliced_nodes
-        last_order_name = sliced.order_values.last.expr.name
 
-        expect(last_order_name).to eq sliced.primary_key
+        order_sql = sliced.order_values.last.to_sql
+
+        expect(order_sql).to end_with(Project.arel_table[:id].desc.to_sql)
       end
     end
 

spec/lib/gitlab/pagination/keyset/simple_order_builder_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Pagination::Keyset::SimpleOrderBuilder do
+  let(:ordered_scope) { described_class.build(scope).first }
+  let(:order_object) { Gitlab::Pagination::Keyset::Order.extract_keyset_order_object(ordered_scope) }
+
+  subject(:sql_with_order) { ordered_scope.to_sql }
+
+  context 'when no order present' do
+    let(:scope) { Project.where(id: [1, 2, 3]) }
+
+    it 'orders by primary key' do
+      expect(sql_with_order).to end_with('ORDER BY "projects"."id" DESC')
+    end
+
+    it 'sets the column definition distinct and not nullable' do
+      column_definition = order_object.column_definitions.first
+
+      expect(column_definition).to be_not_nullable
+      expect(column_definition).to be_distinct
+    end
+  end
+
+  context 'when primary key order present' do
+    let(:scope) { Project.where(id: [1, 2, 3]).order(id: :asc) }
+
+    it 'orders by primary key without altering the direction' do
+      expect(sql_with_order).to end_with('ORDER BY "projects"."id" ASC')
+    end
+  end
+
+  context 'when ordered by other column' do
+    let(:scope) { Project.where(id: [1, 2, 3]).order(created_at: :asc) }
+
+    it 'adds extra primary key order as tie-breaker' do
+      expect(sql_with_order).to end_with('ORDER BY "projects"."created_at" ASC, "projects"."id" DESC')
+    end
+
+    it 'sets the column definition for created_at non-distinct and nullable' do
+      column_definition = order_object.column_definitions.first
+
+      expect(column_definition.attribute_name).to eq('created_at')
+      expect(column_definition.nullable?).to eq(true) # be_nullable calls non_null? method for some reason
+      expect(column_definition).not_to be_distinct
+    end
+  end
+
+  context 'when ordered by two columns where the last one is the tie breaker' do
+    let(:scope) { Project.where(id: [1, 2, 3]).order(created_at: :asc, id: :asc) }
+
+    it 'preserves the order' do
+      expect(sql_with_order).to end_with('ORDER BY "projects"."created_at" ASC, "projects"."id" ASC')
+    end
+  end
+
+  context 'when non-nullable column is given' do
+    let(:scope) { Project.where(id: [1, 2, 3]).order(namespace_id: :asc, id: :asc) }
+
+    it 'sets the column definition for namespace_id non-distinct and non-nullable' do
+      column_definition = order_object.column_definitions.first
+
+      expect(column_definition.attribute_name).to eq('namespace_id')
+      expect(column_definition).to be_not_nullable
+      expect(column_definition).not_to be_distinct
+    end
+  end
+
+  context 'return :unable_to_order symbol when order cannot be built' do
+    subject(:success) { described_class.build(scope).last }
+
+    context 'when raw SQL order is given' do
+      let(:scope) { Project.order('id DESC') }
+
+      it { is_expected.to eq(false) }
+    end
+
+    context 'when NULLS LAST order is given' do
+      let(:scope) { Project.order(::Gitlab::Database.nulls_last_order('created_at', 'ASC')) }
+
+      it { is_expected.to eq(false) }
+    end
+
+    context 'when more than 2 columns are given for the order' do
+      let(:scope) { Project.order(created_at: :asc, updated_at: :desc, id: :asc) }
+
+      it { is_expected.to eq(false) }
+    end
+  end
+end

spec/requests/api/graphql_spec.rb

@@ -283,25 +283,50 @@ RSpec.describe 'GraphQL' do
       )
     end
 
-    it 'paginates datetimes correctly when they have millisecond data' do
-      # let's make sure we're actually querying a timestamp, just in case
-      expect(Gitlab::Graphql::Pagination::Keyset::QueryBuilder)
-        .to receive(:new).with(anything, anything, hash_including('created_at'), anything).and_call_original
+    context 'when new_graphql_keyset_pagination feature flag is off' do
+      before do
+        stub_feature_flags(new_graphql_keyset_pagination: false)
+      end
+
+      it 'paginates datetimes correctly when they have millisecond data' do
+        # let's make sure we're actually querying a timestamp, just in case
+        expect(Gitlab::Graphql::Pagination::Keyset::QueryBuilder)
+          .to receive(:new).with(anything, anything, hash_including('created_at'), anything).and_call_original
+
+        execute_query
+        first_page = graphql_data
+        edges = first_page.dig(*issues_edges)
+        cursor = first_page.dig(*end_cursor)
+
+        expect(edges.count).to eq(6)
+        expect(edges.last['node']['iid']).to eq(issues[4].iid.to_s)
 
-      execute_query
-      first_page = graphql_data
-      edges = first_page.dig(*issues_edges)
-      cursor = first_page.dig(*end_cursor)
+        execute_query(after: cursor)
+        second_page = graphql_data
+        edges = second_page.dig(*issues_edges)
 
-      expect(edges.count).to eq(6)
-      expect(edges.last['node']['iid']).to eq(issues[4].iid.to_s)
+        expect(edges.count).to eq(4)
+        expect(edges.last['node']['iid']).to eq(issues[0].iid.to_s)
+      end
+    end
+
+    context 'when new_graphql_keyset_pagination feature flag is on' do
+      it 'paginates datetimes correctly when they have millisecond data' do
+        execute_query
+        first_page = graphql_data
+        edges = first_page.dig(*issues_edges)
+        cursor = first_page.dig(*end_cursor)
 
-      execute_query(after: cursor)
-      second_page = graphql_data
-      edges = second_page.dig(*issues_edges)
+        expect(edges.count).to eq(6)
+        expect(edges.last['node']['iid']).to eq(issues[4].iid.to_s)
 
-      expect(edges.count).to eq(4)
-      expect(edges.last['node']['iid']).to eq(issues[0].iid.to_s)
+        execute_query(after: cursor)
+        second_page = graphql_data
+        edges = second_page.dig(*issues_edges)
+
+        expect(edges.count).to eq(4)
+        expect(edges.last['node']['iid']).to eq(issues[0].iid.to_s)
+      end
     end
   end
 end