activerecord-multi-tenant-next 2.5.0

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

@@ -0,0 +1,431 @@
+# frozen_string_literal: true
+
+require 'active_record'
+require_relative 'arel_visitors_depth_first' unless Arel::Visitors.const_defined?(:DepthFirst)
+
+# Iterates AST and adds tenant enforcement clauses to all relations
+module MultiTenant
+  class Table
+    attr_reader :arel_table
+
+    def initialize(arel_table)
+      @arel_table = arel_table
+    end
+
+    def eql?(other)
+      self.class == other.class &&
+        equality_fields.eql?(other.equality_fields)
+    end
+
+    def hash
+      equality_fields.hash
+    end
+
+    protected
+
+    def equality_fields
+      [arel_table.name, arel_table.table_alias]
+    end
+  end
+
+  class Context
+    attr_reader :arel_node, :known_relations, :handled_relations
+
+    def initialize(arel_node)
+      @arel_node = arel_node
+      @known_relations = []
+      @handled_relations = []
+      @discovering = false
+    end
+
+    def discover_relations
+      old_discovering = @discovering
+      @discovering = true
+      yield
+      @discovering = old_discovering
+    end
+
+    def visited_relation(relation)
+      return unless @discovering
+
+      @known_relations << Table.new(relation)
+    end
+
+    def visited_handled_relation(relation)
+      @handled_relations << Table.new(relation)
+    end
+
+    def unhandled_relations
+      known_relations.uniq - handled_relations
+    end
+  end
+
+  class ArelTenantVisitor < if Arel::Visitors.const_defined?(:DepthFirst)
+                              Arel::Visitors::DepthFirst
+                            else
+                              ::MultiTenant::ArelVisitorsDepthFirst
+                            end
+    def initialize(arel)
+      super(proc {})
+      @statement_node_id = nil
+
+      @contexts = []
+      @current_context = nil
+      accept(arel.ast)
+    end
+
+    attr_reader :contexts
+
+    # rubocop:disable Naming/MethodName
+    def visit_Arel_Attributes_Attribute(*args)
+      return if @current_context.nil?
+
+      super
+    end
+
+    def visit_Arel_Nodes_Equality(obj, *args)
+      if obj.left.is_a?(Arel::Attributes::Attribute)
+        table_name = MultiTenant::TableNode.table_name(obj.left.relation)
+        model = MultiTenant.multi_tenant_model_for_table(table_name)
+        if model.present? && obj.left.name.to_s == model.partition_key.to_s
+          @current_context.visited_handled_relation(obj.left.relation)
+        end
+      end
+      super
+    end
+
+    def visit_MultiTenant_TenantEnforcementClause(obj, *)
+      @current_context.visited_handled_relation(obj.tenant_attribute.relation)
+    end
+
+    def visit_MultiTenant_TenantJoinEnforcementClause(obj, *)
+      @current_context.visited_handled_relation(obj.tenant_attribute.relation)
+    end
+
+    def visit_Arel_Table(obj, _collector = nil)
+      @current_context.visited_relation(obj) if tenant_relation?(MultiTenant::TableNode.table_name(obj))
+    end
+
+    alias visit_Arel_Nodes_TableAlias visit_Arel_Table
+
+    # A SelectManager can appear inside a condition, e.g. as the right-hand
+    # side of an IN. Neither Arel's DepthFirst (Rails 6.0) nor the gem's copy
+    # of it (Rails 6.1+) knows how to visit one, so descend into its AST here.
+    def visit_Arel_SelectManager(obj, *)
+      visit obj.ast
+    end
+
+    def visit_Arel_Nodes_SelectCore(obj, *_args)
+      nest_context(obj) do
+        @current_context.discover_relations do
+          visit obj.source
+        end
+        visit obj.wheres
+        visit obj.groups
+        visit obj.windows
+        if defined?(obj.having)
+          visit obj.having
+        else
+          visit obj.havings
+        end
+      end
+    end
+
+    # rubocop:enable Naming/MethodName
+
+    # rubocop:disable Naming/MethodName
+    def visit_Arel_Nodes_OuterJoin(obj, _collector = nil)
+      nest_context(obj) do
+        @current_context.discover_relations do
+          visit obj.left
+          visit obj.right
+        end
+      end
+    end
+
+    # rubocop:enable Naming/MethodName
+
+    alias visit_Arel_Nodes_FullOuterJoin visit_Arel_Nodes_OuterJoin
+    alias visit_Arel_Nodes_RightOuterJoin visit_Arel_Nodes_OuterJoin
+
+    alias visit_ActiveModel_Attribute terminal
+
+    private
+
+    def tenant_relation?(table_name)
+      MultiTenant.multi_tenant_model_for_table(table_name).present?
+    end
+
+    DISPATCH = Hash.new do |hash, klass|
+      hash[klass] = "visit_#{(klass.name || '').gsub('::', '_')}"
+    end
+
+    def dispatch
+      DISPATCH
+    end
+
+    # rubocop:disable Naming/AccessorMethodName
+    def get_dispatch_cache
+      dispatch
+    end
+
+    # rubocop:enable Naming/AccessorMethodName
+
+    def nest_context(obj)
+      old_context = @current_context
+      @current_context = Context.new(obj)
+      @contexts << @current_context
+
+      yield
+
+      @current_context = old_context
+    end
+  end
+
+  class BaseTenantEnforcementClause < Arel::Nodes::Node
+    attr_reader :tenant_attribute
+
+    def initialize(tenant_attribute)
+      super()
+      @tenant_attribute = tenant_attribute
+      @tenant_model = MultiTenant.multi_tenant_model_for_table(
+        MultiTenant::TableNode.table_name(tenant_attribute.relation)
+      )
+    end
+
+    def to_s
+      to_sql
+    end
+
+    def to_str
+      to_sql
+    end
+
+    def to_sql(*)
+      collector = Arel::Collectors::SQLString.new
+      collector = @tenant_model.connection.visitor.accept tenant_arel, collector
+      collector.value
+    end
+  end
+
+  class TenantEnforcementClause < BaseTenantEnforcementClause
+    private
+
+    def tenant_arel
+      if defined?(Arel::Nodes::Quoted)
+        @tenant_attribute.eq(Arel::Nodes::Quoted.new(MultiTenant.current_tenant_id))
+      else
+        @tenant_attribute.eq(MultiTenant.current_tenant_id)
+      end
+    end
+  end
+
+  class TenantJoinEnforcementClause < BaseTenantEnforcementClause
+    attr_reader :table_left
+
+    def initialize(tenant_attribute, table_left)
+      super(tenant_attribute)
+      @table_left = table_left
+      @model_left = MultiTenant.multi_tenant_model_for_table(MultiTenant::TableNode.table_name(table_left))
+    end
+
+    private
+
+    def tenant_arel
+      @tenant_attribute.eq(@table_left[@model_left.partition_key])
+    end
+  end
+
+  module TenantValueVisitor
+    # rubocop:disable Naming/MethodName
+    def visit_MultiTenant_TenantEnforcementClause(obj, collector)
+      collector << obj
+    end
+
+    def visit_MultiTenant_TenantJoinEnforcementClause(obj, collector)
+      collector << obj
+    end
+
+    # rubocop:enable Naming/MethodName
+  end
+
+  module DatabaseStatements
+    def join_to_update(update, *args)
+      update = super
+      model = MultiTenant.multi_tenant_model_for_table(MultiTenant::TableNode.table_name(update.ast.relation))
+      if model.present? && !MultiTenant.with_write_only_mode_enabled? && MultiTenant.current_tenant_id.present?
+        update.where(MultiTenant::TenantEnforcementClause.new(model.arel_table[model.partition_key]))
+      end
+      update
+    end
+
+    def join_to_delete(delete, *args)
+      delete = super
+      model = MultiTenant.multi_tenant_model_for_table(MultiTenant::TableNode.table_name(delete.ast.left))
+      if model.present? && !MultiTenant.with_write_only_mode_enabled? && MultiTenant.current_tenant_id.present?
+        delete.where(MultiTenant::TenantEnforcementClause.new(model.arel_table[model.partition_key]))
+      end
+      delete
+    end
+
+    def update(arel, name = nil, binds = [])
+      model = MultiTenant.multi_tenant_model_for_arel(arel)
+      if model.present? && !MultiTenant.with_write_only_mode_enabled? && MultiTenant.current_tenant_id.present?
+        arel.where(MultiTenant::TenantEnforcementClause.new(model.arel_table[model.partition_key]))
+      end
+      super
+    end
+
+    def delete(arel, name = nil, binds = [])
+      model = MultiTenant.multi_tenant_model_for_arel(arel)
+      if model.present? && !MultiTenant.with_write_only_mode_enabled? && MultiTenant.current_tenant_id.present?
+        arel.where(MultiTenant::TenantEnforcementClause.new(model.arel_table[model.partition_key]))
+      end
+      super
+    end
+  end
+end
+
+require 'active_record/connection_adapters/abstract_adapter'
+ActiveRecord::ConnectionAdapters::AbstractAdapter.prepend(MultiTenant::DatabaseStatements)
+
+Arel::Visitors::ToSql.include(MultiTenant::TenantValueVisitor)
+
+module MultiTenant
+  # Rails 8.1 compiles update_all on a joined relation for PostgreSQL as
+  #
+  #   UPDATE t AS __active_record_update_alias SET ...
+  #   FROM t INNER JOIN ... WHERE t.id = __active_record_update_alias.id
+  #
+  # Citus only plans that self-join when it is also on the distribution
+  # column, so for multi-tenant tables add the partition key to it. This also
+  # keeps the rows being updated inside the tenant. Older Rails versions use
+  # an IN (subquery) here and never build the alias, so this is a no-op there.
+  module UpdateJoinPartitionKey
+    private
+
+    def prepare_update_statement(statement)
+      stmt = super
+      return stmt unless stmt.relation.respond_to?(:left) && stmt.relation.left.is_a?(Arel::Nodes::TableAlias)
+
+      table = statement.relation.left
+      model = MultiTenant.multi_tenant_model_for_table(MultiTenant::TableNode.table_name(table))
+      return stmt if model.nil? || model.partition_key.blank? || model.partition_key.to_s == model.primary_key.to_s
+
+      stmt.wheres << stmt.relation.left[model.partition_key].eq(table[model.partition_key])
+      stmt
+    end
+  end
+end
+
+Arel::Visitors::PostgreSQL.prepend(MultiTenant::UpdateJoinPartitionKey)
+
+module MultiTenant
+  module QueryMethodsExtensions
+    def build_arel(*)
+      arel = super
+
+      unless MultiTenant.with_write_only_mode_enabled?
+        visitor = MultiTenant::ArelTenantVisitor.new(arel)
+
+        visitor.contexts.each do |context|
+          node = context.arel_node
+
+          context.unhandled_relations.each do |relation|
+            model = MultiTenant.multi_tenant_model_for_table(MultiTenant::TableNode.table_name(relation.arel_table))
+
+            if MultiTenant.current_tenant_id
+              enforcement_clause = MultiTenant::TenantEnforcementClause.new(relation.arel_table[model.partition_key])
+              case node
+              when Arel::Nodes::Join # Arel::Nodes::OuterJoin, Arel::Nodes::RightOuterJoin, Arel::Nodes::FullOuterJoin
+                node.right.expr = node.right.expr.and(enforcement_clause)
+              when Arel::Nodes::SelectCore
+                if node.wheres.empty?
+                  node.wheres = [enforcement_clause]
+                elsif node.wheres[0].is_a?(Arel::Nodes::And)
+                  node.wheres[0].children << enforcement_clause
+                else
+                  node.wheres[0] = enforcement_clause.and(node.wheres[0])
+                end
+              else
+                raise 'UnknownContext'
+              end
+            end
+
+            next unless node.is_a?(Arel::Nodes::SelectCore) || node.is_a?(Arel::Nodes::Join)
+
+            node_list = if node.is_a? Arel::Nodes::Join
+                          [node]
+                        else
+                          node.source.right
+                        end
+
+            node_list.select { |n| n.is_a? Arel::Nodes::Join }.each do |node_join|
+              next unless node_join.right
+
+              relation_right, relation_left = relations_from_node_join(node_join)
+
+              next unless relation_right && relation_left
+
+              model_right = MultiTenant.multi_tenant_model_for_table(MultiTenant::TableNode.table_name(relation_left))
+              model_left = MultiTenant.multi_tenant_model_for_table(MultiTenant::TableNode.table_name(relation_right))
+              next unless model_right && model_left
+
+              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
+        end
+      end
+
+      arel
+    end
+
+    private
+
+    def relations_from_node_join(node_join)
+      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].flatten
+
+      tenant_applied = children.any? do |c|
+        c.is_a?(MultiTenant::TenantEnforcementClause) || c.is_a?(MultiTenant::TenantJoinEnforcementClause)
+      end
+      return nil, nil if tenant_applied || children.empty?
+
+      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
+
+      [nil, nil]
+    end
+  end
+end
+
+require 'active_record/relation'
+ActiveRecord::QueryMethods.prepend(MultiTenant::QueryMethodsExtensions)
+
+module MultiTenantFindBy
+  if ActiveRecord.gem_version >= Gem::Version.create('7.2.0')
+    def cached_find_by_statement(connection, key, &block)
+      return super unless respond_to?(:scoped_by_tenant?) && scoped_by_tenant?
+
+      super(connection, Array.wrap(key) + [MultiTenant.current_tenant_id.to_s], &block)
+    end
+  else
+    def cached_find_by_statement(key, &block)
+      return super unless respond_to?(:scoped_by_tenant?) && scoped_by_tenant?
+
+      super(Array.wrap(key) + [MultiTenant.current_tenant_id.to_s], &block)
+    end
+  end
+end
+
+ActiveSupport.on_load(:active_record) do |base|
+  base.singleton_class.prepend(MultiTenantFindBy)
+end

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

