#201 | Enhanced tenanted and untenanted associations

lib/active_record/tenanted/base.rb

@@ -34,6 +34,41 @@ def tenanted?
           false
         end
 
+        def belongs_to(name, scope = nil, **options)
+          tenant_key = options.delete(:tenant_key)
+          super(name, scope, **options)
+
+          if tenant_key
+            define_method("#{name}=") do |value|
+              super(value)
+              if value.respond_to?(:tenant)
+                self.send("#{tenant_key}=", value.tenant)
+              end
+            end
+
+            unless tenanted?
+              define_method(name) do
+                tenant_value = send(tenant_key)
+                return nil unless tenant_value
+                target_klass = self.class.reflect_on_association(name).klass
+                if target_klass.tenanted?
+                  tenant_klass = if target_klass.respond_to?(:with_tenant)
+                    target_klass
+                  else
+                    target_klass.tenanted_subtenant_of
+                  end
+
+                  tenant_klass.prohibit_shard_swapping(false) do
+                    tenant_klass.with_tenant(tenant_value) { super() }
+                  end
+                else
+                  super()
+                end
+              end
+            end
+          end
+        end
+
         def table_exists?
           super
         rescue ActiveRecord::Tenanted::NoTenantError

lib/active_record/tenanted/cross_tenant_associations.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+  module Tenanted
+    module CrossTenantAssociations
+      extend ActiveSupport::Concern
+
+      class_methods do
+        def has_one(name, scope = nil, **options)
+          define_enhanced_association(:has_one, name, scope, **options)
+        end
+
+        def has_many(name, scope = nil, **options)
+          define_enhanced_association(:has_many, name, scope, **options)
+        end
+
+        private
+          # For now association methods are identical
+          def define_enhanced_association(association_type, name, scope, **options)
+            tenant_key = options.delete(:tenant_key)
+            class_name = options[:class_name]
+            enhanced_scope = enhance_cross_tenant_association(name, scope, tenant_key: tenant_key || :tenant_id, class_name: class_name)
+            method(association_type).super_method.call(name, enhanced_scope, **options)
+          end
+
+          def enhance_cross_tenant_association(name, scope, tenant_key:, class_name: nil)
+            resolved_class_name = class_name || name.to_s.classify
+
+            ->(record) {
+              target_class = resolved_class_name.constantize
+              base_scope = scope ? target_class.instance_exec(&scope) : target_class.all
+
+              if target_class.tenanted?
+                base_scope
+              else
+                base_scope.where(tenant_key => record.tenant)
+              end
+            }
+          end
+      end
+    end
+  end
+end

lib/active_record/tenanted/subtenant.rb

@@ -6,6 +6,8 @@ module Subtenant
       extend ActiveSupport::Concern
 
       class_methods do
+        include CrossTenantAssociations::ClassMethods
+
         def tenanted?
           true
         end

lib/active_record/tenanted/tenant.rb

@@ -85,6 +85,8 @@ def to_s
       CONNECTION_POOL_CREATION_LOCK = Thread::Mutex.new # :nodoc:
 
       class_methods do
+        include CrossTenantAssociations::ClassMethods
+
         def tenanted?
           true
         end