Rails 6.0 multi-db extensions for migration tasks

db.rake

require "active_record"

databases = ActiveRecord::Tasks::DatabaseTasks.setup_initial_database_yaml

db_namespace = namespace :db do
  if Rails.version.start_with?("6.0")

    # https://github.com/rails/rails/pull/38770/files
    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, spec_name: name)

          ActiveRecord::Base.establish_connection(db_config.config)
          ActiveRecord::Base.connection.migration_context.rollback(step)

          db_namespace["_dump"].invoke
        end
      end
    end

    namespace :migrate do
      ActiveRecord::Tasks::DatabaseTasks.for_each(databases) do |name|
        desc "Migrate #{name} database for current environment"
        task name => :load_config do
          original_db_config = ActiveRecord::Base.connection_config
          db_config = ActiveRecord::Base.configurations.configs_for(env_name: Rails.env, spec_name: name)

          ActiveRecord::Base.establish_connection(db_config.config)
          ActiveRecord::Tasks::DatabaseTasks.migrate
          # https://github.com/rails/rails/pull/38586
          # db_namespace["_dump:#{name}"].invoke
          db_namespace["_dump"].invoke
        ensure
          ActiveRecord::Base.establish_connection(original_db_config)
        end
      end

      # https://github.com/rails/rails/pull/39958
      desc "Rolls back the database one migration and re-migrates up (options: STEP=x, VERSION=x)."
      task redo: :load_config do
        ActiveRecord::Tasks::DatabaseTasks.raise_for_multi_db(command: "db:migrate:redo")

        raise "Empty VERSION provided" if ENV["VERSION"] && ENV["VERSION"].empty?

        if ENV["VERSION"]
          db_namespace["migrate:down"].invoke
          db_namespace["migrate:up"].invoke
        else
          db_namespace["rollback"].invoke
          db_namespace["migrate"].invoke
        end
      end

      namespace :redo do
        ActiveRecord::Tasks::DatabaseTasks.for_each(databases) do |name|
          desc "Rolls back #{name} database one migration and re-migrates up (options: STEP=x, VERSION=x)."
          task name => :load_config do
            raise "Empty VERSION provided" if ENV["VERSION"] && ENV["VERSION"].empty?

            if ENV["VERSION"]
              db_namespace["migrate:down:#{name}"].invoke
              db_namespace["migrate:up:#{name}"].invoke
            else
              db_namespace["rollback:#{name}"].invoke
              db_namespace["migrate:#{name}"].invoke
            end
          end
        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

    # https://github.com/rails/rails/pull/38449/files
    namespace :structure do

      namespace :dump do
        ActiveRecord::Tasks::DatabaseTasks.for_each(databases) do |spec_name|
          desc "Dumps the #{spec_name} database structure to db/structure.sql. Specify another file with SCHEMA=db/my_structure.sql"
          task spec_name => :load_config do
            db_config = ActiveRecord::Base.configurations.configs_for(env_name: ActiveRecord::Tasks::DatabaseTasks.env, spec_name: spec_name)
            ActiveRecord::Base.establish_connection(db_config.config)
            ActiveRecord::Tasks::DatabaseTasks.dump_schema(db_config.config, :sql, db_config.spec_name)
            db_namespace["structure:dump:#{spec_name}"].reenable
          end
        end
      end

      namespace :load do
        ActiveRecord::Tasks::DatabaseTasks.for_each(databases) do |spec_name|
          desc "Recreates the #{spec_name} database from the structure.sql file"
          task spec_name => :load_config do
            db_config = ActiveRecord::Base.configurations.configs_for(env_name: ActiveRecord::Tasks::DatabaseTasks.env, spec_name: spec_name)
            ActiveRecord::Tasks::DatabaseTasks.load_schema(db_config.config, :sql, ENV["SCHEMA"])
          end
        end
      end

    end


    namespace :test do
      ActiveRecord::Tasks::DatabaseTasks.for_each(databases) do |spec_name|
        # desc "Recreate the #{spec_name} test database"
        namespace :load do
          task spec_name => "db:test:purge:#{spec_name}" do
            case ActiveRecord::Base.schema_format
            when :ruby
              db_namespace["test:load_schema:#{spec_name}"].invoke
            when :sql
              db_namespace["test:load_structure:#{spec_name}"].invoke
            end
          end
        end

        # desc "Recreate the #{spec_name} test database from an existent schema.rb file"
        namespace :load_schema do
          task spec_name => "db:test:purge:#{spec_name}" do
            should_reconnect = ActiveRecord::Base.connection_pool.active_connection?
            ActiveRecord::Schema.verbose = false
            filename = ActiveRecord::Tasks::DatabaseTasks.dump_filename(spec_name, :ruby)
            db_config = ActiveRecord::Base.configurations.configs_for(env_name: "test", spec_name: spec_name)
            ActiveRecord::Tasks::DatabaseTasks.load_schema(db_config.config, :ruby, filename)
          ensure
            if should_reconnect
              ActiveRecord::Base.establish_connection(ActiveRecord::Tasks::DatabaseTasks.env.to_sym)
            end
          end
        end

        # desc "Recreate the #{spec_name} test database from an existent structure.sql file"
        namespace :load_structure do
          task spec_name => "db:test:purge:#{spec_name}" do
            filename = ActiveRecord::Tasks::DatabaseTasks.dump_filename(spec_name, :sql)
            db_config = ActiveRecord::Base.configurations.configs_for(env_name: "test", spec_name: spec_name)
            ActiveRecord::Tasks::DatabaseTasks.load_schema(db_config.config, :sql, filename)
          end
        end

        # desc "Empty the #{spec_name} test database"
        namespace :purge do
          task spec_name => %w(load_config check_protected_environments) do
            db_config = ActiveRecord::Base.configurations.configs_for(env_name: "test", spec_name: spec_name)
            ActiveRecord::Tasks::DatabaseTasks.purge(db_config.config)
          end
        end

        # desc 'Load the #{spec_name} database test schema'
        namespace :prepare do
          task spec_name => :load_config do
            db_namespace["test:load:#{spec_name}"].invoke
          end
        end
      end
    end
  else
    raise "Remove this patch: #{__FILE__} #{__LINE__}"
  end
end