Use prepend instead of alias

lib/activerecord-multi-tenant/migrations.rb

@@ -68,22 +68,19 @@ def citus_version
 
 ActiveRecord::Migration.include MultiTenant::MigrationExtensions if defined?(ActiveRecord::Migration)
 
-module ActiveRecord
-  module ConnectionAdapters # :nodoc:
-    module SchemaStatements
-      alias orig_create_table create_table
-
-      def create_table(table_name, options = {}, &block)
-        ret = orig_create_table(table_name, **options.except(:partition_key), &block)
-        if options[:id] != false && options[:partition_key] && options[:partition_key].to_s != 'id'
-          execute "ALTER TABLE #{table_name} DROP CONSTRAINT #{table_name}_pkey"
-          execute "ALTER TABLE #{table_name} ADD PRIMARY KEY(\"#{options[:partition_key]}\", id)"
-        end
-        ret
+module MultiTenant
+  module SchemaStatementsExtensions
+    def create_table(table_name, options = {}, &block)
+      ret = super(table_name, **options.except(:partition_key), &block)
+      if options[:id] != false && options[:partition_key] && options[:partition_key].to_s != 'id'
+        execute "ALTER TABLE #{table_name} DROP CONSTRAINT #{table_name}_pkey"
+        execute "ALTER TABLE #{table_name} ADD PRIMARY KEY(\"#{options[:partition_key]}\", id)"
       end
+      ret
     end
   end
 end
+ActiveRecord::ConnectionAdapters::SchemaStatements.prepend(MultiTenant::SchemaStatementsExtensions)
 
 module ActiveRecord
   class SchemaDumper

lib/activerecord-multi-tenant/model_extensions.rb

@@ -190,21 +190,19 @@ def inherited(subclass)
 end
 
 # skips statement caching for classes that is Multi-tenant or has a multi-tenant relation
-module ActiveRecord
-  module Associations
-    class Association
-      alias skip_statement_cache_orig skip_statement_cache?
-
-      def skip_statement_cache?(*scope)
-        return true if klass.respond_to?(:scoped_by_tenant?) && klass.scoped_by_tenant?
-
-        if reflection.through_reflection
-          through_klass = reflection.through_reflection.klass
-          return true if through_klass.respond_to?(:scoped_by_tenant?) && through_klass.scoped_by_tenant?
-        end
+module MultiTenant
+  module AssociationExtensions
+    def skip_statement_cache?(*scope)
+      return true if klass.respond_to?(:scoped_by_tenant?) && klass.scoped_by_tenant?
 
-        skip_statement_cache_orig(*scope)
+      if reflection.through_reflection
+        through_klass = reflection.through_reflection.klass
+        return true if through_klass.respond_to?(:scoped_by_tenant?) && through_klass.scoped_by_tenant?
       end
+
+      super(*scope)
     end
   end
 end
+
+ActiveRecord::Associations::Association.prepend(MultiTenant::AssociationExtensions)

lib/activerecord-multi-tenant/multi_tenant.rb

@@ -130,39 +130,21 @@ def self.without(&block)
   # Wrap calls to any of `method_names` on an instance Class `klass` with MultiTenant.with
   # when `'owner'` (evaluated in context of the klass instance) is a ActiveRecord model instance that is multi-tenant
   # Instruments the methods provided with previously set Multitenant parameters
-  # In Ruby 2 using splat (*) operator with `&block` is not supported, so we need to use `method(...)` syntax
   # TODO: Could not understand the use of owner here. Need to check
