diff --git a/lib/rls_rails/helpers.rb b/lib/rls_rails/helpers.rb index 250c1ab..5625489 100644 --- a/lib/rls_rails/helpers.rb +++ b/lib/rls_rails/helpers.rb @@ -4,14 +4,12 @@ module RLS # This variable is very problematic and not relieable, since in a # threaded environment of a connection pool this module is changed and # race conditions can occur, de-syncing the module-status with the true db status. - @rls_status = {user_id: '', tenant_id: '', disabled: ''} - def self.disable! return if RLS.status[:disable] === 'true' # do not use disabled? here since it may be blank clear_query_cache execute_sql("SET SESSION rls.disable = TRUE;") - @rls_status.merge!(disabled: 'true') + thread_rls_status.merge!(disabled: 'true') debug_print "WARNING: ROW LEVEL SECURITY DISABLED!\n" end @@ -27,7 +25,7 @@ def self.enable! clear_query_cache debug_print "ROW LEVEL SECURITY ENABLED!\n" execute_sql("SET SESSION rls.disable = FALSE;") - @rls_status.merge!(disabled: 'false') + thread_rls_status.merge!(disabled: 'false') end def self.enabled? @@ -41,7 +39,7 @@ def self.set_tenant tenant clear_query_cache debug_print "Accessing database as #{tenant.name}\n" execute_sql "SET SESSION rls.disable = FALSE; SET SESSION rls.tenant_id = #{tenant.id};" - @rls_status.merge!(tenant_id: tenant.id.to_s) + thread_rls_status.merge!(tenant_id: tenant.id.to_s) end def self.set_user user @@ -51,7 +49,7 @@ def self.set_user user clear_query_cache debug_print "Accessing database as #{user.class}##{user.id}\n" execute_sql "SET SESSION rls.disable = FALSE; SET SESSION rls.user_id = #{user.id};" - @rls_status.merge!(user_id: user.id.to_s) + thread_rls_status.merge!(user_id: user.id.to_s) end def self.current_tenant_id @@ -71,7 +69,7 @@ def self.reset! RESET rls.disable; SQL clear_query_cache - @rls_status.merge!(tenant_id: '', user_id: '', disabled: '') + thread_rls_status.merge!(tenant_id: '', user_id: '', disabled: '') end # Sets the RLS status to the given value in one go. @@ -89,7 +87,7 @@ def self.status= status SET SESSION rls.user_id = '#{user_id}'; SET SESSION rls.tenant_id = '#{tenant_id}'; SQL - @rls_status.merge!(tenant_id: tenant_id, user_id: user_id, disabled: disable) + thread_rls_status.merge!(tenant_id: tenant_id, user_id: user_id, disabled: disable) end # @return [Hash] Values of the current RLS sesssion @@ -171,4 +169,8 @@ def self.execute_sql query def self.debug_print s print s if Railtie.config.rls_rails.verbose end -end \ No newline at end of file + + def self.thread_rls_status + Thread.current["rls_status"] ||= { user_id: '', tenant_id: '', disabled: '' } + end +end diff --git a/lib/rls_rails/railtie.rb b/lib/rls_rails/railtie.rb index 399d574..dc3e1b1 100644 --- a/lib/rls_rails/railtie.rb +++ b/lib/rls_rails/railtie.rb @@ -19,6 +19,20 @@ class Railtie < ::Rails::Railtie ActiveRecord::Migration.include RLS::Statements #ActiveRecord::Migration::CommandRecorder.include Scenic::CommandRecorder #ActiveRecord::SchemaDumper.prepend Scenic::SchemaDumper + + ActiveRecord::Base.connection.class.set_callback :checkout, :after do + # ensure the RLS-related session variables are reset when a + # thread checks out a connection + execute <<~SQL + RESET rls.user_id; + RESET rls.tenant_id; + RESET rls.disable; + SQL + + clear_query_cache + + RLS.thread_rls_status.merge!(tenant_id: '', user_id: '', disabled: '') + end end end