Merge branch 'restrict-update-column-in-batches-for-large-tables' into 'master'

Restrict update column in batches for large tables

See merge request gitlab-org/gitlab-ce!15458

31 files changed, 153 insertions, 67 deletions

diff --git a/db/migrate/20160419122101_add_only_allow_merge_if_build_succeeds_to_projects.rb b/db/migrate/20160419122101_add_only_allow_merge_if_build_succeeds_to_projects.rb
index 22bac46e25c..1716b6e8153 100644
--- a/db/migrate/20160419122101_add_only_allow_merge_if_build_succeeds_to_projects.rb
+++ b/db/migrate/20160419122101_add_only_allow_merge_if_build_succeeds_to_projects.rb
@@ -1,4 +1,4 @@
-# rubocop:disable Migration/AddColumnWithDefaultToLargeTable
+# rubocop:disable Migration/UpdateLargeTable
 class AddOnlyAllowMergeIfBuildSucceedsToProjects < ActiveRecord::Migration
   include Gitlab::Database::MigrationHelpers
   disable_ddl_transaction!
diff --git a/db/migrate/20160608195742_add_repository_storage_to_projects.rb b/db/migrate/20160608195742_add_repository_storage_to_projects.rb
index 0f3664c13ef..e4febd1614d 100644
--- a/db/migrate/20160608195742_add_repository_storage_to_projects.rb
+++ b/db/migrate/20160608195742_add_repository_storage_to_projects.rb
@@ -1,4 +1,4 @@
-# rubocop:disable Migration/AddColumnWithDefaultToLargeTable
+# rubocop:disable Migration/UpdateLargeTable
 class AddRepositoryStorageToProjects < ActiveRecord::Migration
   include Gitlab::Database::MigrationHelpers
   disable_ddl_transaction!
diff --git a/db/migrate/20160615191922_set_missing_stage_on_ci_builds.rb b/db/migrate/20160615191922_set_missing_stage_on_ci_builds.rb
index 5336b036bca..c58cb957df4 100644
--- a/db/migrate/20160615191922_set_missing_stage_on_ci_builds.rb
+++ b/db/migrate/20160615191922_set_missing_stage_on_ci_builds.rb
@@ -1,3 +1,4 @@
+# rubocop:disable Migration/UpdateLargeTable
 # rubocop:disable Migration/UpdateColumnInBatches
 class SetMissingStageOnCiBuilds < ActiveRecord::Migration
   include Gitlab::Database::MigrationHelpers
diff --git a/db/migrate/20160715154212_add_request_access_enabled_to_projects.rb b/db/migrate/20160715154212_add_request_access_enabled_to_projects.rb
index 5dc26f8982a..22c925799a3 100644
--- a/db/migrate/20160715154212_add_request_access_enabled_to_projects.rb
+++ b/db/migrate/20160715154212_add_request_access_enabled_to_projects.rb
@@ -1,4 +1,4 @@
-# rubocop:disable Migration/AddColumnWithDefaultToLargeTable
+# rubocop:disable Migration/UpdateLargeTable
 class AddRequestAccessEnabledToProjects < ActiveRecord::Migration
   include Gitlab::Database::MigrationHelpers
   disable_ddl_transaction!
diff --git a/db/migrate/20160715204316_add_request_access_enabled_to_groups.rb b/db/migrate/20160715204316_add_request_access_enabled_to_groups.rb
index 4a317646788..4fcb29e1325 100644
--- a/db/migrate/20160715204316_add_request_access_enabled_to_groups.rb
+++ b/db/migrate/20160715204316_add_request_access_enabled_to_groups.rb
@@ -1,4 +1,4 @@
-# rubocop:disable Migration/AddColumnWithDefaultToLargeTable
+# rubocop:disable Migration/UpdateLargeTable
 class AddRequestAccessEnabledToGroups < ActiveRecord::Migration
   include Gitlab::Database::MigrationHelpers
   disable_ddl_transaction!
