pg_rls 0.1.2 → 0.1.4

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 5296cf19938df79d340f6da931acf0aee13e967720d964ec979773189634c155
-  data.tar.gz: ebf0278848a707e692198346dc734e8da9024c3b5e1325da760c4d7f7cdc30f2
+  metadata.gz: 1bead0b6e11e78f3aebe2c843c8cceaa178664bad5755f4f6858893f0232979c
+  data.tar.gz: c91f321f93c5d79aebccf09413f98580cc64104a32e350f874dce1c170c1c6a8
 SHA512:
-  metadata.gz: 97465ed5ea998a4c6e6ae39aae82d08ae0663733ef0f442d4716431adb607f13b96fb44139d5c5be83696cbb044cfa778a918a74aa842906ae798869f0fb3145
-  data.tar.gz: c85948bfbcd16f610b2095f2900fee54d3abc3cd27b581eb282144e4b444341a0cfb50c2b1ef98e8cd9d395c22e6218d81fa19d2281df956bb3d8e5cf4e39204
+  metadata.gz: 3baca7c6f170a2ffe07538413374230163d6a907372186bfe3bbab61b5fe8f7eb32402bd4781ef6f0a35b0f8f80e6b6dd2df1115f807810955a606b899f13eb9
+  data.tar.gz: '00866118e32d11bacc622119a0bbd02880c1f231b59fed3f65b58b176b29b96f0a38469b4071ebc3f4d32c3261a5090153683bc48512a556f48f4edbf72d09b3'

data/.rubocop.yml

@@ -24,3 +24,6 @@ Security/MarshalLoad:
 Metrics/ModuleLength:
   Exclude:
   - 'lib/pg_rls.rb'
+Naming/AccessorMethodName:
+  Exclude:
+  - 'lib/pg_rls/tenant.rb'

data/Gemfile.lock

@@ -1,8 +1,7 @@
 PATH
   remote: .
   specs:
-    pg_rls (0.1.2)
-      bundler (~> 2.2)
+    pg_rls (0.1.4)
 
 GEM
   remote: https://rubygems.org/
@@ -284,4 +283,4 @@ RUBY VERSION
    ruby 3.2.2p53
 
 BUNDLED WITH
-   2.3.15
+   2.4.22

data/lib/generators/pg_rls/install_generator.rb

@@ -16,6 +16,9 @@ module PgRls
         end
         super
       end
+      ENVIRONMENT_LINE = 'Rails.application.initialize!'
+      ENVIRONMENT_PATH = 'config/environment.rb'
+
       APPLICATION_LINE = 'class Application < Rails::Application'
       APPLICATION_PATH = 'config/application.rb'
 
@@ -50,6 +53,14 @@ module PgRls
         template 'pg_rls.rb.tt', 'config/initializers/pg_rls.rb'
       end
 
