Module: PgRls::Tenant

Defined in:
lib/pg_rls/tenant.rb

Overview

Tenant Controller

Class Method Summary collapse

Class Method Details

.fetch ⇒ Object 



32
33
34
35
36
 
# File 'lib/pg_rls/tenant.rb', line 32

def fetch
  fetch!
rescue ActiveRecord::StatementInvalid, ActiveRecord::RecordNotFound
  nil
end

.fetch! ⇒ Object 



38
39
40
41
42
43
44
45
46
47
 
# File 'lib/pg_rls/tenant.rb', line 38

def fetch!
  PgRls::Current::Context.tenant ||= PgRls.main_model.connection_pool.with_connection do |connection|
    tenant_id = get_tenant_id(connection)
    if tenant_id.present?
      PgRls.main_model.find_by!(
        tenant_id:
      )
    end
  end
end

.get_tenant_id(connection) ⇒ Object

rubocop:disable Lint/RescueStandardError rubocop:disable Lint/UselessAssignment



51
52
53
54
55
 
# File 'lib/pg_rls/tenant.rb', line 51

def get_tenant_id(connection)
  connection.execute("SELECT current_setting('rls.tenant_id')").getvalue(0, 0)
rescue => e
  nil
end

.on_find_each(ids: [], scope: nil) ⇒ Object 



86
87
88
89
90
91
92
93
94
95
96
97
98
 
# File 'lib/pg_rls/tenant.rb', line 86

def on_find_each(ids: [], scope: nil, &)
  raise 'Invalid Scope' if scope.present? && PgRls.main_model != scope.klass

  result = []

  query = build_on_each_query(ids, scope)

  query.find_each do |tenant|
    result << { tenant_id: tenant.id, result: with_tenant!(tenant, &) }
  end

  result
end

.reset_rls! ⇒ Object

rubocop:enable Lint/RescueStandardError rubocop:enable Lint/UselessAssignment



59
60
61
62
63
64
65
66
67
68
69
70
 
# File 'lib/pg_rls/tenant.rb', line 59

def reset_rls!
  PgRls.execute_rls_in_shards do |connection_class|
    connection_class.connection_pool.with_connection do |connection|
      connection.transaction do
        connection.execute('RESET rls.tenant_id')
      end
    end
  end

  PgRls::Current::Context.clear_all
  nil
end

.set_rls!(tenant) ⇒ Object 



72
73
74
75
76
77
78
79
80
81
82
83
84
 
# File 'lib/pg_rls/tenant.rb', line 72

def set_rls!(tenant)
  tenant_id = tenant.tenant_id
  PgRls.execute_rls_in_shards do |connection_class|
    connection_class.connection_pool.with_connection do |connection|
      connection.transaction do
        connection.execute(format('SET rls.tenant_id = %s',
                                  connection.quote(tenant_id)))
      end
    end
  end
  PgRls::Current::Context.clear_all
  PgRls::Current::Context.tenant = tenant
end

.switch(resource) ⇒ Object 



7
8
9
10
11
 
# File 'lib/pg_rls/tenant.rb', line 7

def switch(resource)
  switch!(resource)
rescue PgRls::Errors::TenantNotFound
  nil
end

.switch!(resource) ⇒ Object 



13
14
15
16
17
18
19
20
 
# File 'lib/pg_rls/tenant.rb', line 13

def switch!(resource)
  tenant = switch_tenant!(resource)

  "RLS changed to '#{tenant.id}'"
rescue StandardError
  Rails.logger.info('connection was not made')
  raise PgRls::Errors::TenantNotFound
end

.with_tenant!(resource) ⇒ Object 



22
23
24
25
26
27
28
29
30
 
# File 'lib/pg_rls/tenant.rb', line 22

def with_tenant!(resource)
  PgRls.main_model.connection_pool.with_connection do
    tenant = switch_tenant!(resource)

    yield(tenant).presence if block_given?
  ensure
    reset_rls! unless PgRls.test_inline_tenant == true || PgRls::Current::Context.tenant.blank?
  end
end