From 8966263e0c738e85d8872aee0bfcf7fbe77b22a6 Mon Sep 17 00:00:00 2001 From: Yorick Peterse Date: Wed, 15 Jun 2016 16:42:52 +0200 Subject: [PATCH 1/3] Customizing of update_column_in_batches queries By passing a block to update_column_in_batches() one can now customize the queries executed. This in turn can be used to only update a specific set of rows instead of simply all the rows in the table. --- lib/gitlab/database/migration_helpers.rb | 95 +++++++++++-------- .../gitlab/database/migration_helpers_spec.rb | 12 +++ 2 files changed, 66 insertions(+), 41 deletions(-) diff --git a/lib/gitlab/database/migration_helpers.rb b/lib/gitlab/database/migration_helpers.rb index dd3ff0ab18..a87c9b038a 100644 --- a/lib/gitlab/database/migration_helpers.rb +++ b/lib/gitlab/database/migration_helpers.rb @@ -28,63 +28,73 @@ module Gitlab # Updates the value of a column in batches. # # This method updates the table in batches of 5% of the total row count. - # Any data inserted while running this method (or after it has finished - # running) is _not_ updated automatically. + # This method will continue updating rows until no rows remain. + # + # When given a block this method will yield to values to the block: + # + # 1. An instance of `Arel::Table` for the table that is being updated. + # 2. The query to run as an Arel object. + # + # By supplying a block one can add extra conditions to the queries being + # executed. Note that the same block is used for _all_ queries. + # + # Example: + # + # update_column_in_batches(:projects, :foo, 10) do |table, query| + # query.where(table[:some_column].eq('hello')) + # end + # + # This would result in this method updating only rows there + # `projects.some_column` equals "hello". # # table - The name of the table. # column - The name of the column to update. # value - The value for the column. def update_column_in_batches(table, column, value) - quoted_table = quote_table_name(table) - quoted_column = quote_column_name(column) - - ## - # Workaround for #17711 - # - # It looks like for MySQL `ActiveRecord::Base.conntection.quote(true)` - # returns correct value (1), but `ActiveRecord::Migration.new.quote` - # returns incorrect value ('true'), which causes migrations to fail. - # - quoted_value = connection.quote(value) + table = Arel::Table.new(table) processed = 0 - total = exec_query("SELECT COUNT(*) AS count FROM #{quoted_table}"). - to_hash. - first['count']. - to_i + 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 # 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_row = exec_query(%Q{ - SELECT id - FROM #{quoted_table} - ORDER BY id ASC - LIMIT 1 OFFSET #{processed} - }).to_hash.first + 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 - stop_row = exec_query(%Q{ - SELECT id - FROM #{quoted_table} - ORDER BY id ASC - LIMIT 1 OFFSET #{processed + batch_size} - }).to_hash.first + stop_arel = table.project(table[:id]). + order(table[:id].asc). + take(1). + skip(processed + batch_size) - query = %Q{ - UPDATE #{quoted_table} - SET #{quoted_column} = #{quoted_value} - WHERE id >= #{start_row['id']} - } + 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). + set([[table[column], value]]). + where(table[:id].gteq(start_row['id'])) + + update_arel = yield table, update_arel if block_given? if stop_row - query += " AND id < #{stop_row['id']}" + update_arel = update_arel.where(table[:id].lt(stop_row['id'])) end - execute(query) + execute(update_arel.to_sql) processed += batch_size end @@ -95,9 +105,9 @@ module Gitlab # This method runs the following steps: # # 1. Add the column with a default value of NULL. - # 2. Update all existing rows in batches. - # 3. Change the default value of the column to the specified value. - # 4. Update any remaining rows. + # 2. Change the default value of the column to the specified value. + # 3. Update all existing rows in batches. + # 4. Set a `NOT NULL` constraint on the column if desired (the default). # # These steps ensure a column can be added to a large and commonly used # table without locking the entire table for the duration of the table @@ -109,7 +119,10 @@ module Gitlab # default - The default value for the column. # allow_null - When set to `true` the column will allow NULL values, the # default is to not allow NULL values. - def add_column_with_default(table, column, type, default:, allow_null: false) + # + # This method can also take a block which is passed directly to the + # `update_column_in_batches` method. + def add_column_with_default(table, column, type, default:, allow_null: false, &block) if transaction_open? raise 'add_column_with_default can not be run inside a transaction, ' \ 'you can disable transactions by calling disable_ddl_transaction! ' \ @@ -126,7 +139,7 @@ module Gitlab begin transaction do - update_column_in_batches(table, column, default) + update_column_in_batches(table, column, default, &block) change_column_null(table, column, false) unless allow_null end diff --git a/spec/lib/gitlab/database/migration_helpers_spec.rb b/spec/lib/gitlab/database/migration_helpers_spec.rb index 1ec539066a..b46f56196e 100644 --- a/spec/lib/gitlab/database/migration_helpers_spec.rb +++ b/spec/lib/gitlab/database/migration_helpers_spec.rb @@ -71,6 +71,18 @@ describe Gitlab::Database::MigrationHelpers, lib: true do expect(Project.where(archived: true).count).to eq(5) end + + context 'when a block is supplied' do + it 'yields an Arel table and query object to the supplied block' do + first_id = Project.first.id + + model.update_column_in_batches(:projects, :archived, true) do |t, query| + query.where(t[:id].eq(first_id)) + end + + expect(Project.where(archived: true).count).to eq(1) + end + end end describe '#add_column_with_default' do From 816c453558c6d25fd8724bd3fe08e6fd221f2887 Mon Sep 17 00:00:00 2001 From: Yorick Peterse Date: Wed, 15 Jun 2016 17:04:07 +0200 Subject: [PATCH 2/3] Don't update columns in batches in a transaction This ensures that whatever locks are acquired aren't held onto until the end of the transaction (= after _all_ rows have been updated). Timing wise there's also no difference between using a transaction and not using one. --- lib/gitlab/database/migration_helpers.rb | 6 ++---- spec/lib/gitlab/database/migration_helpers_spec.rb | 2 +- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/lib/gitlab/database/migration_helpers.rb b/lib/gitlab/database/migration_helpers.rb index a87c9b038a..909ff8677c 100644 --- a/lib/gitlab/database/migration_helpers.rb +++ b/lib/gitlab/database/migration_helpers.rb @@ -138,11 +138,9 @@ module Gitlab end begin - transaction do - update_column_in_batches(table, column, default, &block) + update_column_in_batches(table, column, default, &block) - change_column_null(table, column, false) unless allow_null - end + change_column_null(table, column, false) unless allow_null # We want to rescue _all_ exceptions here, even those that don't inherit # from StandardError. rescue Exception => error # rubocop: disable all diff --git a/spec/lib/gitlab/database/migration_helpers_spec.rb b/spec/lib/gitlab/database/migration_helpers_spec.rb index b46f56196e..9096ad101b 100644 --- a/spec/lib/gitlab/database/migration_helpers_spec.rb +++ b/spec/lib/gitlab/database/migration_helpers_spec.rb @@ -90,7 +90,7 @@ describe Gitlab::Database::MigrationHelpers, lib: true do before do expect(model).to receive(:transaction_open?).and_return(false) - expect(model).to receive(:transaction).twice.and_yield + expect(model).to receive(:transaction).and_yield expect(model).to receive(:add_column). with(:projects, :foo, :integer, default: nil) From c1e756c24263d7c20ae15680ca2ac1d0c2c5451f Mon Sep 17 00:00:00 2001 From: Yorick Peterse Date: Thu, 16 Jun 2016 12:50:11 +0200 Subject: [PATCH 3/3] Fix update_column_in_batches to update all rows 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. --- lib/gitlab/database/migration_helpers.rb | 52 +++++++++++++----------- 1 file changed, 28 insertions(+), 24 deletions(-) diff --git a/lib/gitlab/database/migration_helpers.rb b/lib/gitlab/database/migration_helpers.rb index 909ff8677c..dec20d8659 100644 --- a/lib/gitlab/database/migration_helpers.rb +++ b/lib/gitlab/database/migration_helpers.rb @@ -30,7 +30,7 @@ module Gitlab # This method updates the table in batches of 5% of the total row count. # This method will continue updating rows until no rows remain. # - # When given a block this method will yield to values to the block: + # When given a block this method will yield two values to the block: # # 1. An instance of `Arel::Table` for the table that is being updated. # 2. The query to run as an Arel object. @@ -44,59 +44,63 @@ module Gitlab # query.where(table[:some_column].eq('hello')) # end # - # This would result in this method updating only rows there + # This would result in this method updating only rows where # `projects.some_column` equals "hello". # # 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 + 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 - 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 - 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'])) + where(table[:id].gteq(start_id)) + + 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 update_arel = yield table, update_arel if block_given? - if stop_row - update_arel = update_arel.where(table[:id].lt(stop_row['id'])) - end - execute(update_arel.to_sql) - processed += batch_size + # There are no more rows left to update. + break unless stop_row end end