activerecord-multi-tenant 2.0.0 → 2.4.0

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)
@@ -18,7 +46,8 @@ module MultiTenant
         execute "SELECT citus_run_on_all_workers($$#{sql}$$)" # initial citus_tools.sql with different names
       when nil
         # Do nothing, this is regular Postgres
-      else # 6.1 and newer
+      else
+        # 6.1 and newer
         execute "SELECT run_command_on_workers($$#{sql}$$)"
       end
     end
@@ -28,24 +57,23 @@ 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 +82,52 @@ 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?
+          query_to_execute = <<-SQL.strip
+            SELECT logicalrelid::regclass AS table_name,
+                   column_to_column_name(logicalrelid, partkey) AS dist_col_name
+            FROM pg_dist_partition
+          SQL
+          @connection.execute(query_to_execute).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,18 +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
+      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
@@ -22,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)
@@ -54,43 +67,56 @@ module MultiTenant
         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, :optional).merge(foreign_key: options[:partition_key])
+        if MultiTenant.tenant_klass_defined?(tenant_name, options)
+          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)
+          if MultiTenant.tenant_klass_defined?(tenant_name, options)
             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
@@ -99,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
@@ -127,12 +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
+  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
 
+# 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

@@ -5,27 +5,38 @@ module MultiTenant
     attribute :tenant
   end
 
-  def self.tenant_klass_defined?(tenant_name)
-    !!tenant_name.to_s.classify.safe_constantize
+  def self.tenant_klass_defined?(tenant_name, options = {})
+    class_name = if options[:class_name].present?
+                   options[:class_name]
+                 else
+                   tenant_name.to_s.classify
+                 end
+    !!class_name.safe_constantize
   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,21 +49,23 @@ 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)
+      MultiTenant.multi_tenant_model_for_table(TableNode.table_name(arel.ast.relation.left))
     else
-      MultiTenant.multi_tenant_model_for_table(arel.ast.relation.table_name)
+      MultiTenant.multi_tenant_model_for_table(TableNode.table_name(arel.ast.relation))
     end
   end
 
@@ -74,7 +87,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,35 +96,81 @@ 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
+  # 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 #{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)
 
+  # 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)