From e33075a0efce63658111ab35ce9a9103b8b0d9b4 Mon Sep 17 00:00:00 2001
From: eileencodes <eileencodes@gmail.com>
Date: Thu, 19 Mar 2020 16:41:59 -0400
Subject: [PATCH] Handle db:rollback and db:rollback:[NAME] for multi-db apps

With a multiple database application `db:rollback` becomes problematic.
We can't rollback just the primary, that doesn't match the behavior in
the other tasks. We can't rollback a migration for every database, that
is unexpected.

To solve this I handled `db:rollback` the same way I handled `:up` and
`:down`. If `db:rollback` is called for a multi-db application then it
will raise an error recommending you use `db:rollback:[NAME]` instead.
Calling `db:rollback:primary` or `db:rollback:animals` will rollback
the migration for the number of steps specified.

Closes: #38513
Follow-up to: #34078
---
 activerecord/CHANGELOG.md                     |  6 +++
 .../lib/active_record/railties/databases.rake | 18 +++++++
 .../test/application/rake/multi_dbs_test.rb   | 53 +++++++++++++++++++
 3 files changed, 77 insertions(+)

diff --git a/activerecord/CHANGELOG.md b/activerecord/CHANGELOG.md
index 125f7aa423..025023bbef 100644
--- a/activerecord/CHANGELOG.md
+++ b/activerecord/CHANGELOG.md
@@ -1,3 +1,9 @@
+*   Add support for `db:rollback:name` for multiple database applications.
+
+    Multiple database applications will now raise if `db:rollback` is call and recommend using the `db:rollback:[NAME]` to rollback migrations.
+
+    *Eileen M. Uchitelle*
+
 *   `Relation#pick` now uses already loaded results instead of making another query.
 
     *Eugene Kenny*
diff --git a/activerecord/lib/active_record/railties/databases.rake b/activerecord/lib/active_record/railties/databases.rake
index 098ae4dbd8..bcb9d7f22b 100644
--- a/activerecord/lib/active_record/railties/databases.rake
+++ b/activerecord/lib/active_record/railties/databases.rake
@@ -243,10 +243,28 @@ db_namespace = namespace :db do
     end
   end
 
+  namespace :rollback do
+    ActiveRecord::Tasks::DatabaseTasks.for_each(databases) do |name|
+      task name => :load_config do
+        step = ENV["STEP"] ? ENV["STEP"].to_i : 1
+        db_config = ActiveRecord::Base.configurations.configs_for(env_name: Rails.env, name: name)
+
+        ActiveRecord::Base.establish_connection(db_config)
+        ActiveRecord::Base.connection.migration_context.rollback(step)
+
+        db_namespace["_dump"].invoke
+      end
+    end
+  end
+
   desc "Rolls the schema back to the previous version (specify steps w/ STEP=n)."
   task rollback: :load_config do
+    ActiveRecord::Tasks::DatabaseTasks.raise_for_multi_db(command: "db:migrate:rollback")
+
     step = ENV["STEP"] ? ENV["STEP"].to_i : 1
+
     ActiveRecord::Base.connection.migration_context.rollback(step)
+
     db_namespace["_dump"].invoke
   end
 
diff --git a/railties/test/application/rake/multi_dbs_test.rb b/railties/test/application/rake/multi_dbs_test.rb
index 465dbee350..eed1ce4239 100644
--- a/railties/test/application/rake/multi_dbs_test.rb
+++ b/railties/test/application/rake/multi_dbs_test.rb
@@ -281,6 +281,31 @@ module ApplicationTests
         end
       end
 
+      def db_migrate_and_rollback(namespace = nil)
+        Dir.chdir(app_path) do
+          generate_models_for_animals
+          rails("db:migrate")
+
+          if namespace
+            rollback_output = rails("db:rollback:#{namespace}")
+          else
+            assert_raises RuntimeError, /You're using a multiple database application/ do
+              rollback_output = rails("db:rollback")
+            end
+          end
+
+          case namespace
+          when "primary"
+            assert_no_match(/OneMigration: reverted/, rollback_output)
+            assert_match(/CreateBooks: reverted/, rollback_output)
+          when nil
+          else
+            assert_no_match(/TwoMigration: reverted/, rollback_output)
+            assert_match(/CreateDogs: reverted/, rollback_output)
+          end
+        end
+      end
+
       def db_prepare
         Dir.chdir(app_path) do
           generate_models_for_animals
@@ -480,6 +505,34 @@ module ApplicationTests
         db_up_and_down "02", "animals"
       end
 
+      test "db:rollback raises on a multi-db application" do
+        require "#{app_path}/config/environment"
+
+        app_file "db/migrate/01_one_migration.rb", <<-MIGRATION
+          class OneMigration < ActiveRecord::Migration::Current
+          end
+        MIGRATION
+
+        db_migrate_and_rollback
+      end
+
+      test "db:rollback:namespace works" do
+        require "#{app_path}/config/environment"
+
+        app_file "db/migrate/01_one_migration.rb", <<-MIGRATION
+          class OneMigration < ActiveRecord::Migration::Current
+          end
+        MIGRATION
+
+        app_file "db/animals_migrate/02_two_migration.rb", <<-MIGRATION
+          class TwoMigration < ActiveRecord::Migration::Current
+          end
+        MIGRATION
+
+        db_migrate_and_rollback "primary"
+        db_migrate_and_rollback "animals"
+      end
+
       test "db:migrate:status works on all databases" do
         require "#{app_path}/config/environment"
         db_migrate_and_migrate_status