Merge pull request #32274 from eileencodes/part-1-add-rake-tasks-for-multi-db

Part 1 Easy Multi db in Rails: Add basic rake tasks for multi db setup

6 files changed, 271 insertions, 22 deletions

diff --git a/activerecord/lib/active_record.rb b/activerecord/lib/active_record.rb
index d43378c64f..d198466dbf 100644
--- a/activerecord/lib/active_record.rb
+++ b/activerecord/lib/active_record.rb
@@ -40,6 +40,7 @@ module ActiveRecord
   autoload :Core
   autoload :ConnectionHandling
   autoload :CounterCache
+  autoload :DatabaseConfigurations
   autoload :DynamicMatchers
   autoload :Enum
   autoload :InternalMetadata
diff --git a/activerecord/lib/active_record/base.rb b/activerecord/lib/active_record/base.rb
index cc99401390..7ab9160265 100644
--- a/activerecord/lib/active_record/base.rb
+++ b/activerecord/lib/active_record/base.rb
@@ -290,6 +290,7 @@ module ActiveRecord #:nodoc:
     extend CollectionCacheKey
 
     include Core
+    include DatabaseConfigurations
     include Persistence
     include ReadonlyAttributes
     include ModelSchema
diff --git a/activerecord/lib/active_record/database_configurations.rb b/activerecord/lib/active_record/database_configurations.rb
new file mode 100644
index 0000000000..86624a41c9
--- /dev/null
+++ b/activerecord/lib/active_record/database_configurations.rb
@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+  module DatabaseConfigurations # :nodoc:
+    class DatabaseConfig
+      attr_reader :env_name, :spec_name, :config
+
+      def initialize(env_name, spec_name, config)
+        @env_name = env_name
+        @spec_name = spec_name
+        @config = config
+      end
+    end
+
+    # Selects the config for the specified environment and specification name
+    #
+    # For example if passed :development, and :animals it will select the database
+    # under the :development and :animals configuration level
+    def self.config_for_env_and_spec(environment, specification_name, configs = ActiveRecord::Base.configurations) # :nodoc:
+      configs_for(environment, configs).find do |db_config|
+        db_config.spec_name == specification_name
+      end
+    end
+
+    # Collects the configs for the environment passed in.
+    #
+    # If a block is given returns the specification name and configuration
+    # otherwise returns an array of DatabaseConfig structs for the environment.
+    def self.configs_for(env, configs = ActiveRecord::Base.configurations, &blk) # :nodoc:
+      env_with_configs = db_configs(configs).select do |db_config|
+        db_config.env_name == env
+      end
+
+      if block_given?
+        env_with_configs.each do |env_with_config|
+          yield env_with_config.spec_name, env_with_config.config
+        end
+      else
+        env_with_configs
+      end
+    end
+
+    # Given an env, spec and config creates DatabaseConfig structs with
+    # each attribute set.
+    def self.walk_configs(env_name, spec_name, config) # :nodoc:
+      if config["database"] || env_name == "default"
+        DatabaseConfig.new(env_name, spec_name, config)
+      else
+        config.each_pair.map do |spec_name, sub_config|
+          walk_configs(env_name, spec_name, sub_config)
+        end
+      end
+    end
+
+    # Walks all the configs passed in and returns an array
+    # of DatabaseConfig structs for each configuration.
+    def self.db_configs(configs = ActiveRecord::Base.configurations) # :nodoc:
+      configs.each_pair.flat_map do |env_name, config|
+        walk_configs(env_name, "primary", config)
+      end
+    end
+  end
+end
diff --git a/activerecord/lib/active_record/railties/databases.rake b/activerecord/lib/active_record/railties/databases.rake
index 662a8bc720..76dbcafffb 100644
--- a/activerecord/lib/active_record/railties/databases.rake
+++ b/activerecord/lib/active_record/railties/databases.rake
@@ -22,6 +22,14 @@ db_namespace = namespace :db do
     task all: :load_config do
       ActiveRecord::Tasks::DatabaseTasks.create_all
     end
+
+    ActiveRecord::Tasks::DatabaseTasks.for_each do |spec_name|
+      desc "Create #{spec_name} database for current environment"
+      task spec_name => :load_config do
+        db_config = ActiveRecord::DatabaseConfigurations.config_for_env_and_spec(Rails.env, spec_name)
+        ActiveRecord::Tasks::DatabaseTasks.create(db_config.config)
+      end
+    end
   end
 
   desc "Creates the database from DATABASE_URL or config/database.yml for the current RAILS_ENV (use db:create:all to create all databases in the config). Without RAILS_ENV or when RAILS_ENV is development, it defaults to creating the development and test databases."
