Extract `execute_in_batches` migrations helper

2 files changed, 60 insertions, 22 deletions

diff --git a/lib/gitlab/database/migration_helpers.rb b/lib/gitlab/database/migration_helpers.rb
index 29e812c3c78..5acc96331c0 100644
--- a/lib/gitlab/database/migration_helpers.rb
+++ b/lib/gitlab/database/migration_helpers.rb
@@ -248,7 +248,7 @@ module Gitlab
         # rows for GitLab.com.
         batch_size = max_size if batch_size > max_size
 
-        walk_table_in_batches(table, of: batch_size, scope: scope) do
+        execute_in_batches(table, of: batch_size, scope: scope) do
           Arel::UpdateManager.new(ActiveRecord::Base)
             .table(table_arel)
             .set([[table_arel[column], value]])
@@ -281,23 +281,32 @@ module Gitlab
           stop_id = exec_query(stop_arel.to_sql)
             .to_hash.first.to_h['id'].to_i
 
-          action = yield(batch, start_id, stop_id)
+          yield batch, start_id, stop_id
 
-          if action.is_a?(Arel::TreeManager)
-            exec_arel = action.where(table[:id].gteq(start_id))
-            exec_arel = exec_arel.where(table[:id].lt(stop_id)) if stop_id.nonzero?
-            exec_arel = scope.call(table, exec_arel) if scope
+          stop_id.zero? ? break : start_id = stop_id
+        end
+      end
 
-            execute(exec_arel.to_sql)
-          end
+      def execute_in_batches(table, of: 1000, scope: nil)
+        if transaction_open?
+          raise <<-MSG
+            execute_in_batches helper can not be run inside a transaction.
+            You can disable transactions by calling `disable_ddl_transaction!`
+            method in the body of your migration class.
+          MSG
+        end
 
-          if stop_id.zero?
-            # there are no more rows left to update
-            break
-          else
-            # next loop
-            start_id = stop_id
-          end
+       # raise 'This method requires a block!' unless block_given?
+
+        table_arel = Arel::Table.new(table)
+
+        walk_table_in_batches(table, of: of, scope: scope) do |_batch, start_id, stop_id|
+          exec_arel = yield table_arel
+          exec_arel = exec_arel.where(table_arel[:id].gteq(start_id))
+          exec_arel = exec_arel.where(table_arel[:id].lt(stop_id)) if stop_id.nonzero?
+          exec_arel = scope.call(table_arel, exec_arel) if scope
+
+          execute(exec_arel.to_sql)
         end
       end
 
diff --git a/spec/lib/gitlab/database/migration_helpers_spec.rb b/spec/lib/gitlab/database/migration_helpers_spec.rb
index 65af2e1cb52..f4a66b7e2a2 100644
--- a/spec/lib/gitlab/database/migration_helpers_spec.rb
+++ b/spec/lib/gitlab/database/migration_helpers_spec.rb
@@ -262,7 +262,8 @@ describe Gitlab::Database::MigrationHelpers, lib: true do
   describe '#update_column_in_batches' do
     context 'when running outside of a transaction' do
       before do
-        expect(model).to receive(:transaction_open?).twice.and_return(false)
+        expect(model).to receive(:transaction_open?)
+          .at_least(:once).and_return(false)
 
         create_list(:empty_project, 5)
       end
@@ -336,10 +337,39 @@ describe Gitlab::Database::MigrationHelpers, lib: true do
           first_id = Project.first.id
           scope = ->(table, query) { query.where(table[:id].eq(first_id)) }
 
-          model.walk_table_in_batches(:projects, scope: scope) do
+          expect { |b| model.walk_table_in_batches(:projects, of: 1, scope: scope, &b) }
+            .to yield_control.exactly(:once)
+        end
+      end
+    end
+
+    context 'when running inside the transaction' do
+      it 'raises RuntimeError' do
+        expect(model).to receive(:transaction_open?).and_return(true)
+
+        expect { model.walk_table_in_batches(:projects, of: 2) }
+          .to raise_error(RuntimeError)
+      end
+    end
+  end
+
+  describe '#execute_in_batches' do
+    context 'when running outside of a transaction' do
+      before do
+        expect(model).to receive(:transaction_open?)
+          .at_least(:once).and_return(false)
+
+        create_list(:empty_project, 6)
+      end
+
+      context 'when a scope is provided' do
+        it 'limits the scope of the statement provided inside the block' do
+          first_id = Project.first.id
+          scope = ->(table, query) { query.where(table[:id].eq(first_id)) }
+
+          model.execute_in_batches(:projects, scope: scope) do |table|
             Arel::UpdateManager.new(ActiveRecord::Base)
-              .table(Arel::Table.new(:projects))
-              .set([[Arel::Table.new(:projects)[:archived], true]])
+              .table(table).set([[table[:archived], true]])
           end
 
           expect(Project.where(archived: true).count).to eq(1)
@@ -351,9 +381,8 @@ describe Gitlab::Database::MigrationHelpers, lib: true do
       it 'raises RuntimeError' do
         expect(model).to receive(:transaction_open?).and_return(true)
 
-        expect do
-          model.update_column_in_batches(:projects, :star_count, Arel.sql('1+1'))
-        end.to raise_error(RuntimeError)
+        expect { model.execute_in_batches(:projects)}
+          .to raise_error(RuntimeError)
       end
     end
   end