|
1 | 1 | # frozen_string_literal: true |
2 | 2 |
|
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 |
60 | 18 | end |
61 | 19 | end |
62 | 20 |
|
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 | | - |
67 | | - if primary_key.is_a?(Array) |
68 | | - # For composite primary keys, project all primary key columns |
69 | | - primary_key_columns = primary_key.map { |pk| table[pk] } |
70 | | - subquery = subquery.project(*primary_key_columns) |
71 | | - |
72 | | - # Create IN condition using composite primary key columns |
73 | | - Arel::Nodes::In.new( |
74 | | - Arel::Nodes::Grouping.new(primary_key_columns), |
75 | | - subquery.ast |
76 | | - ) |
77 | | - else |
78 | | - subquery = subquery.project(table[primary_key]) |
79 | | - Arel::Nodes::In.new(table[primary_key], subquery.ast) |
80 | | - end |
| 21 | + alias prepare_delete_statement prepare_update_statement |
81 | 22 | end |
82 | 23 | end |
83 | 24 | end |
84 | 25 |
|
85 | | -# Patch ActiveRecord::Relation with the extension module |
86 | | -ActiveRecord::Relation.prepend(Arel::ActiveRecordRelationExtension) |
| 26 | +Arel::Visitors::ToSql.prepend(Arel::Visitors::ToSqlPatch) |
0 commit comments