Split keyset pagination in two steps

Keyset pagination steps:
1. Apply LIMIT clause
1. Add header with link to next page

(2) is based on the last record in the current page, to derive the
filter values for the next page.

With splitting this into two explicit steps, we make sure that (2)
always operates on the actual data rather than - possibly - an
intermediate state of the database query.

lib/api/helpers.rb

@@ -4,6 +4,7 @@ module API
   module Helpers
     include Gitlab::Utils
     include Helpers::Pagination
+    include Helpers::PaginationStrategies
 
     SUDO_HEADER = "HTTP_SUDO"
     GITLAB_SHARED_SECRET_HEADER = "Gitlab-Shared-Secret"

lib/api/helpers/pagination.rb

@@ -4,30 +4,8 @@ module API
   module Helpers
     module Pagination
       def paginate(relation)
-        return paginate_with_offset(relation) unless keyset_pagination_enabled?
-
-        paginate_with_key(relation)
-      end
-
-      private
-
-      def paginate_with_offset(relation)
         Gitlab::Pagination::OffsetPagination.new(self).paginate(relation)
       end
-
-      def paginate_with_key(relation)
-        request_context = Gitlab::Pagination::Keyset::RequestContext.new(self)
-
-        unless Gitlab::Pagination::Keyset.available?(request_context, relation)
-          return error!('Keyset pagination is not yet available for this type of request', 405)
-        end
-
-        Gitlab::Pagination::Keyset.paginate(request_context, relation)
-      end
-
-      def keyset_pagination_enabled?
-        params[:pagination] == 'keyset' && Feature.enabled?(:api_keyset_pagination, default_enabled: true)
-      end
     end
   end
 end

lib/api/helpers/pagination_strategies.rb

+# frozen_string_literal: true
+
+module API
+  module Helpers
+    module PaginationStrategies
+      def paginate_with_strategies(relation)
+        paginator = paginator(relation)
+
+        yield(paginator.paginate(relation)).tap do |records, _|
+          paginator.finalize(records)
+        end
+      end
+
+      def paginator(relation)
+        return Gitlab::Pagination::OffsetPagination.new(self) unless keyset_pagination_enabled?
+
+        request_context = Gitlab::Pagination::Keyset::RequestContext.new(self)
+
+        unless Gitlab::Pagination::Keyset.available?(request_context, relation)
+          return error!('Keyset pagination is not yet available for this type of request', 405)
+        end
+
+        Gitlab::Pagination::Keyset::Pager.new(request_context)
+      end
+
+      private
+
+      def keyset_pagination_enabled?
+        params[:pagination] == 'keyset' && Feature.enabled?(:api_keyset_pagination, default_enabled: true)
+      end
+    end
+  end
+end

lib/api/projects.rb

@@ -90,18 +90,22 @@ module API
       def present_projects(projects, options = {})
         projects = reorder_projects(projects)
         projects = apply_filters(projects)
-        projects = paginate(projects)
-        projects, options = with_custom_attributes(projects, options)
 
-        options = options.reverse_merge(
-          with: current_user ? Entities::ProjectWithAccess : Entities::BasicProjectDetails,
-          statistics: params[:statistics],
-          current_user: current_user,
-          license: false
-        )
-        options[:with] = Entities::BasicProjectDetails if params[:simple]
+        records, options = paginate_with_strategies(projects) do |projects|
+          projects, options = with_custom_attributes(projects, options)
+
+          options = options.reverse_merge(
+            with: current_user ? Entities::ProjectWithAccess : Entities::BasicProjectDetails,
+            statistics: params[:statistics],
+            current_user: current_user,
+            license: false
+          )
+          options[:with] = Entities::BasicProjectDetails if params[:simple]
+
+          [options[:with].prepare_relation(projects, options), options]
+        end
 
-        present options[:with].prepare_relation(projects, options), options
+        present records, options
       end
 
       def translate_params_for_compatibility(params)

lib/gitlab/pagination/keyset.rb

@@ -3,10 +3,6 @@
 module Gitlab
   module Pagination
     module Keyset
-      def self.paginate(request_context, relation)
-        Gitlab::Pagination::Keyset::Pager.new(request_context).paginate(relation)
-      end
-
       def self.available?(request_context, relation)
         order_by = request_context.page.order_by
 

lib/gitlab/pagination/keyset/pager.rb

@@ -14,10 +14,11 @@ module Gitlab
           # Validate assumption: The last two columns must match the page order_by
           validate_order!(relation)
 