@@ -33,6 +41,14 @@ db_namespace = namespace :db do
     task all: [:load_config, :check_protected_environments] do
       ActiveRecord::Tasks::DatabaseTasks.drop_all
     end
+
+    ActiveRecord::Tasks::DatabaseTasks.for_each do |spec_name|
+      desc "Drop #{spec_name} database for current environment"
+      task spec_name => [:load_config, :check_protected_environments] do
+        db_config = ActiveRecord::DatabaseConfigurations.config_for_env_and_spec(Rails.env, spec_name)
+        ActiveRecord::Tasks::DatabaseTasks.drop(db_config.config)
+      end
+    end
   end
 
   desc "Drops the database from DATABASE_URL or config/database.yml for the current RAILS_ENV (use db:drop:all to drop all databases in the config). Without RAILS_ENV or when RAILS_ENV is development, it defaults to dropping the development and test databases."
@@ -57,7 +73,10 @@ db_namespace = namespace :db do
 
   desc "Migrate the database (options: VERSION=x, VERBOSE=false, SCOPE=blog)."
   task migrate: :load_config do
-    ActiveRecord::Tasks::DatabaseTasks.migrate
+    ActiveRecord::DatabaseConfigurations.configs_for(Rails.env) do |spec_name, config|
+      ActiveRecord::Base.establish_connection(config)
+      ActiveRecord::Tasks::DatabaseTasks.migrate
+    end
     db_namespace["_dump"].invoke
   end
 
@@ -77,6 +96,15 @@ db_namespace = namespace :db do
   end
 
   namespace :migrate do
+    ActiveRecord::Tasks::DatabaseTasks.for_each do |spec_name|
+      desc "Migrate #{spec_name} database for current environment"
+      task spec_name => :load_config do
+        db_config = ActiveRecord::DatabaseConfigurations.config_for_env_and_spec(Rails.env, spec_name)
+        ActiveRecord::Base.establish_connection(db_config.config)
+        ActiveRecord::Tasks::DatabaseTasks.migrate
+      end
+    end
+
     # desc  'Rollbacks the database one migration and re migrate up (options: STEP=x, VERSION=x).'
     task redo: :load_config do
       raise "Empty VERSION provided" if ENV["VERSION"] && ENV["VERSION"].empty?
@@ -246,10 +274,15 @@ db_namespace = namespace :db do
     desc "Creates a db/schema.rb file that is portable against any DB supported by Active Record"
     task dump: :load_config do
       require "active_record/schema_dumper"
-      filename = ENV["SCHEMA"] || File.join(ActiveRecord::Tasks::DatabaseTasks.db_dir, "schema.rb")
-      File.open(filename, "w:utf-8") do |file|
-        ActiveRecord::SchemaDumper.dump(ActiveRecord::Base.connection, file)
+
+      ActiveRecord::DatabaseConfigurations.configs_for(Rails.env) do |spec_name, config|
+        filename = ActiveRecord::Tasks::DatabaseTasks.dump_filename(spec_name, :ruby)
+        File.open(filename, "w:utf-8") do |file|
+          ActiveRecord::Base.establish_connection(config)
+          ActiveRecord::SchemaDumper.dump(ActiveRecord::Base.connection, file)
+        end
       end
+
       db_namespace["schema:dump"].reenable
     end
 
@@ -276,22 +309,25 @@ db_namespace = namespace :db do
         rm_f filename, verbose: false
       end
     end
-
   end
 
   namespace :structure do
     desc "Dumps the database structure to db/structure.sql. Specify another file with SCHEMA=db/my_structure.sql"
     task dump: :load_config do
-      filename = ENV["SCHEMA"] || File.join(ActiveRecord::Tasks::DatabaseTasks.db_dir, "structure.sql")
-      current_config = ActiveRecord::Tasks::DatabaseTasks.current_config
-      ActiveRecord::Tasks::DatabaseTasks.structure_dump(current_config, filename)
-
-      if ActiveRecord::SchemaMigration.table_exists?
-        File.open(filename, "a") do |f|
-          f.puts ActiveRecord::Base.connection.dump_schema_information
-          f.print "\n"
+      ActiveRecord::DatabaseConfigurations.configs_for(Rails.env) do |spec_name, config|
+        ActiveRecord::Base.establish_connection(config)
+        filename = ActiveRecord::Tasks::DatabaseTasks.dump_filename(spec_name, :sql)
+        current_config = ActiveRecord::Tasks::DatabaseTasks.current_config
+        ActiveRecord::Tasks::DatabaseTasks.structure_dump(config, filename)
+
+        if ActiveRecord::SchemaMigration.table_exists?
+          File.open(filename, "a") do |f|
+            f.puts ActiveRecord::Base.connection.dump_schema_information
+            f.print "\n"
+          end
         end
       end