@@ -0,0 +1,90 @@
+# frozen_string_literal: true
+
+module Arel
+  module ActiveRecordRelationExtension
+    # Overrides the delete_all method to include tenant scoping
+    def delete_all
+      model = MultiTenant.multi_tenant_model_for_table(table_name)
+
+      # Call the original delete_all method if the current tenant is identified by an ID
+      return super if model.nil? || MultiTenant.current_tenant_is_id? || MultiTenant.current_tenant.nil?
+
+      stmt = Arel::DeleteManager.new.from(table)
+      stmt.wheres = [generate_in_condition_subquery]
+
+      # Execute the delete statement using the connection and return the result
+      klass.connection.delete(stmt, "#{klass} Delete All").tap { reset }
+    end
+
+    # Overrides the update_all method to include tenant scoping
+    def update_all(updates)
+      model = MultiTenant.multi_tenant_model_for_table(table_name)
+
+      # Call the original update_all method if the current tenant is identified by an ID
+      return super if model.nil? || MultiTenant.current_tenant_is_id? || MultiTenant.current_tenant.nil?
+
+      if updates.is_a?(Hash)
+        if klass.locking_enabled? &&
+           !updates.key?(klass.locking_column) &&
+           !updates.key?(klass.locking_column.to_sym)
+          attr = table[klass.locking_column]
+          updates[attr.name] = _increment_attribute(attr)
+        end
+        values = _substitute_values(updates)
+      else
+        values = Arel.sql(klass.sanitize_sql_for_assignment(updates, table.name))
+      end
+
+      stmt = Arel::UpdateManager.new
+      stmt.table(table)
+      stmt.set values
+      stmt.wheres = [generate_in_condition_subquery]
+
+      klass.connection.update(stmt, "#{klass} Update All").tap { reset }
+    end
+
+    private
+
+    # The generate_in_condition_subquery method generates a subquery that selects
+    # records associated with the current tenant.
+    def generate_in_condition_subquery
+      # Get the tenant key and tenant ID based on the current tenant
+      tenant_key = MultiTenant.partition_key(MultiTenant.current_tenant_class)
+      tenant_id = MultiTenant.current_tenant_id
+
+      # Build an Arel query. Only Active Record 7.2 and 8.0 take the
+      # connection as the first argument of build_arel; every other version
+      # takes only aliases, which is required (without a default) in 6.0 and 8.1.
+      arel = if eager_loading?
+               apply_join_dependency.arel
+             elsif ActiveRecord.gem_version >= Gem::Version.create('7.2.0') &&
+                   ActiveRecord.gem_version < Gem::Version.create('8.1.0')
+               build_arel(klass.connection)
+             else
+               build_arel(nil)
+             end
+
+      arel.source.left = table
+
+      # If the tenant ID is present and the tenant key is a column in the model,
+      # add a condition to only include records where the tenant key equals the tenant ID
+      if tenant_id && klass.column_names.include?(tenant_key)
+        tenant_condition = table[tenant_key].eq(tenant_id)
+        unless arel.constraints.any? { |node| node.to_sql.include?(tenant_condition.to_sql) }
+          arel = arel.where(tenant_condition)
+        end
+      end
+
+      # Clone the query, clear its projections, and set its projection to the primary key of the table
+      subquery = arel.clone
+      subquery.projections.clear
+      subquery = subquery.project(table[primary_key])
+
+      # Create an IN condition node with the primary key of the table and the subquery
+      Arel::Nodes::In.new(table[primary_key], subquery.ast)
+    end
+  end
+end
+
+# Patch ActiveRecord::Relation with the extension module
+ActiveRecord::Relation.prepend(Arel::ActiveRecordRelationExtension)

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

