activerecord-multi-tenant 1.2.0 → 2.2.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: fa2e595cc76e33a877613504121106791680118da0a0c4957ff1190e75e99c3a
-  data.tar.gz: 8a232e3431462d5b8b3972f0825664c210cc51729b2e94d08c5529851365d7f8
+  metadata.gz: 15ad19bf20781129dc1bd57d4e4b8eb8731044cea3b89f29a62789d0f5a45aee
+  data.tar.gz: 0f76290b00a7d540495972ab5a9de09d3abbb1dde11758d58b08493bf743b5da
 SHA512:
-  metadata.gz: a2df2315ca041de29d2e82786a33904a9ac1dedd3f1ee4a2e470822cff977e9f4548ab0d33bd6eefec8055b9fc74ff8db2fdcf4c93cba4992aaaa681dcd1a695
-  data.tar.gz: 9da5eec93a1bcd33bff29ab1e0035022592802f512d27139608e92988a17365d3fe4cff0a73fc0f83e9513cea9dd2bf6233bb4deac9533320fb0884f7ef8b219
+  metadata.gz: 17f7b912ecc314e5462affd74daf2e1b79239f8f65bfa30c8896f4acb9e1d787c240aeefe8cb09045b34ab2772d46aab4c42bb8cdb3374df247f1131a08758d0
+  data.tar.gz: 713df8fdefb6cd50b2c02408a9763da39e532b957d176fe5a8a6e1fbe265d2cf79457ae4d196d0e6ddfdf8effe0ca370c0297ea7186a5ce145717e4a93772be1

data/.github/workflows/CI.yml

@@ -13,40 +13,24 @@ jobs:
       fail-fast: false
       matrix:
         ruby:
-          - '2.5'
           - '2.6'
           - '2.7'
           - '3.0'
           - '3.1'
         gemfile:
-          - rails_5.2
           - rails_6.0
           - rails_6.1
           - rails_7.0
-          - active_record_5.2
           - active_record_6.0
           - active_record_6.1
           - active_record_7.0
         prepared_statements: [true, false]
         exclude:
           # activesupport-7.0.0 requires ruby version >= 2.7.0
-          - ruby: '2.5'
-            gemfile: 'rails_7.0'
-          - ruby: '2.5'
-            gemfile: 'active_record_7.0'
           - ruby: '2.6'
             gemfile: 'rails_7.0'
           - ruby: '2.6'
             gemfile: 'active_record_7.0'
-          # ruby >3 and activesupport 5.2 are not compatible
-          - ruby: '3.0'
-            gemfile: 'rails_5.2'
-          - ruby: '3.0'
-            gemfile: 'active_record_5.2'
-          - ruby: '3.1'
-            gemfile: 'rails_5.2'
-          - ruby: '3.1'
-            gemfile: 'active_record_5.2'
     name: Ruby ${{ matrix.ruby }} / ${{ matrix.gemfile }} ${{ (matrix.prepared_statements && 'w/ prepared statements') || '' }}
     env:
        BUNDLE_GEMFILE: gemfiles/${{ matrix.gemfile }}.gemfile

data/Appraisals

@@ -1,15 +1,3 @@
-appraise 'rails-5.2' do
-  gem 'rails', '~> 5.2.0'
-  gem 'i18n',  '~> 0.9.5'
-  gem 'nokogiri', '~> 1.7.1'
-  gem 'nio4r', '~> 2.3.1'
-  gem 'sprockets', '~> 3.7.1'
-  gem 'byebug', '~> 11.0'
-  gem 'rake', '12.0.0'
-  gem 'redis', '3.3.3'
-  gem 'pry-byebug', '3.9.0'
-end
-
 appraise 'rails-6.0' do
   gem 'rails', '~> 6.0.3'
 end
@@ -22,18 +10,6 @@ appraise 'rails-7.0' do
   gem 'rails', '~> 7.0.0'
 end
 
-appraise 'active-record-5.2' do
-  gem 'activerecord', '~> 5.2.0'
-  gem 'i18n',  '~> 0.9.5'
-  gem 'nokogiri', '~> 1.7.1'
-  gem 'nio4r', '~> 2.3.1'
-  gem 'sprockets', '~> 3.7.1'
-  gem 'byebug', '~> 11.0'
-  gem 'rake', '12.0.0'
-  gem 'redis', '3.3.3'
-  gem 'pry-byebug', '3.9.0'
-end
-
 appraise 'active-record-6.0' do
   gem 'activerecord', '~> 6.0.3'
 end

data/CHANGELOG.md

@@ -1,5 +1,43 @@
 # Changelog
 
