Enhance BulkChangeTable to support block wrappers

lib/rubocop/cop/rails/bulk_change_table.rb

@@ -124,16 +124,8 @@ def on_def(node)
           return unless MIGRATION_METHODS.include?(node.method_name)
           return unless node.body
 
-          recorder = AlterMethodsRecorder.new
-
-          node.body.child_nodes.each do |child_node|
-            if call_to_combinable_alter_method? child_node
-              recorder.process(child_node)
-            else
-              recorder.flush
-            end
-          end
-
+          recorder = AlterMethodsRecorder.new(combinable_alter_methods)
+          recorder.process_migration_body(node.body)
           recorder.offensive_nodes.each { |n| add_offense_for_alter_methods(n) }
         end
 
@@ -190,10 +182,6 @@ def support_bulk_alter?
           end
         end
 
-        def call_to_combinable_alter_method?(child_node)
-          child_node.send_type? && combinable_alter_methods.include?(child_node.method_name)
-        end
-
         def combinable_alter_methods
           case database
           when MYSQL
@@ -233,11 +221,22 @@ def add_offense_for_change_table(node)
 
         # Record combinable alter methods and register offensive nodes.
         class AlterMethodsRecorder
-          def initialize
+          def initialize(combinable_methods)
+            @combinable_methods = combinable_methods
             @nodes = []
             @offensive_nodes = []
           end
 
+          def process_migration_body(body_node)
+            if body_node.block_type?
+              process_block_body(body_node)
+            else
+              body_node.child_nodes.each do |child_node|
+                process_node_for_alter_methods(child_node)
+              end
+            end
+          end
+
           # @param new_node [RuboCop::AST::SendNode]
           def process(new_node)
             # arguments: [{(sym :table)(str "table")} ...]
@@ -261,6 +260,60 @@ def offensive_nodes
             flush
             @offensive_nodes
           end
+
+          private
+
+          def process_node_for_alter_methods(node)
+            if call_to_combinable_alter_method?(node)
+              process(node)
+            elsif node.block_type?
+              process_block_body(node)
+            else
+              flush
+            end
+          end
+
+          def process_block_body(block_node)
+            body = block_node.children[2]
+            return unless body
+
+            case body.type
+            when :begin
+              process_begin_block_body(body)
+            when :send, :block
+              process_single_node_block_body(body)
+            else
+              process_other_block_body(body)
+            end
+          end
+
+          def process_begin_block_body(body)
+            body.child_nodes.each do |child_node|
+              if child_node.block_type?
+                flush
+                process_node_for_alter_methods(child_node)
+                flush
+              else
+                process_node_for_alter_methods(child_node)
+              end
+            end
+          end
+
+          def process_single_node_block_body(body)
+            process_node_for_alter_methods(body)
+          end
+
+          def process_other_block_body(body)
+            return unless body.respond_to?(:child_nodes)
+
+            body.child_nodes.each do |child_node|
+              process_node_for_alter_methods(child_node)
+            end
+          end
+
+          def call_to_combinable_alter_method?(child_node)
+            child_node.send_type? && @combinable_methods.include?(child_node.method_name)
+          end
         end
       end
     end

spec/rubocop/cop/rails/bulk_change_table_spec.rb

@@ -461,6 +461,205 @@ def change
         end
       RUBY
     end