+      def inject_include_to_environment
+        return if environment_already_included?
+
+        gsub_file(ENVIRONMENT_PATH, /(#{Regexp.escape(ENVIRONMENT_LINE)})/mio) do |match|
+          "require_relative 'initializers/pg_rls'\n#{match}"
+        end
+      end
+
       def inject_include_to_application
         return if aplication_already_included?
 
@@ -74,6 +85,10 @@ module PgRls
         File.readlines(APPLICATION_PATH).grep(/config.active_record.schema_format = :sql/).any?
       end
 
+      def environment_already_included?
+        File.readlines(ENVIRONMENT_PATH).grep(%r{require_relative 'initializers/pg_rls'}).any?
+      end
+
       def initialize_error_text
         <<~ERROR
           TO DO

data/lib/pg_rls/database/admin_statements.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+module PgRls
+  module Admin
+    module ActiveRecord
+      module Migrator
+        def initialize(*args)
+          PgRls.instance_variable_set(:@as_db_admin, true)
+          super
+        end
+      end
+
+      module Tasks
+        module DatabaseTasks
+          def resolve_configuration(configuration)
+            PgRls.instance_variable_set(:@as_db_admin, true) unless PgRls.as_db_admin?
+            super
+          end
+
+          def migration_class
+            PgRls.instance_variable_set(:@as_db_admin, true) unless PgRls.as_db_admin?
+            super
+          end
+        end
+      end
+    end
+  end
+end

data/lib/pg_rls/database/configurations.rb

@@ -13,6 +13,8 @@ module ActiveRecord
       end
 
       def configuration_hash
+        reset_pg_rls_configuration if db_changed?
+
         return admin_configuration_hash if PgRls.as_db_admin?
 
         rls_configuration_hash
@@ -28,6 +30,15 @@ module ActiveRecord
           config[:password] = PgRls.password
         end.freeze
       end
+
+      def db_changed?
+        admin_configuration_hash[:database] != @configuration_hash[:database]
+      end
+
+      def reset_pg_rls_configuration
+        @rls_configuration_hash = nil
+        @admin_configuration_hash = nil
+      end
     end
   end
 end

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

@@ -23,130 +23,19 @@ namespace :db do
   include PgRls::Schema::UpStatements
 
   override_task grant_usage: :load_config do
-    PgRls.admin_tasks_execute do
-      create_rls_user
-    end
-  end
-
-  override_task create: :load_config do
-    PgRls.admin_tasks_execute do
-      Rake::Task['db:create:original'].invoke
-    end
-  end
-
-  override_task drop: :load_config do
-    PgRls.admin_tasks_execute do
-      Rake::Task['db:drop:original'].invoke
-    end
-  end
-
-  override_task migrate: :load_config do
-    PgRls.admin_tasks_execute do
-      Rake::Task['db:migrate:original'].invoke
-    end
-  end
-
-  override_task rollback: :load_config do
-    PgRls.admin_tasks_execute do
-      Rake::Task['db:rollback:original'].invoke
-    end
-  end
-
-  override_task prepare: :load_config do
-    PgRls.admin_tasks_execute do
-      Rake::Task['db:prepare:original'].invoke
-    end
-  end
-
-  override_task setup: :load_config do
-    PgRls.admin_tasks_execute do
-      Rake::Task['db:setup:original'].invoke
-    end
-  end
-
-  override_task prepare: :load_config do
-    PgRls.admin_tasks_execute do
-      Rake::Task['db:reset:original'].invoke
-    end
-  end
-
-  override_task purge: :load_config do
-    PgRls.admin_tasks_execute do
-      Rake::Task['db:purge:original'].invoke
-    end
+    PgRls.instance_variable_set(:@as_db_admin, true)
+    create_rls_user
   end
 
   override_task abort_if_pending_migrations: :load_config do
-    PgRls.admin_tasks_execute do
-      Rake::Task['db:abort_if_pending_migrations:original'].invoke
-    end
+    PgRls.instance_variable_set(:@as_db_admin, true)
+    Rake::Task['db:abort_if_pending_migrations:original'].invoke
   end
 
   namespace :test do
     override_task grant_usage: :load_config do
-      PgRls.admin_tasks_execute do
-        create_rls_user
-      end
-    end
-
-    override_task create: :load_config do
-      PgRls.admin_tasks_execute do
-        Rake::Task['db:test:create:original'].invoke
-      end
-    end
-
-    override_task drop: :load_config do
-      PgRls.admin_tasks_execute do
-        Rake::Task['db:test:drop:original'].invoke
-      end
-    end
-
-    override_task prepare: :load_config do
-      PgRls.admin_tasks_execute do
-        Rake::Task['db:test:prepare:original'].invoke
-      end
-    end
-
-    override_task setup: :load_config do
-      PgRls.admin_tasks_execute do
-        Rake::Task['db:test:setup:original'].invoke
-      end
-    end
-
-    override_task purge: :load_config do
-      PgRls.admin_tasks_execute do
-        Rake::Task['db:test:purge:original'].invoke
-      end
-    end
-
-    override_task load_schema: :load_config do
-      PgRls.admin_tasks_execute do
-        Rake::Task['db:test:load_schema:original'].invoke
-      end
-    end
-  end
-
-  namespace :environment do
-    override_task set: :load_config do
-      PgRls.admin_tasks_execute do
-        Rake::Task['db:environment:set:original'].invoke
-      end
-    end
-  end
-
-  namespace :schema do
-    override_task load: :load_config 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
-      end
-    end
-
-    override_task dump: :load_config do
-      PgRls.admin_tasks_execute do
-        Rake::Task['db:schema:dump:original'].invoke
-      end
+      PgRls.instance_variable_set(:@as_db_admin, true)
+      create_rls_user
     end
   end
 end

data/lib/pg_rls/errors/index.rb

@@ -2,4 +2,3 @@
 
 require_relative 'rake_only_error'
 require_relative 'tenant_not_found'
-require_relative 'admin_username'

data/lib/pg_rls/multi_tenancy.rb

@@ -28,19 +28,5 @@ module PgRls
       session[:_tenant] = nil
       raise PgRls::Errors::TenantNotFound, 'No tenant was found'
     end
-
-    def switch_tenant_by_resource!(resource = nil)
-      Tenant.switch!(resource)
-      session[:_tenant] = resource
-    rescue PgRls::Errors::TenantNotFound
-      Tenant.switch(session[:_tenant])
-    rescue NoMethodError
-      session[:tenant] = nil
-      redirect_to '/'
-    end
-
-    def tenant_match_session_cookies?
-      session[:_tenant] == request.subdomain
-    end
   end
 end

data/lib/pg_rls/tenant.rb

@@ -4,15 +4,9 @@ module PgRls
   # Tenant Controller
   module Tenant
     class << self
-      attr_reader :tenant
-
       def switch(resource)
-        tenant = switch_tenant!(resource)
-
-        "RLS changed to '#{tenant.id}'"
-      rescue StandardError => e
-        Rails.logger.info('connection was not made')
-        Rails.logger.info(e)
+        switch!(resource)
+      rescue PgRls::Errors::TenantNotFound
         nil
       end
 
@@ -20,17 +14,19 @@ module PgRls
         tenant = switch_tenant!(resource)
 
         "RLS changed to '#{tenant.id}'"
-      rescue StandardError => e
+      rescue StandardError
         Rails.logger.info('connection was not made')
-        raise e
+        raise PgRls::Errors::TenantNotFound
       end
 
       def with_tenant!(resource)
-        tenant = switch_tenant!(resource)
+        PgRls.main_model.connection_pool.with_connection do
+          tenant = switch_tenant!(resource)
 
-        yield(tenant) if block_given?
-      ensure
-        reset_rls! unless PgRls.test_inline_tenant == true
+          yield(tenant).presence if block_given?
+        ensure
+          reset_rls! unless PgRls.test_inline_tenant == true
+        end
       end
 
       def fetch
@@ -39,19 +35,15 @@ module PgRls
         nil
       end
 
-      def tenant!
-        @tenant ||= PgRls.main_model.find_by!(
+      def fetch!
+        PgRls.main_model.find_by!(
           tenant_id: PgRls.connection_class.connection.execute(
             "SELECT current_setting('rls.tenant_id')"
           ).getvalue(0, 0)
         )
       end
-      alias fetch! tenant!
 
       def reset_rls!
-        return if @tenant.blank?
-
-        @tenant = nil
         PgRls.execute_rls_in_shards do |connection_class|
           connection_class.transaction do
             connection_class.connection.execute('RESET rls.tenant_id')
@@ -61,6 +53,15 @@ module PgRls
         nil
       end
 
+      def set_rls!(tenant_id)
+        PgRls.execute_rls_in_shards do |connection_class|
+          connection_class.transaction do
+            connection_class.connection.execute(format('SET rls.tenant_id = %s',
+                                                       connection_class.connection.quote(tenant_id)))
+          end
+        end
+      end
+
       private
 
       def switch_tenant!(resource)
@@ -68,31 +69,29 @@ module PgRls
         PgRls.main_model.ignored_columns = []
         # rubocop: enable Rails/IgnoredColumnsAssignment
 
-        find_tenant(resource)
+        tenant = find_tenant(resource)
 
-        PgRls.execute_rls_in_shards do |connection_class|
-          connection_class.transaction do
-            connection_class.connection.execute(format('SET rls.tenant_id = %s',
-                                                       connection_class.connection.quote(tenant.tenant_id)))
-          end
-        end
+        set_rls!(tenant.tenant_id)
 
         tenant
+      rescue NoMethodError
+        raise PgRls::Errors::TenantNotFound
       end
 
       def find_tenant(resource)
-        raise PgRls::Errors::AdminUsername if PgRls.admin_connection?
-
         reset_rls!
 
+        tenant = nil
+
         PgRls.search_methods.each do |method|
-          break if @tenant.present?
+          break if tenant.present?
 
-          @method = method
-          @tenant = find_tenant_by_method(resource, method)
+          tenant = find_tenant_by_method(resource, method)
         end
 
         raise PgRls::Errors::TenantNotFound if tenant.blank?
+
+        tenant
       end
 
       def find_tenant_by_method(resource, method)

data/lib/pg_rls/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module PgRls
-  VERSION = '0.1.2'
+  VERSION = '0.1.4'
 end

data/lib/pg_rls.rb

@@ -6,27 +6,24 @@ require_relative 'pg_rls/version'
 require_relative 'pg_rls/database/prepared'
 require_relative 'pg_rls/schema/statements'
 require_relative 'pg_rls/database/configurations'
+require_relative 'pg_rls/database/admin_statements'
 require_relative 'pg_rls/tenant'
 require_relative 'pg_rls/multi_tenancy'
 require_relative 'pg_rls/railtie' if defined?(Rails)
 require_relative 'pg_rls/errors/index'
 
+ActiveRecord::Migrator.prepend PgRls::Admin::ActiveRecord::Migrator
+ActiveRecord::Tasks::DatabaseTasks.prepend PgRls::Admin::ActiveRecord::Tasks::DatabaseTasks
+ActiveRecord::ConnectionAdapters::AbstractAdapter.include PgRls::Schema::Statements
 # PostgreSQL Row Level Security
 module PgRls
   class Error < StandardError; end
-  SECURE_USERNAME = 'app_user'
-
   class << self
     extend Forwardable
 
     WRITER_METHODS = %i[table_name class_name search_methods].freeze
-    READER_METHODS = %i[
-      connection_class execute table_name class_name search_methods
-    ].freeze
-    DELEGATORS_METHODS = %i[
-      connection_class execute table_name search_methods
-      class_name main_model
-    ].freeze
+    READER_METHODS = %i[connection_class execute table_name class_name search_methods].freeze
+    DELEGATORS_METHODS = %i[connection_class execute table_name search_methods class_name main_model].freeze
 
     attr_writer(*WRITER_METHODS)
     attr_reader(*READER_METHODS)
@@ -34,7 +31,6 @@ module PgRls
     def_delegators(*DELEGATORS_METHODS)
 
     def setup
-      ActiveRecord::ConnectionAdapters::AbstractAdapter.include PgRls::Schema::Statements
       ActiveRecord::Base.ignored_columns += %w[tenant_id]
 
       yield self
@@ -44,32 +40,11 @@ module PgRls
       @connection_class ||= ActiveRecord::Base
     end
 
-    def rake_tasks?
-      Rake.application.top_level_tasks.present? || ARGV.any? { |arg| arg =~ /rake|dsl/ }
-    rescue NoMethodError
-      false
-    end
-
-    def admin_tasks_execute
-      raise PgRls::Errors::RakeOnlyError unless rake_tasks?
-
-      self.as_db_admin = true
-
-      yield
-    ensure
-      self.as_db_admin = false
-    end
-
     def admin_execute(query = nil, &)
-      current_tenant = PgRls::Tenant.fetch
-      establish_new_connection!(admin: true)
-
-      return ensure_block_execution(&) if block_given?
-
-      execute(query)
+      current_tenant, reset_rls_connection = establish_admin_connection
+      execute_query_or_block(query, &)
     ensure
-      establish_new_connection!
-      PgRls::Tenant.switch(current_tenant) if current_tenant.present?
+      reset_connection_if_needed(current_tenant, reset_rls_connection)
     end
 
     def establish_new_connection!(admin: false)
@@ -142,9 +117,37 @@ module PgRls
       establish_new_connection!(admin: true) if reset_connection
     end
 
+    def establish_admin_connection
+      reset_rls_connection = false
+      current_tenant = nil
+
+      unless admin_connection?
+        reset_rls_connection = true
+        current_tenant = PgRls::Tenant.fetch
+        establish_new_connection!(admin: true)
+      end
+
+      [current_tenant, reset_rls_connection]
+    end
+
     def ensure_block_execution(*, **)
       yield(*, **).presence
     end
+
+    def execute_query_or_block(query = nil, &)
+      if block_given?
+        ensure_block_execution(&)
+      else
+        execute(query)
+      end
+    end
+
+    def reset_connection_if_needed(current_tenant, reset_rls_connection)
+      return unless reset_rls_connection
+
+      establish_new_connection!
+      PgRls::Tenant.switch(current_tenant) if current_tenant.present?
+    end
   end
 
   mattr_accessor :table_name

metadata

@@ -1,34 +1,20 @@
 --- !ruby/object:Gem::Specification
 name: pg_rls
 version: !ruby/object:Gem::Version
-  version: 0.1.2
+  version: 0.1.4
 platform: ruby
 authors:
 - Daniel Laloush
-autorequire: 
+autorequire:
 bindir: exe
 cert_chain: []
-date: 2023-10-14 00:00:00.000000000 Z
-dependencies:
-- !ruby/object:Gem::Dependency
-  name: bundler
-  requirement: !ruby/object:Gem::Requirement
-    requirements:
-    - - "~>"
-      - !ruby/object:Gem::Version
-        version: '2.2'
-  type: :runtime
-  prerelease: false
-  version_requirements: !ruby/object:Gem::Requirement
-    requirements:
-    - - "~>"
-      - !ruby/object:Gem::Version
-        version: '2.2'
+date: 2023-11-13 00:00:00.000000000 Z
+dependencies: []
 description: |2
       This gem will help you to integrate PostgreSQL RLS to help you develop a great multitenancy application
       checkout the repository at https://github.com/Dandush03/pg_rls
 email:
-- daniel.laloush@influitive.com
+- d.laloush@outlook.com
 executables: []
 extensions: []
 extra_rdoc_files: []
@@ -61,10 +47,10 @@ files:
 - lib/generators/templates/pg_rls.rb.tt
 - lib/pg_rls.rb
 - lib/pg_rls/Rakefile
+- lib/pg_rls/database/admin_statements.rb
 - lib/pg_rls/database/configurations.rb
 - lib/pg_rls/database/prepared.rb
 - lib/pg_rls/database/tasks/admin_database.rake
-- lib/pg_rls/errors/admin_username.rb
 - lib/pg_rls/errors/index.rb
 - lib/pg_rls/errors/rake_only_error.rb
 - lib/pg_rls/errors/tenant_not_found.rb
@@ -85,7 +71,7 @@ licenses:
 - MIT
 metadata:
   rubygems_mfa_required: 'true'
-post_install_message: 
+post_install_message:
 rdoc_options: []
 require_paths:
 - lib
@@ -100,8 +86,8 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.4.19
-signing_key: 
+rubygems_version: 3.4.22
+signing_key:
 specification_version: 4
 summary: Write a short summary, because RubyGems requires one.
 test_files: []

data/lib/pg_rls/errors/admin_username.rb

@@ -1,12 +0,0 @@
-# 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