Name rake tasks with the database name (#223)

Blocking #222 

This PR is mostly a refactoring, but brings some breaking changes that I
think are worth the improvements in usability and extensibility.

Specifically, the problem I'm facing right now is updating the rake
tasks to be specific to a database, so for example `db:migrate:tenant`
should be named `db:migrate:primary` (or more generally,
`db:migrate:<DBNAME>`) so that if I have a secondary database it will
have its own rake tasks like `db:migrate:secondary`.

So the breaking change is around task names; however everything is still
a dependency of the Rails tasks `db:migrate`, `db:prepare`, `db:drop`,
and `db:reset` so it should be a huge deal for anybody (I hope).

- db:migrate:DBNAME replaces db:migrate:tenant and db:migrate:tenant:all
  - it operates on all tenants by default
- if there are no tenants it will create a database for the default
tenant
- the ARTENANT env var can be specified to run against a specific tenant
- db:drop:DBNAME replaces db:drop:tenant
  - it operates on all tenants by default
- NEW: the ARTENANT env var can be specified to run against a specific
tenant
- db:reset:DBNAME replaces db:reset:tenant
  - it operates on all tenants by default
- NEW: the ARTENANT env var can be specified to run against a specific
tenant
- Tenanted::DatabaseTasks.base_config has been removed

Some additional changes:

- Tenanted::DatabaseTasks is now a class that takes a tenanted base
config as a constructor parameter. Yay for encapsulation.
- ActiveRecord::Tenanted.base_configs is a new method that returns all
the tenanted base configs for the current environment.

Author: flavorjones

GUIDE.md

@@ -466,75 +466,17 @@ TODO:
   - [x] install a load hook
 
 - database tasks
-  - [x] make `db:migrate:tenant:all` iterate over all the tenants on disk
-  - [x] make `db:migrate:tenant ARTENANT=asdf` run migrations on just that tenant
-  - [x] make `db:migrate:tenant` run migrations on `development-tenant` in dev
-  - [x] make `db:migrate` run `db:migrate:tenant` in dev
-  - [x] make `db:prepare` run `db:migrate:tenant` in dev
+  - [x] make `db:migrate:__dbname__` migrate all the existing tenants
+  - [x] make `db:migrate:__dbname__ ARTENANT=asdf` run migrations on just that tenant
+  - [x] make `db:drop:__dbname__` drop all the existing tenants
+  - [x] make `db:drop:__dbname__ ARTENANT=asdf` drop just that tenant
+  - [x] make `db:migrate` run `db:migrate:__dbname__`
+  - [x] make `db:prepare` run `db:migrate:__dbname__`
+  - [x] make `db:drop` run `db:drop:__dbname__`
   - [x] make a decision on what output tasks should emit, and whether we need a separate verbose setting
+  - [x] use the database name instead of "tenant", e.g. "db:migrate:primary"
   - [ ] make the implicit migration opt-in
-  - [ ] use the database name instead of "tenant", e.g. "db:migrate:primary"
-  - [ ] fully implement all the relevant database tasks:
-    - [ ] `db:_dump`
-    - [ ] `db:_dump:__name__`
-    - [ ] `db:abort_if_pending_migrations`
-    - [ ] `db:abort_if_pending_migrations:__name__`
-    - [ ] `db:charset`
-    - [ ] `db:check_protected_environments`
-    - [ ] `db:collation`
-    - [ ] `db:create`
-    - [ ] `db:create:all`
-    - [ ] `db:create:__name__`
-    - [ ] `db:drop`
-    - [ ] `db:drop:_unsafe`
-    - [ ] `db:drop:all`
-    - [ ] `db:drop:__name__`
-    - [ ] `db:encryption:init`
-    - [ ] `db:environment:set`
-    - [ ] `db:fixtures:identify`
-    - [ ] `db:fixtures:load`
-    - [ ] `db:forward`
-    - [ ] `db:install:migrations`
-    - [ ] `db:load_config`
-    - [ ] `db:migrate` with support for VERSION
-    - [ ] `db:migrate:down` with support for VERSION
-    - [ ] `db:migrate:down:__name__`
-    - [ ] `db:migrate:__name__`
-    - [ ] `db:migrate:redo` with support for STEP and VERSION
-    - [ ] `db:migrate:redo:__name__`
-    - [ ] `db:migrate:reset`
-    - [ ] `db:migrate:status`
-    - [ ] `db:migrate:status:__name__`
-    - [ ] `db:migrate:up` with support for VERSION
-    - [ ] `db:migrate:up:__name__`
-    - [ ] `db:prepare`
-    - [ ] `db:purge` (see Known Issues below)
-    - [ ] `db:purge:all` (see Known Issues below)
-    - [ ] `db:reset`
-    - [ ] `db:reset:all`
-    - [ ] `db:reset:__name__`
-    - [ ] `db:rollback` with support for STEP
-    - [ ] `db:rollback:__name__`
-    - [ ] `db:schema:cache:clear`
-    - [ ] `db:schema:cache:dump`
-    - [ ] `db:schema:dump`
-    - [ ] `db:schema:dump:__name__`
-    - [ ] `db:schema:load`
-    - [ ] `db:schema:load:__name__`
-    - [ ] `db:seed`
-    - [ ] `db:seed:replant`
-    - [ ] `db:setup`
-    - [ ] `db:setup:all`
-    - [ ] `db:setup:__name__`
-    - [ ] `db:test:load_schema`
-    - [ ] `db:test:load_schema:__name__`
-    - [ ] `db:test:prepare`
-    - [ ] `db:test:prepare:__name__`
-    - [ ] `db:test:purge`
-    - [ ] `db:test:purge:__name__`
-    - [ ] `db:truncate_all`
-    - [ ] `db:version`
-    - [ ] `db:version:__name__`
+  - [ ] fully implement all the relevant database tasks - see https://github.com/basecamp/activerecord-tenanted/issues/222
 
 - installation
   - [ ] install a variation on the default database.yml with primary tenanted and non-primary "global" untenanted

lib/active_record/tenanted.rb

@@ -41,10 +41,19 @@ class IntegrationNotConfiguredError < Error; end
     # Raised when an unsupported database adapter is used.
     class UnsupportedDatabaseError < Error; end
 
+    # Return the constantized connection class configured in `config.active_record_tenanted.connection_class`,
+    # or nil if none is configured.
     def self.connection_class
       # TODO: cache this / speed this up
       Rails.application.config.active_record_tenanted.connection_class&.constantize
     end
+
+    # Return an Array of the tenanted database configurations.
+    def self.base_configs(configurations = ActiveRecord::Base.configurations)
+      configurations
+        .configs_for(env_name: ActiveRecord::Tasks::DatabaseTasks.env, include_hidden: true)
+        .select { |c| c.configuration_hash[:tenanted] }
+    end
   end
 end
 

lib/active_record/tenanted/console.rb

@@ -5,7 +5,11 @@ module Tenanted
     module Console # :nodoc:
       module IRBConsole
         def start
-          ActiveRecord::Tenanted::DatabaseTasks.set_current_tenant if Rails.env.local?
+          # TODO: we could be setting the current tenant for all tenanted configs.
+          if Rails.env.local? && ActiveRecord::Tenanted.connection_class
+            config = ActiveRecord::Tenanted.connection_class.connection_pool.db_config
+            ActiveRecord::Tenanted::DatabaseTasks.new(config).set_current_tenant
+          end
           super
         end
       end

lib/active_record/tenanted/database_tasks.rb

@@ -1,54 +1,61 @@
 # frozen_string_literal: true
 
+require "rake"
+
 module ActiveRecord
   module Tenanted
-    module DatabaseTasks # :nodoc:
-      extend self
-
-      def migrate_all
-        raise ArgumentError, "Could not find a tenanted database" unless config = base_config
+    class DatabaseTasks # :nodoc:
+      include Rake::DSL
 
-        tenants = config.tenants.presence || [ get_current_tenant ].compact
-        tenants.each do |tenant|
-          tenant_config = config.new_tenant_config(tenant)
-          migrate(tenant_config)
+      class << self
+        def verbose?
+          ActiveRecord::Tasks::DatabaseTasks.send(:verbose?)
         end
       end
 
-      def migrate_tenant(tenant_name = set_current_tenant)
-        raise ArgumentError, "Could not find a tenanted database" unless config = base_config
+      attr_reader :config
 
-        tenant_config = config.new_tenant_config(tenant_name)
+      def initialize(config)
+        unless config.is_a?(ActiveRecord::Tenanted::DatabaseConfigurations::BaseConfig)
+          raise TypeError, "Argument must be an instance of ActiveRecord::Tenanted::DatabaseConfigurations::BaseConfig"
+        end
+        @config = config
+      end
 
-        migrate(tenant_config)
+      def migrate_all
+        tenants = config.tenants.presence || [ get_default_tenant ].compact
+        tenants.each do |tenant|
+          migrate_tenant(tenant)
+        end
       end
 
-      def drop_all
-        raise ArgumentError, "Could not find a tenanted database" unless config = base_config
+      def migrate_tenant(tenant = set_current_tenant)
+        db_config = config.new_tenant_config(tenant)
+        migrate(db_config)
+        $stdout.puts "Migrated database '#{db_config.database}'" if verbose?
+      end
 
+      def drop_all
         config.tenants.each do |tenant|
-          db_config = config.new_tenant_config(tenant)
-          db_config.config_adapter.drop_database
-          $stdout.puts "Dropped database '#{db_config.database}'" if verbose?
+          drop_tenant(tenant)
         end
       end
 
-      def base_config
-        db_configs = ActiveRecord::Base.configurations.configs_for(
-          env_name: ActiveRecord::Tasks::DatabaseTasks.env,
-          include_hidden: true
-        )
-        db_configs.detect { |c| c.configuration_hash[:tenanted] }
+      def drop_tenant(tenant = set_current_tenant)
+        db_config = config.new_tenant_config(tenant)
+        db_config.config_adapter.drop_database
+        $stdout.puts "Dropped database '#{db_config.database}'" if verbose?
       end
 
-      def get_current_tenant
+      def get_default_tenant
+        # TODO: needs to work with multiple tenanted configs, maybe using ENV["ARTENANT_#{config.name}"]
         tenant = ENV["ARTENANT"]
 
         if tenant.present?
           $stdout.puts "Setting current tenant to #{tenant.inspect}" if verbose?
         elsif Rails.env.local?
           tenant = Rails.application.config.active_record_tenanted.default_tenant
-          $stdout.puts "Defaulting current tenant to #{tenant.inspect}" if verbose?
+          $stdout.puts "Defaulting current tenant for #{config.name.inspect} to #{tenant.inspect}" if verbose?
         else
           tenant = nil
           $stdout.puts "Cannot determine an implicit tenant: ARTENANT not set, and Rails.env is not local." if verbose?
@@ -64,7 +71,7 @@ def set_current_tenant
         end
 
         if connection_class.current_tenant.nil?
-          connection_class.current_tenant = get_current_tenant
+          connection_class.current_tenant = get_default_tenant
         else
           connection_class.current_tenant
         end
@@ -107,7 +114,49 @@ def migrate(config)
       end
 
       def verbose?
-        ActiveRecord::Tasks::DatabaseTasks.send(:verbose?)
+        self.class.verbose?
+      end
+
+      def register_rake_tasks
+        name = config.name
+
+        desc "Migrate tenanted #{name} databases for current environment"
+        task "db:migrate:#{name}" => "load_config" do
+          verbose_was = ActiveRecord::Migration.verbose
+          ActiveRecord::Migration.verbose = ActiveRecord::Tenanted::DatabaseTasks.verbose?
+
+          tenant = ENV["ARTENANT"]
+          if tenant.present?
+            migrate_tenant(tenant)
+          else
+            migrate_all
+          end
+        ensure
+          ActiveRecord::Migration.verbose = verbose_was
+        end
+        task "db:migrate" => "db:migrate:#{name}"
+        task "db:prepare" => "db:migrate:#{name}"
+
+        desc "Drop tenanted #{name} databases for current environment"
+        task "db:drop:#{name}" => "load_config" do
+          verbose_was = ActiveRecord::Migration.verbose
+          ActiveRecord::Migration.verbose = ActiveRecord::Tenanted::DatabaseTasks.verbose?
+
+          tenant = ENV["ARTENANT"]
+          if tenant.present?
+            drop_tenant(tenant)
+          else
+            drop_all
+          end
+        ensure
+          ActiveRecord::Migration.verbose = verbose_was
+        end
+        task "db:drop" => "db:drop:#{name}"
+
+        # TODO: Rails' database tasks include "db:seed" in the tasks that "db:reset" runs.
+        desc "Drop and recreate tenanted #{name} database from its schema for the current environment"
+        task "db:reset:#{name}" => [ "db:drop:#{name}", "db:migrate:#{name}" ]
+        task "db:reset" => "db:reset:#{name}"
       end
     end
   end

lib/active_record/tenanted/tenant.rb

@@ -120,15 +120,16 @@ def with_tenant(tenant_name, prohibit_shard_swapping: true, &block)
 
         def create_tenant(tenant_name, if_not_exists: false, &block)
           created_db = false
-          adapter = tenanted_root_config.new_tenant_config(tenant_name).config_adapter
+          base_config = tenanted_root_config
+          adapter = base_config.new_tenant_config(tenant_name).config_adapter
 
           adapter.acquire_ready_lock do
             unless adapter.database_exist?
               adapter.create_database
 
               with_tenant(tenant_name) do
                 connection_pool(schema_version_check: false)
-                ActiveRecord::Tenanted::DatabaseTasks.migrate_tenant(tenant_name)
+                ActiveRecord::Tenanted::DatabaseTasks.new(base_config).migrate_tenant(tenant_name)
               end
 
               created_db = true

lib/tasks/active_record/tenanted_tasks.rake

@@ -1,78 +1,20 @@
 # frozen_string_literal: true
 
-namespace :db do
-  desc "Migrate the database for tenant ARTENANT"
-  task "migrate:tenant" => "load_config" do
-    unless ActiveRecord::Tenanted.connection_class
-      warn "ActiveRecord::Tenanted integration is not configured via connection_class"
-      next
-    end
-
-    unless ActiveRecord::Tenanted::DatabaseTasks.base_config
-      warn "WARNING: No tenanted database found, skipping tenanted migration"
-    else
-      begin
-        verbose_was = ActiveRecord::Migration.verbose
-        ActiveRecord::Migration.verbose = ActiveRecord::Tenanted::DatabaseTasks.verbose?
-
-        ActiveRecord::Tenanted::DatabaseTasks.migrate_tenant
-      ensure
-        ActiveRecord::Migration.verbose = verbose_was
-      end
-    end
-  end
-
-  desc "Migrate the database for all existing tenants"
-  task "migrate:tenant:all" => "load_config" do
-    unless ActiveRecord::Tenanted.connection_class
-      warn "ActiveRecord::Tenanted integration is not configured via connection_class"
-      next
-    end
-
-    verbose_was = ActiveRecord::Migration.verbose
-    ActiveRecord::Migration.verbose = ActiveRecord::Tenanted::DatabaseTasks.verbose?
-
-    ActiveRecord::Tenanted::DatabaseTasks.migrate_all
-  ensure
-    ActiveRecord::Migration.verbose = verbose_was
-  end
-
-  desc "Drop and recreate all tenant databases from their schema for the current environment"
-  task "reset:tenant" => [ "db:drop:tenant", "db:migrate:tenant" ]
-
-  desc "Drop all tenanted databases for the current environment"
-  task "drop:tenant" => "load_config" do
-    unless ActiveRecord::Tenanted::DatabaseTasks.base_config
-      warn "WARNING: No tenanted database found, skipping tenanted reset"
-    else
-      begin
-        verbose_was = ActiveRecord::Migration.verbose
-        ActiveRecord::Migration.verbose = ActiveRecord::Tenanted::DatabaseTasks.verbose?
-
-        ActiveRecord::Tenanted::DatabaseTasks.drop_all
-      ensure
-        ActiveRecord::Migration.verbose = verbose_was
-      end
-    end
+# Ensure a default tenant is set for database tasks that may need it.
+desc "Set the current tenant to ARTENANT if present, else the environment default"
+task "db:tenant" => "load_config" do
+  unless ActiveRecord::Tenanted.connection_class
+    warn "ActiveRecord::Tenanted integration is not configured via connection_class"
+    next
   end
 
-  desc "Set the current tenant to ARTENANT if present, else the environment default"
-  task "tenant" => "load_config" do
-    unless ActiveRecord::Tenanted.connection_class
-      warn "ActiveRecord::Tenanted integration is not configured via connection_class"
-      next
-    end
-
-    ActiveRecord::Tenanted::DatabaseTasks.set_current_tenant
-  end
+  config = ActiveRecord::Tenanted.connection_class.connection_pool.db_config
+  ActiveRecord::Tenanted::DatabaseTasks.new(config).set_current_tenant
 end
-
-# Decorate database tasks with the tenanted version.
-task "db:migrate" => "db:migrate:tenant:all"
-task "db:prepare" => "db:migrate:tenant:all"
-task "db:reset"   => "db:reset:tenant"
-task "db:drop"    => "db:drop:tenant"
-
-# Ensure a default tenant is set for database tasks that may need it.
 task "db:fixtures:load" => "db:tenant"
 task "db:seed"          => "db:tenant"
+
+# Create tenanted rake tasks
+ActiveRecord::Tenanted.base_configs(ActiveRecord::DatabaseConfigurations.new(ActiveRecord::Tasks::DatabaseTasks.setup_initial_database_yaml)).each do |config|
+  ActiveRecord::Tenanted::DatabaseTasks.new(config).register_rake_tasks
+end

test/test_helper.rb

@@ -51,6 +51,12 @@ def test(name, &block)
           super
         end
 
+        def for_each_db_scenario(s = all_scenarios, &block)
+          s.each_key do |db_scenario|
+            with_db_scenario(db_scenario, &block)
+          end
+        end
+
         def for_each_scenario(s = all_scenarios, except: {}, &block)
           s.each do |db_scenario, model_scenarios|
             with_db_scenario(db_scenario) do