diff --git a/db/migrate/20160721081015_drop_and_readd_has_external_wiki_in_projects.rb b/db/migrate/20160721081015_drop_and_readd_has_external_wiki_in_projects.rb
index abe8e701e23..58f7f2a2841 100644
--- a/db/migrate/20160721081015_drop_and_readd_has_external_wiki_in_projects.rb
+++ b/db/migrate/20160721081015_drop_and_readd_has_external_wiki_in_projects.rb
@@ -1,3 +1,4 @@
+# rubocop:disable Migration/UpdateLargeTable
 # rubocop:disable Migration/UpdateColumnInBatches
 class DropAndReaddHasExternalWikiInProjects < ActiveRecord::Migration
   include Gitlab::Database::MigrationHelpers
diff --git a/db/migrate/20160831223750_remove_features_enabled_from_projects.rb b/db/migrate/20160831223750_remove_features_enabled_from_projects.rb
index 7414a28ac97..aec709aaf59 100644
--- a/db/migrate/20160831223750_remove_features_enabled_from_projects.rb
+++ b/db/migrate/20160831223750_remove_features_enabled_from_projects.rb
@@ -1,7 +1,7 @@
 # See http://doc.gitlab.com/ce/development/migration_style_guide.html
 # for more information on how to write migrations for GitLab.
 
-# rubocop:disable Migration/AddColumnWithDefaultToLargeTable
+# rubocop:disable Migration/UpdateLargeTable
 class RemoveFeaturesEnabledFromProjects < ActiveRecord::Migration
   include Gitlab::Database::MigrationHelpers
   disable_ddl_transaction!
diff --git a/db/migrate/20160913162434_remove_projects_pushes_since_gc.rb b/db/migrate/20160913162434_remove_projects_pushes_since_gc.rb
index 0100e30a733..df7d922b816 100644
--- a/db/migrate/20160913162434_remove_projects_pushes_since_gc.rb
+++ b/db/migrate/20160913162434_remove_projects_pushes_since_gc.rb
@@ -1,7 +1,7 @@
 # See http://doc.gitlab.com/ce/development/migration_style_guide.html
 # for more information on how to write migrations for GitLab.
 
-# rubocop:disable Migration/AddColumnWithDefaultToLargeTable
+# rubocop:disable Migration/UpdateLargeTable
 class RemoveProjectsPushesSinceGc < ActiveRecord::Migration
   include Gitlab::Database::MigrationHelpers
 
diff --git a/db/migrate/20170124193147_add_two_factor_columns_to_namespaces.rb b/db/migrate/20170124193147_add_two_factor_columns_to_namespaces.rb
index ae37da275fd..27ebe0af33b 100644
--- a/db/migrate/20170124193147_add_two_factor_columns_to_namespaces.rb
+++ b/db/migrate/20170124193147_add_two_factor_columns_to_namespaces.rb
@@ -1,4 +1,4 @@
-# rubocop:disable Migration/AddColumnWithDefaultToLargeTable
+# rubocop:disable Migration/UpdateLargeTable
 class AddTwoFactorColumnsToNamespaces < ActiveRecord::Migration
   include Gitlab::Database::MigrationHelpers
 
diff --git a/db/migrate/20170124193205_add_two_factor_columns_to_users.rb b/db/migrate/20170124193205_add_two_factor_columns_to_users.rb
index 8d4aefa4365..558a1837c79 100644
--- a/db/migrate/20170124193205_add_two_factor_columns_to_users.rb
+++ b/db/migrate/20170124193205_add_two_factor_columns_to_users.rb
@@ -1,4 +1,4 @@
-# rubocop:disable Migration/AddColumnWithDefaultToLargeTable
+# rubocop:disable Migration/UpdateLargeTable
 class AddTwoFactorColumnsToUsers < ActiveRecord::Migration
   include Gitlab::Database::MigrationHelpers
 
