Cop to restrict methods within with_lock_retries

- Eliminates WithLockRetriesWithoutDdlTransaction cop.
- Introduce a new cop to restrict methods within
`with_lock_retries`

db/migrate/20200215222507_drop_forked_project_links_fk.rb

@@ -8,20 +8,16 @@ class DropForkedProjectLinksFk < ActiveRecord::Migration[6.0]
   disable_ddl_transaction!
 
   def up
-    # rubocop: disable Migration/WithLockRetriesWithoutDdlTransaction
     with_lock_retries do
       remove_foreign_key_if_exists :forked_project_links, column: :forked_to_project_id
     end
-    # rubocop: enable Migration/WithLockRetriesWithoutDdlTransaction
   end
 
   def down
     unless foreign_key_exists?(:forked_project_links, :projects, column: :forked_to_project_id)
-      # rubocop: disable Migration/WithLockRetriesWithoutDdlTransaction
       with_lock_retries do
         add_foreign_key :forked_project_links, :projects, column: :forked_to_project_id, on_delete: :cascade, validate: false
       end
-      # rubocop: enable Migration/WithLockRetriesWithoutDdlTransaction
     end
 
     fk_name = concurrent_foreign_key_name(:forked_project_links, :forked_to_project_id, prefix: 'fk_rails_')

db/migrate/20200304024025_add_sprint_id_index_to_issues.rb

@@ -13,7 +13,7 @@ class AddSprintIdIndexToIssues < ActiveRecord::Migration[6.0]
   end
 
   def down
-    with_lock_retries do # rubocop:disable Migration/WithLockRetriesWithoutDdlTransaction
+    with_lock_retries do
       remove_foreign_key :issues, column: :sprint_id
     end
     remove_concurrent_index :issues, :sprint_id

db/migrate/20200304024042_add_sprint_id_index_to_merge_requests.rb

@@ -13,7 +13,7 @@ class AddSprintIdIndexToMergeRequests < ActiveRecord::Migration[6.0]
   end
 
   def down
-    with_lock_retries do # rubocop:disable Migration/WithLockRetriesWithoutDdlTransaction
+    with_lock_retries do
       remove_foreign_key :merge_requests, column: :sprint_id
     end
     remove_concurrent_index :merge_requests, :sprint_id

db/migrate/20200407171133_add_protected_tag_create_access_levels_user_id_foreign_key.rb

 # frozen_string_literal: true
 
-# rubocop: disable Migration/WithLockRetriesWithoutDdlTransaction
 class AddProtectedTagCreateAccessLevelsUserIdForeignKey < ActiveRecord::Migration[6.0]
   include Gitlab::Database::MigrationHelpers
 

db/migrate/20200408154331_add_protected_branch_merge_access_levels_user_id_foreign_key.rb

 # frozen_string_literal: true
 
-# rubocop: disable Migration/WithLockRetriesWithoutDdlTransaction
 class AddProtectedBranchMergeAccessLevelsUserIdForeignKey < ActiveRecord::Migration[6.0]
   include Gitlab::Database::MigrationHelpers
 

db/migrate/20200408154411_add_path_locks_user_id_foreign_key.rb

 # frozen_string_literal: true
 
-# rubocop: disable Migration/WithLockRetriesWithoutDdlTransaction
 class AddPathLocksUserIdForeignKey < ActiveRecord::Migration[6.0]
   include Gitlab::Database::MigrationHelpers
 

db/migrate/20200408154455_add_protected_branch_push_access_levels_user_id_foreign_key.rb

 # frozen_string_literal: true
 
-# rubocop: disable Migration/WithLockRetriesWithoutDdlTransaction
 class AddProtectedBranchPushAccessLevelsUserIdForeignKey < ActiveRecord::Migration[6.0]
   include Gitlab::Database::MigrationHelpers
 

db/migrate/20200408154604_add_u2f_registrations_user_id_foreign_key.rb

 # frozen_string_literal: true
 
-# rubocop: disable Migration/WithLockRetriesWithoutDdlTransaction
 class AddU2fRegistrationsUserIdForeignKey < ActiveRecord::Migration[6.0]
   include Gitlab::Database::MigrationHelpers
 

db/migrate/20200420172752_add_sprints_foreign_key_to_projects.rb

@@ -12,7 +12,7 @@ class AddSprintsForeignKeyToProjects < ActiveRecord::Migration[6.0]
   end
 
   def down
-    with_lock_retries do # rubocop:disable Migration/WithLockRetriesWithoutDdlTransaction
+    with_lock_retries do
       remove_foreign_key :sprints, column: :project_id
     end
   end

db/migrate/20200420172927_add_sprints_foreign_key_to_groups.rb

