pg_rls 0.1.3 → 0.1.4

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 0e0133ebb0867d0bfeeea122f6ddef47f861c6b17e2b8b65a07acdca4768c66a
-  data.tar.gz: 5bd700b9dd72150b01e72086f07795772cbb3f3b512e7106dbf58b2850713414
+  metadata.gz: 1bead0b6e11e78f3aebe2c843c8cceaa178664bad5755f4f6858893f0232979c
+  data.tar.gz: c91f321f93c5d79aebccf09413f98580cc64104a32e350f874dce1c170c1c6a8
 SHA512:
-  metadata.gz: af0c6529cd940354cbaf1432c7086e1ae336f0d160df036108820873f31699c580833b7b153956e1b599b26946de15e8bfd5eda2dc94e10fa8b0b757fc94bcf1
-  data.tar.gz: 04feda60ca3f7fa0e98f096246db9ba2d4eecac158af08f719631e7f078b224e5dc33b660e92129a6eb70e217bb7e550fdf74843015b1df988272954d756192c
+  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.1)
-      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,7 @@ module PgRls
         end
         super
       end
+      ENVIRONMENT_LINE = 'Rails.application.initialize!'
       ENVIRONMENT_PATH = 'config/environment.rb'
 
       APPLICATION_LINE = 'class Application < Rails::Application'
@@ -55,7 +56,9 @@ module PgRls
       def inject_include_to_environment
         return if environment_already_included?
 
-        prepend_to_file(ENVIRONMENT_PATH, "\nrequire_relative 'initializers/pg_rls'\n")
+        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
@@ -83,7 +86,7 @@ module PgRls
       end
 
       def environment_already_included?
-        File.readlines(ENVIRONMENT_PATH).grep(/require_relative 'initializers\/pg_rls'/).any?
+        File.readlines(ENVIRONMENT_PATH).grep(%r{require_relative 'initializers/pg_rls'}).any?
       end
 
       def initialize_error_text

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/tenant.rb

@@ -14,12 +14,11 @@ module PgRls
         tenant = switch_tenant!(resource)
 
         "RLS changed to '#{tenant.id}'"
-      rescue StandardError => e
+      rescue StandardError
         Rails.logger.info('connection was not made')
         raise PgRls::Errors::TenantNotFound
       end
 
-
       def with_tenant!(resource)
         PgRls.main_model.connection_pool.with_connection do
           tenant = switch_tenant!(resource)
@@ -54,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)
@@ -63,12 +71,7 @@ module PgRls
 
         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
@@ -76,8 +79,6 @@ module PgRls
       end
 
       def find_tenant(resource)
-        raise PgRls::Errors::AdminUsername if PgRls.admin_connection?
-
         reset_rls!
 
         tenant = nil

data/lib/pg_rls/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module PgRls
-  VERSION = '0.1.3'
+  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.3
+  version: 0.1.4
 platform: ruby
 authors:
 - Daniel Laloush
 autorequire:
 bindir: exe
 cert_chain: []
-date: 2023-11-02 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
@@ -100,7 +86,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.3.7
+rubygems_version: 3.4.22
 signing_key:
 specification_version: 4
 summary: Write a short summary, because RubyGems requires one.

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