diff --git a/db/migrate/20170301125302_add_printing_merge_request_link_enabled_to_project.rb b/db/migrate/20170301125302_add_printing_merge_request_link_enabled_to_project.rb
index 7ad01a04815..6d43f346d4f 100644
--- a/db/migrate/20170301125302_add_printing_merge_request_link_enabled_to_project.rb
+++ b/db/migrate/20170301125302_add_printing_merge_request_link_enabled_to_project.rb
@@ -1,7 +1,7 @@
 # See http://doc.gitlab.com/ce/development/migration_style_guide.html
 # for more information on how to write migrations for GitLab.
 
-# rubocop:disable Migration/AddColumnWithDefaultToLargeTable
+# rubocop:disable Migration/UpdateLargeTable
 class AddPrintingMergeRequestLinkEnabledToProject < ActiveRecord::Migration
   include Gitlab::Database::MigrationHelpers
   disable_ddl_transaction!
diff --git a/db/migrate/20170305180853_add_auto_cancel_pending_pipelines_to_project.rb b/db/migrate/20170305180853_add_auto_cancel_pending_pipelines_to_project.rb
index f335e77fb5e..3c5cd95726a 100644
--- a/db/migrate/20170305180853_add_auto_cancel_pending_pipelines_to_project.rb
+++ b/db/migrate/20170305180853_add_auto_cancel_pending_pipelines_to_project.rb
@@ -1,4 +1,4 @@
-# rubocop:disable Migration/AddColumnWithDefaultToLargeTable
+# rubocop:disable Migration/UpdateLargeTable
 class AddAutoCancelPendingPipelinesToProject < ActiveRecord::Migration
   include Gitlab::Database::MigrationHelpers
 
diff --git a/db/migrate/20170315174634_revert_add_notified_of_own_activity_to_users.rb b/db/migrate/20170315174634_revert_add_notified_of_own_activity_to_users.rb
index 6c9fe19ca34..807dfcb385d 100644
--- a/db/migrate/20170315174634_revert_add_notified_of_own_activity_to_users.rb
+++ b/db/migrate/20170315174634_revert_add_notified_of_own_activity_to_users.rb
@@ -1,4 +1,4 @@
-# rubocop:disable Migration/AddColumnWithDefaultToLargeTable
+# rubocop:disable Migration/UpdateLargeTable
 class RevertAddNotifiedOfOwnActivityToUsers < ActiveRecord::Migration
   include Gitlab::Database::MigrationHelpers
   disable_ddl_transaction!
diff --git a/db/migrate/20170320173259_migrate_assignees.rb b/db/migrate/20170320173259_migrate_assignees.rb
index 7b61e811317..255b5e9c4db 100644
--- a/db/migrate/20170320173259_migrate_assignees.rb
+++ b/db/migrate/20170320173259_migrate_assignees.rb
@@ -1,3 +1,4 @@
+# rubocop:disable Migration/UpdateLargeTable
 # rubocop:disable Migration/UpdateColumnInBatches
 class MigrateAssignees < ActiveRecord::Migration
   include Gitlab::Database::MigrationHelpers
diff --git a/db/post_migrate/20170131214021_reset_users_authorized_projects_populated.rb b/db/post_migrate/20170131214021_reset_users_authorized_projects_populated.rb
index 82f8147547e..f1f81691f81 100644
--- a/db/post_migrate/20170131214021_reset_users_authorized_projects_populated.rb
+++ b/db/post_migrate/20170131214021_reset_users_authorized_projects_populated.rb
@@ -1,3 +1,4 @@
+# rubocop:disable Migration/UpdateLargeTable
 # rubocop:disable Migration/UpdateColumnInBatches
 class ResetUsersAuthorizedProjectsPopulated < ActiveRecord::Migration
   include Gitlab::Database::MigrationHelpers
diff --git a/db/post_migrate/20170309171644_reset_relative_position_for_issue.rb b/db/post_migrate/20170309171644_reset_relative_position_for_issue.rb
index 01fff680183..49fd46b0262 100644
--- a/db/post_migrate/20170309171644_reset_relative_position_for_issue.rb
+++ b/db/post_migrate/20170309171644_reset_relative_position_for_issue.rb
@@ -1,3 +1,4 @@
+# rubocop:disable Migration/UpdateLargeTable
 # rubocop:disable Migration/UpdateColumnInBatches
 class ResetRelativePositionForIssue < ActiveRecord::Migration
   include Gitlab::Database::MigrationHelpers
