Merge branch 'keyset-iterator-pass-cursor' into 'master'

Allow passing the cursor to Keyset::Iterator

See merge request gitlab-org/gitlab!71292

lib/gitlab/pagination/keyset/iterator.rb

@@ -6,10 +6,11 @@ module Gitlab
       class Iterator
         UnsupportedScopeOrder = Class.new(StandardError)
 
-        def initialize(scope:, use_union_optimization: true, in_operator_optimization_options: nil)
+        def initialize(scope:, cursor: {}, use_union_optimization: true, in_operator_optimization_options: nil)
           @scope, success = Gitlab::Pagination::Keyset::SimpleOrderBuilder.build(scope)
           raise(UnsupportedScopeOrder, 'The order on the scope does not support keyset pagination') unless success
 
+          @cursor = cursor
           @order = Gitlab::Pagination::Keyset::Order.extract_keyset_order_object(scope)
           @use_union_optimization = in_operator_optimization_options ? false : use_union_optimization
           @in_operator_optimization_options = in_operator_optimization_options
@@ -17,11 +18,9 @@ module Gitlab
 
         # rubocop: disable CodeReuse/ActiveRecord
         def each_batch(of: 1000)
-          cursor_attributes = {}
-
           loop do
             current_scope = scope.dup
-            relation = order.apply_cursor_conditions(current_scope, cursor_attributes, keyset_options)
+            relation = order.apply_cursor_conditions(current_scope, cursor, keyset_options)
             relation = relation.reorder(order) unless @in_operator_optimization_options
             relation = relation.limit(of)
 
@@ -30,14 +29,14 @@ module Gitlab
             last_record = relation.last
             break unless last_record
 
-            cursor_attributes = order.cursor_attributes_for_node(last_record)
+            @cursor = order.cursor_attributes_for_node(last_record)
           end
         end
         # rubocop: enable CodeReuse/ActiveRecord
 
         private
 
-        attr_reader :scope, :order
+        attr_reader :scope, :cursor, :order
 
         def keyset_options
           {

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

@@ -32,8 +32,11 @@ RSpec.describe Gitlab::Pagination::Keyset::Iterator do
     ])
   end
 
+  let(:iterator_params) { nil }
   let(:scope) { project.issues.reorder(custom_reorder) }
 
+  subject(:iterator) { described_class.new(**iterator_params) }
+
   shared_examples 'iterator examples' do
     describe '.each_batch' do
       it 'yields an ActiveRecord::Relation when a block is given' do
@@ -56,6 +59,30 @@ RSpec.describe Gitlab::Pagination::Keyset::Iterator do
         expect(count).to eq(9)
       end
 
+      it 'continues after the cursor' do
+        loaded_records = []
+        cursor = nil
+
+        # stopping the iterator after the first batch and storing the cursor
+        iterator.each_batch(of: 2) do |relation| # rubocop: disable Lint/UnreachableLoop
+          loaded_records.concat(relation.to_a)
+          record = loaded_records.last
+
+          cursor = custom_reorder.cursor_attributes_for_node(record)
+          break
+        end
+
+        expect(loaded_records).to eq(project.issues.order(custom_reorder).take(2))
+
+        # continuing the iteration
+        new_iterator = described_class.new(**iterator_params.merge(cursor: cursor))
+        new_iterator.each_batch(of: 2) do |relation|
+          loaded_records.concat(relation.to_a)
+        end
+
+        expect(loaded_records).to eq(project.issues.order(custom_reorder))
+      end
+
       it 'allows updating of the yielded relations' do
         time = Time.current
 
@@ -131,13 +158,13 @@ RSpec.describe Gitlab::Pagination::Keyset::Iterator do
   end
 
   context 'when use_union_optimization is used' do
-    subject(:iterator) { described_class.new(scope: scope, use_union_optimization: true) }
+    let(:iterator_params) { { scope: scope, use_union_optimization: true } }
 
     include_examples 'iterator examples'
   end
 
   context 'when use_union_optimization is not used' do
-    subject(:iterator) { described_class.new(scope: scope, use_union_optimization: false) }
+    let(:iterator_params) { { scope: scope, use_union_optimization: false } }
 
     include_examples 'iterator examples'
   end