+
+    context 'with methods wrapped in blocks' do
+      it 'registers an offense when alter methods are wrapped in safety_assured block with do..end syntax' do
+        expect_offense(<<~RUBY)
+          def change
+            safety_assured do
+              add_column :users, :name, :string, null: false
+              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ You can use `change_table :users, bulk: true` to combine alter queries.
+              add_column :users, :email, :string
+            end
+          end
+        RUBY
+      end
+
+      it 'registers an offense when alter methods are wrapped in safety_assured block with {} syntax' do
+        expect_offense(<<~RUBY)
+          def change
+            safety_assured {
+              add_column :users, :name, :string, null: false
+              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ You can use `change_table :users, bulk: true` to combine alter queries.
+              remove_column :users, :nickname
+            }
+          end
+        RUBY
+      end
+
+      it 'registers an offense when alter methods are wrapped in safety_assured with mixed table operations' do
+        expect_offense(<<~RUBY)
+          def change
+            safety_assured do
+              add_column :users, :name, :string, null: false
+              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ You can use `change_table :users, bulk: true` to combine alter queries.
+              add_column :users, :email, :string
+              remove_column :users, :nickname
+              add_column :profiles, :bio, :text
+              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ You can use `change_table :profiles, bulk: true` to combine alter queries.
+              add_column :profiles, :avatar_url, :string
+            end
+          end
+        RUBY
+      end
+
+      it 'registers an offense when alter methods are wrapped in reversible block' do
+        expect_offense(<<~RUBY)
+          def change
+            reversible do |dir|
+              dir.up do
+                add_column :users, :name, :string, null: false
+                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ You can use `change_table :users, bulk: true` to combine alter queries.
+                add_column :users, :email, :string
+              end
+              dir.down do
+                remove_column :users, :email
+                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ You can use `change_table :users, bulk: true` to combine alter queries.
+                remove_column :users, :name
+              end
+            end
+          end
+        RUBY
+      end
+
+      it 'registers an offense when alter methods are wrapped in up_only block' do
+        expect_offense(<<~RUBY)
+          def change
+            up_only do
+              add_column :users, :name, :string, null: false
+              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ You can use `change_table :users, bulk: true` to combine alter queries.
+              add_column :users, :email, :string
+              remove_column :users, :old_field
+            end
+          end
+        RUBY
+      end
+
+      it 'registers an offense when alter methods are wrapped in custom block' do
+        expect_offense(<<~RUBY)
+          def change
+            with_lock_retries do
+              add_column :users, :name, :string, null: false
+              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ You can use `change_table :users, bulk: true` to combine alter queries.
+              add_column :users, :email, :string
+            end
+          end
+        RUBY
+      end
+
+      it 'registers an offense when alter methods are wrapped in nested blocks' do
+        expect_offense(<<~RUBY)
+          def change
+            safety_assured do
+              with_lock_retries do
+                add_column :users, :name, :string, null: false
+                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ You can use `change_table :users, bulk: true` to combine alter queries.
+                add_column :users, :email, :string
+              end
+            end
+          end
+        RUBY
+      end
+
+      it 'registers an offense when some alter methods are in blocks and others are not' do
+        expect_offense(<<~RUBY)
+          def change
+            add_column :users, :first_name, :string
+            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ You can use `change_table :users, bulk: true` to combine alter queries.
+            safety_assured do
+              add_column :users, :last_name, :string
+            end
+            add_column :users, :email, :string
+          end
+        RUBY
+      end
+
+      it 'registers an offense when alter methods are wrapped in transaction block' do
+        expect_offense(<<~RUBY)
+          def change
+            transaction do
+              add_column :users, :name, :string, null: false
+              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ You can use `change_table :users, bulk: true` to combine alter queries.
+              add_column :users, :email, :string
+            end
+          end
+        RUBY
+      end
+
+      it 'registers an offense when alter methods are wrapped in disable_ddl_transaction block' do
+        expect_offense(<<~RUBY)
+          def change
+            disable_ddl_transaction!
+
+            safety_assured do
+              add_index :users, :email, algorithm: :concurrently
+              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ You can use `change_table :users, bulk: true` to combine alter queries.
+              add_index :users, :name, algorithm: :concurrently
+            end
+          end
+        RUBY
+      end
+
+      it 'does not register an offense when alter methods target different tables in the same block' do
+        expect_no_offenses(<<~RUBY)
+          def change
+            safety_assured do
+              add_column :users, :name, :string
+              add_column :profiles, :bio, :text
+            end
+          end
+        RUBY
+      end
+
+      it 'does not register an offense when only one alter method is in a block' do
+        expect_no_offenses(<<~RUBY)
+          def change
+            safety_assured do
+              add_column :users, :name, :string
+            end
+          end
+        RUBY
+      end
+
+      it 'does not register an offense when blocks contain non-combinable operations' do
+        expect_no_offenses(<<~RUBY)
+          def change
+            safety_assured do
+              add_column :users, :name, :string
+              add_reference :users, :team
+              remove_column :users, :old_field
+            end
+          end
+        RUBY
+      end
+
+      it 'registers an offense when change_table transformations are wrapped in safety_assured' do
+        expect_offense(<<~RUBY)
+          def change
+            safety_assured do
+              change_table :users do |t|
+              ^^^^^^^^^^^^^^^^^^^ You can combine alter queries using `bulk: true` options.
+                t.string :name, null: false
+                t.string :email
+              end
+            end
+          end
+        RUBY
+      end
+
+      it 'does not register an offense when change_table with bulk: true is wrapped in safety_assured' do
+        expect_no_offenses(<<~RUBY)
+          def change
+            safety_assured do
+              change_table :users, bulk: true do |t|
+                t.string :name, null: false
+                t.string :email
+              end
+            end
+          end
+        RUBY
+      end
+    end
   end
 
   context 'when database is PostgreSQL' do