diff --git a/db/post_migrate/20170324160416_migrate_user_activities_to_users_last_activity_on.rb b/db/post_migrate/20170324160416_migrate_user_activities_to_users_last_activity_on.rb
index cb1b4f1855d..78413a608f1 100644
--- a/db/post_migrate/20170324160416_migrate_user_activities_to_users_last_activity_on.rb
+++ b/db/post_migrate/20170324160416_migrate_user_activities_to_users_last_activity_on.rb
@@ -1,3 +1,4 @@
+# rubocop:disable Migration/UpdateLargeTable
 class MigrateUserActivitiesToUsersLastActivityOn < ActiveRecord::Migration
   include Gitlab::Database::MigrationHelpers
 
diff --git a/db/post_migrate/20170406142253_migrate_user_project_view.rb b/db/post_migrate/20170406142253_migrate_user_project_view.rb
index c4e910b3b44..d6061dd416d 100644
--- a/db/post_migrate/20170406142253_migrate_user_project_view.rb
+++ b/db/post_migrate/20170406142253_migrate_user_project_view.rb
@@ -1,3 +1,4 @@
+# rubocop:disable Migration/UpdateLargeTable
 # See http://doc.gitlab.com/ce/development/migration_style_guide.html
 # for more information on how to write migrations for GitLab.
 
diff --git a/db/post_migrate/20170502070007_enable_auto_cancel_pending_pipelines_for_all.rb b/db/post_migrate/20170502070007_enable_auto_cancel_pending_pipelines_for_all.rb
index 765daa0a347..bba37e32c01 100644
--- a/db/post_migrate/20170502070007_enable_auto_cancel_pending_pipelines_for_all.rb
+++ b/db/post_migrate/20170502070007_enable_auto_cancel_pending_pipelines_for_all.rb
@@ -1,3 +1,4 @@
+# rubocop:disable Migration/UpdateLargeTable
 # rubocop:disable Migration/UpdateColumnInBatches
 class EnableAutoCancelPendingPipelinesForAll < ActiveRecord::Migration
   include Gitlab::Database::MigrationHelpers
diff --git a/db/post_migrate/20170503004427_update_retried_for_ci_build.rb b/db/post_migrate/20170503004427_update_retried_for_ci_build.rb
index 9d9f36550e7..b0b58ab3011 100644
--- a/db/post_migrate/20170503004427_update_retried_for_ci_build.rb
+++ b/db/post_migrate/20170503004427_update_retried_for_ci_build.rb
@@ -1,3 +1,4 @@
+# rubocop:disable Migration/UpdateLargeTable
 class UpdateRetriedForCiBuild < ActiveRecord::Migration
   include Gitlab::Database::MigrationHelpers
 
diff --git a/db/post_migrate/20170508170547_add_head_pipeline_for_each_merge_request.rb b/db/post_migrate/20170508170547_add_head_pipeline_for_each_merge_request.rb
index f77078ddd70..81e9d050668 100644
--- a/db/post_migrate/20170508170547_add_head_pipeline_for_each_merge_request.rb
+++ b/db/post_migrate/20170508170547_add_head_pipeline_for_each_merge_request.rb
@@ -1,3 +1,4 @@
+# rubocop:disable Migration/UpdateLargeTable
 class AddHeadPipelineForEachMergeRequest < ActiveRecord::Migration
   include Gitlab::Database::MigrationHelpers
 
diff --git a/db/post_migrate/20170518231126_fix_wrongly_renamed_routes.rb b/db/post_migrate/20170518231126_fix_wrongly_renamed_routes.rb
index c78beda9d21..3e952980866 100644
--- a/db/post_migrate/20170518231126_fix_wrongly_renamed_routes.rb
+++ b/db/post_migrate/20170518231126_fix_wrongly_renamed_routes.rb
@@ -1,3 +1,4 @@
+# rubocop:disable Migration/UpdateLargeTable
 # See http://doc.gitlab.com/ce/development/migration_style_guide.html
 # for more information on how to write migrations for GitLab.
 
