pg_rls 0.0.2.6.11 → 0.1.0

data/lib/generators/pg_rls/active_record/active_record_generator.rb

@@ -20,11 +20,11 @@ module PgRls
       end
 
       def create_tenant_migration_file
-        if creating?
-          migration_template(create_migration_template_path,
-                             "#{migration_path}/#{create_file_sub_name}_#{table_name}.rb",
-                             migration_version:)
-        end
+        return unless creating?
+
+        migration_template(create_migration_template_path,
+                           "#{migration_path}/#{create_file_sub_name}_#{table_name}.rb",
+                           migration_version:)
       end
 
       def convert_tenant_migration_file
@@ -34,11 +34,11 @@ module PgRls
                              migration_version:)
         end
 
-        if installation_in_progress?
-          migration_template('convert_migration_backport.rb.tt',
-                             "#{migration_path}/pg_rls_backport_#{table_name}.rb",
-                             migration_version:)
-        end
+        return unless installation_in_progress?
+
+        migration_template('convert_migration_backport.rb.tt',
+                           "#{migration_path}/pg_rls_backport_#{table_name}.rb",
+                           migration_version:)
       end
 
       def create_model_file

data/lib/generators/pg_rls/active_record/templates/convert_migration_backport.rb.tt

@@ -3,7 +3,7 @@
 class PgRlsBackport<%= table_name.camelize %> < ActiveRecord::Migration<%= migration_version %>
   def up
     # Suggested Code:
-    # PgRls.all_tenants do |tenant|
+    # PgRls.on_each_tenant do |tenant|
     #   tenant.<%= table_name %>.in_batches(of: 100) do |<%= table_name %>|
     #     <%= table_name %>.each { |<%= table_name.singularize %>| <%= table_name.singularize %>.update_attribute('tenant_id', tenant.tenant_id) }
     #   end

data/lib/generators/pg_rls/install_generator.rb

@@ -32,13 +32,13 @@ module PgRls
       hook_for :orm, required: true
 
       def orm_error_message
-        <<-ERROR.strip_heredoc
-        An ORM must be set to install PgRls in your application.
-        Be sure to have an ORM like Active Record or loaded in your
-        app or configure your own at `config/application.rb`.
-          config.generators do |g|
-            g.orm :your_orm_gem
-          end
+        <<~ERROR
+          An ORM must be set to install PgRls in your application.
+          Be sure to have an ORM like Active Record or loaded in your
+          app or configure your own at `config/application.rb`.
+            config.generators do |g|
+              g.orm :your_orm_gem
+            end
         ERROR
       end
 
@@ -46,7 +46,6 @@ module PgRls
         raise MissingORMError, orm_error_message unless options[:orm]
 
         inject_include_to_application
-        inject_include_to_application_record
         inject_include_to_application_controller
         template 'pg_rls.rb.tt', 'config/initializers/pg_rls.rb'
       end
@@ -54,23 +53,15 @@ module PgRls
       def inject_include_to_application
         return if aplication_already_included?
 
