diff --git a/lib/annotate_rb/tasks/annotate_models_migrate.rake b/lib/annotate_rb/tasks/annotate_models_migrate.rake index 4071d4c8..f12aed52 100644 --- a/lib/annotate_rb/tasks/annotate_models_migrate.rake +++ b/lib/annotate_rb/tasks/annotate_models_migrate.rake @@ -28,8 +28,13 @@ migration_tasks.each do |task| next unless Rake::Task.task_defined?(task) next if config[:skip_on_db_migrate] - Rake::Task[task].enhance do # This block is ran after `task` completes - task_name = Rake.application.top_level_tasks.last # The name of the task that was run, e.g. "db:migrate" + Rake::Task[task].enhance do |current_task| # This block is ran after `task` completes + # Prefer the top-level task (the one invoked from the CLI, e.g. "db:migrate") so that + # when sub-tasks chain (e.g. db:migrate:redo invokes db:rollback then db:migrate), we + # defer the annotation run to after everything completes. Fall back to the currently + # enhanced task when there is no top-level task (e.g. when the task is invoked + # programmatically from application code rather than from the Rake CLI). + task_name = Rake.application.top_level_tasks.last || current_task.name Rake::Task[task_name].enhance do ::AnnotateRb::Runner.run_after_migration diff --git a/spec/lib/tasks/annotate_models_migrate_spec.rb b/spec/lib/tasks/annotate_models_migrate_spec.rb index c05d04d5..5617e0ef 100644 --- a/spec/lib/tasks/annotate_models_migrate_spec.rb +++ b/spec/lib/tasks/annotate_models_migrate_spec.rb @@ -68,6 +68,18 @@ Rake.application.top_level end end + + describe "programmatic invocation (no top-level task)" do + it "should still annotate model files" do + # Simulate programmatic invocation from application code — e.g. Rails' + # in-browser "Run pending migrations" button, which calls + # Rake::Task["db:migrate"].invoke directly, bypassing Rake's CLI argv parser. + Rake.application.instance_variable_set(:@top_level_tasks, []) + + expect(AnnotateRb::Runner).to receive(:run_after_migration) + expect { Rake::Task["db:migrate"].invoke }.not_to raise_error + end + end end context "when skip_on_db_migrate is enabled" do