+## 2.2.0      2022-12-06
+* Handle changing tenant from `nil` to a value [#173](https://github.com/citusdata/activerecord-multi-tenant/pull/173)
+* Allow Partitioned tables to be created without a primary key [#172](https://github.com/citusdata/activerecord-multi-tenant/pull/172)
+* Only attempt to reload with MultiTenant when parition_key is present [#175](https://github.com/citusdata/activerecord-multi-tenant/pull/175)
+* Remove support for Ruby 2.5 & ActiveRecord 5.2
+
+## 2.1.6      2022-11-23
+* Fix undefined wrap_methods error & wrap_methods version check [#170](https://github.com/citusdata/activerecord-multi-tenant/pull/170)
+
+## 2.1.5      2022-11-20
+* Fix `MultiTenant.without` codegen bug in Rails 6.1+ [#168](https://github.com/citusdata/activerecord-multi-tenant/pull/168)
+
+## 2.1.4      2022-11-03
+* Fixes #166 where db:schema:dump is broken when using this gem with MySQL [#167](https://github.com/citusdata/activerecord-multi-tenant/pull/167)
+
+## 2.1.3      2022-10-27
+* Error when calling a method that takes keyword arguments with MultiTenant.wrap_methods [#164](https://github.com/citusdata/activerecord-multi-tenant/pull/164)
+
+## 2.1.2      2022-10-26
+* Fixes issue when wraping methods that require a block [#162](https://github.com/citusdata/activerecord-multi-tenant/pull/162)
+
+## 2.1.1      2022-10-20
+* Fix query building for models with mismatched partition_keys [#150](https://github.com/citusdata/activerecord-multi-tenant/pull/150)
+* Identify tenant even if class name is nonstandard [#152](https://github.com/citusdata/activerecord-multi-tenant/pull/152)
+* Add current_tenant_id to WHERE clauses when calling methods on activerecord instance or its associations [#154](https://github.com/citusdata/activerecord-multi-tenant/pull/154)
+* Make create_distributed_table, create_reference_table reversible & add ruby wrapper for rebalance_table_shards [#155](https://github.com/citusdata/activerecord-multi-tenant/pull/155)
+* Support create_distributed_table, create_reference_table in schema.rb [#156](https://github.com/citusdata/activerecord-multi-tenant/pull/156)
+* Add client and server sidekiq middleware to sidekiq middleware chain [#158](https://github.com/citusdata/activerecord-multi-tenant/pull/158)
+
+## 2.0.0      2022-05-19
+
+* Replace RequestStore with CurrentAttributes [#139](https://github.com/citusdata/activerecord-multi-tenant/pull/139)
+* Support changing table_name after calling multi_tenant [#128](https://github.com/citusdata/activerecord-multi-tenant/pull/128)
+* Allow to use uuid as primary key on partition table [#112](https://github.com/citusdata/activerecord-multi-tenant/pull/112)
+* Support latest Rails 5.2 [#145](https://github.com/citusdata/activerecord-multi-tenant/pull/145)
+* Support optional: true for belongs_to [#147](https://github.com/citusdata/activerecord-multi-tenant/pull/147)
+
+
 ## 1.2.0      2022-03-29
 
 * Test Rails 7 & Ruby 3

data/README.md

@@ -16,7 +16,7 @@ gem 'activerecord-multi-tenant'
 
 ## Supported Rails versions
 
-All Ruby on Rails versions starting with 4.2 or newer (up to 6.0) are supported.
+All Ruby on Rails versions starting with 6.0 or newer (up to 7.0) are supported.
 
 This gem only supports ActiveRecord (the Rails default ORM), and not alternative ORMs like Sequel.
 

data/activerecord-multi-tenant.gemspec

@@ -15,8 +15,7 @@ Gem::Specification.new do |s|
   s.homepage      = 'https://github.com/citusdata/activerecord-multi-tenant'
   s.license       = 'MIT'
 
-  s.add_runtime_dependency('request_store', '>= 1.0.5')
-  s.add_dependency('rails','>= 4.2')
+  s.add_dependency 'rails', '>= 6'
 
   s.add_development_dependency 'rspec', '>= 3.0'
   s.add_development_dependency 'rspec-rails'

data/lib/activerecord-multi-tenant/migrations.rb

@@ -2,12 +2,40 @@ module MultiTenant
   module MigrationExtensions
     def create_distributed_table(table_name, partition_key)
       return unless citus_version.present?
-      execute "SELECT create_distributed_table($$#{table_name}$$, $$#{partition_key}$$)"
+
+      reversible do |dir|
+        dir.up do
+          execute "SELECT create_distributed_table($$#{table_name}$$, $$#{partition_key}$$)"
+        end
+        dir.down do
+          undistribute_table(table_name)
+        end
+      end
     end
 
     def create_reference_table(table_name)
       return unless citus_version.present?
-      execute "SELECT create_reference_table($$#{table_name}$$)"
+
+      reversible do |dir|
+        dir.up do
+          execute "SELECT create_reference_table($$#{table_name}$$)"
+        end
+        dir.down do
+          undistribute_table(table_name)
+        end
+      end
+    end
+
+    def undistribute_table(table_name)
+      return unless citus_version.present?
+
+      execute "SELECT undistribute_table($$#{table_name}$$))"
+    end
+
+    def rebalance_table_shards
+      return unless citus_version.present?
+
+      execute 'SELECT rebalance_table_shards()'
     end
 
     def execute_on_all_nodes(sql)
@@ -28,24 +56,22 @@ module MultiTenant
     end
 
     def citus_version
-      execute("SELECT extversion FROM pg_extension WHERE extname = 'citus'").getvalue(0,0).try(:split, '-').try(:first)
+      execute("SELECT extversion FROM pg_extension WHERE extname = 'citus'").getvalue(0, 0).try(:split, '-').try(:first)
     rescue ArgumentError => e
-      raise unless e.message == "invalid tuple number 0"
+      raise unless e.message == 'invalid tuple number 0'
     end
   end
 end
 
-if defined?(ActiveRecord::Migration)
-  ActiveRecord::Migration.send(:include, MultiTenant::MigrationExtensions)
-end
+ActiveRecord::Migration.include MultiTenant::MigrationExtensions if defined?(ActiveRecord::Migration)
 
 module ActiveRecord
   module ConnectionAdapters # :nodoc:
     module SchemaStatements
-      alias :orig_create_table :create_table
+      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[:partition_key] && options[:partition_key].to_s != 'id'
+        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
@@ -54,3 +80,44 @@ module ActiveRecord
     end
   end
 end
+
+module ActiveRecord
+  class SchemaDumper
+    private
+
+    alias initialize_without_citus initialize
+    def initialize(connection, options = {})
+      initialize_without_citus(connection, options)
+
+      citus_version = begin
+        ActiveRecord::Migration.citus_version
+      rescue StandardError
+        # Handle the case where this gem is used with MySQL https://github.com/citusdata/activerecord-multi-tenant/issues/166
+        nil
+      end
+      @distribution_columns =
+        if citus_version.present?
+          @connection.execute('SELECT logicalrelid::regclass AS table_name, column_to_column_name(logicalrelid, partkey) AS dist_col_name FROM pg_dist_partition').to_h do |v|
+            [v['table_name'], v['dist_col_name']]
+          end
+        else
+          {}
+        end
+    end
+
+    # Support for create_distributed_table & create_reference_table
+    alias table_without_citus table
+    def table(table, stream)
+      table_without_citus(table, stream)
+      table_name = remove_prefix_and_suffix(table)
+      distribution_column = @distribution_columns[table_name]
+      if distribution_column
+        stream.puts "  create_distributed_table(#{table_name.inspect}, #{distribution_column.inspect})"
+        stream.puts
+      elsif @distribution_columns.key?(table_name)
+        stream.puts "  create_reference_table(#{table_name.inspect})"
+        stream.puts
+      end
+    end
+  end
+end

data/lib/activerecord-multi-tenant/model_extensions.rb

@@ -1,12 +1,20 @@
+require_relative './multi_tenant'
+
 module MultiTenant
   module ModelExtensionsClassMethods
     DEFAULT_ID_FIELD = 'id'.freeze
 
     def multi_tenant(tenant_name, options = {})
-      if to_s.underscore.to_sym == tenant_name
+      if to_s.underscore.to_sym == tenant_name || (!table_name.nil? && table_name.singularize.to_sym == tenant_name)
         unless MultiTenant.with_write_only_mode_enabled?
           # This is the tenant model itself. Workaround for https://github.com/citusdata/citus/issues/687
-          before_create -> { self.id ||= self.class.connection.select_value("SELECT nextval('" + [self.class.table_name, self.class.primary_key, 'seq'].join('_') + "'::regclass)") }
+          before_create -> do
+           if self.class.columns_hash[self.class.primary_key].type == :uuid
+             self.id ||= SecureRandom.uuid
+           else
+             self.id ||= self.class.connection.select_value("SELECT nextval('#{self.class.table_name}_#{self.class.primary_key}_seq'::regclass)")
+           end
+          end
         end
       else
         class << self
@@ -38,18 +46,18 @@ module MultiTenant
 
           def inherited(subclass)
             super
-            MultiTenant.register_multi_tenant_model(subclass.table_name, subclass) if subclass.table_name
+            MultiTenant.register_multi_tenant_model(subclass)
           end
         end
 
-        MultiTenant.register_multi_tenant_model(table_name, self) if table_name
+        MultiTenant.register_multi_tenant_model(self)
 
         @partition_key = options[:partition_key] || MultiTenant.partition_key(tenant_name)
         partition_key = @partition_key
 
         # Create an implicit belongs_to association only if tenant class exists
         if MultiTenant.tenant_klass_defined?(tenant_name)
-          belongs_to tenant_name, **options.slice(:class_name, :inverse_of).merge(foreign_key: options[:partition_key])
+          belongs_to tenant_name, **options.slice(:class_name, :inverse_of, :optional).merge(foreign_key: options[:partition_key])
         end
 
         # New instances should have the tenant set
@@ -67,7 +75,6 @@ module MultiTenant
             # Rails 5 `attribute_will_change!` uses the attribute-method-call rather than `read_attribute`
             # and will raise ActiveModel::MissingAttributeError if that column was not selected.
             # This is rescued as NoMethodError and in MRI attribute_was is assigned an arbitrary Object
-            # This is still true after the Rails 5.2 refactor
             was = send("#{partition_key}_was")
             was_nil_or_skipped = was.nil? || was.class == Object
 
@@ -94,7 +101,8 @@ module MultiTenant
         include to_include
 
         around_save -> (record, block) {
-          if persisted? && MultiTenant.current_tenant_id.nil?
+          record_tenant = record.attribute_was(partition_key)
+          if persisted? && MultiTenant.current_tenant_id.nil? && !record_tenant.nil?
             MultiTenant.with(record.public_send(partition_key)) { block.call }
           else
             block.call
@@ -102,7 +110,8 @@ module MultiTenant
         }
 
         around_update -> (record, block) {
-          if MultiTenant.current_tenant_id.nil?
+          record_tenant = record.attribute_was(partition_key)
+          if MultiTenant.current_tenant_id.nil? && !record_tenant.nil?
             MultiTenant.with(record.public_send(partition_key)) { block.call }
           else
             block.call
@@ -123,6 +132,20 @@ end
 
 ActiveSupport.on_load(:active_record) do |base|
   base.extend MultiTenant::ModelExtensionsClassMethods
+
+  # Ensure we have current_tenant_id in where clause when a cached ActiveRecord instance is being reloaded, or update_columns without callbacks is called
+  MultiTenant.wrap_methods(ActiveRecord::Base, 'self', :delete, :reload, :update_columns)
+
+  # Any queuries fired for fetching a singular association have the correct current_tenant_id in WHERE clause
+  # reload is called anytime any record's association is accessed
+  MultiTenant.wrap_methods(ActiveRecord::Associations::Association, 'owner', :reload)
+
+  # For collection associations, we need to wrap multiple methods in returned proxy so that any queries have the correct current_tenant_id in WHERE clause
+  ActiveRecord::Associations::CollectionProxy.alias_method :equals_mt, :== # Hack to prevent syntax error due to invalid method name
+  ActiveRecord::Associations::CollectionProxy.alias_method :append_mt, :<< # Hack to prevent syntax error due to invalid method name
+  MultiTenant.wrap_methods(ActiveRecord::Associations::CollectionProxy, '@association.owner', :find, :last, :take, :build, :create, :create!, :replace, :delete_all, :destroy_all, :delete, :destroy, :calculate, :pluck, :size, :empty?, :include?, :equals_mt, :records, :append_mt, :find_nth_with_limit, :find_nth_from_last, :null_scope?, :find_from_target?, :exec_queries)
+  ActiveRecord::Associations::CollectionProxy.alias_method :==, :equals_mt
+  ActiveRecord::Associations::CollectionProxy.alias_method :<<, :append_mt
 end
 
 class ActiveRecord::Associations::Association

data/lib/activerecord-multi-tenant/multi_tenant.rb

@@ -1,6 +1,10 @@
-require 'request_store'
+require 'active_support/current_attributes'
 
 module MultiTenant
+  class Current < ::ActiveSupport::CurrentAttributes
+    attribute :tenant
+  end
+
   def self.tenant_klass_defined?(tenant_name)
     !!tenant_name.to_s.classify.safe_constantize
   end
@@ -24,13 +28,23 @@ module MultiTenant
   def self.with_lock_workaround_enabled?; @@enable_with_lock_workaround; end
 
   # Registry that maps table names to models (used by the query rewriter)
-  def self.register_multi_tenant_model(table_name, model_klass)
-    @@multi_tenant_models ||= {}
-    @@multi_tenant_models[table_name.to_s] = model_klass
+  def self.register_multi_tenant_model(model_klass)
+    @@multi_tenant_models ||= []
+    @@multi_tenant_models.push(model_klass)
+
+    remove_class_variable(:@@multi_tenant_model_table_names) if defined?(@@multi_tenant_model_table_names)
   end
+
   def self.multi_tenant_model_for_table(table_name)
-    @@multi_tenant_models ||= {}
-    @@multi_tenant_models[table_name.to_s]
+    @@multi_tenant_models ||= []
+
+    if !defined?(@@multi_tenant_model_table_names)
+      @@multi_tenant_model_table_names = @@multi_tenant_models.map { |model|
+        [model.table_name, model] if model.table_name
+      }.compact.to_h
+    end
+
+    @@multi_tenant_model_table_names[table_name.to_s]
   end
 
   def self.multi_tenant_model_for_arel(arel)
@@ -43,11 +57,11 @@ module MultiTenant
   end
 
   def self.current_tenant=(tenant)
-    RequestStore.store[:current_tenant] = tenant
+    Current.tenant = tenant
   end
 
   def self.current_tenant
-    RequestStore.store[:current_tenant]
+    Current.tenant
   end
 
   def self.current_tenant_id
@@ -95,6 +109,41 @@ module MultiTenant
     end
   end
 
+  # 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
+  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 #{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}(...) }
+          else
+            #{original_method_name}(...)
+          end
+        end
+        CODE
+      end
+    end
+  end
+
   # Preserve backward compatibility for people using .with_id
   singleton_class.send(:alias_method, :with_id, :with)
 

data/lib/activerecord-multi-tenant/query_rewriter.rb

@@ -295,12 +295,9 @@ module ActiveRecord
               end
 
               node_list.select{ |n| n.is_a? Arel::Nodes::Join }.each do |node_join|
-                if (!node_join.right ||
-                    (ActiveRecord::VERSION::MAJOR == 5 &&
-                     !node_join.right.expr.right.is_a?(Arel::Attributes::Attribute)))
+                if !node_join.right
                   next
                 end
-
                 relation_right, relation_left = relations_from_node_join(node_join)
 
                 next unless relation_right && relation_left
@@ -308,7 +305,7 @@ module ActiveRecord
                 model_right = MultiTenant.multi_tenant_model_for_table(relation_left.table_name)
                 model_left = MultiTenant.multi_tenant_model_for_table(relation_right.table_name)
                 if model_right && model_left
-                  join_enforcement_clause = MultiTenant::TenantJoinEnforcementClause.new(relation_left[model_left.partition_key], relation_right)
+                  join_enforcement_clause = MultiTenant::TenantJoinEnforcementClause.new(relation_right[model_right.partition_key], relation_left)
                   node_join.right.expr = node_join.right.expr.and(join_enforcement_clause)
                 end
               end
@@ -322,19 +319,20 @@ module ActiveRecord
 
     private
     def relations_from_node_join(node_join)
-      if ActiveRecord::VERSION::MAJOR == 5 || node_join.right.expr.is_a?(Arel::Nodes::Equality)
+      if node_join.right.expr.is_a?(Arel::Nodes::Equality)
         return node_join.right.expr.right.relation, node_join.right.expr.left.relation
       end
 
-      children = node_join.right.expr.children
+      children = [node_join.right.expr.children].flatten
 
-      tenant_applied = children.any?(MultiTenant::TenantEnforcementClause) || children.any?(MultiTenant::TenantJoinEnforcementClause)
+      tenant_applied = children.any?{|c| c.is_a?(MultiTenant::TenantEnforcementClause) || c.is_a?(MultiTenant::TenantJoinEnforcementClause)}
       if tenant_applied || children.empty?
         return nil, nil
       end
 
-      if children[0].right.respond_to?('relation') && children[0].left.respond_to?('relation')
-        return children[0].right.relation, children[0].left.relation
+      child = children.first.respond_to?(:children) ? children.first.children.first : children.first
+      if child.right.respond_to?(:relation) && child.left.respond_to?(:relation)
+        return child.right.relation, child.left.relation
       end
 
       return nil, nil

data/lib/activerecord-multi-tenant/sidekiq.rb

@@ -33,6 +33,22 @@ module Sidekiq::Middleware::MultiTenant
   end
 end
 
+Sidekiq.configure_server do |config|
+  config.server_middleware do |chain|
+    chain.add Sidekiq::Middleware::MultiTenant::Server
+  end
+  config.client_middleware do |chain|
+    chain.add Sidekiq::Middleware::MultiTenant::Client
+  end
+end
+
+Sidekiq.configure_client do |config|
+  config.client_middleware do |chain|
+    chain.add Sidekiq::Middleware::MultiTenant::Client
+  end
+end
+
+
 module Sidekiq
   class Client
     def push_bulk_with_tenants(items)

data/lib/activerecord-multi-tenant/version.rb

@@ -1,3 +1,3 @@
 module MultiTenant
-  VERSION = '1.2.0'
+  VERSION = '2.2.0'
 end

data/lib/activerecord-multi-tenant.rb

@@ -10,4 +10,3 @@ require_relative 'activerecord-multi-tenant/query_rewriter'
 require_relative 'activerecord-multi-tenant/query_monitor'
 require_relative 'activerecord-multi-tenant/version'
 require_relative 'activerecord-multi-tenant/with_lock'
-require_relative 'activerecord-multi-tenant/persistence_extension'

data/spec/activerecord-multi-tenant/associations_spec.rb

@@ -0,0 +1,21 @@
+require 'spec_helper'
+
+describe MultiTenant, 'Association methods' do
+  let(:account1) { Account.create! name: 'test1' }
+  let(:account2) { Account.create! name: 'test2' }
+  let(:project1) { Project.create! name: 'something1', account: account1 }
+  let(:project2) { Project.create! name: 'something2', account: account2, id: project1.id }
+  let(:task1) { Task.create! name: 'task1', project: project1, account: account1 }
+  let(:task2) { Task.create! name: 'task2', project: project2, account: account2, id: task1.id }
+
+  context 'include the tenant_id in queries and' do
+    it 'creates a task with correct account_id' do
+      expect(project2.tasks.create(name: 'task3').account_id).to eq(account2.id)
+    end
+    it 'return correct account_id' do
+      expect(task1.project.account_id).to_not eq(task2.project.account_id) # belongs_to
+      expect(project2.tasks.count).to eq(1)
+      expect(project2.tasks.first.account_id).to eq(account2.id) # has_many
+    end
+  end
+end

data/spec/activerecord-multi-tenant/model_extensions_spec.rb

@@ -70,6 +70,65 @@ describe MultiTenant do
     it { expect(@partition_key_not_model_task.non_model_id).to be 77 }
   end
 
+
+  describe 'Tenant model with a nonstandard class name' do
+    let(:account_klass) do
+      Class.new(ActiveRecord::Base) do
+        self.table_name = 'account'
+        def self.name
+          'UserAccount'
+        end
+
+        multi_tenant(:account)
+      end
+    end
+    it "does not register the tenant model" do
+      expect(MultiTenant).not_to receive(:register_multi_tenant_model)
+      account_klass
+    end
+  end
+
+  describe 'Changes table_name after multi_tenant called' do
+    before do
+      account_klass.has_many(:posts, anonymous_class: post_klass)
+      post_klass.belongs_to(:account, anonymous_class: account_klass)
+
+      @account1 = account_klass.create! name: 'foo'
+      @account2 = account_klass.create! name: 'bar'
+
+      @post1 = @account1.posts.create! name: 'foobar'
+      @post2 = @account2.posts.create! name: 'baz'
+
+      MultiTenant.current_tenant = @account1
+      @posts = post_klass.all
+    end
+
+    let(:account_klass) do
+      Class.new(Account) do
+        def self.name
+          'Account'
+        end
+      end
+    end
+
+    let(:post_klass) do
+      Class.new(ActiveRecord::Base) do
+        self.table_name = 'unknown'
+
+        multi_tenant(:account)
+
+        self.table_name = 'posts'
+
+        def self.name
+          'Post'
+        end
+      end
+    end
+
+    it { expect(@posts.length).to eq(1) }
+    it { expect(@posts).to eq([@post1]) }
+  end
+
   # Scoping models
   describe 'Project.all should be scoped to the current tenant if set' do
     before do
@@ -144,6 +203,22 @@ describe MultiTenant do
       end
     end
 
+    it 'handles belongs_to with optional: true' do
+      record = OptionalSubTask.create(sub_task_id: sub_task.id)
+      expect(record.reload.sub_task).to eq(sub_task)
+      expect(record.account_id).to eq(nil)
+    end
+
+    it 'handles changing tenant from nil to a value' do
+      record = OptionalSubTask.create(sub_task_id: sub_task.id)
+      expect(record.reload.sub_task).to eq(sub_task)
+      expect(record.account_id).to eq(nil)
+
+      record.account = account
+      record.save!
+      expect(record.reload.account_id).to eq(account.id)
+    end
+
     it 'handles has_many through' do
       MultiTenant.with(account) do
         expect(project.sub_tasks).to eq [sub_task]
@@ -234,6 +309,55 @@ describe MultiTenant do
     end
   end
 
+  # Joins
+  describe 'joins for models' do
+    context 'for models with where condition in associations' do
+      let(:account) { Account.create!(name: 'Account 1') }
+
+      it 'should add tenant condition to the queries when tenant is set' do
+        expected_join_sql = <<-SQL.strip
+          SELECT "comments".* FROM "comments" INNER JOIN "tasks" ON "tasks"."id" = "comments"."commentable_id" AND "comments"."commentable_type" = 'Task' AND "tasks"."account_id" = 1 WHERE "comments"."account_id" = 1
+        SQL
+
+        MultiTenant.with(account) do
+          expect(Comment.joins(:task).to_sql).to eq(expected_join_sql)
+        end
+      end
+
+      it 'should add tenant condition to the queries when tenant is not set' do
+        MultiTenant.without do
+          expected_join_sql = <<-SQL.strip
+            SELECT "comments".* FROM "comments" INNER JOIN "tasks" ON "tasks"."id" = "comments"."commentable_id" AND "comments"."commentable_type" = 'Task' AND "comments"."account_id" = "tasks"."account_id"
+          SQL
+          expect(Comment.joins(:task).to_sql).to eq(expected_join_sql)
+        end
+      end
+    end
+
+    context 'for models with default associations' do
+      let(:account) { Account.create!(name: 'Account 1') }
+
+      it 'should add tenant condition to the queries when tenant is set' do
+        expected_join_sql = <<-SQL.strip
+          SELECT "projects".* FROM "projects" INNER JOIN "tasks" ON "tasks"."project_id" = "projects"."id" AND "tasks"."account_id" = 1 WHERE "projects"."account_id" = 1
+        SQL
+
+        MultiTenant.with(account) do
+          expect(Project.joins(:tasks).to_sql).to eq(expected_join_sql)
+        end
+      end
+
+      it 'should add tenant condition to the queries when tenant is not set' do
+        MultiTenant.without do
+          expected_join_sql = <<-SQL.strip
+            SELECT "projects".* FROM "projects" INNER JOIN "tasks" ON "tasks"."project_id" = "projects"."id" AND "projects"."account_id" = "tasks"."account_id"
+          SQL
+          expect(Project.joins(:tasks).to_sql).to eq(expected_join_sql)
+        end
+      end
+    end
+  end
+
   # ::with
   describe "::with" do
     it "should set current_tenant to the specified tenant inside the block" do
@@ -349,10 +473,10 @@ describe MultiTenant do
 
   it "applies the team_id conditions in the where clause" do
     option1 = <<-sql.strip
-      SELECT "sub_tasks".* FROM "sub_tasks" INNER JOIN "tasks" ON "sub_tasks"."task_id" = "tasks"."id" AND "sub_tasks"."account_id" = "tasks"."account_id" WHERE "tasks"."project_id" = 1 AND "sub_tasks"."account_id" = 1 AND "tasks"."account_id" = 1
+      SELECT "sub_tasks".* FROM "sub_tasks" INNER JOIN "tasks" ON "sub_tasks"."task_id" = "tasks"."id" AND "tasks"."account_id" = "sub_tasks"."account_id" WHERE "tasks"."project_id" = 1 AND "sub_tasks"."account_id" = 1 AND "tasks"."account_id" = 1
     sql
     option2 = <<-sql.strip
-      SELECT "sub_tasks".* FROM "sub_tasks" INNER JOIN "tasks" ON "sub_tasks"."task_id" = "tasks"."id" AND "sub_tasks"."account_id" = "tasks"."account_id" WHERE "sub_tasks"."account_id" = 1 AND "tasks"."project_id" = 1 AND "tasks"."account_id" = 1
+      SELECT "sub_tasks".* FROM "sub_tasks" INNER JOIN "tasks" ON "sub_tasks"."task_id" = "tasks"."id" AND "tasks"."account_id" = "sub_tasks"."account_id" WHERE "sub_tasks"."account_id" = 1 AND "tasks"."project_id" = 1 AND "tasks"."account_id" = 1
     sql
 
     account1 = Account.create! name: 'Account 1'
@@ -367,7 +491,7 @@ describe MultiTenant do
 
     MultiTenant.without do
       expected_sql = <<-sql
-                        SELECT "sub_tasks".* FROM "sub_tasks" INNER JOIN "tasks" ON "sub_tasks"."task_id" = "tasks"."id" AND "sub_tasks"."account_id" = "tasks"."account_id" WHERE "tasks"."project_id" = 1
+                        SELECT "sub_tasks".* FROM "sub_tasks" INNER JOIN "tasks" ON "sub_tasks"."task_id" = "tasks"."id" AND "tasks"."account_id" = "sub_tasks"."account_id" WHERE "tasks"."project_id" = 1
                      sql
 
       project = Project.first
@@ -405,7 +529,7 @@ describe MultiTenant do
       expect(project.categories).to include(category1)
 
       expected_sql = <<-sql
-                        SELECT "projects".* FROM "projects" INNER JOIN "project_categories" ON "project_categories"."project_id" = "projects"."id" AND "project_categories"."account_id" = "projects"."account_id" INNER JOIN "categories" ON "categories"."id" = "project_categories"."category_id" WHERE "projects"."account_id" = 1
+                        SELECT "projects".* FROM "projects" INNER JOIN "project_categories" ON "project_categories"."project_id" = "projects"."id" AND "projects"."account_id" = "project_categories"."account_id" INNER JOIN "categories" ON "categories"."id" = "project_categories"."category_id" WHERE "projects"."account_id" = 1
                     sql
 
       expect(Project.where(account_id: 1).joins(:categories).to_sql).to eq(expected_sql.strip)
@@ -439,7 +563,7 @@ describe MultiTenant do
 
     MultiTenant.without do
       expected_sql = <<-sql
-                     SELECT "projects"."id" AS t0_r0, "projects"."account_id" AS t0_r1, "projects"."name" AS t0_r2, "categories"."id" AS t1_r0, "categories"."name" AS t1_r1 FROM "projects" LEFT OUTER JOIN "project_categories" ON "project_categories"."project_id" = "projects"."id" AND "project_categories"."account_id" = "projects"."account_id" LEFT OUTER JOIN "categories" ON "categories"."id" = "project_categories"."category_id" WHERE "projects"."account_id" = 1
+                     SELECT "projects"."id" AS t0_r0, "projects"."account_id" AS t0_r1, "projects"."name" AS t0_r2, "categories"."id" AS t1_r0, "categories"."name" AS t1_r1 FROM "projects" LEFT OUTER JOIN "project_categories" ON "project_categories"."project_id" = "projects"."id" AND "projects"."account_id" = "project_categories"."account_id" LEFT OUTER JOIN "categories" ON "categories"."id" = "project_categories"."category_id" WHERE "projects"."account_id" = 1
                      sql
 
       expect(Project.where(account_id: 1).eager_load(:categories).to_sql).to eq(expected_sql.strip)
@@ -471,7 +595,7 @@ describe MultiTenant do
 
     MultiTenant.without do
       expected_sql = <<-sql
-                     SELECT "tasks".* FROM "tasks" INNER JOIN "projects" ON "projects"."id" = "tasks"."project_id" AND "projects"."account_id" = "tasks"."account_id" LEFT JOIN project_categories pc ON project.category_id = pc.id WHERE "tasks"."account_id" = 1
+                     SELECT "tasks".* FROM "tasks" INNER JOIN "projects" ON "projects"."id" = "tasks"."project_id" AND "tasks"."account_id" = "projects"."account_id" LEFT JOIN project_categories pc ON project.category_id = pc.id WHERE "tasks"."account_id" = 1
                      sql
 
       expect(Task.where(account_id: 1).joins(:project).joins('LEFT JOIN project_categories pc ON project.category_id = pc.id').to_sql).to eq(expected_sql.strip)

data/spec/activerecord-multi-tenant/record_modifications_spec.rb

@@ -54,6 +54,25 @@ describe MultiTenant, 'Record modifications' do
     end
   end
 
+  it 'should not update other objects with same id when calling object.update_columns' do
+    # When two records with same id but different account_id are updated, it should only update the current one
+    expect(project.account).to eq(account)
+    expect(project2.account).to eq(account2)
+    expect(project.id).to eq(project2.id)
+
+    MultiTenant.without do
+      project2.update_columns(name: 'newthing2')
+      expect(project.reload.name).to eq('something')
+      expect(project2.reload.name).to eq('newthing2')
+    end
+  end
+
+  it 'should return the same object when calling object.reload' do
+    # When two records with same id but different account_id are updated, it should not return the other object
+    expect(project.reload.account_id).to eq(account.id)
+    expect(project2.reload.account_id).to eq(account2.id)
+  end
+
   it 'test delete for reference tables' do
     category1 = Category.create! name: 'Category 1'
     expect(Category.count).to eq(1)

data/spec/schema.rb

@@ -34,6 +34,13 @@ ARGV.grep(/\w+_spec\.rb/).empty? && ActiveRecord::Schema.define(version: 1) do
     t.column :type, :string
   end
 
+  create_table :optional_sub_tasks, force: true do |t|
+    t.references :account, :integer
+    t.column :sub_task_id, :integer
+    t.column :name, :string
+    t.column :type, :string
+  end
+
   create_table :countries, force: true do |t|
     t.column :name, :string
   end
@@ -106,6 +113,11 @@ ARGV.grep(/\w+_spec\.rb/).empty? && ActiveRecord::Schema.define(version: 1) do
     t.column :domain_id, :integer
   end
 
+  create_table :posts, force: true, partition_key: :account_id do |t|
+    t.column :account_id, :integer
+    t.column :name, :string
+  end
+
   create_distributed_table :accounts, :id
   create_distributed_table :projects, :account_id
   create_distributed_table :managers, :account_id
@@ -121,6 +133,7 @@ ARGV.grep(/\w+_spec\.rb/).empty? && ActiveRecord::Schema.define(version: 1) do
   create_distributed_table :allowed_places, :account_id
   create_distributed_table :domains, :account_id
   create_distributed_table :pages, :account_id
+  create_distributed_table :posts, :account_id
   create_reference_table :categories
 end
 
@@ -128,6 +141,7 @@ class Account < ActiveRecord::Base
   multi_tenant :account
   has_many :projects
   has_one :manager, inverse_of: :account
+  has_many :optional_sub_tasks
 end
 
 class Project < ActiveRecord::Base
@@ -159,6 +173,14 @@ class SubTask < ActiveRecord::Base
   multi_tenant :account
   belongs_to :task
   has_one :project, through: :task
+  has_many :optional_sub_tasks
+end
+
+with_belongs_to_required_by_default do
+  class OptionalSubTask < ActiveRecord::Base
+    multi_tenant :account, optional: true
+    belongs_to :sub_task
+  end
 end
 
 class StiSubTask < SubTask
@@ -198,6 +220,7 @@ class Comment < ActiveRecord::Base
 end
 
 class Organization < ActiveRecord::Base
+  multi_tenant :organization
   has_many :uuid_records
 end
 

data/spec/spec_helper.rb

@@ -46,4 +46,11 @@ def uses_prepared_statements?
   ActiveRecord::Base.connection.prepared_statements
 end
 
+def with_belongs_to_required_by_default(&block)
+  default_value = ActiveRecord::Base.belongs_to_required_by_default
+  ActiveRecord::Base.belongs_to_required_by_default = true
+  yield
+ensure
+  ActiveRecord::Base.belongs_to_required_by_default = default_value
+end
 require 'schema'

metadata

@@ -1,43 +1,29 @@
 --- !ruby/object:Gem::Specification
 name: activerecord-multi-tenant
 version: !ruby/object:Gem::Version
-  version: 1.2.0
+  version: 2.2.0
 platform: ruby
 authors:
 - Citus Data
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2022-04-11 00:00:00.000000000 Z
+date: 2022-12-06 00:00:00.000000000 Z
 dependencies:
-- !ruby/object:Gem::Dependency
-  name: request_store
-  requirement: !ruby/object:Gem::Requirement
-    requirements:
-    - - ">="
-      - !ruby/object:Gem::Version
-        version: 1.0.5
-  type: :runtime
-  prerelease: false
-  version_requirements: !ruby/object:Gem::Requirement
-    requirements:
-    - - ">="
-      - !ruby/object:Gem::Version
-        version: 1.0.5
 - !ruby/object:Gem::Dependency
   name: rails
   requirement: !ruby/object:Gem::Requirement
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '4.2'
+        version: '6'
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '4.2'
+        version: '6'
 - !ruby/object:Gem::Dependency
   name: rspec
   requirement: !ruby/object:Gem::Requirement
@@ -168,11 +154,9 @@ files:
 - activerecord-multi-tenant.gemspec
 - docker-compose.yml
 - gemfiles/.bundle/config
-- gemfiles/active_record_5.2.gemfile
 - gemfiles/active_record_6.0.gemfile
 - gemfiles/active_record_6.1.gemfile
 - gemfiles/active_record_7.0.gemfile
-- gemfiles/rails_5.2.gemfile
 - gemfiles/rails_6.0.gemfile
 - gemfiles/rails_6.1.gemfile
 - gemfiles/rails_7.0.gemfile
@@ -184,12 +168,12 @@ files:
 - lib/activerecord-multi-tenant/migrations.rb
 - lib/activerecord-multi-tenant/model_extensions.rb
 - lib/activerecord-multi-tenant/multi_tenant.rb
-- lib/activerecord-multi-tenant/persistence_extension.rb
 - lib/activerecord-multi-tenant/query_monitor.rb
 - lib/activerecord-multi-tenant/query_rewriter.rb
 - lib/activerecord-multi-tenant/sidekiq.rb
 - lib/activerecord-multi-tenant/version.rb
 - lib/activerecord-multi-tenant/with_lock.rb
+- spec/activerecord-multi-tenant/associations_spec.rb
 - spec/activerecord-multi-tenant/controller_extensions_spec.rb
 - spec/activerecord-multi-tenant/fast_truncate_spec.rb
 - spec/activerecord-multi-tenant/model_extensions_spec.rb

data/gemfiles/active_record_5.2.gemfile

@@ -1,16 +0,0 @@
-# This file was generated by Appraisal
-
-source "https://rubygems.org"
-
-gem "appraisal"
-gem "activerecord", "~> 5.2.0", "< 5.2.4" # FIXME
-gem "i18n", "~> 0.9.5"
-gem "nokogiri", "~> 1.7.1"
-gem "nio4r", "~> 2.3.1"
-gem "sprockets", "~> 3.7.1"
-gem "byebug", "~> 11.0"
-gem "rake", "12.0.0"
-gem "redis", "3.3.3"
-gem "pry-byebug", "3.9.0"
-
-gemspec path: "../"

data/gemfiles/rails_5.2.gemfile

@@ -1,16 +0,0 @@
-# This file was generated by Appraisal
-
-source "https://rubygems.org"
-
-gem "appraisal"
-gem "rails", "~> 5.2.0", "< 5.2.4" # FIXME
-gem "i18n", "~> 0.9.5"
-gem "nokogiri", "~> 1.7.1"
-gem "nio4r", "~> 2.3.1"
-gem "sprockets", "~> 3.7.1"
-gem "byebug", "~> 11.0"
-gem "rake", "12.0.0"
-gem "redis", "3.3.3"
-gem "pry-byebug", "3.9.0"
-
-gemspec path: "../"

data/lib/activerecord-multi-tenant/persistence_extension.rb

@@ -1,13 +0,0 @@
-module ActiveRecord
-  module Persistence
-    alias :delete_orig :delete
-
-    def delete
-      if MultiTenant.multi_tenant_model_for_table(self.class.table_name).present? && persisted? && MultiTenant.current_tenant_id.nil?
-        MultiTenant.with(self.public_send(self.class.partition_key)) { delete_orig }
-      else
-        delete_orig
-      end
-    end
-  end
-end