@@ -12,7 +12,7 @@ class AddSprintsForeignKeyToGroups < ActiveRecord::Migration[6.0]
   end
 
   def down
-    with_lock_retries do # rubocop:disable Migration/WithLockRetriesWithoutDdlTransaction
+    with_lock_retries do
       remove_foreign_key :sprints, column: :group_id
     end
   end

db/migrate/20200423081441_add_foreign_key_from_users_to_metrics_users_starred_dashboars.rb

@@ -11,7 +11,7 @@ class AddForeignKeyFromUsersToMetricsUsersStarredDashboars < ActiveRecord::Migra
   end
 
   def down
-    with_lock_retries do # rubocop:disable Migration/WithLockRetriesWithoutDdlTransaction
+    with_lock_retries do
       remove_foreign_key_if_exists :metrics_users_starred_dashboards, column: :user_id
     end
   end

db/migrate/20200423081519_add_foreign_key_from_projects_to_metrics_users_starred_dashboars.rb

@@ -11,7 +11,7 @@ class AddForeignKeyFromProjectsToMetricsUsersStarredDashboars < ActiveRecord::Mi
   end
 
   def down
-    with_lock_retries do # rubocop:disable Migration/WithLockRetriesWithoutDdlTransaction
+    with_lock_retries do
       remove_foreign_key_if_exists :metrics_users_starred_dashboards, column: :project_id
     end
   end

db/post_migrate/20200424043515_drop_namespaces_plan_id.rb

@@ -8,14 +8,14 @@ class DropNamespacesPlanId < ActiveRecord::Migration[6.0]
   disable_ddl_transaction!
 
   def up
-    with_lock_retries do # rubocop: disable Migration/WithLockRetriesWithoutDdlTransaction
+    with_lock_retries do
       remove_column :namespaces, :plan_id
     end
   end
 
   def down
     unless column_exists?(:namespaces, :plan_id)
-      with_lock_retries do # rubocop: disable Migration/WithLockRetriesWithoutDdlTransaction
+      with_lock_retries do
         add_column :namespaces, :plan_id, :integer
       end
     end

doc/development/migration_style_guide.md