-        gsub_file(APPLICATION_PATH, /(#{Regexp.escape(APPLICATION_LINE)})/mi) do |match|
+        gsub_file(APPLICATION_PATH, /(#{Regexp.escape(APPLICATION_LINE)})/mio) do |match|
           "#{match}\n  config.active_record.schema_format = :sql\n"
         end
       end
 
-      def inject_include_to_application_record
-        return if aplication_record_already_included?
-
-        gsub_file(APPLICATION_RECORD_PATH, /(#{Regexp.escape(APPLICATION_RECORD_LINE)})/mi) do |match|
-          "#{match}\n  include PgRls::SecureConnection\n"
-        end
-      end
-
       def inject_include_to_application_controller
         return if aplication_controller_already_included?
 
-        gsub_file(APPLICATION_CONTROLLER_PATH, /(#{Regexp.escape(APPLICATION_CONTROLLER_LINE)})/mi) do |match|
+        gsub_file(APPLICATION_CONTROLLER_PATH, /(#{Regexp.escape(APPLICATION_CONTROLLER_LINE)})/mio) do |match|
           "#{match}\n  include PgRls::MultiTenancy\n"
         end
       end
@@ -79,16 +70,12 @@ module PgRls
         File.readlines(APPLICATION_CONTROLLER_PATH).grep(/include PgRls::MultiTenancy/).any?
       end
 
-      def aplication_record_already_included?
-        File.readlines(APPLICATION_RECORD_PATH).grep(/include PgRls::SecureConnection/).any?
-      end
-
       def aplication_already_included?
         File.readlines(APPLICATION_PATH).grep(/config.active_record.schema_format = :sql/).any?
       end
 
       def initialize_error_text
-        <<-ERROR.strip_heredoc
+        <<~ERROR
           TO DO
         ERROR
       end

data/lib/generators/templates/README

@@ -2,12 +2,10 @@ README
 ===============================================================================
 WARNING!!
 
-PgRls required that ActiveRecord format migration as SQL 
+PgRls required that ActiveRecord format migration as SQL
 
 Once you remove a tenant all of his data would be removed as well
 
-PgRls::SecureConnection was included to your ApplicationRecord do not remove it
-
 If you're setting a custom user, make sure to regenerate the structure.sql since
 Postgresql policies are created on each user
 

data/lib/generators/templates/pg_rls.rb.tt

@@ -20,25 +20,18 @@ PgRls.setup do |config|
   # config.username = Rails.application.credentials.dig(:database, :username)
   # config.password = Rails.application.credentials.dig(:database, :password)
 
+  ## ------------------------------ Middleware SetResetConnection -----------------------------
+  ## Uncomment this lines if you're using SetResetConnection Middleware
+  #
+  # config.session_store_server = Rails.application.config_for(:redis).session
+  #
+  ## Uncomment this line if you're not using warden as your authentication system or if you
+  ## changed the default warden key. Devise, uses warden authentication.
+  # config.session_store_default_warden_key = '2'
+  #
+  ## Uncomment this line if you're setting a diferent session key than stablished under your
+  ## redis server configuration
+  # config.session_key_prefix = '_hub_session'
   ##
-  ## Uncomment this lines in order to enable solo mode
-  ## Solo mode is made for API mode where we don't want to repeate the same
-  ## data structure across many project, Solo mode create a hidden tenant table
-  ## which is autopopulated on each request
-  # config.solo_mode = true
-
-  ##
-  ## After installing the PgRls gem, you can add the `PgRls::Middleware::SetResetConnection`
-  ## middleware to your Rails application to secure all database connections using Row Level Security.
-  ## To add the middleware using the `middleware.use` method, add the following
-  ## lines to your `config/application.rb` file:
-  ## require 'pg_rls/middleware/set_reset_connection'
-  ## config.middleware.use PgRls::Middleware::SetResetConnection
-  ## Note: Be sure to add the `PgRls::Middleware::SetResetConnection` middleware after any
-  ## middleware that sets up the Rails session, since the RLS middleware depends
-  ## on the presence of the session to work correctly.
-  ## Additionally, you will need to manually require the `PgRls::Middleware::SetResetConnection`
-  ## file in your application, as shown in the first line of the code snippet above.
-  # config.session_key = '_hub_sessions'
-  # config.session_prefix = '_session_id:2::'
+  ## ------------------------------ Middleware SetResetConnection -----------------------------
 end

data/lib/pg_rls/database/configurations.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+require 'active_record/database_configurations'
+
+module ActiveRecord
+  # ActiveRecord::DatabaseConfigurations
+  class DatabaseConfigurations
+    class HashConfig
+      def initialize(env_name, name, configuration_hash)
+        @env_name = env_name
+        @name = name
+        @configuration_hash = configuration_hash
+      end
+
+      def configuration_hash
+        return admin_configuration_hash if PgRls.as_db_admin?
+
+        rls_configuration_hash
+      end
+
+      def admin_configuration_hash
+        @admin_configuration_hash ||= @configuration_hash
+      end
+
+      def rls_configuration_hash
+        @rls_configuration_hash ||= @configuration_hash.deep_dup.tap do |config|
+          config[:username] = PgRls.username
+          config[:password] = PgRls.password
+        end.freeze
+      end
+    end
+  end
+end

data/lib/pg_rls/database/prepared.rb

@@ -6,7 +6,7 @@ module PgRls
     module Prepared
       class << self
         def grant_user_credentials(name: PgRls.username, schema: 'public')
-          PgRls.admin_execute <<-SQL
+          PgRls.admin_execute <<-SQL.squish
             DO
             $do$
             BEGIN

data/lib/pg_rls/database/tasks/admin_database.rake

@@ -22,128 +22,105 @@ end
 namespace :db do
   include PgRls::Schema::UpStatements
 
-  def admin_connection
-    PgRls.as_db_admin = true
-    PgRls.establish_new_connection
-
-    yield
-  ensure
-    PgRls.as_db_admin = false
-    ActiveRecord::Base.connection.disconnect!
-    PgRls.establish_new_connection
-  end
-
   override_task grant_usage: :load_config do
-    admin_connection do
+    PgRls.admin_tasks_execute do
       create_rls_user
     end
   end
 
   override_task create: :load_config do
-    admin_connection do
+    PgRls.admin_tasks_execute do
       Rake::Task['db:create:original'].invoke
     end
   end
 
   override_task drop: :load_config do
-    admin_connection do
+    PgRls.admin_tasks_execute do
       Rake::Task['db:drop:original'].invoke
     end
   end
 
   override_task migrate: :load_config do
-    admin_connection do
+    PgRls.admin_tasks_execute do
       Rake::Task['db:migrate:original'].invoke
     end
   end
 
   override_task rollback: :load_config do
-    admin_connection do
+    PgRls.admin_tasks_execute do
       Rake::Task['db:rollback:original'].invoke
     end
   end
 
   override_task prepare: :load_config do
-    admin_connection do
+    PgRls.admin_tasks_execute do
       Rake::Task['db:prepare:original'].invoke
     end
   end
 
   override_task setup: :load_config do
-    admin_connection do
+    PgRls.admin_tasks_execute do
       Rake::Task['db:setup:original'].invoke
     end
   end
 
   override_task prepare: :load_config do
-    admin_connection do
+    PgRls.admin_tasks_execute do
       Rake::Task['db:reset:original'].invoke
     end
   end
 
   override_task purge: :load_config do
-    admin_connection do
+    PgRls.admin_tasks_execute do
       Rake::Task['db:purge:original'].invoke
     end
   end
 
   override_task abort_if_pending_migrations: :load_config do
-    admin_connection do
+    PgRls.admin_tasks_execute do
       Rake::Task['db:abort_if_pending_migrations:original'].invoke
     end
   end
 
   namespace :test do
-    def admin_connection_test_db
-      Rails.env = 'test'
-      PgRls.as_db_admin = true
-      PgRls.establish_new_connection
-
-      yield
-    ensure
-      PgRls.as_db_admin = false
-      ActiveRecord::Base.connection.disconnect!
-      PgRls.establish_new_connection
-    end
-
     override_task grant_usage: :load_config do
-      admin_connection_test_db do
+      PgRls.admin_tasks_execute do
         create_rls_user
       end
     end
 
     override_task create: :load_config do
-      admin_connection_test_db do
+      PgRls.admin_tasks_execute do
         Rake::Task['db:test:create:original'].invoke
       end
     end
 
     override_task drop: :load_config do
-      admin_connection_test_db do
+      PgRls.admin_tasks_execute do
         Rake::Task['db:test:drop:original'].invoke
       end
     end
 
     override_task prepare: :load_config do
-      admin_connection_test_db do
+      PgRls.admin_tasks_execute do
         Rake::Task['db:test:prepare:original'].invoke
       end
     end
 
     override_task setup: :load_config do
-      admin_connection_test_db do
+      PgRls.admin_tasks_execute do
         Rake::Task['db:test:setup:original'].invoke
       end
     end
 
     override_task purge: :load_config do
-      admin_connection_test_db do
+      PgRls.admin_tasks_execute do
         Rake::Task['db:test:purge:original'].invoke
       end
     end
 
     override_task load_schema: :load_config do
-      admin_connection_test_db do
+      PgRls.admin_tasks_execute do
         Rake::Task['db:test:load_schema:original'].invoke
       end
     end
@@ -151,7 +128,7 @@ namespace :db do
 
   namespace :environment do
     override_task set: :load_config do
-      admin_connection do
+      PgRls.admin_tasks_execute do
         Rake::Task['db:environment:set:original'].invoke
       end
     end
@@ -159,7 +136,7 @@ namespace :db do
 
   namespace :schema do
     override_task load: :load_config do
-      admin_connection do
+      PgRls.admin_tasks_execute do
         Rake::Task['db:schema:load:original'].invoke
         Rake::Task['db:grant_usage'].invoke
         Rake::Task['db:test:grant_usage'].invoke
@@ -167,7 +144,7 @@ namespace :db do
     end
 
     override_task dump: :load_config do
-      admin_connection do
+      PgRls.admin_tasks_execute do
         Rake::Task['db:schema:dump:original'].invoke
       end
     end

data/lib/pg_rls/errors/admin_username.rb

@@ -0,0 +1,12 @@
+# frozen_string_literal: true
+
+module PgRls
+  module Errors
+    class AdminUsername < StandardError
+      def initialize(msg = nil)
+        msg ||= 'Cannot set or reset tenant for admin user'
+        super(msg)
+      end
+    end
+  end
+end

data/lib/pg_rls/errors/index.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+require_relative 'rake_only_error'
+require_relative 'tenant_not_found'
+require_relative 'admin_username'

data/lib/pg_rls/errors/rake_only_error.rb

@@ -0,0 +1,12 @@
+# frozen_string_literal: true
+
+module PgRls
+  module Errors
+    class RakeOnlyError < StandardError
+      def initialize(msg = nil)
+        msg ||= 'This method can only be executed through rake tasks'
+        super(msg)
+      end
+    end
+  end
+end

data/lib/pg_rls/errors/tenant_not_found.rb

@@ -5,18 +5,10 @@ module PgRls
     # Raise Tenant Not found and ensure that the tenant is resetted
     class TenantNotFound < StandardError
       def initialize(msg = nil)
-        reset_tenant_id
-        @msg = msg
+        PgRls::Tenant.reset_rls!
+        msg ||= "Tenant Doesn't exist"
         super(msg)
       end
-
-      def message
-        @msg || "Tenant Doesn't exist"
-      end
-
-      def reset_tenant_id
-        PgRls.connection_class.connection.execute('RESET rls.tenant_id')
-      end
     end
   end
 end

data/lib/pg_rls/middleware/set_reset_connection.rb

@@ -1,5 +1,43 @@
 # frozen_string_literal: true
 
+module PgRls
+  class FrozenConfiguration < StandardError; end
+
+  def self.sessions
+    @sessions ||= Redis.new(@session_store_server)
+  end
+
+  def self.session_prefix
+    @session_prefix ||= begin
+      store_default_warden_key = @session_store_default_warden_key || '2'
+
+      "#{session_key_prefix}:#{store_default_warden_key}::"
+    end
+  end
+
+  def self.session_store_server=(opts = {})
+    raise Errors::FrozenConfiguration unless @sessions.nil?
+
+    @session_store_server = opts.deep_symbolize_keys
+  end
+
+  def self.session_store_default_warden_key=(val)
+    raise Errors::FrozenConfiguration unless @sessions.nil?
+
+    @session_store_default_warden_key = val
+  end
+
+  def self.session_key_prefix
+    @session_key_prefix ||= @session_key_prefix || @session_store_server[:key_prefix]
+  end
+
+  def self.session_key_prefix=(val)
+    raise Errors::FrozenConfiguration unless @sessions.nil?
+
+    @session_key_prefix = val
+  end
+end
+
 module PgRls
   module Middleware
     # Set RLS if sessions present.
@@ -26,7 +64,7 @@ module PgRls
         cookie_string = env['HTTP_COOKIE']
         return if cookie_string.nil?
 
-        cookie_regex = /#{PgRls.session_key}=([^;]+)/
+        cookie_regex = /#{PgRls.session_key_prefix}=([^;]+)/
         match = cookie_regex.match(cookie_string)
         match[1] if match
       end
@@ -36,8 +74,15 @@ module PgRls
 
         return if cookie.blank?
 
-        sessions = Rails.cache.read("#{PgRls.session_prefix}#{Digest::SHA256.hexdigest(cookie)}")
-        sessions['_tenant'] if sessions.present?
+        redis_session_key = "#{PgRls.session_prefix}#{Digest::SHA256.hexdigest(cookie)}"
+        tenant_session = Marshal.load(PgRls.sessions.get(redis_session_key))
+
+        return if tenant_session.blank?
+        return if tenant_session['_tenant'].blank?
+
+        tenant_session['_tenant'] if tenant_session.present?
+      rescue TypeError
+        nil
       end
 
       def rails_active_storage_request?(env)

data/lib/pg_rls/middleware/sidekiq/client.rb

@@ -7,9 +7,18 @@ module PgRls
       # Set PgRls Policies
       class Client
         def call(_job_class, msg, _queue, _redis_pool)
-          msg['pg_rls'] = PgRls::Tenant.fetch.id
+          load_tenant_attribute!(msg)
           yield
         end
+
+        def load_tenant_attribute!(msg)
+          if PgRls.username == PgRls.current_db_username
+            tenant = PgRls::Tenant.fetch!
+            msg['pg_rls'] = tenant.id
+          else
+            msg['admin'] = true
+          end
+        end
       end
     end
   end

data/lib/pg_rls/middleware/sidekiq/server.rb

@@ -7,7 +7,11 @@ module PgRls
       # Set PgRls Policies
       class Server
         def call(_job_instance, msg, _queue, &)
-          PgRls::Tenant.with_tenant!(msg['pg_rls'], &)
+          if msg['admin']
+            PgRls.admin_execute(&)
+          else
+            PgRls::Tenant.with_tenant!(msg['pg_rls'], &)
+          end
         end
       end
     end

data/lib/pg_rls/multi_tenancy.rb

@@ -16,13 +16,13 @@ module PgRls
 
     private
 
-    def switch_tenant!(&block)
+    def switch_tenant!
       fetched_tenant = session[:_tenant] || current_tenant
-      return block.call if PgRls::Tenant.fetch.present?
+      return yield if PgRls::Tenant.fetch.present?
 
       Tenant.with_tenant!(fetched_tenant) do |tenant|
         session[:_tenant] = tenant
-        block.call(tenant)
+        yield(tenant)
       end
     rescue NoMethodError
       session[:_tenant] = nil
@@ -32,7 +32,7 @@ module PgRls
     def switch_tenant_by_resource!(resource = nil)
       Tenant.switch!(resource)
       session[:_tenant] = resource
-    rescue PgRls::Errors::TenantNotFound, ActiveRecord::RecordNotFound
+    rescue PgRls::Errors::TenantNotFound
       Tenant.switch(session[:_tenant])
     rescue NoMethodError
       session[:tenant] = nil

data/lib/pg_rls/schema/down_statements.rb

@@ -5,7 +5,7 @@ module PgRls
     # Down Schema Statements
     module DownStatements
       def drop_rls_user
-        ActiveRecord::Migration.execute <<~SQL
+        ActiveRecord::Migration.execute <<~SQL.squish
           DROP OWNED BY #{PgRls.username};
           REVOKE ALL PRIVILEGES ON ALL TABLES IN SCHEMA public FROM #{PgRls.username};
           REVOKE ALL PRIVILEGES ON ALL SEQUENCES IN SCHEMA public FROM #{PgRls.username};
@@ -23,28 +23,28 @@ module PgRls
       end
 
       def detach_blocking_function(table_name)
-        ActiveRecord::Migration.execute <<-SQL
+        ActiveRecord::Migration.execute <<-SQL.squish
           DROP TRIGGER IF EXISTS id_safe_guard
             ON #{table_name};
         SQL
       end
 
       def detach_trigger_function(table_name)
-        ActiveRecord::Migration.execute <<-SQL
+        ActiveRecord::Migration.execute <<-SQL.squish
           DROP TRIGGER IF EXISTS tenant_id_setter
             ON #{table_name};
         SQL
       end
 
       def drop_rls_column(table_name)
-        ActiveRecord::Migration.execute <<-SQL
+        ActiveRecord::Migration.execute <<-SQL.squish
           ALTER TABLE #{table_name}
             DROP COLUMN IF EXISTS tenant_id;
         SQL
       end
 
       def drop_rls_policy(table_name)
-        ActiveRecord::Migration.execute <<-SQL
+        ActiveRecord::Migration.execute <<-SQL.squish
           DROP POLICY #{table_name}_#{PgRls.username} ON #{table_name};
           ALTER TABLE #{table_name} DISABLE ROW LEVEL SECURITY;
         SQL

data/lib/pg_rls/schema/statements.rb

@@ -10,17 +10,17 @@ module PgRls
       include UpStatements
       include DownStatements
 
-      def create_rls_tenant_table(table_name, **options, &)
+      def create_rls_tenant_table(table_name, **, &)
         create_rls_user
         create_rls_setter_function
         create_rls_blocking_function
-        create_table(table_name, **options, &)
+        create_table(table_name, **, &)
         add_rls_column_to_tenant_table(table_name)
         append_blocking_function(table_name)
       end
 
-      def create_rls_table(table_name, **options, &)
-        create_table(table_name, **options, &)
+      def create_rls_table(table_name, **, &)
+        create_table(table_name, **, &)
         add_rls_column(table_name)
         create_rls_policy(table_name)
         append_trigger_function(table_name)