activerecord-multi-tenant 2.1.6 → 2.3.0

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

@@ -1,20 +1,32 @@
 require_relative './multi_tenant'
 
 module MultiTenant
+  # Extension to the model to allow scoping of models to the current tenant. This is done by adding
+  # the multitenant method to the models that need to be scoped. This method is called in the
+  # model declaration.
+  # Adds scoped_by_tenant? partition_key, primary_key and inherited methods to the model
   module ModelExtensionsClassMethods
     DEFAULT_ID_FIELD = 'id'.freeze
-
+    # executes when multi_tenant method is called in the model. This method adds the following
+    # methods to the model that calls it.
+    # scoped_by_tenant? - returns true if the model is scoped by tenant
+    # partition_key - returns the partition key for the model
+    # primary_key - returns the primary key for the model
+    #
     def multi_tenant(tenant_name, options = {})
       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 -> 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
+          before_create lambda {
+            id = if self.class.columns_hash[self.class.primary_key].type == :uuid
+                   SecureRandom.uuid
+                 else
+                   self.class.connection.select_value(
+                     "SELECT nextval('#{self.class.table_name}_#{self.class.primary_key}_seq'::regclass)"
+                   )
+                 end
+            self.id ||= id
+          }
         end
       else
         class << self
@@ -24,24 +36,23 @@ module MultiTenant
 
           # Allow partition_key to be set from a superclass if not already set in this class
           def partition_key
-            @partition_key ||= ancestors.detect{ |k| k.instance_variable_get(:@partition_key) }
-                                 .try(:instance_variable_get, :@partition_key)
+            @partition_key ||= ancestors.detect { |k| k.instance_variable_get(:@partition_key) }
+                                        .try(:instance_variable_get, :@partition_key)
           end
 
           # Avoid primary_key errors when using composite primary keys (e.g. id, tenant_id)
           def primary_key
-            return @primary_key if @primary_key
+            if defined?(PRIMARY_KEY_NOT_SET) ? !PRIMARY_KEY_NOT_SET.equal?(@primary_key) : @primary_key
+              return @primary_key
+            end
 
             primary_object_keys = Array.wrap(connection.schema_cache.primary_keys(table_name)) - [partition_key]
 
-            if primary_object_keys.size == 1
-              @primary_key = primary_object_keys.first
-            elsif connection.schema_cache.columns_hash(table_name).include? DEFAULT_ID_FIELD
-              @primary_key = DEFAULT_ID_FIELD
-            else
-              # table without a primary key and DEFAULT_ID_FIELD is not present in the table
-              @primary_key = nil
-            end
+            @primary_key = if primary_object_keys.size == 1
+                             primary_object_keys.first
+                           elsif connection.schema_cache.columns_hash(table_name).include? DEFAULT_ID_FIELD
+                             DEFAULT_ID_FIELD
+                           end
           end
 
           def inherited(subclass)
@@ -57,42 +68,55 @@ module MultiTenant
 
         # 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, :optional).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
