Revert recent changes in migration helpers

lib/gitlab/database/migration_helpers.rb

@@ -221,19 +221,17 @@ module Gitlab
       # make things _more_ complex).
       #
       # rubocop: disable Metrics/AbcSize
-      def update_column_in_batches(table, column, value, &scope)
+      def update_column_in_batches(table, column, value)
         if transaction_open?
-          raise <<-MSG
-            update_column_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
+          raise 'update_column_in_batches can not be run inside a transaction, ' \
+            'you can disable transactions by calling disable_ddl_transaction! ' \
+            'in the body of your migration class'
         end
 
-        table_arel = Arel::Table.new(table)
+        table = Arel::Table.new(table)
 
-        count_arel = table_arel.project(Arel.star.count.as('count'))
-        count_arel = yield table_arel, count_arel if block_given?
+        count_arel = table.project(Arel.star.count.as('count'))
+        count_arel = yield table, count_arel if block_given?
 
         total = exec_query(count_arel.to_sql).to_hash.first['count'].to_i
 
@@ -248,103 +246,37 @@ module Gitlab
         # rows for GitLab.com.
         batch_size = max_size if batch_size > max_size
 
-        execute_in_batches(table, of: batch_size, scope: scope) do
-          Arel::UpdateManager.new(ActiveRecord::Base)
-            .table(table_arel)
-            .set([[table_arel[column], value]])
-        end
-      end
-
-      ##
-      # Iterates a table and executes a block for given range.
-      #
-      # Yields batch index, start and stop ids.
-      #
-      # Optional `scope` keyword argument is a closure that is meant to limit
-      # the scope the statement is going to be applied onto.
-      #
-      # Arel statement this helper will execute must be defined inside the
-      # block.
-      #
-      # Example:
-      #
-      # scope = ->(table, query) { query.where(table[:id].gt(100) }
-      #
-      # walk_table_in_batches(:table, of: 10, scope: scope) do |index, start, stop|
-      #   # do something here
-      # end
-      #
-      def walk_table_in_batches(table, of: 1000, scope: nil)
-        if transaction_open?
-          raise <<-MSG
-            walk_table_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
-
-        table = Arel::Table.new(table)
-
         start_arel = table.project(table[:id]).order(table[:id].asc).take(1)
-        start_arel = scope.call(table, start_arel) if scope
-        start_id = exec_query(start_arel.to_sql).to_hash.first.to_h['id'].to_i
+        start_arel = yield table, start_arel if block_given?
+        start_id = exec_query(start_arel.to_sql).to_hash.first['id'].to_i
 
-        1.step do |batch|
+        loop do
           stop_arel = table.project(table[:id])
             .where(table[:id].gteq(start_id))
             .order(table[:id].asc)
             .take(1)
-            .skip(of)
-
-          stop_arel = scope.call(table, stop_arel) if scope
-          stop_id = exec_query(stop_arel.to_sql)
-            .to_hash.first.to_h['id'].to_i
+            .skip(batch_size)
 
-          yield batch, start_id, stop_id
+          stop_arel = yield table, stop_arel if block_given?
+          stop_row = exec_query(stop_arel.to_sql).to_hash.first
 
-          stop_id.zero? ? break : start_id = stop_id
-        end
-      end
-
-      ##
-      # Executes an SQL statement in batches, created by Arel manager.
-      #
-      # Optional `scope` keyword argument is a closure that is meant to limit
-      # the scope the statement is going to be applied onto.
-      #
-      # Arel statement this helper will execute must be defined inside the
-      # block.
-      #
-      # Example:
-      #
-      # scope = ->(table, query) { query.where(table[:id].gt(100) }
-      #
-      # execute_in_batches(:table, of: 10000, scope: scope) do |table|
-      #   Arel::UpdateManager.new(ActiveRecord::Base)
-      #     .table(table)
-      #     .set([[table[:field], 101]])
-      # 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
+          update_arel = Arel::UpdateManager.new(ActiveRecord::Base)
+            .table(table)
+            .set([[table[column], value]])
+            .where(table[:id].gteq(start_id))
 
-        raise ArgumentError, 'This method requires a block!' unless block_given?
+          if stop_row
+            stop_id = stop_row['id'].to_i
+            start_id = stop_id
+            update_arel = update_arel.where(table[:id].lt(stop_id))
+          end
 
-        table_arel = Arel::Table.new(table)
+          update_arel = yield table, update_arel if block_given?
 
-        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(update_arel.to_sql)
 
-          execute(exec_arel.to_sql)
+          # There are no more rows left to update.
+          break unless stop_row
         end
       end
 

spec/lib/gitlab/database/migration_helpers_spec.rb

@@ -2,7 +2,9 @@ require 'spec_helper'
 
 describe Gitlab::Database::MigrationHelpers, lib: true do
   let(:model) do
-    ActiveRecord::Migration.new.extend(described_class)
+    ActiveRecord::Migration.new.extend(
+      Gitlab::Database::MigrationHelpers
+    )
   end
 
   before do
@@ -262,8 +264,7 @@ 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?)
-          .at_least(:once).and_return(false)
+        expect(model).to receive(:transaction_open?).and_return(false)
 
         create_list(:empty_project, 5)
       end
@@ -312,81 +313,6 @@ describe Gitlab::Database::MigrationHelpers, lib: true do
     end
   end
 
-  describe '#walk_table_in_batches' do
-    context 'when running outside of a transaction' do
-      before do
-        expect(model).to receive(:transaction_open?).and_return(false)
-
-        create_list(:empty_project, 6)
-      end
-
-      it 'yields for each batch' do
-        expect { |b| model.walk_table_in_batches(:projects, of: 2, &b) }
-          .to yield_control.exactly(3).times
-      end
-
-      it 'yields successive ranges' do
-        expect { |b| model.walk_table_in_batches(:projects, of: 2, &b) }
-          .to yield_successive_args([1, Integer, Integer],
-                                    [2, Integer, Integer],
-                                    [3, Integer, 0])
-      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)) }
-
-          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(table).set([[table[:archived], true]])
-          end
-
-          expect(Project.where(archived: true).count).to eq(1)
-        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.execute_in_batches(:projects)}
-          .to raise_error(RuntimeError)
-      end
-    end
-  end
-
   describe '#add_column_with_default' do
     context 'outside of a transaction' do
       context 'when a column limit is not set' do