Skip to content

Commit e938a43

Browse files
committed
Refactor tenant enforcement to use Arel visitor pattern
1 parent 272e95c commit e938a43

File tree

3 files changed

+43
-76
lines changed

3 files changed

+43
-76
lines changed

lib/activerecord-multi-tenant/query_rewriter.rb

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -265,19 +265,33 @@ def join_to_delete(delete, *args)
265265

266266
def update(arel, name = nil, binds = [])
267267
model = MultiTenant.multi_tenant_model_for_arel(arel)
268-
if model.present? && !MultiTenant.with_write_only_mode_enabled? && MultiTenant.current_tenant_id.present?
268+
if model.present? &&
269+
!MultiTenant.with_write_only_mode_enabled? &&
270+
MultiTenant.current_tenant_id.present? &&
271+
!already_has_tenant_enforcement_clause?(arel)
269272
arel.where(MultiTenant::TenantEnforcementClause.new(model.arel_table[model.partition_key]))
270273
end
271274
super
272275
end
273276

274277
def delete(arel, name = nil, binds = [])
275278
model = MultiTenant.multi_tenant_model_for_arel(arel)
276-
if model.present? && !MultiTenant.with_write_only_mode_enabled? && MultiTenant.current_tenant_id.present?
279+
if model.present? &&
280+
!MultiTenant.with_write_only_mode_enabled? &&
281+
MultiTenant.current_tenant_id.present? &&
282+
!already_has_tenant_enforcement_clause?(arel)
277283
arel.where(MultiTenant::TenantEnforcementClause.new(model.arel_table[model.partition_key]))
278284
end
279285
super
280286
end
287+
288+
private
289+
290+
def already_has_tenant_enforcement_clause?(arel)
291+
arel.try(:ast).try(:wheres).to_a.any? do |where|
292+
where.is_a?(MultiTenant::BaseTenantEnforcementClause)
293+
end
294+
end
281295
end
282296
end
283297

Lines changed: 17 additions & 67 deletions
Original file line numberDiff line numberDiff line change
@@ -1,75 +1,25 @@
11
# frozen_string_literal: true
22

3-
module Arel
4-
module ActiveRecordRelationExtension
5-
# Overrides the delete_all method to include tenant scoping
6-
def delete_all
7-
model = MultiTenant.multi_tenant_model_for_table(table_name)
8-
9-
# Call the original delete_all method if the current tenant is identified by an ID
10-
return super if model.nil? || MultiTenant.current_tenant_is_id? || MultiTenant.current_tenant.nil?
11-
12-
stmt = Arel::DeleteManager.new.from(table)
13-
stmt.wheres = [generate_in_condition_subquery]
14-
15-
# Execute the delete statement using the connection and return the result
16-
klass.connection.delete(stmt, "#{klass} Delete All").tap { reset }
17-
end
18-
19-
# Overrides the update_all method to include tenant scoping
20-
def update_all(updates)
21-
model = MultiTenant.multi_tenant_model_for_table(table_name)
22-
23-
# Call the original update_all method if the current tenant is identified by an ID
24-
return super if model.nil? || MultiTenant.current_tenant_is_id? || MultiTenant.current_tenant.nil?
25-
26-
stmt = Arel::UpdateManager.new
27-
stmt.table(table)
28-
stmt.set Arel.sql(klass.send(:sanitize_sql_for_assignment, updates))
29-
stmt.wheres = [generate_in_condition_subquery]
30-
31-
klass.connection.update(stmt, "#{klass} Update All").tap { reset }
32-
end
33-
34-
private
35-
36-
# The generate_in_condition_subquery method generates a subquery that selects
37-
# records associated with the current tenant.
38-
def generate_in_condition_subquery
39-
# Get the tenant key and tenant ID based on the current tenant
40-
tenant_key = MultiTenant.partition_key(MultiTenant.current_tenant_class)
41-
tenant_id = MultiTenant.current_tenant_id
42-
43-
# Build an Arel query
44-
arel = if eager_loading?
45-
apply_join_dependency.arel
46-
elsif ActiveRecord.gem_version >= Gem::Version.create('7.2.0')
47-
build_arel(klass.connection)
48-
else
49-
build_arel
50-
end
51-
52-
arel.source.left = table
53-
54-
# If the tenant ID is present and the tenant key is a column in the model,
55-
# add a condition to only include records where the tenant key equals the tenant ID
56-
if tenant_id && klass.column_names.include?(tenant_key)
57-
tenant_condition = table[tenant_key].eq(tenant_id)
58-
unless arel.constraints.any? { |node| node.to_sql.include?(tenant_condition.to_sql) }
59-
arel = arel.where(tenant_condition)
3+
module Arel # :nodoc: all
4+
module Visitors
5+
module ToSqlPatch
6+
def prepare_update_statement(object)
7+
if object.key && (has_limit_or_offset_or_orders?(object) || has_join_sources?(object))
8+
stmt = super
9+
10+
model = MultiTenant.multi_tenant_model_for_table(MultiTenant::TableNode.table_name(object.relation.left))
11+
if model.present? && !MultiTenant.with_write_only_mode_enabled? && MultiTenant.current_tenant_id.present?
12+
stmt.wheres << MultiTenant::TenantEnforcementClause.new(model.arel_table[model.partition_key])
13+
end
14+
15+
stmt
16+
else
17+
super
6018
end
6119
end
62-
63-
# Clone the query, clear its projections, and set its projection to the primary key of the table
64-
subquery = arel.clone
65-
subquery.projections.clear
66-
subquery = subquery.project(table[primary_key])
67-
68-
# Create an IN condition node with the primary key of the table and the subquery
69-
Arel::Nodes::In.new(table[primary_key], subquery.ast)
20+
alias prepare_delete_statement prepare_update_statement
7021
end
7122
end
7223
end
7324