diff --git a/db/post_migrate/20170526190000_migrate_build_stage_reference_again.rb b/db/post_migrate/20170526190000_migrate_build_stage_reference_again.rb
index 97cb242415d..31a73bb3b27 100644
--- a/db/post_migrate/20170526190000_migrate_build_stage_reference_again.rb
+++ b/db/post_migrate/20170526190000_migrate_build_stage_reference_again.rb
@@ -1,3 +1,4 @@
+# rubocop:disable Migration/UpdateLargeTable
 class MigrateBuildStageReferenceAgain < ActiveRecord::Migration
   include Gitlab::Database::MigrationHelpers
 
diff --git a/db/post_migrate/20170927112318_update_legacy_diff_notes_type_for_import.rb b/db/post_migrate/20170927112318_update_legacy_diff_notes_type_for_import.rb
index a238216253b..b040c81b316 100644
--- a/db/post_migrate/20170927112318_update_legacy_diff_notes_type_for_import.rb
+++ b/db/post_migrate/20170927112318_update_legacy_diff_notes_type_for_import.rb
@@ -1,3 +1,4 @@
+# rubocop:disable Migration/UpdateLargeTable
 class UpdateLegacyDiffNotesTypeForImport < ActiveRecord::Migration
   include Gitlab::Database::MigrationHelpers
 
diff --git a/db/post_migrate/20170927112319_update_notes_type_for_import.rb b/db/post_migrate/20170927112319_update_notes_type_for_import.rb
index 1e70acd9868..5a400c71b02 100644
--- a/db/post_migrate/20170927112319_update_notes_type_for_import.rb
+++ b/db/post_migrate/20170927112319_update_notes_type_for_import.rb
@@ -1,3 +1,4 @@
+# rubocop:disable Migration/UpdateLargeTable
 class UpdateNotesTypeForImport < ActiveRecord::Migration
   include Gitlab::Database::MigrationHelpers
 
diff --git a/doc/development/migration_style_guide.md b/doc/development/migration_style_guide.md
index 9b8ab5da74e..a235dd74909 100644
--- a/doc/development/migration_style_guide.md
+++ b/doc/development/migration_style_guide.md
@@ -198,7 +198,43 @@ end
 
 Keep in mind that this operation can easily take 10-15 minutes to complete on
 larger installations (e.g. GitLab.com). As a result you should only add default
-values if absolutely necessary.
+values if absolutely necessary. There is a RuboCop cop that will fail if this
+method is used on some tables that are very large on GitLab.com, which would
+cause other issues.
+
+## Updating an existing column
+
+To update an existing column to a particular value, you can use
+`update_column_in_batches` (`add_column_with_default` uses this internally to
+fill in the default value). This will split the updates into batches, so we
+don't update too many rows at in a single statement.
+
+This updates the column `foo` in the `projects` table to 10, where `some_column`
+is `'hello'`:
+
+```ruby
+update_column_in_batches(:projects, :foo, 10) do |table, query|
+  query.where(table[:some_column].eq('hello'))
+end
+```
+
+To perform a computed update, the value can be wrapped in `Arel.sql`, so Arel
+treats it as an SQL literal. The below example is the same as the one above, but
+the value is set to the product of the `bar` and `baz` columns:
+
+```ruby
+update_value = Arel.sql('bar * baz')
+
+update_column_in_batches(:projects, :foo, update_value) do |table, query|
+  query.where(table[:some_column].eq('hello'))
+end
+```
+
+Like `add_column_with_default`, there is a RuboCop cop to detect usage of this
+on large tables. In the case of `update_column_in_batches`, it may be acceptable
+to run on a large table, as long as it is only updating a small subset of the
+rows in the table, but do not ignore that without validating on the GitLab.com
+staging environment - or asking someone else to do so for you - beforehand.
 
 ## Integer column type
 
