Fix update_column_in_batches to update all rowsupdate-column-in-batches-where

This changes update_column_in_batches to ensure it always updates all
rows now. These changes also allow for an extra SELECT query to be
removed, nor does it use the row count for determining offsets and the
likes; instead it's only used to determine the batch size.

1 files changed, 25 insertions, 21 deletions

diff --git a/lib/gitlab/database/migration_helpers.rb b/lib/gitlab/database/migration_helpers.rb
index 909ff8677cb..03ccc278d8e 100644
--- a/lib/gitlab/database/migration_helpers.rb
+++ b/lib/gitlab/database/migration_helpers.rb
@@ -50,53 +50,57 @@ module Gitlab
       # table - The name of the table.
       # column - The name of the column to update.
       # value - The value for the column.
+      #
+      # Rubocop's Metrics/AbcSize metric is disabled for this method as Rubocop
+      # determines this method to be too complex while there's no way to make it
+      # less "complex" without introducing extra methods (which actually will
+      # make things _more_ complex).
+      #
+      # rubocop: disable Metrics/AbcSize
       def update_column_in_batches(table, column, value)
         table = Arel::Table.new(table)
-        processed = 0
 
         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
 
+        return if total == 0
+
         # Update in batches of 5% until we run out of any rows to update.
         batch_size = ((total / 100.0) * 5.0).ceil
 
-        loop do
-          start_arel = table.project(table[:id]).
-            order(table[:id].asc).
-            take(1).
-            skip(processed)
-
-          start_arel = yield table, start_arel if block_given?
-          start_row = exec_query(start_arel.to_sql).to_hash.first
-
-          # There are no more rows to process
-          break unless start_row
+        start_arel = table.project(table[:id]).order(table[:id].asc).take(1)
+        start_arel = yield table, start_arel if block_given?
+        start_id = exec_query(start_arel.to_sql).to_hash.first['id'].to_i
 
+        loop do
           stop_arel = table.project(table[:id]).
+            where(table[:id].gteq(start_id)).
             order(table[:id].asc).
             take(1).
-            skip(processed + batch_size)
+            skip(batch_size)
 
           stop_arel = yield table, stop_arel if block_given?
           stop_row = exec_query(stop_arel.to_sql).to_hash.first
 
-          update_manager = Arel::UpdateManager.new(ActiveRecord::Base)
-
-          update_arel = update_manager.table(table).
+          update_arel = Arel::UpdateManager.new(ActiveRecord::Base).
+            table(table).
             set([[table[column], value]]).
-            where(table[:id].gteq(start_row['id']))
-
-          update_arel = yield table, update_arel if block_given?
+            where(table[:id].gteq(start_id))
 
           if stop_row
-            update_arel = update_arel.where(table[:id].lt(stop_row['id']))
+            stop_id = stop_row['id'].to_i
+            start_id = stop_id
+            update_arel = update_arel.where(table[:id].lt(stop_id))
           end
 
+          update_arel = yield table, update_arel if block_given?
+
           execute(update_arel.to_sql)
 
-          processed += batch_size
+          # There are no more rows left to update.
+          break unless stop_row
         end
       end