74-
# Patch ActiveRecord::Relation with the extension module
75-
ActiveRecord::Relation.prepend(Arel::ActiveRecordRelationExtension)
25+
Arel::Visitors::ToSql.prepend(Arel::Visitors::ToSqlPatch)

spec/activerecord-multi-tenant/query_rewriter_spec.rb

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@
4545

4646
it 'update_all the records with expected query' do
4747
expected_query = <<-SQL.strip
48-
UPDATE "projects" SET "name" = 'New Name' WHERE "projects"."id" IN
48+
UPDATE "projects" SET "name" = 'New Name' WHERE ("projects"."id") IN
4949
(SELECT "projects"."id" FROM "projects"
5050
INNER JOIN "managers" ON "managers"."project_id" = "projects"."id"
5151
and "managers"."account_id" = :account_id
@@ -63,7 +63,9 @@
6363
@queries.each do |actual_query|
6464
next unless actual_query.include?('UPDATE "projects" SET "name"')
6565

66-
expect(format_sql(actual_query)).to eq(format_sql(expected_query.gsub(':account_id', account.id.to_s)))
66+
expect(format_sql(actual_query.gsub('$1', "'New Name'"))).to eq(format_sql(expected_query.gsub(
67+
':account_id', account.id.to_s
68+
)))
6769
end
6870
end
6971

@@ -79,7 +81,7 @@
7981
SET
8082
"name" = 'New Name'
8183
WHERE
82-
"projects"."id" IN (
84+
("projects"."id") IN (
8385
SELECT
8486
"projects"."id"
8587
FROM
@@ -99,8 +101,9 @@
99101
@queries.each do |actual_query|
100102
next unless actual_query.include?('UPDATE "projects" SET "name"')
101103

102-
expect(format_sql(actual_query.gsub('$1',
103-
limit.to_s)).strip).to eq(format_sql(expected_query).strip)
104+
expect(
105+
format_sql(actual_query.gsub('$1', "'#{new_name}'").gsub('$2', limit.to_s)).strip
106+
).to eq(format_sql(expected_query).strip)
104107
end
105108
end
106109
end
@@ -126,7 +129,7 @@
126129

127130
it 'delete_all the records' do
128131
expected_query = <<-SQL.strip
129-
DELETE FROM "projects" WHERE "projects"."id" IN
132+
DELETE FROM "projects" WHERE ("projects"."id") IN
130133
(SELECT "projects"."id" FROM "projects"
131134
INNER JOIN "managers" ON "managers"."project_id" = "projects"."id"
132135
and "managers"."account_id" = :account_id
@@ -179,7 +182,7 @@
179182
DELETE FROM
180183
"projects"
181184
WHERE
182-
"projects"."id" IN (
185+
("projects"."id") IN (
183186
SELECT
184187
"projects"."id"
185188
FROM

0 commit comments

Comments
 (0)