diff --git a/lib/gitlab/database/migration_helpers.rb b/lib/gitlab/database/migration_helpers.rb
index 2c35da8f1aa..c276c3566b4 100644
--- a/lib/gitlab/database/migration_helpers.rb
+++ b/lib/gitlab/database/migration_helpers.rb
@@ -220,6 +220,15 @@ module Gitlab
       # column - The name of the column to update.
       # value - The value for the column.
       #
+      # The `value` argument is typically a literal. To perform a computed
+      # update, an Arel literal can be used instead:
+      #
+      #     update_value = Arel.sql('bar * baz')
+      #
+      #     update_column_in_batches(:projects, :foo, update_value) do |table, query|
+      #       query.where(table[:some_column].eq('hello'))
+      #     end
+      #
       # 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
diff --git a/rubocop/cop/migration/add_column_with_default_to_large_table.rb b/rubocop/cop/migration/update_large_table.rb
index fb363f95b56..3ae3fb1b68e 100644
--- a/rubocop/cop/migration/add_column_with_default_to_large_table.rb
+++ b/rubocop/cop/migration/update_large_table.rb
@@ -12,12 +12,12 @@ module RuboCop
       #
       # See https://gitlab.com/gitlab-com/infrastructure/issues/1602 for more
       # information.
-      class AddColumnWithDefaultToLargeTable < RuboCop::Cop::Cop
+      class UpdateLargeTable < RuboCop::Cop::Cop
         include MigrationHelpers
 
-        MSG = 'Using `add_column_with_default` on the `%s` table will take a ' \
-          'long time to complete, and should be avoided unless absolutely ' \
-          'necessary'.freeze
+        MSG = 'Using `%s` on the `%s` table will take a long time to ' \
+              'complete, and should be avoided unless absolutely ' \
+              'necessary'.freeze
 
         LARGE_TABLES = %i[
           ci_pipelines
@@ -34,20 +34,22 @@ module RuboCop
           users
         ].freeze
 
-        def_node_matcher :add_column_with_default?, <<~PATTERN
-          (send nil :add_column_with_default $(sym ...) ...)
+        def_node_matcher :batch_update?, <<~PATTERN
+          (send nil ${:add_column_with_default :update_column_in_batches} $(sym ...) ...)
         PATTERN
 
         def on_send(node)
           return unless in_migration?(node)
 
-          matched = add_column_with_default?(node)
-          return unless matched
+          matches = batch_update?(node)
+          return unless matches
+
+          update_method = matches.first
+          table = matches.last.to_a.first
 
-          table = matched.to_a.first
           return unless LARGE_TABLES.include?(table)
 
-          add_offense(node, :expression, format(MSG, table))
+          add_offense(node, :expression, format(MSG, update_method, table))
         end
       end
     end
diff --git a/rubocop/rubocop.rb b/rubocop/rubocop.rb
index 4ebbe010e90..a1668749dd9 100644
--- a/rubocop/rubocop.rb
+++ b/rubocop/rubocop.rb
@@ -7,7 +7,6 @@ require_relative 'cop/polymorphic_associations'
 require_relative 'cop/project_path_helper'
 require_relative 'cop/redirect_with_status'
 require_relative 'cop/migration/add_column'
-require_relative 'cop/migration/add_column_with_default_to_large_table'
 require_relative 'cop/migration/add_concurrent_foreign_key'
 require_relative 'cop/migration/add_concurrent_index'
 require_relative 'cop/migration/add_index'
@@ -20,6 +19,7 @@ require_relative 'cop/migration/reversible_add_column_with_default'
 require_relative 'cop/migration/safer_boolean_column'
 require_relative 'cop/migration/timestamps'
 require_relative 'cop/migration/update_column_in_batches'
+require_relative 'cop/migration/update_large_table'
 require_relative 'cop/rspec/env_assignment'
 require_relative 'cop/rspec/single_line_hook'
 require_relative 'cop/rspec/verbose_include_metadata'
