Adds rollback task

CHANGELOG.md

@@ -5,6 +5,11 @@
 ### Fixed
 
 - `.current_tenant = nil` now clears the tenant context, properly setting the shard to `UNTENANTED_SENTINEL` instead of `""` @flavorjones
+- Prevent `db:rollback` from raising `ActiveRecord::Tenanted::NoTenantError` when only tenanted databases are configured.
+
+### Added
+
+- Adds tenanted database rollback support via `ActiveRecord::Tenanted::DatabaseTasks#rollback_*` and the `db:rollback:DBNAME` rake task.
 
 
 ## v0.6.0 / 2025-11-05

GUIDE.md

@@ -420,6 +420,11 @@ Documentation outline:
     - 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:rollback:DBNAME
+    - dependency of db:rollback
+    - it operates on all tenants by default
+    - the ARTENANT env var can be specified to run against a specific tenant
+    - the STEP env var determines how many migrations to revert (defaults to 1)
   - db:drop:DBNAME replaces db:drop:tenant
     - dependency of db:drop
     - it operates on all tenants by default
@@ -492,6 +497,7 @@ TODO:
   - [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:rollback:__dbname__` rollback all the existing tenants
   - [x] make `db:migrate` run `db:migrate:__dbname__`
   - [x] make `db:prepare` run `db:migrate:__dbname__`
   - [x] make `db:drop` run `db:drop:__dbname__`

lib/active_record/tenanted/database_tasks.rb

@@ -11,6 +11,38 @@ class << self
         def verbose?
           ActiveRecord::Tasks::DatabaseTasks.send(:verbose?)
         end
+
+        # When we only have tenanted databases configured, we want to
+        # prevent the default Rails db tasks from running, since they
+        # will blow up with NoTenantError.
+        def wrap_rails_task(task_name, &should_run_original)
+          wrapped_rails_tasks[task_name] ||= begin
+            task = Rake::Task[task_name]
+            original_prerequisites = task.prerequisites.dup
+            original_actions = task.actions.dup
+            run_original = should_run_original || -> { true }
+
+            task.clear
+            task.enhance(original_prerequisites)
+            task.enhance do |t, args|
+              if run_original.call
+                original_actions.each { |action| action.call(t, args) }
+              end
+            end
+
+            true
+          end
+        end
+
+        private
+          def wrapped_rails_tasks
+            @wrapped_rails_tasks ||= {}
+          end
+
+          def default_database_tasks_present?
+            configs = ActiveRecord::Base.configurations
+            configs && configs.configs_for(env_name: Rails.env).any?
+          end
       end
 
       attr_reader :config
@@ -45,6 +77,17 @@ def drop_tenant(tenant = set_current_tenant)
         $stdout.puts "Dropped database '#{db_config.database}'" if verbose?
       end
 
+      def rollback_all(step = 1)
+        tenants.each do |tenant|
+          rollback_tenant(tenant, step: step)
+        end
+      end
+
+      def rollback_tenant(tenant = set_current_tenant, step: 1)
+        db_config = config.new_tenant_config(tenant)
+        rollback(db_config, step: step)
+      end
+
       def tenants
         config.tenants.presence || [ get_default_tenant ].compact
       end
@@ -115,13 +158,38 @@ def migrate(config)
         end
       end
 
+      def rollback(config, step:)
+        ActiveRecord::Tasks::DatabaseTasks.with_temporary_connection(config) do |conn|
+          pool = conn.pool
+
+          pool.migration_context.rollback(step)
+          pool.schema_cache.clear!
+
+          if Rails.env.development? || ENV["ARTENANT_SCHEMA_DUMP"].present?
+            ActiveRecord::Tasks::DatabaseTasks.dump_schema(config)
+
+            cache_dump = ActiveRecord::Tasks::DatabaseTasks.cache_dump_filename(config)
+            ActiveRecord::Tasks::DatabaseTasks.dump_schema_cache(pool, cache_dump)
+          end
+        end
+      end
+
       def verbose?
         self.class.verbose?
       end
 
       def register_rake_tasks
         name = config.name
 
+        self.class.wrap_rails_task("db:rollback") do
+          self.class.send(:default_database_tasks_present?)
+        end
+
+        step_from_env = -> {
+          step = ENV["STEP"]
+          step.present? ? step.to_i : 1
+        }
+
         desc "Migrate tenanted #{name} databases for current environment"
         task "db:migrate:#{name}" => "load_config" do
           verbose_was = ActiveRecord::Migration.verbose
@@ -139,6 +207,24 @@ def register_rake_tasks
         task "db:migrate" => "db:migrate:#{name}"
         task "db:prepare" => "db:migrate:#{name}"
 
+        desc "Rollback tenanted #{name} databases for current environment"
+        task "db:rollback:#{name}" => "load_config" do
+          verbose_was = ActiveRecord::Migration.verbose
+          ActiveRecord::Migration.verbose = ActiveRecord::Tenanted::DatabaseTasks.verbose?
+
+          step = step_from_env.call
+          tenant = ENV["ARTENANT"]
+
+          if tenant.present?
+            rollback_tenant(tenant, step: step)
+          else
+            rollback_all(step)
+          end
+        ensure
+          ActiveRecord::Migration.verbose = verbose_was
+        end
+        task "db:rollback" => "db:rollback:#{name}"
+
         desc "Drop tenanted #{name} databases for current environment"
         task "db:drop:#{name}" => "load_config" do
           verbose_was = ActiveRecord::Migration.verbose

test/unit/database_tasks_test.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require "test_helper"
+require "rake"
 
 describe ActiveRecord::Tenanted::DatabaseTasks do
   describe ".migrate_tenant" do
@@ -132,4 +133,108 @@
       end
     end
   end
+
+  describe ".rollback_tenant" do
+    for_each_scenario do
+      setup do
+        with_new_migration_file
+        ActiveRecord::Tenanted::DatabaseTasks.new(base_config).migrate_tenant("foo")
+      end
+
+      test "rolls back the most recent migration" do
+        ActiveRecord::Tenanted::DatabaseTasks.new(base_config).rollback_tenant("foo")
+
+        config = base_config.new_tenant_config("foo")
+        ActiveRecord::Tasks::DatabaseTasks.with_temporary_connection(config) do |conn|
+          assert_equal(20250203191115, conn.pool.migration_context.current_version)
+        end
+      end
+
+      test "accepts a custom step count" do
+        ActiveRecord::Tenanted::DatabaseTasks.new(base_config).rollback_tenant("foo", step: 2)
+
+        config = base_config.new_tenant_config("foo")
+        ActiveRecord::Tasks::DatabaseTasks.with_temporary_connection(config) do |conn|
+          assert_equal(0, conn.pool.migration_context.current_version)
+        end
+      end
+    end
+  end
+
+  describe ".rollback_all" do
+    for_each_scenario do
+      let(:tenants) { %w[foo bar baz] }
+
+      setup do
+        tenants.each do |tenant|
+          TenantedApplicationRecord.create_tenant(tenant)
+        end
+
+        with_new_migration_file
+        ActiveRecord::Tenanted::DatabaseTasks.new(base_config).migrate_all
+      end
+
+      test "rolls back all existing tenants" do
+        ActiveRecord::Tenanted::DatabaseTasks.new(base_config).rollback_all
+
+        tenants.each do |tenant|
+          config = base_config.new_tenant_config(tenant)
+          ActiveRecord::Tasks::DatabaseTasks.with_temporary_connection(config) do |conn|
+            assert_equal(20250203191115, conn.pool.migration_context.current_version)
+          end
+        end
+      end
+    end
+  end
+
+  describe ".wrap_rails_task" do
+    setup do
+      @original_rake_application = Rake.application
+      Rake.application = Rake::Application.new
+      ActiveRecord::Tenanted::DatabaseTasks.instance_variable_set(:@wrapped_rails_tasks, nil)
+    end
+
+    teardown do
+      ActiveRecord::Tenanted::DatabaseTasks.instance_variable_set(:@wrapped_rails_tasks, nil)
+      Rake.application = @original_rake_application
+    end
+
+    test "skips the original task when the guard returns false" do
+      Rake::Task.define_task("db:rollback") { raise "should not run" }
+
+      ActiveRecord::Tenanted::DatabaseTasks.send(:wrap_rails_task, "db:rollback") { false }
+
+      Rake::Task["db:rollback"].invoke
+      assert(true)
+    end
+
+    test "runs the original task when the guard returns true" do
+      Rake::Task.define_task("db:rollback") { raise "original task ran" }
+
+      ActiveRecord::Tenanted::DatabaseTasks.send(:wrap_rails_task, "db:rollback") { true }
+
+      error = assert_raises(RuntimeError) { Rake::Task["db:rollback"].invoke }
+      assert_equal("original task ran", error.message)
+    end
+
+    test "guard skips original db:rollback when no default configurations exist" do
+      Rake::Task.define_task("db:rollback") { raise "should not run" }
+
+      fake_configs = Class.new do
+        def configs_for(...)
+          []
+        end
+      end.new
+
+      ActiveRecord::Base.stub(:configurations, fake_configs) do
+        ActiveRecord::Tenanted::DatabaseTasks.send(:wrap_rails_task, "db:rollback") do
+          ActiveRecord::Tenanted::DatabaseTasks.send(:default_database_tasks_present?)
+        end
+
+        Rake::Task["db:rollback"].invoke
+      end
+
+      assert(true)
+    end
+  end
 end