-          # This performs the query and retrieves records
-          relation.limit(page.per_page).load.tap do |records| # rubocop: disable CodeReuse/ActiveRecord
-            apply_headers(records.last)
-          end
+          relation.limit(page.per_page) # rubocop: disable CodeReuse/ActiveRecord
+        end
+
+        def finalize(records)
+          apply_headers(records.last)
         end
 
         private

lib/gitlab/pagination/offset_pagination.rb

@@ -16,6 +16,10 @@ module Gitlab
         end
       end
 
+      def finalize(records)
+        # nothing to do
+      end
+
       private
 
       def paginate_with_limit_optimization(relation)

spec/lib/api/helpers/pagination_spec.rb

@@ -5,56 +5,14 @@ require 'spec_helper'
 describe API::Helpers::Pagination do
   subject { Class.new.include(described_class).new }
 
-  let(:expected_result) { double("result") }
-  let(:relation) { double("relation") }
-  let(:params) { {} }
+  let(:paginator) { double('paginator') }
+  let(:relation) { double('relation') }
+  let(:expected_result) { double('expected result') }
 
-  before do
-    allow(subject).to receive(:params).and_return(params)
-  end
-
-  describe '#paginate' do
-    let(:offset_pagination) { double("offset pagination") }
-
-    context 'for offset pagination' do
-      it 'delegates to OffsetPagination' do
-        expect(::Gitlab::Pagination::OffsetPagination).to receive(:new).with(subject).and_return(offset_pagination)
-        expect(offset_pagination).to receive(:paginate).with(relation).and_return(expected_result)
-
-        result = subject.paginate(relation)
-
-        expect(result).to eq(expected_result)
-      end
-    end
-
-    context 'for keyset pagination' do
-      let(:params) { { pagination: 'keyset' } }
-      let(:request_context) { double('request context') }
-
-      before do
-        allow(Gitlab::Pagination::Keyset::RequestContext).to receive(:new).with(subject).and_return(request_context)
-      end
-
-      context 'when keyset pagination is available' do
-        it 'delegates to KeysetPagination' do
-          expect(Gitlab::Pagination::Keyset).to receive(:available?).and_return(true)
-          expect(Gitlab::Pagination::Keyset).to receive(:paginate).with(request_context, relation).and_return(expected_result)
-
-          result = subject.paginate(relation)
-
-          expect(result).to eq(expected_result)
-        end
-      end
-
-      context 'when keyset pagination is not available' do
-        it 'renders a 501 error if keyset pagination isnt available yet' do
-          expect(Gitlab::Pagination::Keyset).to receive(:available?).with(request_context, relation).and_return(false)
-          expect(Gitlab::Pagination::Keyset).not_to receive(:paginate)
-          expect(subject).to receive(:error!).with(/not yet available/, 405)
+  it 'delegates to OffsetPagination' do
+    expect(Gitlab::Pagination::OffsetPagination).to receive(:new).with(subject).and_return(paginator)
+    expect(paginator).to receive(:paginate).with(relation).and_return(expected_result)
 
-          subject.paginate(relation)
-        end
-      end
-    end
+    expect(subject.paginate(relation)).to eq(expected_result)
   end
 end

