From 3a1c6787d6dfaaf843e3e2206ecbbadf448091a7 Mon Sep 17 00:00:00 2001
From: Ikraam Ghoor <ikraamghoor@gmail.com>
Date: Thu, 16 Jan 2025 13:27:08 +0000
Subject: [PATCH] Pass the ar_relation for context to `should_require_tenant?`

This is a backwards compatible addition that allows finer control over when
`ActsAsTenant::Errors::NoTenantSet` is raised in the default scope check.
---
 lib/acts_as_tenant.rb                     | 11 +++++------
 lib/acts_as_tenant/model_extensions.rb    | 12 ++++++------
 spec/acts_as_tenant/configuration_spec.rb | 15 +++++++++++++++
 3 files changed, 26 insertions(+), 12 deletions(-)

diff --git a/lib/acts_as_tenant.rb b/lib/acts_as_tenant.rb
index f4074015..3a8dec55 100644
--- a/lib/acts_as_tenant.rb
+++ b/lib/acts_as_tenant.rb
@@ -150,12 +150,11 @@ def self.with_mutable_tenant(&block)
     ActsAsTenant.mutable_tenant!(false)
   end
 
-  def self.should_require_tenant?
-    if configuration.require_tenant.respond_to?(:call)
-      !!configuration.require_tenant.call
-    else
-      !!configuration.require_tenant
-    end
+  def self.should_require_tenant?(context = nil)
+    config = configuration.require_tenant
+    return !!config unless config.respond_to?(:call)
+
+    config.arity.zero? ? !!config.call : !!config.call(context)
   end
 end
 
diff --git a/lib/acts_as_tenant/model_extensions.rb b/lib/acts_as_tenant/model_extensions.rb
index 053b9b85..8dc3dc9b 100644
--- a/lib/acts_as_tenant/model_extensions.rb
+++ b/lib/acts_as_tenant/model_extensions.rb
@@ -17,7 +17,7 @@ def acts_as_tenant(tenant = :account, scope = nil, **options)
         belongs_to tenant, scope, **valid_options
 
         default_scope lambda {
-          if ActsAsTenant.should_require_tenant? && ActsAsTenant.current_tenant.nil? && !ActsAsTenant.unscoped?
+          if ActsAsTenant.should_require_tenant?(self) && ActsAsTenant.current_tenant.nil? && !ActsAsTenant.unscoped?
             raise ActsAsTenant::Errors::NoTenantSet
           end
 
@@ -46,10 +46,10 @@ def acts_as_tenant(tenant = :account, scope = nil, **options)
         before_validation proc { |m|
           if ActsAsTenant.current_tenant
             if options[:polymorphic]
-              m.send("#{fkey}=".to_sym, ActsAsTenant.current_tenant.class.to_s) if m.send(fkey.to_s).nil?
-              m.send("#{polymorphic_type}=".to_sym, ActsAsTenant.current_tenant.class.to_s) if m.send(polymorphic_type.to_s).nil?
+              m.send(:"#{fkey}=", ActsAsTenant.current_tenant.class.to_s) if m.send(fkey.to_s).nil?
+              m.send(:"#{polymorphic_type}=", ActsAsTenant.current_tenant.class.to_s) if m.send(polymorphic_type.to_s).nil?
             else
-              m.send "#{fkey}=".to_sym, ActsAsTenant.current_tenant.send(pkey)
+              m.send :"#{fkey}=", ActsAsTenant.current_tenant.send(pkey)
             end
           end
         }, on: :create
@@ -80,13 +80,13 @@ def acts_as_tenant(tenant = :account, scope = nil, **options)
         # - Add an override to prevent unnecessary db hits
         # - Add a helper method to verify if a model has been scoped by AaT
         to_include = Module.new {
-          define_method "#{fkey}=" do |integer|
+          define_method :"#{fkey}=" do |integer|
             write_attribute(fkey.to_s, integer)
             raise ActsAsTenant::Errors::TenantIsImmutable if !ActsAsTenant.mutable_tenant? && tenant_modified?
             integer
           end
 
-          define_method "#{ActsAsTenant.tenant_klass}=" do |model|
+          define_method :"#{ActsAsTenant.tenant_klass}=" do |model|
             super(model)
             raise ActsAsTenant::Errors::TenantIsImmutable if !ActsAsTenant.mutable_tenant? && tenant_modified?
             model
diff --git a/spec/acts_as_tenant/configuration_spec.rb b/spec/acts_as_tenant/configuration_spec.rb
index c12edd4e..d8b0500e 100644
--- a/spec/acts_as_tenant/configuration_spec.rb
+++ b/spec/acts_as_tenant/configuration_spec.rb
@@ -30,6 +30,21 @@
       expect(ActsAsTenant.should_require_tenant?).to eq(false)
     end
 
+    it "evaluates lambda with context" do
+      account = accounts(:foo)
+      test_context = nil
+
+      ActsAsTenant.configure do |config|
+        config.require_tenant = ->(context) {
+          test_context = context
+          context.klass.name == "Project"
+        }
+      end
+
+      expect { account.projects.create!(name: "foobar") }.to raise_error(ActsAsTenant::Errors::NoTenantSet)
+      expect(test_context.klass.name).to eq("Project")
+    end
+
     it "evaluates boolean" do
       ActsAsTenant.configure do |config|
         config.require_tenant = true