+
       db_namespace["structure:dump"].reenable
     end
 
diff --git a/activerecord/lib/active_record/tasks/database_tasks.rb b/activerecord/lib/active_record/tasks/database_tasks.rb
index af1bbc7e93..5787660148 100644
--- a/activerecord/lib/active_record/tasks/database_tasks.rb
+++ b/activerecord/lib/active_record/tasks/database_tasks.rb
@@ -134,6 +134,13 @@ module ActiveRecord
         end
       end
 
+      def for_each
+        databases = Rails.application.config.load_database_yaml
+        ActiveRecord::DatabaseConfigurations.configs_for(Rails.env, databases) do |spec_name, _|
+          yield spec_name
+        end
+      end
+
       def create_current(environment = env)
         each_current_configuration(environment) { |configuration|
           create configuration
@@ -252,17 +259,31 @@ module ActiveRecord
       end
 
       def schema_file(format = ActiveRecord::Base.schema_format)
+        File.join(db_dir, schema_file_type(format))
+      end
+
+      def schema_file_type(format = ActiveRecord::Base.schema_format)
         case format
         when :ruby
-          File.join(db_dir, "schema.rb")
+          "schema.rb"
         when :sql
-          File.join(db_dir, "structure.sql")
+          "structure.sql"
         end
       end
 
+      def dump_filename(namespace, format = ActiveRecord::Base.schema_format)
+        filename = if namespace == "primary"
+          schema_file_type(format)
+        else
+          "#{namespace}_#{schema_file_type(format)}"
+        end
+
+        ENV["SCHEMA"] || File.join(ActiveRecord::Tasks::DatabaseTasks.db_dir, filename)
+      end
+
       def load_schema_current(format = ActiveRecord::Base.schema_format, file = nil, environment = env)
-        each_current_configuration(environment) { |configuration, configuration_environment|
-          load_schema configuration, format, file, configuration_environment
+        each_current_configuration(environment) { |configuration, spec_name, env|
+          load_schema configuration, format, file, env
         }
         ActiveRecord::Base.establish_connection(environment.to_sym)
       end
@@ -312,10 +333,10 @@ module ActiveRecord
           environments = [environment]
           environments << "test" if environment == "development"
 
-          ActiveRecord::Base.configurations.slice(*environments).each do |configuration_environment, configuration|
-            next unless configuration["database"]
-
-            yield configuration, configuration_environment
+          environments.each do |env|
+            ActiveRecord::DatabaseConfigurations.configs_for(env) do |spec_name, configuration|
+              yield configuration, spec_name, env
+            end
           end
         end
 
diff --git a/activerecord/test/cases/tasks/database_tasks_test.rb b/activerecord/test/cases/tasks/database_tasks_test.rb
index 21226352ff..38c2c8b2f3 100644
--- a/activerecord/test/cases/tasks/database_tasks_test.rb
+++ b/activerecord/test/cases/tasks/database_tasks_test.rb
@@ -221,7 +221,76 @@ module ActiveRecord
       ENV["RAILS_ENV"] = old_env
     end
 
-    def test_establishes_connection_for_the_given_environment
+    def test_establishes_connection_for_the_given_environments
+      ActiveRecord::Tasks::DatabaseTasks.stubs(:create).returns true
+
+      ActiveRecord::Base.expects(:establish_connection).with(:development)
+
+      ActiveRecord::Tasks::DatabaseTasks.create_current(
+        ActiveSupport::StringInquirer.new("development")
+      )
+    end
+  end
+
+  class DatabaseTasksCreateCurrentThreeTierTest < ActiveRecord::TestCase
+    def setup
+      @configurations = {
+        "development" => { "primary" => { "database" => "dev-db" }, "secondary" => { "database" => "secondary-dev-db" } },
+        "test" => { "primary" => { "database" => "test-db" }, "secondary" => { "database" => "secondary-test-db" } },
+        "production" => { "primary" => { "database" => "prod-db" }, "secondary" => { "database" => "secondary-prod-db" } }
+      }
+
+      ActiveRecord::Base.stubs(:configurations).returns(@configurations)
+      ActiveRecord::Base.stubs(:establish_connection).returns(true)
+    end
+
+    def test_creates_current_environment_database
+      ActiveRecord::Tasks::DatabaseTasks.expects(:create).
+        with("database" => "prod-db")
+
+      ActiveRecord::Tasks::DatabaseTasks.expects(:create).
+        with("database" => "secondary-prod-db")
+
+      ActiveRecord::Tasks::DatabaseTasks.create_current(
+        ActiveSupport::StringInquirer.new("production")
+      )
+    end
+
+    def test_creates_test_and_development_databases_when_env_was_not_specified
+      ActiveRecord::Tasks::DatabaseTasks.expects(:create).
+        with("database" => "dev-db")
+      ActiveRecord::Tasks::DatabaseTasks.expects(:create).
+        with("database" => "secondary-dev-db")
+      ActiveRecord::Tasks::DatabaseTasks.expects(:create).
+        with("database" => "test-db")
+      ActiveRecord::Tasks::DatabaseTasks.expects(:create).
+        with("database" => "secondary-test-db")
+
+      ActiveRecord::Tasks::DatabaseTasks.create_current(
+        ActiveSupport::StringInquirer.new("development")
+      )
+    end
+
+    def test_creates_test_and_development_databases_when_rails_env_is_development
+      old_env = ENV["RAILS_ENV"]
+      ENV["RAILS_ENV"] = "development"
+      ActiveRecord::Tasks::DatabaseTasks.expects(:create).
+        with("database" => "dev-db")
+      ActiveRecord::Tasks::DatabaseTasks.expects(:create).
+        with("database" => "secondary-dev-db")
+      ActiveRecord::Tasks::DatabaseTasks.expects(:create).
+        with("database" => "test-db")
+      ActiveRecord::Tasks::DatabaseTasks.expects(:create).
+        with("database" => "secondary-test-db")
+
+      ActiveRecord::Tasks::DatabaseTasks.create_current(
+        ActiveSupport::StringInquirer.new("development")
+      )
+    ensure
+      ENV["RAILS_ENV"] = old_env
+    end
+
+    def test_establishes_connection_for_the_given_environments_config
       ActiveRecord::Tasks::DatabaseTasks.stubs(:create).returns true
 
       ActiveRecord::Base.expects(:establish_connection).with(:development)
@@ -347,6 +416,64 @@ module ActiveRecord
     end
   end
 
+  class DatabaseTasksDropCurrentThreeTierTest < ActiveRecord::TestCase
+    def setup
+      @configurations = {
+        "development" => { "primary" => { "database" => "dev-db" }, "secondary" => { "database" => "secondary-dev-db" } },
+        "test" => { "primary" => { "database" => "test-db" }, "secondary" => { "database" => "secondary-test-db" } },
+        "production" => { "primary" => { "database" => "prod-db" }, "secondary" => { "database" => "secondary-prod-db" } }
+      }
+
+      ActiveRecord::Base.stubs(:configurations).returns(@configurations)
+    end
+
+    def test_drops_current_environment_database
+      ActiveRecord::Tasks::DatabaseTasks.expects(:drop).
+        with("database" => "prod-db")
+
+      ActiveRecord::Tasks::DatabaseTasks.expects(:drop).
+        with("database" => "secondary-prod-db")
+
+      ActiveRecord::Tasks::DatabaseTasks.drop_current(
+        ActiveSupport::StringInquirer.new("production")
+      )
+    end
+
+    def test_drops_test_and_development_databases_when_env_was_not_specified
+      ActiveRecord::Tasks::DatabaseTasks.expects(:drop).
+        with("database" => "dev-db")
+      ActiveRecord::Tasks::DatabaseTasks.expects(:drop).
+        with("database" => "secondary-dev-db")
+      ActiveRecord::Tasks::DatabaseTasks.expects(:drop).
+        with("database" => "test-db")
+      ActiveRecord::Tasks::DatabaseTasks.expects(:drop).
+        with("database" => "secondary-test-db")
+
+      ActiveRecord::Tasks::DatabaseTasks.drop_current(
+        ActiveSupport::StringInquirer.new("development")
+      )
+    end
+
+    def test_drops_testand_development_databases_when_rails_env_is_development
+      old_env = ENV["RAILS_ENV"]
+      ENV["RAILS_ENV"] = "development"
+      ActiveRecord::Tasks::DatabaseTasks.expects(:drop).
+        with("database" => "dev-db")
+      ActiveRecord::Tasks::DatabaseTasks.expects(:drop).
+        with("database" => "secondary-dev-db")
+      ActiveRecord::Tasks::DatabaseTasks.expects(:drop).
+        with("database" => "test-db")
+      ActiveRecord::Tasks::DatabaseTasks.expects(:drop).
+        with("database" => "secondary-test-db")
+
+      ActiveRecord::Tasks::DatabaseTasks.drop_current(
+        ActiveSupport::StringInquirer.new("development")
+      )
+    ensure
+      ENV["RAILS_ENV"] = old_env
+    end
+  end
+
   if current_adapter?(:SQLite3Adapter) && !in_memory_db?
     class DatabaseTasksMigrateTest < ActiveRecord::TestCase
       self.use_transactional_tests = false