diff --git a/spec/rubocop/cop/migration/add_column_with_default_to_large_table_spec.rb b/spec/rubocop/cop/migration/add_column_with_default_to_large_table_spec.rb
deleted file mode 100644
index 07cb3fc4a2e..00000000000
--- a/spec/rubocop/cop/migration/add_column_with_default_to_large_table_spec.rb
+++ /dev/null
@@ -1,44 +0,0 @@
-require 'spec_helper'
-
-require 'rubocop'
-require 'rubocop/rspec/support'
-
-require_relative '../../../../rubocop/cop/migration/add_column_with_default_to_large_table'
-
-describe RuboCop::Cop::Migration::AddColumnWithDefaultToLargeTable do
-  include CopHelper
-
-  subject(:cop) { described_class.new }
-
-  context 'in migration' do
-    before do
-      allow(cop).to receive(:in_migration?).and_return(true)
-    end
-
-    described_class::LARGE_TABLES.each do |table|
-      it "registers an offense for the #{table} table" do
-        inspect_source(cop, "add_column_with_default :#{table}, :column, default: true")
-
-        aggregate_failures do
-          expect(cop.offenses.size).to eq(1)
-          expect(cop.offenses.map(&:line)).to eq([1])
-        end
-      end
-    end
-
-    it 'registers no offense for non-blacklisted tables' do
-      inspect_source(cop, "add_column_with_default :table, :column, default: true")
-
-      expect(cop.offenses).to be_empty
-    end
-  end
-
-  context 'outside of migration' do
-    it 'registers no offense' do
-      table = described_class::LARGE_TABLES.sample
-      inspect_source(cop, "add_column_with_default :#{table}, :column, default: true")
-
-      expect(cop.offenses).to be_empty
-    end
-  end
-end
diff --git a/spec/rubocop/cop/migration/update_large_table_spec.rb b/spec/rubocop/cop/migration/update_large_table_spec.rb
new file mode 100644
index 00000000000..17b19e139e4
--- /dev/null
+++ b/spec/rubocop/cop/migration/update_large_table_spec.rb
@@ -0,0 +1,69 @@
+require 'spec_helper'
+
+require 'rubocop'
+require 'rubocop/rspec/support'
+
+require_relative '../../../../rubocop/cop/migration/update_large_table'
+
+describe RuboCop::Cop::Migration::UpdateLargeTable do
+  include CopHelper
+
+  subject(:cop) { described_class.new }
+
+  context 'in migration' do
+    before do
+      allow(cop).to receive(:in_migration?).and_return(true)
+    end
+
+    shared_examples 'large tables' do |update_method|
+      described_class::LARGE_TABLES.each do |table|
+        it "registers an offense for the #{table} table" do
+          inspect_source(cop, "#{update_method} :#{table}, :column, default: true")
+
+          aggregate_failures do
+            expect(cop.offenses.size).to eq(1)
+            expect(cop.offenses.map(&:line)).to eq([1])
+          end
+        end
+      end
+    end
+
+    context 'for the add_column_with_default method' do
+      include_examples 'large tables', 'add_column_with_default'
+    end
+
+    context 'for the update_column_in_batches method' do
+      include_examples 'large tables', 'update_column_in_batches'
+    end
+
+    it 'registers no offense for non-blacklisted tables' do
+      inspect_source(cop, "add_column_with_default :table, :column, default: true")
+
+      expect(cop.offenses).to be_empty
+    end
+
+    it 'registers no offense for non-blacklisted methods' do
+      table = described_class::LARGE_TABLES.sample
+
+      inspect_source(cop, "some_other_method :#{table}, :column, default: true")
+
+      expect(cop.offenses).to be_empty
+    end
+  end
+
+  context 'outside of migration' do
+    let(:table) { described_class::LARGE_TABLES.sample }
+
+    it 'registers no offense for add_column_with_default' do
+      inspect_source(cop, "add_column_with_default :#{table}, :column, default: true")
+
+      expect(cop.offenses).to be_empty
+    end
+
+    it 'registers no offense for update_column_in_batches' do
+      inspect_source(cop, "add_column_with_default :#{table}, :column, default: true")
+
+      expect(cop.offenses).to be_empty
+    end
+  end
+end