@@ -327,6 +327,34 @@ def down
 end
 ```
 
+**Usage with `disable_ddl_transaction!`**
+
+Generally the `with_lock_retries` helper should work with `disabled_ddl_transaction!`. A custom RuboCop rule ensures that only allowed methods can be placed within the lock retries block.
+
+```ruby
+disable_ddl_transaction!
+
+def up
+  with_lock_retries do
+    add_column :users, :name, :text
+  end
+
+  add_text_limit :users, :name, 255 # Includes constraint validation (full table scan)
+end
+```
+
+The RuboCop rule generally allows standard Rails migration methods, listed below. This example will cause a rubocop offense:
+
+```ruby
+disabled_ddl_transaction!
+
+def up
+  with_lock_retries do
+    add_concurrent_index :users, :name
+  end
+end
+```
+
 ### When to use the helper method
 
 The `with_lock_retries` helper method can be used when you normally use
@@ -350,8 +378,6 @@ Example changes:
 - `change_column_default`
 - `create_table` / `drop_table`
 
-**Note:** `with_lock_retries` method **cannot** be used with `disable_ddl_transaction!`.
-
 **Note:** `with_lock_retries` method **cannot** be used within the `change` method, you must manually define the `up` and `down` methods to make the migration reversible.
 
 ### How the helper method works

rubocop/cop/migration/with_lock_retries_disallowed_method.rb

+# frozen_string_literal: true
+
+require_relative '../../migration_helpers'
+
+module RuboCop
+  module Cop
+    module Migration
+      class WithLockRetriesDisallowedMethod < RuboCop::Cop::Cop
+        include MigrationHelpers
+
+        ALLOWED_MIGRATION_METHODS = %i[
+          create_table
+          drop_table
+          add_foreign_key
+          remove_foreign_key
+          add_column
+          remove_column
+          execute
+          change_column_default
+          remove_foreign_key_if_exists
+          remove_foreign_key_without_error
+          table_exists?
+          index_exists_by_name?
+          foreign_key_exists?
+          index_exists?
+          column_exists?
+        ].sort.freeze
+
+        MSG = "The method is not allowed to be called within the `with_lock_retries` block, the only allowed methods are: #{ALLOWED_MIGRATION_METHODS.join(', ')}"
+
+        def_node_matcher :send_node?, <<~PATTERN
+        send
+        PATTERN
+
+        def on_block(node)
+          block_body = node.body
+
+          return unless in_migration?(node)
+          return unless block_body
+          return unless node.method_name == :with_lock_retries
+
+          if send_node?(block_body)
+            check_node(block_body)
+          else
+            block_body.children.each { |n| check_node(n) }
+          end
+        end
+
+        def check_node(node)
+          return unless send_node?(node)
+
+          name = node.children[1]
+          add_offense(node, location: :expression) unless ALLOWED_MIGRATION_METHODS.include?(name)
+        end
+      end
+    end
+  end
+end

rubocop/cop/migration/with_lock_retries_without_ddl_transaction.rb

-# frozen_string_literal: true
-
-require_relative '../../migration_helpers'
-
-module RuboCop
-  module Cop
-    module Migration
-      # Cop that prevents usage of `with_lock_retries` with `disable_ddl_transaction!`
-      class WithLockRetriesWithoutDdlTransaction < RuboCop::Cop::Cop
-        include MigrationHelpers
-
-        MSG = '`with_lock_retries` cannot be used with disabled DDL transactions (`disable_ddl_transaction!`). ' \
-          'Please remove the `disable_ddl_transaction!` call from your migration.'.freeze
-
-        def_node_matcher :disable_ddl_transaction?, <<~PATTERN
-        (send _ :disable_ddl_transaction!)
-        PATTERN
-
-        def_node_matcher :with_lock_retries?, <<~PATTERN
-          (send _ :with_lock_retries)
-        PATTERN
-
-        def on_send(node)
-          return unless in_migration?(node)
-          return unless with_lock_retries?(node)
-
-          node.each_ancestor(:begin) do |begin_node|
-            disable_ddl_transaction_node = begin_node.children.find { |n| disable_ddl_transaction?(n) }
-
-            add_offense(node, location: :expression) if disable_ddl_transaction_node
-          end
-        end
-      end
-    end
-  end
-end

spec/rubocop/cop/migration/with_lock_retries_without_ddl_transaction_spec.rb → spec/rubocop/cop/migration/with_lock_retries_disallowed_method_spec.rb

@@ -5,14 +5,11 @@ require 'spec_helper'
 require 'rubocop'
 require 'rubocop/rspec/support'
 
-require_relative '../../../../rubocop/cop/migration/with_lock_retries_without_ddl_transaction'
+require_relative '../../../../rubocop/cop/migration/with_lock_retries_disallowed_method'
 
-describe RuboCop::Cop::Migration::WithLockRetriesWithoutDdlTransaction do
+describe RuboCop::Cop::Migration::WithLockRetriesDisallowedMethod do
   include CopHelper
 
-  let(:valid_source) { 'class MigrationClass < ActiveRecord::Migration[6.0]; def up; with_lock_retries {}; end; end' }
-  let(:invalid_source) { 'class MigrationClass < ActiveRecord::Migration[6.0]; disable_ddl_transaction!; def up; with_lock_retries {}; end; end' }
-
   subject(:cop) { described_class.new }
 
   context 'in migration' do
@@ -20,8 +17,8 @@ describe RuboCop::Cop::Migration::WithLockRetriesWithoutDdlTransaction do
       allow(cop).to receive(:in_migration?).and_return(true)
     end
 
-    it 'registers an offense when `with_lock_retries` is used with `disable_ddl_transaction!` method' do
-      inspect_source(invalid_source)
+    it 'registers an offense when `with_lock_retries` block has disallowed method' do
+      inspect_source('def change; with_lock_retries { disallowed_method }; end')
 
       aggregate_failures do
         expect(cop.offenses.size).to eq(1)
@@ -29,8 +26,33 @@ describe RuboCop::Cop::Migration::WithLockRetriesWithoutDdlTransaction do
       end
     end
 
-    it 'registers no offense when `with_lock_retries` is used inside an `up` method' do
-      inspect_source(valid_source)
+    it 'registers an offense when `with_lock_retries` block has disallowed methods' do
+      source = <<~HEREDOC
+      def change
+        with_lock_retries do
+          disallowed_method
+
+          create_table do |t|
+            t.text :text
+          end
+
+          other_disallowed_method
+
+          add_column :users, :name
+        end
+      end
+      HEREDOC
+
+      inspect_source(source)
+
+      aggregate_failures do
+        expect(cop.offenses.size).to eq(2)
+        expect(cop.offenses.map(&:line)).to eq([3, 9])
+      end
+    end
+
+    it 'registers no offense when `with_lock_retries` has only allowed method' do
+      inspect_source('def up; with_lock_retries { add_foreign_key :foo, :bar }; end')
 
       expect(cop.offenses.size).to eq(0)
     end
@@ -38,7 +60,7 @@ describe RuboCop::Cop::Migration::WithLockRetriesWithoutDdlTransaction do
 
   context 'outside of migration' do
     it 'registers no offense' do
-      inspect_source(invalid_source)
+      inspect_source('def change; with_lock_retries { disallowed_method }; end')
 
       expect(cop.offenses.size).to eq(0)
     end