-        after_initialize Proc.new { |record|
+        after_initialize proc { |record|
           if MultiTenant.current_tenant_id &&
-              (!record.attribute_present?(partition_key) || record.public_send(partition_key.to_sym).nil?)
+             (!record.attribute_present?(partition_key) || record.public_send(partition_key.to_sym).nil?)
             record.public_send("#{partition_key}=".to_sym, MultiTenant.current_tenant_id)
           end
         }
 
+        # Below block adds the following methods to the model that calls it.
+        # partition_key= - returns the partition key for the model.class << self 'partition' method defined above
+        # is the getter method. Here, there is additional check to assure that the tenant id is not changed once set
+        # tenant_name- returns the name of the tenant model. Its setter and getter methods defined separately
+        # Getter checks for the tenant association and if it is not loaded, returns the current tenant id set
+        # in the MultiTenant module
         to_include = Module.new do
           define_method "#{partition_key}=" do |tenant_id|
-            write_attribute("#{partition_key}", tenant_id)
+            write_attribute(partition_key.to_s, tenant_id)
 
             # 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
+            was_nil_or_skipped = was.nil? || was.instance_of?(Object)
+
+            if send("#{partition_key}_changed?") && persisted? && !was_nil_or_skipped
+              raise MultiTenant::TenantIsImmutable
+            end
 
-            raise MultiTenant::TenantIsImmutable if send("#{partition_key}_changed?") && persisted? && !was_nil_or_skipped
             tenant_id
           end
 
           if MultiTenant.tenant_klass_defined?(tenant_name)
             define_method "#{tenant_name}=" do |model|
               super(model)
-              raise MultiTenant::TenantIsImmutable if send("#{partition_key}_changed?") && persisted? && !send("#{partition_key}_was").nil?
+              if send("#{partition_key}_changed?") && persisted? && !send("#{partition_key}_was").nil?
+                raise MultiTenant::TenantIsImmutable
+              end
+
               model
             end
 
-            define_method "#{tenant_name}" do
-              if !association(tenant_name.to_sym).loaded? && !MultiTenant.current_tenant_is_id? && MultiTenant.current_tenant_id && public_send(partition_key) == MultiTenant.current_tenant_id
-                return MultiTenant.current_tenant
+            define_method tenant_name.to_s do
+              if !association(tenant_name.to_sym).loaded? && !MultiTenant.current_tenant_is_id? &&
+                 MultiTenant.current_tenant_id && public_send(partition_key) == MultiTenant.current_tenant_id
+                MultiTenant.current_tenant
               else
                 super()
               end
@@ -101,23 +125,28 @@ module MultiTenant
         end
         include to_include
 
-        around_save -> (record, block) {
-          if persisted? && MultiTenant.current_tenant_id.nil?
+        # Below blocks sets tenant_id for the current session with the tenant_id of the record
+        # If the tenant is not set for the `session.After` the save operation current session tenant is set to nil
+        # If tenant is set for the session, save operation is performed as it is
+        around_save lambda { |record, block|
+          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
           end
         }
 
-        around_update -> (record, block) {
-          if MultiTenant.current_tenant_id.nil?
+        around_update lambda { |record, block|
+          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
           end
         }
 
-        around_destroy -> (record, block) {
+        around_destroy lambda { |record, block|
           if MultiTenant.current_tenant_id.nil?
             MultiTenant.with(record.public_send(partition_key)) { block.call }
           else
@@ -129,26 +158,39 @@ module MultiTenant
   end
 end
 
+# Below code block is executed on Model, Associations and CollectionProxy objects
+# when ActiveRecord is loaded and decorates defined methods with MultiTenant.with function.
+# Additionally, adds aliases for some operators.
 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
+  # 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)
+  # 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
 
+# skips statement caching for classes that is Multi-tenant or has a multi-tenant relation
 class ActiveRecord::Associations::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?
 

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

@@ -10,22 +10,28 @@ module MultiTenant
   end
 
   def self.partition_key(tenant_name)
-    "#{tenant_name.to_s}_id"
+    "#{tenant_name}_id"
   end
 
+  # rubocop:disable Style/ClassVars
   # In some cases we only have an ID - if defined we'll return the default tenant class in such cases
-  def self.default_tenant_class=(tenant_class); @@default_tenant_class = tenant_class; end
-  def self.default_tenant_class; @@default_tenant_class ||= nil; end
+  def self.default_tenant_class=(tenant_class)
+    @@default_tenant_class = tenant_class
+  end
+
+  def self.default_tenant_class
+    @@default_tenant_class ||= nil
+  end
 
   # Write-only Mode - this only adds the tenant_id to new records, but doesn't
   # require its presence for SELECTs/UPDATEs/DELETEs
-  def self.enable_write_only_mode; @@enable_write_only_mode = true; end
-  def self.with_write_only_mode_enabled?; @@enable_write_only_mode ||= false; end
+  def self.enable_write_only_mode
+    @@enable_write_only_mode = true
+  end
 
-  # Workaroud to make "with_lock" work until https://github.com/citusdata/citus/issues/1236 is fixed
-  @@enable_with_lock_workaround = false
-  def self.enable_with_lock_workaround; @@enable_with_lock_workaround = true; end
-  def self.with_lock_workaround_enabled?; @@enable_with_lock_workaround; end
+  def self.with_write_only_mode_enabled?
+    @@enable_write_only_mode ||= false
+  end
 
   # Registry that maps table names to models (used by the query rewriter)
   def self.register_multi_tenant_model(model_klass)
@@ -38,17 +44,19 @@ module MultiTenant
   def self.multi_tenant_model_for_table(table_name)
     @@multi_tenant_models ||= []
 
-    if !defined?(@@multi_tenant_model_table_names)
-      @@multi_tenant_model_table_names = @@multi_tenant_models.map { |model|
+    unless defined?(@@multi_tenant_model_table_names)
+      @@multi_tenant_model_table_names = @@multi_tenant_models.map do |model|
         [model.table_name, model] if model.table_name
-      }.compact.to_h
+      end.compact.to_h
     end
 
     @@multi_tenant_model_table_names[table_name.to_s]
+    # rubocop:enable Style/ClassVars
   end
 
   def self.multi_tenant_model_for_arel(arel)
     return nil unless arel.respond_to?(:ast)
+
     if arel.ast.relation.is_a? Arel::Nodes::JoinSource
       MultiTenant.multi_tenant_model_for_table(arel.ast.relation.left.table_name)
     else
@@ -74,7 +82,7 @@ module MultiTenant
 
   def self.current_tenant_class
     if current_tenant_is_id?
-      MultiTenant.default_tenant_class || fail('Only have tenant id, and no default tenant class set')
+      MultiTenant.default_tenant_class || raise('Only have tenant id, and no default tenant class set')
     elsif current_tenant
       MultiTenant.current_tenant.class.name
     end
@@ -83,33 +91,40 @@ module MultiTenant
   def self.load_current_tenant!
     return MultiTenant.current_tenant if MultiTenant.current_tenant && !current_tenant_is_id?
     raise 'MultiTenant.current_tenant must be set to load' if MultiTenant.current_tenant.nil?
-    klass = MultiTenant.default_tenant_class || fail('Only have tenant id, and no default tenant class set')
+
+    klass = MultiTenant.default_tenant_class || raise('Only have tenant id, and no default tenant class set')
     self.current_tenant = klass.find(MultiTenant.current_tenant_id)
   end
 
   def self.with(tenant, &block)
-    return block.call if self.current_tenant == tenant
-    old_tenant = self.current_tenant
+    return block.call if current_tenant == tenant
+
+    old_tenant = current_tenant
     begin
       self.current_tenant = tenant
-      return block.call
+      block.call
     ensure
       self.current_tenant = old_tenant
     end
   end
 
   def self.without(&block)
-    return block.call if self.current_tenant.nil?
-    old_tenant = self.current_tenant
+    return block.call if current_tenant.nil?
+
+    old_tenant = current_tenant
     begin
       self.current_tenant = nil
-      return block.call
+      block.call
     ensure
       self.current_tenant = old_tenant
     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
+  # 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|
@@ -117,7 +132,7 @@ module MultiTenant
         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?
+            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)
@@ -133,7 +148,7 @@ module MultiTenant
         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?
+          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}(...)
@@ -147,6 +162,10 @@ module MultiTenant
   # Preserve backward compatibility for people using .with_id
   singleton_class.send(:alias_method, :with_id, :with)
 
+  # This exception is raised when a there is an attempt to change tenant
   class TenantIsImmutable < StandardError
   end
+
+  class MissingTenantError < StandardError
+  end
 end

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

@@ -1,18 +1,34 @@
 # Add generic warning when queries fail and there is no tenant set
+# To handle this case, a QueryMonitor hook is created and registered
+# to sql.active_record. This hook will log a warning when a query fails
+# This hook is executed after the query is executed.
 module MultiTenant
+  # rubocop:disable Style/ClassVars
   # Option to enable query monitor
   @@enable_query_monitor = false
-  def self.enable_query_monitor; @@enable_query_monitor = true; end
-  def self.query_monitor_enabled?; @@enable_query_monitor; end
 
+  def self.enable_query_monitor
+    @@enable_query_monitor = true
+  end
+
+  def self.query_monitor_enabled?
+    @@enable_query_monitor
+  end
+
+  # rubocop:enable Style/ClassVars
+  # QueryMonitor class to log a warning when a query fails and there is no tenant set
+  # start and finish methods are required to be register sql.active_record hook
   class QueryMonitor
-    def start(name, id, payload); end
-    def finish(name, id, payload)
+    def start(_name, _id, _payload) end
+
+    def finish(_name, _id, payload)
       return unless MultiTenant.query_monitor_enabled?
+
       return unless payload[:exception].present? && MultiTenant.current_tenant_id.nil?
+
       Rails.logger.info 'WARNING: Tenant not present - make sure to add MultiTenant.with(tenant) { ... }'
     end
   end
 end
-
+# Actual code to register the hook.
 ActiveSupport::Notifications.subscribe('sql.active_record', MultiTenant::QueryMonitor.new)