activerecord-multi-tenant 2.0.0 → 2.1.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 457cfac4996b93bb7b0b25cd99c2687f2742a8bc580888008baa4f903255cd74
-  data.tar.gz: 8c369e04a906752deda52d9b7b78a2377551e4b85f74e3b6d4b0742d30147783
+  metadata.gz: e198970d8d5a4db04c8e203e28021668c765fdde78f9bb9a7cd0e6c577bfbe0e
+  data.tar.gz: 6abb30aac520063dbef7ddc5689714ef1dd9eb691aa382498b30c8bbcf03844b
 SHA512:
-  metadata.gz: 7715dadcbd9cd13d86135518d3fc62283d9ab85e802016ab1be15fe39d815ec344c1d3e4699f96870b6e50b6d3b66b4a38f780d07ad6d3d4b8e21668deacca18
-  data.tar.gz: d1d364ebde87012cc5b91a6cf812a1a8cc0f89bb5ee8859b3de01ab0392dc5a62b74a3f93ac721d3e583343b703fc872d73c1f301cff9716e686492f06922d36
+  metadata.gz: 5010dbbd2a8036ad5a3adae79a49567a279a318931ffe832663dd2430f455d1895e26cca6d96456484b894cc241cca21fbab31aa301f9ab4ef0538a4266b4005
+  data.tar.gz: e94b7ee0f619b0763879efefbae3a69792e00e8b9f2a2661a82666d4c129decea132f20b44999177bdc0162dec0e9b3cac4f402d70a76fb14518e283f4441d6c

data/CHANGELOG.md

@@ -1,5 +1,13 @@
 # Changelog
 
