Add annotation hooks for db-specific migration tasks (#686)

lib/tasks/annotate_models_migrate.rake

@@ -4,7 +4,18 @@
 # Append annotations to Rake tasks for ActiveRecord, so annotate automatically gets
 # run after doing db:migrate.
 
-%w(db:migrate db:migrate:up db:migrate:down db:migrate:reset db:migrate:redo db:rollback).each do |task|
+migration_tasks = %w(db:migrate db:migrate:up db:migrate:down db:migrate:reset db:migrate:redo db:rollback)
+if defined?(Rails::Application) && Rails.version.split('.').first.to_i >= 6
+  require 'active_record'
+
+  databases = ActiveRecord::Tasks::DatabaseTasks.setup_initial_database_yaml
+
+  ActiveRecord::Tasks::DatabaseTasks.for_each(databases) do |spec_name|
+    migration_tasks.concat(%w(db:migrate db:migrate:up db:migrate:down).map { |task| "#{task}:#{spec_name}" })
+  end
+end
+
+migration_tasks.each do |task|
   Rake::Task[task].enhance do
     Rake::Task[Rake.application.top_level_tasks.last].enhance do
       annotation_options_task = if Rake::Task.task_defined?('app:set_annotation_options')

spec/integration/rails_6.0.2.1/config/application.rb

@@ -28,5 +28,6 @@ module Rails6021
     # Application configuration can go into files in config/initializers
     # -- all .rb files in that directory are automatically loaded after loading
     # the framework and any gems in your application.
+    self.paths['config/database'] = 'config/multi-database.yml' if ENV['MULTI_DB']
   end
 end

spec/integration/rails_6.0.2.1/config/multi-database.yml

+# SQLite. Versions 3.8.0 and up are supported.
+#   gem install sqlite3
+#
+#   Ensure the SQLite 3 gem is defined in your Gemfile
+#   gem 'sqlite3'
+#
+default: &default
+  adapter: sqlite3
+  pool: <%= ENV.fetch("RAILS_MAX_THREADS") { 5 } %>
+  timeout: 5000
+
+development:
+  primary:
+    <<: *default
+    database: db/development.sqlite3
+  secondary:
+    <<: *default
+    database: db/development-secondary.sqlite3
+
+# Warning: The database defined as "test" will be erased and
+# re-generated from your development database when you run "rake".
+# Do not set this db to the same as development or production.
+test:
+  primary:
+    <<: *default
+    database: db/test.sqlite3
+  secondary:
+    <<: *default
+    database: db/test-secondary.sqlite3
+
+production:
+  primary:
+    <<: *default
+    database: db/production.sqlite3
+  secondary:
+    <<: *default
+    database: db/production-secondary.sqlite3

spec/integration/rails_6.0.2.1_spec.rb

@@ -9,23 +9,6 @@ describe 'Integration testing on Rails 6.0.2.1', if: IntegrationHelper.able_to_r
   ::RAILS_6_0_APP_PATH = File.expand_path(RAILS_6_0_APP_NAME, __dir__).freeze
 
   let!(:git) { Git.open(RAILS_6_0_PROJECT_PATH) }
-
-  before(:all) do
-    Bundler.with_clean_env do
-      Dir.chdir RAILS_6_0_APP_PATH do
-        puts `bundle install`
-        puts `bin/rails db:migrate`
-      end
-    end
-  end
-
-  after(:each) do
-    git.reset_hard
-  end
-
-  describe 'annotate --models' do
-    let(:command) { 'bundle exec annotate --models' }
-
   let(:task_model) do
     patch = <<~PATCH
       +# == Schema Information
@@ -90,9 +73,31 @@ describe 'Integration testing on Rails 6.0.2.1', if: IntegrationHelper.able_to_r
     }
   end
 
-    it 'annotate models' do
+  before(:all) do
+    Bundler.with_clean_env do
+      Dir.chdir RAILS_6_0_APP_PATH do
+        puts `bundle install`
+        puts `bin/rails db:migrate`
+      end
+    end
+  end
+
+  around(:each) do |example|
     Bundler.with_clean_env do
       Dir.chdir RAILS_6_0_APP_PATH do
+        example.run
+      end
+    end
+  end
+
+  after(:each) do
+    git.reset_hard
+  end
+
+  describe 'annotate --models' do
+    let(:command) { 'bundle exec annotate --models' }
+
+    it 'annotate models' do
       expect(git.diff.any?).to be_falsy
 
       puts `#{command}`
@@ -104,8 +109,6 @@ describe 'Integration testing on Rails 6.0.2.1', if: IntegrationHelper.able_to_r
                                   )
     end
   end
-    end
-  end
 
   describe 'annotate --routes' do
     let(:command) { 'bundle exec annotate --routes' }
@@ -156,8 +159,6 @@ describe 'Integration testing on Rails 6.0.2.1', if: IntegrationHelper.able_to_r
     end
 
     it 'annotate routes.rb' do
-      Bundler.with_clean_env do
-        Dir.chdir RAILS_6_0_APP_PATH do
       expect(git.diff.any?).to be_falsy
 
       puts `#{command}`
@@ -165,21 +166,34 @@ describe 'Integration testing on Rails 6.0.2.1', if: IntegrationHelper.able_to_r
       expect(git.diff.entries).to contain_exactly(an_object_having_attributes(task_routes))
     end
   end
-    end
-  end
 
   describe 'rails g annotate:install' do
     let(:command) { 'bin/rails g annotate:install' }
     let(:rake_file_path) { 'lib/tasks/auto_annotate_models.rake' }
+    let(:full_path) { File.expand_path(rake_file_path) }
+
+    after(:each) do
+      File.delete(full_path)
+    end
 
     it 'generates the rake file' do
-      Bundler.with_clean_env do
-        Dir.chdir RAILS_6_0_APP_PATH do
-          full_path = File.expand_path(rake_file_path)
       expect { `#{command}` }.to change { File.exist?(rake_file_path) }.from(false).to(true)
-
-          File.delete(full_path)
     end
+
+    context 'with multi-db environment' do
+      let(:migrate_command) { 'bin/rails db:migrate:primary' }
+
+      it 'hooks database-specific commands and annotates models' do
+        expect(git.diff.any?).to be_falsy
+
+        system({ 'MULTI_DB' => 'true' }, command)
+        system({ 'MULTI_DB' => 'true' }, migrate_command)
+
+        expect(git.diff.entries).to contain_exactly(
+                                        an_object_having_attributes(task_model),
+                                        an_object_having_attributes(task_test),
+                                        an_object_having_attributes(task_fixture)
+                                    )
       end
     end
   end