@@ -0,0 +1,91 @@
+# frozen_string_literal: true
+
+require 'sidekiq/client'
+
+# Adds methods to handle tenant information both in the client and server.
+module Sidekiq::Middleware::MultiTenant
+  # Get the current tenant and store in the message to be sent to Sidekiq.
+  class Client
+    def call(_worker_class, msg, _queue, _redis_pool)
+      if MultiTenant.current_tenant.present?
+        msg['multi_tenant'] ||=
+          {
+            'class' => MultiTenant.current_tenant_class,
+            'id' => MultiTenant.current_tenant_id
+          }
+      end
+
+      yield
+    end
+  end
+
+  # Pull the tenant out and run the current thread with it.
+  class Server
+    def call(_worker_class, msg, _queue, &block)
+      if msg.key?('multi_tenant')
+        tenant = begin
+          msg['multi_tenant']['class'].constantize.find(msg['multi_tenant']['id'])
+        rescue ActiveRecord::RecordNotFound
+          msg['multi_tenant']['id']
+        end
+        MultiTenant.with(tenant, &block)
+      else
+        yield
+      end
+    end
+  end
+end
+
+# Configure Sidekiq to use the multi-tenant client and server middleware to add (client/server)/process(server)
+# tenant information.
+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
+
+# Bulk push support for Sidekiq while setting multi-tenant information.
+# This is a copy of the Sidekiq::Client#push_bulk method with the addition of
+# setting the multi-tenant information for each job.
+module Sidekiq
+  class Client
+    # Allows the caller to enqueue multiple Sidekiq jobs with
+    # tenant information in a single call. It ensures that each job is processed
+    # within the correct tenant context and returns an array of job IDs for the enqueued jobs
+    def push_bulk_with_tenants(items)
+      first_job = items['jobs'].first
+      return [] unless first_job # no jobs to push
+      unless first_job.is_a?(Hash)
+        raise ArgumentError, "Bulk arguments must be an Array of Hashes: [{ 'args' => [1], 'tenant_id' => 1 }, ...]"
+      end
+
+      normed = normalize_item(items.except('jobs').merge('args' => []))
+      payloads = items['jobs'].map do |job|
+        MultiTenant.with(job['tenant_id']) do
+          copy = normed.merge('args' => job['args'], 'jid' => SecureRandom.hex(12), 'enqueued_at' => Time.now.to_f)
+          result = process_single(items['class'], copy)
+          result || nil
+        end
+      end.compact
+
+      raw_push(payloads) unless payloads.empty?
+      payloads.collect { |payload| payload['jid'] }
+    end
+
+    # Enabling the push_bulk_with_tenants method to be called directly on the Sidekiq::Client class
+    class << self
+      def push_bulk_with_tenants(items)
+        new.push_bulk_with_tenants(items)
+      end
+    end
+  end
+end

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

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+module MultiTenant
+  module TableNode
+    # Return table name
+    def self.table_name(node)
+      # NOTE: Arel::Nodes::Table#table_name is removed in Rails 7.1
+      if node.is_a?(Arel::Nodes::TableAlias)
+        node.table_name
+      else
+        node.name
+      end
+    end
+  end
+end

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

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module MultiTenant
+  VERSION = '2.5.0'
+end

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

@@ -0,0 +1,4 @@
+# frozen_string_literal: true
+
+# Lets Bundler auto-require the gem under its published name.
+require_relative 'activerecord-multi-tenant'

data/lib/activerecord-multi-tenant.rb

@@ -0,0 +1,3 @@
+# frozen_string_literal: true
+
+require 'activerecord_multi_tenant'

data/lib/activerecord_multi_tenant.rb

@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+require_relative 'activerecord-multi-tenant/controller_extensions' if Object.const_defined?(:ActionController)
+require_relative 'activerecord-multi-tenant/copy_from_client'
+require_relative 'activerecord-multi-tenant/fast_truncate'
+require_relative 'activerecord-multi-tenant/migrations'
+require_relative 'activerecord-multi-tenant/model_extensions'
+require_relative 'activerecord-multi-tenant/multi_tenant'
+require_relative 'activerecord-multi-tenant/query_rewriter'
+require_relative 'activerecord-multi-tenant/query_monitor'
+require_relative 'activerecord-multi-tenant/table_node'
+require_relative 'activerecord-multi-tenant/version'
+require_relative 'activerecord-multi-tenant/habtm'
+require_relative 'activerecord-multi-tenant/relation_extension'