+## 2.1.0      2022-10-20
+* Fix query building for models with mismatched partition_keys [#150](https://github.com/citusdata/activerecord-multi-tenant/pull/150)
+* Identify tenant even if class name is nonstandard [#152](https://github.com/citusdata/activerecord-multi-tenant/pull/152)
+* Add current_tenant_id to WHERE clauses when calling methods on activerecord instance or its associations [#154](https://github.com/citusdata/activerecord-multi-tenant/pull/154)
+* Make create_distributed_table, create_reference_table reversible & add ruby wrapper for rebalance_table_shards [#155](https://github.com/citusdata/activerecord-multi-tenant/pull/155)
+* Support create_distributed_table, create_reference_table in schema.rb [#156](https://github.com/citusdata/activerecord-multi-tenant/pull/156)
+* Add client and server sidekiq middleware to sidekiq middleware chain [#158](https://github.com/citusdata/activerecord-multi-tenant/pull/158)
+
 ## 2.0.0      2022-05-19
 
 * Replace RequestStore with CurrentAttributes [#139](https://github.com/citusdata/activerecord-multi-tenant/pull/139)

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

@@ -2,12 +2,40 @@ module MultiTenant
   module MigrationExtensions
     def create_distributed_table(table_name, partition_key)
       return unless citus_version.present?
-      execute "SELECT create_distributed_table($$#{table_name}$$, $$#{partition_key}$$)"
+
+      reversible do |dir|
+        dir.up do
+          execute "SELECT create_distributed_table($$#{table_name}$$, $$#{partition_key}$$)"
+        end
+        dir.down do
+          undistribute_table(table_name)
+        end
+      end
     end
 
     def create_reference_table(table_name)
       return unless citus_version.present?
-      execute "SELECT create_reference_table($$#{table_name}$$)"
+
+      reversible do |dir|
+        dir.up do
+          execute "SELECT create_reference_table($$#{table_name}$$)"
+        end
+        dir.down do
+          undistribute_table(table_name)
+        end
+      end
+    end
+
+    def undistribute_table(table_name)
+      return unless citus_version.present?
+
+      execute "SELECT undistribute_table($$#{table_name}$$))"
+    end
+
+    def rebalance_table_shards
+      return unless citus_version.present?
+
+      execute 'SELECT rebalance_table_shards()'
     end
 
     def execute_on_all_nodes(sql)
@@ -28,21 +56,19 @@ 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'
@@ -54,3 +80,38 @@ module ActiveRecord
     end
   end
 end
+
+module ActiveRecord
+  class SchemaDumper
+    private
+
+    alias initialize_without_citus initialize
+    def initialize(connection, options = {})
+      initialize_without_citus(connection, options)
+
+      @distribution_columns =
+        if ActiveRecord::Migration.citus_version.present?
+          @connection.execute('SELECT logicalrelid::regclass AS table_name, column_to_column_name(logicalrelid, partkey) AS dist_col_name FROM pg_dist_partition').to_h do |v|
+            [v['table_name'], v['dist_col_name']]
+          end
+        else
+          {}
+        end
+    end
+
+    # Support for create_distributed_table & create_reference_table
+    alias table_without_citus table
+    def table(table, stream)
+      table_without_citus(table, stream)
+      table_name = remove_prefix_and_suffix(table)
+      distribution_column = @distribution_columns[table_name]
+      if distribution_column
+        stream.puts "  create_distributed_table(#{table_name.inspect}, #{distribution_column.inspect})"
+        stream.puts
+      elsif @distribution_columns.key?(table_name)
+        stream.puts "  create_reference_table(#{table_name.inspect})"
+        stream.puts
+      end
+    end
+  end
+end

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

@@ -3,7 +3,7 @@ module MultiTenant
     DEFAULT_ID_FIELD = 'id'.freeze
 
     def multi_tenant(tenant_name, options = {})
-      if to_s.underscore.to_sym == tenant_name
+      if to_s.underscore.to_sym == tenant_name || (!table_name.nil? && table_name.singularize.to_sym == tenant_name)
         unless MultiTenant.with_write_only_mode_enabled?
           # This is the tenant model itself. Workaround for https://github.com/citusdata/citus/issues/687
           before_create -> do
@@ -129,6 +129,20 @@ end
 
 ActiveSupport.on_load(:active_record) do |base|
   base.extend MultiTenant::ModelExtensionsClassMethods
+
+  # Ensure we have current_tenant_id in where clause when a cached ActiveRecord instance is being reloaded, or update_columns without callbacks is called
+  MultiTenant.wrap_methods(ActiveRecord::Base, 'self', :delete, :reload, :update_columns)
+
+  # Any queuries fired for fetching a singular association have the correct current_tenant_id in WHERE clause
+  # reload is called anytime any record's association is accessed
+  MultiTenant.wrap_methods(ActiveRecord::Associations::Association, 'owner', :reload)
+
+  # For collection associations, we need to wrap multiple methods in returned proxy so that any queries have the correct current_tenant_id in WHERE clause
+  ActiveRecord::Associations::CollectionProxy.alias_method :equals_mt, :== # Hack to prevent syntax error due to invalid method name
+  ActiveRecord::Associations::CollectionProxy.alias_method :append_mt, :<< # Hack to prevent syntax error due to invalid method name
+  MultiTenant.wrap_methods(ActiveRecord::Associations::CollectionProxy, '@association.owner', :find, :last, :take, :build, :create, :create!, :replace, :delete_all, :destroy_all, :delete, :destroy, :calculate, :pluck, :size, :empty?, :include?, :equals_mt, :records, :append_mt, :find_nth_with_limit, :find_nth_from_last, :null_scope?, :find_from_target?, :exec_queries)
+  ActiveRecord::Associations::CollectionProxy.alias_method :==, :equals_mt
+  ActiveRecord::Associations::CollectionProxy.alias_method :<<, :append_mt
 end
 
 class ActiveRecord::Associations::Association

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

@@ -109,6 +109,23 @@ module MultiTenant
     end
   end
 
+  # Wrap calls to any of `method_names` on an instance Class `klass` with MultiTenant.with when `'owner'` (evaluated in context of the klass instance) is a ActiveRecord model instance that is multi-tenant
+  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?
+            MultiTenant.with(#{owner}.public_send(#{owner}.class.partition_key)) { #{original_method_name}(*args) }
+          else
+            #{original_method_name}(*args)
+          end
+        end
+      CODE
+    end
+  end
+
   # Preserve backward compatibility for people using .with_id
   singleton_class.send(:alias_method, :with_id, :with)
 

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

@@ -305,7 +305,7 @@ module ActiveRecord
                 model_right = MultiTenant.multi_tenant_model_for_table(relation_left.table_name)
                 model_left = MultiTenant.multi_tenant_model_for_table(relation_right.table_name)
                 if model_right && model_left
-                  join_enforcement_clause = MultiTenant::TenantJoinEnforcementClause.new(relation_left[model_left.partition_key], relation_right)
+                  join_enforcement_clause = MultiTenant::TenantJoinEnforcementClause.new(relation_right[model_right.partition_key], relation_left)
                   node_join.right.expr = node_join.right.expr.and(join_enforcement_clause)
                 end
               end

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

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

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

@@ -1,3 +1,3 @@
 module MultiTenant
-  VERSION = '2.0.0'
+  VERSION = '2.1.1'
 end

data/lib/activerecord-multi-tenant.rb

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

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

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

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

@@ -70,6 +70,24 @@ describe MultiTenant do
     it { expect(@partition_key_not_model_task.non_model_id).to be 77 }
   end
 
+
+  describe 'Tenant model with a nonstandard class name' do
+    let(:account_klass) do
+      Class.new(ActiveRecord::Base) do
+        self.table_name = 'account'
+        def self.name
+          'UserAccount'
+        end
+
+        multi_tenant(:account)
+      end
+    end
+    it "does not register the tenant model" do
+      expect(MultiTenant).not_to receive(:register_multi_tenant_model)
+      account_klass
+    end
+  end
+
   describe 'Changes table_name after multi_tenant called' do
     before do
       account_klass.has_many(:posts, anonymous_class: post_klass)
@@ -186,11 +204,7 @@ describe MultiTenant do
     end
 
     it 'handles belongs_to with optional: true' do
-      MultiTenant.with(account) do
-        sub_task
-      end
-
-      record = sub_task.optional_sub_tasks.create!
+      record = OptionalSubTask.create(sub_task_id: sub_task.id)
       expect(record.reload.sub_task).to eq(sub_task)
       expect(record.account_id).to eq(nil)
     end
@@ -400,10 +414,10 @@ describe MultiTenant do
 
   it "applies the team_id conditions in the where clause" do
     option1 = <<-sql.strip
-      SELECT "sub_tasks".* FROM "sub_tasks" INNER JOIN "tasks" ON "sub_tasks"."task_id" = "tasks"."id" AND "sub_tasks"."account_id" = "tasks"."account_id" WHERE "tasks"."project_id" = 1 AND "sub_tasks"."account_id" = 1 AND "tasks"."account_id" = 1
+      SELECT "sub_tasks".* FROM "sub_tasks" INNER JOIN "tasks" ON "sub_tasks"."task_id" = "tasks"."id" AND "tasks"."account_id" = "sub_tasks"."account_id" WHERE "tasks"."project_id" = 1 AND "sub_tasks"."account_id" = 1 AND "tasks"."account_id" = 1
     sql
     option2 = <<-sql.strip
-      SELECT "sub_tasks".* FROM "sub_tasks" INNER JOIN "tasks" ON "sub_tasks"."task_id" = "tasks"."id" AND "sub_tasks"."account_id" = "tasks"."account_id" WHERE "sub_tasks"."account_id" = 1 AND "tasks"."project_id" = 1 AND "tasks"."account_id" = 1
+      SELECT "sub_tasks".* FROM "sub_tasks" INNER JOIN "tasks" ON "sub_tasks"."task_id" = "tasks"."id" AND "tasks"."account_id" = "sub_tasks"."account_id" WHERE "sub_tasks"."account_id" = 1 AND "tasks"."project_id" = 1 AND "tasks"."account_id" = 1
     sql
 
     account1 = Account.create! name: 'Account 1'
@@ -418,7 +432,7 @@ describe MultiTenant do
 
     MultiTenant.without do
       expected_sql = <<-sql
-                        SELECT "sub_tasks".* FROM "sub_tasks" INNER JOIN "tasks" ON "sub_tasks"."task_id" = "tasks"."id" AND "sub_tasks"."account_id" = "tasks"."account_id" WHERE "tasks"."project_id" = 1
+                        SELECT "sub_tasks".* FROM "sub_tasks" INNER JOIN "tasks" ON "sub_tasks"."task_id" = "tasks"."id" AND "tasks"."account_id" = "sub_tasks"."account_id" WHERE "tasks"."project_id" = 1
                      sql
 
       project = Project.first
@@ -456,7 +470,7 @@ describe MultiTenant do
       expect(project.categories).to include(category1)
 
       expected_sql = <<-sql
-                        SELECT "projects".* FROM "projects" INNER JOIN "project_categories" ON "project_categories"."project_id" = "projects"."id" AND "project_categories"."account_id" = "projects"."account_id" INNER JOIN "categories" ON "categories"."id" = "project_categories"."category_id" WHERE "projects"."account_id" = 1
+                        SELECT "projects".* FROM "projects" INNER JOIN "project_categories" ON "project_categories"."project_id" = "projects"."id" AND "projects"."account_id" = "project_categories"."account_id" INNER JOIN "categories" ON "categories"."id" = "project_categories"."category_id" WHERE "projects"."account_id" = 1
                     sql
 
       expect(Project.where(account_id: 1).joins(:categories).to_sql).to eq(expected_sql.strip)
@@ -490,7 +504,7 @@ describe MultiTenant do
 
     MultiTenant.without do
       expected_sql = <<-sql
-                     SELECT "projects"."id" AS t0_r0, "projects"."account_id" AS t0_r1, "projects"."name" AS t0_r2, "categories"."id" AS t1_r0, "categories"."name" AS t1_r1 FROM "projects" LEFT OUTER JOIN "project_categories" ON "project_categories"."project_id" = "projects"."id" AND "project_categories"."account_id" = "projects"."account_id" LEFT OUTER JOIN "categories" ON "categories"."id" = "project_categories"."category_id" WHERE "projects"."account_id" = 1
+                     SELECT "projects"."id" AS t0_r0, "projects"."account_id" AS t0_r1, "projects"."name" AS t0_r2, "categories"."id" AS t1_r0, "categories"."name" AS t1_r1 FROM "projects" LEFT OUTER JOIN "project_categories" ON "project_categories"."project_id" = "projects"."id" AND "projects"."account_id" = "project_categories"."account_id" LEFT OUTER JOIN "categories" ON "categories"."id" = "project_categories"."category_id" WHERE "projects"."account_id" = 1
                      sql
 
       expect(Project.where(account_id: 1).eager_load(:categories).to_sql).to eq(expected_sql.strip)
@@ -522,7 +536,7 @@ describe MultiTenant do
 
     MultiTenant.without do
       expected_sql = <<-sql
-                     SELECT "tasks".* FROM "tasks" INNER JOIN "projects" ON "projects"."id" = "tasks"."project_id" AND "projects"."account_id" = "tasks"."account_id" LEFT JOIN project_categories pc ON project.category_id = pc.id WHERE "tasks"."account_id" = 1
+                     SELECT "tasks".* FROM "tasks" INNER JOIN "projects" ON "projects"."id" = "tasks"."project_id" AND "tasks"."account_id" = "projects"."account_id" LEFT JOIN project_categories pc ON project.category_id = pc.id WHERE "tasks"."account_id" = 1
                      sql
 
       expect(Task.where(account_id: 1).joins(:project).joins('LEFT JOIN project_categories pc ON project.category_id = pc.id').to_sql).to eq(expected_sql.strip)

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

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

data/spec/schema.rb

@@ -179,7 +179,6 @@ end
 with_belongs_to_required_by_default do
   class OptionalSubTask < ActiveRecord::Base
     multi_tenant :account, optional: true
-    belongs_to :account, optional: true
     belongs_to :sub_task
   end
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: activerecord-multi-tenant
 version: !ruby/object:Gem::Version
-  version: 2.0.0
+  version: 2.1.1
 platform: ruby
 authors:
 - Citus Data
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2022-05-19 00:00:00.000000000 Z
+date: 2022-10-20 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: rails
@@ -172,12 +172,12 @@ files:
 - lib/activerecord-multi-tenant/migrations.rb
 - lib/activerecord-multi-tenant/model_extensions.rb
 - lib/activerecord-multi-tenant/multi_tenant.rb
-- lib/activerecord-multi-tenant/persistence_extension.rb
 - lib/activerecord-multi-tenant/query_monitor.rb
 - lib/activerecord-multi-tenant/query_rewriter.rb
 - lib/activerecord-multi-tenant/sidekiq.rb
 - lib/activerecord-multi-tenant/version.rb
 - lib/activerecord-multi-tenant/with_lock.rb
+- spec/activerecord-multi-tenant/associations_spec.rb
 - spec/activerecord-multi-tenant/controller_extensions_spec.rb
 - spec/activerecord-multi-tenant/fast_truncate_spec.rb
 - spec/activerecord-multi-tenant/model_extensions_spec.rb

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

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