-  if Gem::Version.create(RUBY_VERSION) < Gem::Version.new('3.0.0')
-    def self.wrap_methods(klass, owner, *method_names)
-      method_names.each do |method_name|
-        original_method_name = :"_mt_original_#{method_name}"
-        klass.class_eval <<-CODE, __FILE__, __LINE__ + 1
-          alias_method :#{original_method_name}, :#{method_name}
-          def #{method_name}(*args, &block)
-            if MultiTenant.multi_tenant_model_for_table(#{owner}.class.table_name).present? && #{owner}.persisted? && MultiTenant.current_tenant_id.nil? && #{owner}.class.respond_to?(:partition_key) && #{owner}.attributes.include?(#{owner}.class.partition_key)
-              MultiTenant.with(#{owner}.public_send(#{owner}.class.partition_key)) { #{original_method_name}(*args, &block) }
-            else
-              #{original_method_name}(*args, &block)
-            end
-          end
-        CODE
-      end
-    end
-  else
-    def self.wrap_methods(klass, owner, *method_names)
-      method_names.each do |method_name|
-        original_method_name = :"_mt_original_#{method_name}"
-        klass.class_eval <<-CODE, __FILE__, __LINE__ + 1
-        alias_method :#{original_method_name}, :#{method_name}
+  def self.wrap_methods(klass, owner, *method_names)
+    mod = Module.new
+    klass.prepend(mod)
+
+    method_names.each do |method_name|
+      mod.module_eval <<-CODE, __FILE__, __LINE__ + 1
         def #{method_name}(...)
           if MultiTenant.multi_tenant_model_for_table(#{owner}.class.table_name).present? && #{owner}.persisted? && MultiTenant.current_tenant_id.nil? && #{owner}.class.respond_to?(:partition_key) && #{owner}.attributes.include?(#{owner}.class.partition_key)
-            MultiTenant.with(#{owner}.public_send(#{owner}.class.partition_key)) { #{original_method_name}(...) }
+            MultiTenant.with(#{owner}.public_send(#{owner}.class.partition_key)) { super }
           else
-            #{original_method_name}(...)
+            super
           end
         end
-        CODE
-      end
+      CODE
     end
   end
 

lib/activerecord-multi-tenant/query_rewriter.rb

@@ -284,13 +284,10 @@ def delete(arel, name = nil, binds = [])
 
 Arel::Visitors::ToSql.include(MultiTenant::TenantValueVisitor)
 
-require 'active_record/relation'
-module ActiveRecord
-  module QueryMethods
-    alias build_arel_orig build_arel
-
-    def build_arel(*args)
-      arel = build_arel_orig(*args)
+module MultiTenant
+  module QueryMethodsExtensions
+    def build_arel(*)
+      arel = super
 
       unless MultiTenant.with_write_only_mode_enabled?
         visitor = MultiTenant::ArelTenantVisitor.new(arel)
@@ -374,6 +371,9 @@ def relations_from_node_join(node_join)
   end
 end
 
+require 'active_record/relation'
+ActiveRecord::QueryMethods.prepend(MultiTenant::QueryMethodsExtensions)
+
 module MultiTenantFindBy
   def cached_find_by_statement(key, &block)
     return super unless respond_to?(:scoped_by_tenant?) && scoped_by_tenant?

spec/activerecord-multi-tenant/multi_tenant_spec.rb

@@ -104,7 +104,7 @@ class SampleNestedTenant < ActiveRecord::Base
             end
             # rubocop:disable Layout/TrailingWhitespace
             # Trailing whitespace is intentionally left here
-            
+
             class AnotherTenant < ActiveRecord::Base
             end
             # rubocop:enable Layout/TrailingWhitespace
@@ -118,4 +118,32 @@ class AnotherTenant < ActiveRecord::Base
       end
     end
   end
+
+  describe '.wrap_methods' do
+    context 'when method is already prepended' do
+      it 'is not an stack error' do
+        klass = Class.new do
+          def hello
+            'hello'
+          end
+        end
+
+        klass.prepend(Module.new do
+          def hello
+            "#{super} world"
+          end
+
+          def owner
+            Class.new(ActiveRecord::Base) do
+              self.table_name = 'accounts'
+            end.new
+          end
+        end)
+
+        MultiTenant.wrap_methods(klass, :owner, :hello)
+
+        expect(klass.new.hello).to eq('hello world')
+      end
+    end
+  end
 end