activerecord-multi-tenant-next 2.5.0

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

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+module MultiTenant
+  # Designed to be mixed into an ActiveRecord model to provide
+  # a copy_from_client method that allows for efficient bulk insertion of
+  # data into a PostgreSQL database using the COPY command
+  class CopyFromClientHelper
+    attr_reader :count
+
+    def initialize(conn, column_types)
+      @count = 0
+      @conn = conn
+      @column_types = column_types
+    end
+
+    def <<(row)
+      row = row.map.with_index { |val, idx| @column_types[idx].serialize(val) }
+      @conn.put_copy_data(row)
+      @count += 1
+    end
+  end
+
+  module CopyFromClient
+    def copy_from_client(columns, &block)
+      conn         = connection.raw_connection
+      column_types = columns.map { |c| type_for_attribute(c.to_s) }
+      helper = MultiTenant::CopyFromClientHelper.new(conn, column_types)
+      conn.copy_data %{COPY #{quoted_table_name}("#{columns.join('","')}") FROM STDIN}, PG::TextEncoder::CopyRow.new do
+        block.call helper
+      end
+      helper.count
+    end
+  end
+end
+
+# Add copy_from_client to ActiveRecord::Base
+ActiveSupport.on_load(:active_record) do |base|
+  base.extend(MultiTenant::CopyFromClient)
+end

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

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+# Truncates only the tables that have been modified, according to sequence
+# values
+# Faster alternative to DatabaseCleaner.clean_with(:truncation, pre_count: true)
+module MultiTenant
+  module FastTruncate
+    def self.run(exclude: ['schema_migrations'])
+      # This is a slightly faster version of DatabaseCleaner.clean_with(:truncation, pre_count: true)
+      ActiveRecord::Base.connection.execute format(%(
+      DO LANGUAGE plpgsql $$
+      DECLARE
+        t record;
+        tables text[];
+        seq_exists boolean;
+        needs_truncate boolean;
+      BEGIN
+        FOR t IN SELECT schemaname, tablename FROM pg_tables WHERE schemaname = 'public' AND tablename NOT IN (%s) LOOP
+          EXECUTE 'SELECT EXISTS (SELECT * from pg_class c WHERE c.relkind = ''S''
+           AND c.relname=''' || t.tablename || '_id_seq'')' into seq_exists;
+          IF seq_exists THEN
+            EXECUTE 'SELECT is_called FROM ' || t.tablename || '_id_seq' INTO needs_truncate;
+          ELSE
+            needs_truncate := true;
+          END IF;
+
+          IF needs_truncate THEN
+            tables := array_append(tables, quote_ident(t.schemaname) || '.' || quote_ident(t.tablename));
+          END IF;
+        END LOOP;
+
+        IF array_length(tables, 1) > 0 THEN
+          EXECUTE 'TRUNCATE TABLE ' || array_to_string(tables, ', ') || ' RESTART IDENTITY CASCADE';
+        END IF;
+      END$$;), exclude.map { |t| "'#{t}'" }.join('\n'))
+    end
+  end
+end

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

@@ -0,0 +1,49 @@
+# frozen_string_literal: true
+
+# This module extension is a monkey patch to the ActiveRecord::Associations::ClassMethods module.
+# It overrides the has_and_belongs_to_many method to add the tenant_id to the join table if the
+# tenant_enabled option is set to true.
+
+module ActiveRecord
+  module Associations
+    module ClassMethods
+      # rubocop:disable Naming/PredicateName
+      def has_and_belongs_to_many_with_tenant(name, scope = nil, **options, &extension)
+        # rubocop:enable Naming/PredicateName
+        has_and_belongs_to_many_without_tenant(name, scope, **options, &extension)
+
+        middle_reflection = _reflect_on_association(name.to_s).through_reflection
+        join_model = middle_reflection.klass
+
+        # get tenant_enabled from options and if it is not set, set it to false
+        tenant_enabled = options[:tenant_enabled] || false
+
+        return unless tenant_enabled
+
+        tenant_class_name = options[:tenant_class_name]
+        tenant_column = options[:tenant_column]
+
+        match = tenant_column.match(/(\w+)_id/)
+        tenant_field_name = match ? match[1] : 'tenant'
+
+        join_model.class_eval do
+          belongs_to tenant_field_name.to_sym, class_name: tenant_class_name
+          before_create :tenant_set
+
+          private
+
+          # This method sets the tenant_id on the join table and executes before creation of the join table record.
+          define_method :tenant_set do
+            return unless tenant_enabled
+            raise MultiTenant::MissingTenantError, 'Tenant Id is not set' unless MultiTenant.current_tenant_id
+
+            send("#{tenant_column}=", MultiTenant.current_tenant_id)
+          end
+        end
+      end
+
+      alias has_and_belongs_to_many_without_tenant has_and_belongs_to_many
+      alias has_and_belongs_to_many has_and_belongs_to_many_with_tenant
+    end
+  end
+end

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

@@ -0,0 +1,134 @@
+# frozen_string_literal: true
+
+module MultiTenant
+  module MigrationExtensions
+    def create_distributed_table(table_name, partition_key)
+      return unless citus_version.present?
+
+      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?
+
+      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)
+      execute sql
+
+      case citus_version
+      when '6.0'
+        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
+        execute "SELECT run_command_on_workers($$#{sql}$$)"
+      end
+    end
+
+    def enable_extension_on_all_nodes(extension)
+      execute_on_all_nodes "CREATE EXTENSION IF NOT EXISTS \"#{extension}\""
+    end
+
+    def citus_version
+      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'
+    end
+  end
+end
+
+ActiveRecord::Migration.include MultiTenant::MigrationExtensions if defined?(ActiveRecord::Migration)
+
+module MultiTenant
+  module SchemaStatementsExtensions
+    def create_table(table_name, options = {}, &block)
+      ret = super(table_name, **options.except(:partition_key), &block)
+      if options[:id] != false && options[:partition_key] && options[:partition_key].to_s != 'id'
+        execute "ALTER TABLE #{table_name} DROP CONSTRAINT #{table_name}_pkey"
+        execute "ALTER TABLE #{table_name} ADD PRIMARY KEY(\"#{options[:partition_key]}\", id)"
+      end
+      ret
+    end
+  end
+end
+ActiveRecord::ConnectionAdapters::SchemaStatements.prepend(MultiTenant::SchemaStatementsExtensions)
+
+module MultiTenant
+  module SchemaDumperExtensions
+    private
+
+    def initialize(connection, options = {})
+      super
+
+      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
+    def table(table, stream)
+      super
+      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
+
+module ActiveRecord
+  class SchemaDumper
+    prepend MultiTenant::SchemaDumperExtensions
+  end
+end

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

@@ -0,0 +1,207 @@
+# frozen_string_literal: true
+
+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'
+    # 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 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
+          def scoped_by_tenant?
+            true
+          end
+
+          # 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)
+          end
+
+          def reset_primary_key
+            primary_object_keys = Array.wrap(connection.schema_cache.primary_keys(table_name)) - [partition_key]
+
+            self.primary_key = if primary_object_keys.size == 1
+                                 primary_object_keys.first
+                               elsif table_name &&
+                                     connection.schema_cache.columns_hash(table_name).include?(DEFAULT_ID_FIELD)
+                                 DEFAULT_ID_FIELD
+                               end
+          end
+
+          def inherited(subclass)
+            super
+            MultiTenant.register_multi_tenant_model(subclass)
+          end
+        end
+
+        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, options)
+          belongs_to(
+            tenant_name,
+            **options.slice(:class_name, :inverse_of, :optional),
+            foreign_key: options[:partition_key]
+          )
+        end
+
+        # New instances should have the tenant set
+        after_initialize proc { |record|
+          if MultiTenant.current_tenant_id &&
+             (!record.attribute_present?(partition_key) || record.public_send(partition_key.to_sym).nil?)
+            record.public_send(:"#{partition_key}=", 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.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
+            was = send("#{partition_key}_was")
+            was_nil_or_skipped = was.nil? || was.instance_of?(Object)
+
+            if send("#{partition_key}_changed?") && persisted? && !was_nil_or_skipped
+              raise MultiTenant::TenantIsImmutable
+            end
+
+            tenant_id
+          end
+
+          if MultiTenant.tenant_klass_defined?(tenant_name, options)
+            define_method "#{tenant_name}=" do |model|
+              super(model)
+              if send("#{partition_key}_changed?") && persisted? && !send("#{partition_key}_was").nil?
+                raise MultiTenant::TenantIsImmutable
+              end
+
+              model
+            end
+
+            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
+            end
+          end
+        end
+        include to_include
+
+        # 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 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 lambda { |record, block|
+          if MultiTenant.current_tenant_id.nil?
+            MultiTenant.with(record.public_send(partition_key)) { block.call }
+          else
+            block.call
+          end
+        }
+      end
+    end
+  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
+module MultiTenant
+  module AssociationExtensions
+    def skip_statement_cache?(*scope)
+      return true if klass.respond_to?(:scoped_by_tenant?) && klass.scoped_by_tenant?
+
+      if reflection.through_reflection
+        through_klass = reflection.through_reflection.klass
+        return true if through_klass.respond_to?(:scoped_by_tenant?) && through_klass.scoped_by_tenant?
+      end
+
+      super
+    end
+  end
+end
+
+ActiveRecord::Associations::Association.prepend(MultiTenant::AssociationExtensions)

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

@@ -0,0 +1,160 @@
+# frozen_string_literal: true
+
+require 'active_support/current_attributes'
+
+module MultiTenant
+  class Current < ::ActiveSupport::CurrentAttributes
+    attribute :tenant
+  end
+
+  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.underscore}_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
+
+  # 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
+
+  # Registry that maps table names to models (used by the query rewriter)
+  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 ||= []
+
+    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
+      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(TableNode.table_name(arel.ast.relation.left))
+    else
+      MultiTenant.multi_tenant_model_for_table(TableNode.table_name(arel.ast.relation))
+    end
+  end
+
+  def self.current_tenant=(tenant)
+    Current.tenant = tenant
+  end
+
+  def self.current_tenant
+    Current.tenant
+  end
+
+  def self.current_tenant_id
+    current_tenant_is_id? ? current_tenant : current_tenant.try(:id)
+  end
+
+  def self.current_tenant_is_id?
+    current_tenant.is_a?(String) || current_tenant.is_a?(Integer)
+  end
+
+  def self.current_tenant_class
+    if current_tenant_is_id?
+      MultiTenant.default_tenant_class || raise('Only have tenant id, and no default tenant class set')
+    elsif current_tenant
+      MultiTenant.current_tenant.class.name
+    end
+  end
+
+  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 || 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 current_tenant == tenant
+
+    old_tenant = current_tenant
+    begin
+      self.current_tenant = tenant
+      block.call
+    ensure
+      self.current_tenant = old_tenant
+    end
+  end
+
+  def self.without(&block)
+    return block.call if current_tenant.nil?
+
+    old_tenant = current_tenant
+    begin
+      self.current_tenant = nil
+      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
+  # TODO: Could not understand the use of owner here. Need to check
+  def self.wrap_methods(klass, owner, *method_names)
+    mod = Module.new
+    klass.prepend(mod)
+
+    method_names.each do |method_name|
+      mod.module_eval <<-CODE, __FILE__, __LINE__ + 1
+        def #{method_name}(...)
+          if MultiTenant.multi_tenant_model_for_table(#{owner}.class.table_name).present? && #{owner}.persisted? && MultiTenant.current_tenant_id.nil? && #{owner}.class.respond_to?(:partition_key) && #{owner}.attributes.include?(#{owner}.class.partition_key)
+            MultiTenant.with(#{owner}.public_send(#{owner}.class.partition_key)) { super }
+          else
+            super
+          end
+        end
+      CODE
+    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

@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+# 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
+
+  # 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)
+      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)