spec/lib/api/helpers/pagination_strategies_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe API::Helpers::PaginationStrategies do
+  subject { Class.new.include(described_class).new }
+
+  let(:expected_result) { double("result") }
+  let(:relation) { double("relation") }
+  let(:params) { {} }
+
+  before do
+    allow(subject).to receive(:params).and_return(params)
+  end
+
+  describe '#paginate_with_strategies' do
+    let(:paginator) { double("paginator", paginate: expected_result, finalize: nil) }
+
+    before do
+      allow(subject).to receive(:paginator).with(relation).and_return(paginator)
+    end
+
+    it 'yields paginated relation' do
+      expect { |b| subject.paginate_with_strategies(relation, &b) }.to yield_with_args(expected_result)
+    end
+
+    it 'calls #finalize with first value returned from block' do
+      return_value = double
+      expect(paginator).to receive(:finalize).with(return_value)
+
+      subject.paginate_with_strategies(relation) do |records|
+        some_options = {}
+        [return_value, some_options]
+      end
+    end
+
+    it 'returns whatever the block returns' do
+      return_value = [double, double]
+
+      result = subject.paginate_with_strategies(relation) do |records|
+        return_value
+      end
+
+      expect(result).to eq(return_value)
+    end
+  end
+
+  describe '#paginator' do
+    context 'offset pagination' do
+      let(:paginator) { double("paginator") }
+
+      before do
+        allow(subject).to receive(:keyset_pagination_enabled?).and_return(false)
+      end
+
+      it 'delegates to OffsetPagination' do
+        expect(Gitlab::Pagination::OffsetPagination).to receive(:new).with(subject).and_return(paginator)
+
+        expect(subject.paginator(relation)).to eq(paginator)
+      end
+    end
+
+    context 'for keyset pagination' do
+      let(:params) { { pagination: 'keyset' } }
+      let(:request_context) { double('request context') }
+      let(:pager) { double('pager') }
+
+      before do
+        allow(subject).to receive(:keyset_pagination_enabled?).and_return(true)
+        allow(Gitlab::Pagination::Keyset::RequestContext).to receive(:new).with(subject).and_return(request_context)
+      end
+
+      context 'when keyset pagination is available' do
+        before do
+          allow(Gitlab::Pagination::Keyset).to receive(:available?).and_return(true)
+          allow(Gitlab::Pagination::Keyset::Pager).to receive(:new).with(request_context).and_return(pager)
+        end
+
+        it 'delegates to Pager' do
+          expect(subject.paginator(relation)).to eq(pager)
+        end
+      end
+
+      context 'when keyset pagination is not available' do
+        before do
+          allow(Gitlab::Pagination::Keyset).to receive(:available?).with(request_context, relation).and_return(false)
+        end
+
+        it 'renders a 501 error' do
+          expect(subject).to receive(:error!).with(/not yet available/, 405)
+
+          subject.paginator(relation)
+        end
+      end
+    end
+  end
+end

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

@@ -15,14 +15,32 @@ describe Gitlab::Pagination::Keyset::Pager do
   describe '#paginate' do
     subject { described_class.new(request).paginate(relation) }
 
-    it 'loads the result relation only once' do
+    it 'does not execute a query' do
       expect do
         subject
-      end.not_to exceed_query_limit(1)
+      end.not_to exceed_query_limit(0)
     end
 
+    it 'returns a limited relation' do
+      expect(subject).to eq(relation.limit(page.per_page))
+    end
+
+    context 'validating the order clause' do
+      let(:page) { Gitlab::Pagination::Keyset::Page.new(order_by: { created_at: :asc }, per_page: 3) }
+
+      it 'raises an error if has a different order clause than the page' do
+        expect { subject }.to raise_error(ArgumentError, /order_by does not match/)
+      end
+    end
+  end
+
+  describe '#finalize' do
+    let(:records) { relation.limit(page.per_page).load }
+
+    subject { described_class.new(request).finalize(records) }
+
     it 'passes information about next page to request' do
-      lower_bounds = relation.limit(page.per_page).last.slice(:id)
+      lower_bounds = records.last.slice(:id)
       expect(page).to receive(:next).with(lower_bounds, false).and_return(next_page)
       expect(request).to receive(:apply_headers).with(next_page)
 
@@ -52,17 +70,5 @@ describe Gitlab::Pagination::Keyset::Pager do
         subject
       end
     end
-
-    it 'returns a limited relation' do
-      expect(subject).to eq(relation.limit(page.per_page))
-    end
-
-    context 'validating the order clause' do
-      let(:page) { Gitlab::Pagination::Keyset::Page.new(order_by: { created_at: :asc }, per_page: 3) }
-
-      it 'raises an error if has a different order clause than the page' do
-        expect { subject }.to raise_error(ArgumentError, /order_by does not match/)
-      end
-    end
   end
 end

spec/lib/gitlab/pagination/keyset_spec.rb

@@ -3,22 +3,6 @@
 require 'spec_helper'
 
 describe Gitlab::Pagination::Keyset do
-  describe '.paginate' do
-    subject { described_class.paginate(request_context, relation) }
-
-    let(:request_context) { double }
-    let(:relation) { double }
-    let(:pager) { double }
-    let(:result) { double }
-
-    it 'uses Pager to paginate the relation' do
-      expect(Gitlab::Pagination::Keyset::Pager).to receive(:new).with(request_context).and_return(pager)
-      expect(pager).to receive(:paginate).with(relation).and_return(result)
-
-      expect(subject).to eq(result)
-    end
-  end
-
   describe '.available?' do
     subject { described_class }