Add annotation hooks for db-specific migration tasks

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 @@ class Application < Rails::Application
     # 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

@@ -0,0 +1,37 @@
+# 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,6 +9,69 @@
   ::RAILS_6_0_APP_PATH = File.expand_path(RAILS_6_0_APP_NAME, __dir__).freeze
 
   let!(:git) { Git.open(RAILS_6_0_PROJECT_PATH) }
+  let(:task_model) do
+    patch = <<~PATCH
+      +# == Schema Information
+      +#
+      +# Table name: tasks
+      +#
+      +#  id         :integer          not null, primary key
+      +#  content    :string
+      +#  count      :integer          default(0)
+      +#  status     :boolean          default(FALSE)
+      +#  created_at :datetime         not null
+      +#  updated_at :datetime         not null
+      +#
+    PATCH
+
+    path = 'app/models/task.rb'
+    {
+        path: include(path),
+        patch: include(patch)
+    }
+  end
+  let(:task_test) do
+    patch = <<~PATCH
+      +# == Schema Information
+      +#
+      +# Table name: tasks
+      +#
+      +#  id         :integer          not null, primary key
+      +#  content    :string
+      +#  count      :integer          default(0)
+      +#  status     :boolean          default(FALSE)
+      +#  created_at :datetime         not null
+      +#  updated_at :datetime         not null
+      +#
+    PATCH
+
+    path = 'test/models/task_test.rb'
+    {
+        path: include(path),
+        patch: include(patch)
+    }
+  end
+  let(:task_fixture) do
+    patch = <<~PATCH
+      +# == Schema Information
+      +#
+      +# Table name: tasks
+      +#
+      +#  id         :integer          not null, primary key
+      +#  content    :string
+      +#  count      :integer          default(0)
+      +#  status     :boolean          default(FALSE)
+      +#  created_at :datetime         not null
+      +#  updated_at :datetime         not null
+      +#
+    PATCH
+
+    path = 'test/fixtures/tasks.yml'
+    {
+        path: include(path),
+        patch: include(patch)
+    }
+  end
 
   before(:all) do
     Bundler.with_clean_env do
@@ -19,91 +82,31 @@
     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' }
 
-    let(:task_model) do
-      patch = <<~PATCH
-        +# == Schema Information
-        +#
-        +# Table name: tasks
-        +#
-        +#  id         :integer          not null, primary key
-        +#  content    :string
-        +#  count      :integer          default(0)
-        +#  status     :boolean          default(FALSE)
-        +#  created_at :datetime         not null
-        +#  updated_at :datetime         not null
-        +#
-      PATCH
-
-      path = 'app/models/task.rb'
-      {
-          path: include(path),
-          patch: include(patch)
-      }
-    end
-    let(:task_test) do
-      patch = <<~PATCH
-        +# == Schema Information
-        +#
-        +# Table name: tasks
-        +#
-        +#  id         :integer          not null, primary key
-        +#  content    :string
-        +#  count      :integer          default(0)
-        +#  status     :boolean          default(FALSE)
-        +#  created_at :datetime         not null
-        +#  updated_at :datetime         not null
-        +#
-      PATCH
-
-      path = 'test/models/task_test.rb'
-      {
-          path: include(path),
-          patch: include(patch)
-      }
-    end
-    let(:task_fixture) do
-      patch = <<~PATCH
-        +# == Schema Information
-        +#
-        +# Table name: tasks
-        +#
-        +#  id         :integer          not null, primary key
-        +#  content    :string
-        +#  count      :integer          default(0)
-        +#  status     :boolean          default(FALSE)
-        +#  created_at :datetime         not null
-        +#  updated_at :datetime         not null
-        +#
-      PATCH
-
-      path = 'test/fixtures/tasks.yml'
-      {
-          path: include(path),
-          patch: include(patch)
-      }
-    end
-
     it 'annotate models' do
-      Bundler.with_clean_env do
-        Dir.chdir RAILS_6_0_APP_PATH do
-          expect(git.diff.any?).to be_falsy
-
-          puts `#{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
+      expect(git.diff.any?).to be_falsy
+
+      puts `#{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
 
@@ -156,30 +159,41 @@
     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
+      expect(git.diff.any?).to be_falsy
 
-          puts `#{command}`
+      puts `#{command}`
 
-          expect(git.diff.entries).to contain_exactly(an_object_having_attributes(task_routes))
-        end
-      end
+      expect(git.diff.entries).to contain_exactly(an_object_having_attributes(task_routes))
     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)
+      expect { `#{command}` }.to change { File.exist?(rake_file_path) }.from(false).to(true)
+    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)
 
-          File.delete(full_path)
-        end
+        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