Introduce callbacks for tenant change methods

CHANGELOG.md

@@ -59,6 +59,8 @@ Read the [Rails Guide documentation on `config.active_record.query_log_tags`](ht
 
 ### Added
 
+- Add callbacks for `:with_tenant` which are invoked when `.with_tenant` is called.
+- Add callbacks for `:set_current_tenant` which are invoked when `.current_tenant=` is called.
 - `UntenantedConnectionPool#size` returns the database configuration's `max_connections` value, so that code (like Solid Queue) can inspect config params without a tenant context.
 
 

GUIDE.md

@@ -396,7 +396,11 @@ TODO:
 
 Documentation outline:
 
-  - invalid characters in a tenant name (which is database-dependent)
+- setting the tenant
+  - `.with_tenant` and `.current_tenant=`
+    - and the callbacks for each, `:with_tenant` and `:set_current_tenant`
+  - validation
+    - invalid characters in a tenant name (which is database-dependent)
     - and how the application may want to do additional validation (e.g. ICANN subdomain restrictions)
   - `#tenant` is a readonly attribute on all tenanted model instances
   - `.current_tenant` returns the execution context for the model connection class

lib/active_record/tenanted/tenant.rb

@@ -104,7 +104,9 @@ def current_tenant=(tenant_name)
             tenant_name = tenant_name.to_s
           end
 
-          connection_class_for_self.connecting_to(shard: tenant_name, role: ActiveRecord.writing_role)
+          run_callbacks :set_current_tenant do
+            connection_class_for_self.connecting_to(shard: tenant_name, role: ActiveRecord.writing_role)
+          end
         end
 
         def tenant_exist?(tenant_name)
@@ -115,11 +117,13 @@ def with_tenant(tenant_name, prohibit_shard_swapping: true, &block)
           tenant_name = tenant_name.to_s unless tenant_name == UNTENANTED_SENTINEL
 
           if tenant_name == current_tenant
-            yield
+            run_callbacks :with_tenant, &block
           else
             connection_class_for_self.connected_to(shard: tenant_name, role: ActiveRecord.writing_role) do
-              prohibit_shard_swapping(prohibit_shard_swapping) do
-                log_tenant_tag(tenant_name, &block)
+              run_callbacks :with_tenant do
+                prohibit_shard_swapping(prohibit_shard_swapping) do
+                  log_tenant_tag(tenant_name, &block)
+                end
               end
             end
           end
@@ -262,9 +266,13 @@ def log_tenant_tag(tenant_name, &block)
         self.default_shard = ActiveRecord::Tenanted::Tenant::UNTENANTED_SENTINEL
 
         prepend TenantCommon
+        extend ActiveSupport::Callbacks
 
         cattr_accessor :tenanted_config_name
         cattr_accessor(:tenanted_connection_pools) { LRU.new }
+
+        define_callbacks :with_tenant
+        define_callbacks :set_current_tenant
       end
 
       def tenanted?

test/unit/tenant_test.rb

@@ -115,6 +115,22 @@
         assert_equal("bar", TenantedApplicationRecord.current_tenant)
       end
 
+      test ".current_tenant= fires callbacks" do
+        before_callback_fired = false
+        after_callback_fired = false
+        TenantedApplicationRecord.set_callback :set_current_tenant, :before do
+          before_callback_fired = true
+        end
+        TenantedApplicationRecord.set_callback :set_current_tenant, :after do
+          after_callback_fired = true
+        end
+
+        TenantedApplicationRecord.current_tenant = "foo"
+
+        assert(before_callback_fired, "Before callback should be fired")
+        assert(after_callback_fired, "After callback should be fired")
+      end
+
       test "using a record after changing tenant raises WrongTenantError" do
         TenantedApplicationRecord.create_tenant("foo")
         TenantedApplicationRecord.create_tenant("bar")
@@ -281,6 +297,40 @@
           TenantedApplicationRecord.with_tenant("baz") { User.count }
         end
       end
+
+      test ".with_tenant fires callbacks" do
+        around_callback_fired = false
+        block_called = false
+        TenantedApplicationRecord.set_callback :with_tenant, :around do |_, block|
+          around_callback_fired = true
+          block.call
+        end
+
+        TenantedApplicationRecord.with_tenant("foo") do
+          block_called = true
+        end
+
+        assert(around_callback_fired, "Around callback should be fired")
+        assert(block_called, "Block should be called")
+      end
+
+      test ".with_tenant fires callbacks even when tenant doesn't change" do
+        around_callback_fired = 0
+        block_called = false
+        TenantedApplicationRecord.set_callback :with_tenant, :around do |_, block|
+          around_callback_fired += 1
+          block.call
+        end
+
+        TenantedApplicationRecord.with_tenant("foo") do
+          TenantedApplicationRecord.with_tenant("foo") do
+            block_called = true
+          end
+        end
+
+        assert_equal(2, around_callback_fired, "Around callback should be fired for each call")
+        assert(block_called, "Block should be called")
+      end
     end
 
     for_each_scenario(except: { primary_db: [ :subtenant_record ] }) do