pg_rls 0.1.10 → 0.2.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 492256f8c775835d6ed888cc64e93bff4c0a3f9583eb9c459e1ed9a5520d7fab
-  data.tar.gz: f05611afd74ad000316df9622663de94d8feda84538de06cf5c548e749a60302
+  metadata.gz: e81a178c9c3c7d4a42c2dff85faa720c8340a2c3373eab5c38edb50e862bdfb9
+  data.tar.gz: 8a8d78bd15baca991c1f50eb3adc956f5b467a640c6b6b629d8260e19b608ea4
 SHA512:
-  metadata.gz: cabd8e2479dbf5432a26b34680932a6ff2a4b1d2f6cbd6c1c369d8f4b954a69baa4cfefe8a81f86995b07f2a966a8129d664adf9115de23c0dbb3a91623cc52a
-  data.tar.gz: 4f66a3686927071d5aa5f66374a63450cda26c9be6ebc86d1cf085fc373f90033e2e7d2fe50f4b1945452088122df3884ca5cd97746c43905615858cccf5fed6
+  metadata.gz: c461c7b3263d1e936a04a0693509cef55d4a699169a8d4addd7e5818b02ef608cf1cadc39aee287b5b3d6cc11705e3bf10eefeafbc8d62e3db413b0010ebcb36
+  data.tar.gz: f216cec76a8d3a5703abd0913decc36ef1344d6327331555e6e97168c5b4cf7aeda46254cf04daaadcca159adf56a3747c54c95be8d16b1f32d5de8dc9e5bd31

data/Gemfile

@@ -17,4 +17,5 @@ gem 'rubocop-performance'
 gem 'rubocop-rails'
 gem 'rubocop-rake'
 gem 'rubocop-rspec'
+gem 'sidekiq'
 gem 'solargraph'

data/Gemfile.lock

@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    pg_rls (0.1.6)
+    pg_rls (0.1.10)
 
 GEM
   remote: https://rubygems.org/
@@ -109,6 +109,7 @@ GEM
     kramdown-parser-gfm (1.1.0)
       kramdown (~> 2.0)
     language_server-protocol (3.17.0.3)
+    logger (1.6.0)
     loofah (2.21.4)
       crass (~> 1.0.2)
       nokogiri (>= 1.12.0)
@@ -184,6 +185,8 @@ GEM
     rbs (2.8.4)
     rdoc (6.5.0)
       psych (>= 4.0.0)
+    redis-client (0.22.2)
+      connection_pool
     regexp_parser (2.8.2)
     reline (0.3.9)
       io-console (~> 0.5)
@@ -236,6 +239,12 @@ GEM
       rubocop-factory_bot (~> 2.22)
     ruby-progressbar (1.13.0)
     ruby2_keywords (0.0.5)
+    sidekiq (7.3.1)
+      concurrent-ruby (< 2)
+      connection_pool (>= 2.3.0)
+      logger
+      rack (>= 2.2.4)
+      redis-client (>= 0.22.2)
     solargraph (0.49.0)
       backport (~> 1.2)
       benchmark
@@ -268,6 +277,7 @@ GEM
 
 PLATFORMS
   arm64-darwin-22
+  arm64-darwin-23
   x86_64-linux
 
 DEPENDENCIES
@@ -280,6 +290,7 @@ DEPENDENCIES
   rubocop-rails
   rubocop-rake
   rubocop-rspec
+  sidekiq
   solargraph
 
 RUBY VERSION

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

@@ -2,9 +2,6 @@
 
 <% module_namespacing do -%>
 class <%= PgRls.class_name.camelize %> < <%= parent_class_name.classify %>
-  def self.current
-    PgRls::Tenant.fetch
-  end
 <% attributes.select(&:reference?).each do |attribute| -%>
   belongs_to :<%= attribute.name %><%= ", polymorphic: true" if attribute.polymorphic? %>
 <% end -%>

data/lib/generators/pg_rls/install_generator.rb

@@ -48,7 +48,6 @@ module PgRls
       def copy_initializer
         raise MissingORMError, orm_error_message unless options[:orm]
 
-        inject_include_to_application
         inject_include_to_application_controller
         template 'pg_rls.rb.tt', 'config/initializers/pg_rls.rb'
       end
@@ -61,14 +60,6 @@ module PgRls
         end
       end
 
-      def inject_include_to_application
-        return if aplication_already_included?
-
-        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_controller
         return if aplication_controller_already_included?
 
@@ -81,10 +72,6 @@ module PgRls
         File.readlines(APPLICATION_CONTROLLER_PATH).grep(/include PgRls::MultiTenancy/).any?
       end
 
-      def aplication_already_included?
-        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

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

@@ -9,6 +9,17 @@ PgRls.setup do |config|
   config.class_name = :<%= PgRls.class_name %>
   config.table_name = :<%= PgRls.table_name %>
   config.search_methods = <%= PgRls.search_methods %>
+  # If you are using `solid_queue`, `solid_cache`, or `solid_cable` with a sharding configuration,
+  # we recommend excluding these shards from Row-Level Security (RLS) to avoid the need to reset
+  # RLS on each shard.
+  #
+  # By default, RLS will be enabled for all shards.
+  # You can specify which shards to exclude from RLS using the `config.excluded_shards` option:
+  #
+  #   config.excluded_shards = []
+  #
+  # Note: While it's technically possible to leave `solid_cache` and `solid_cable` under RLS,
+  # it is generally unnecessary and may introduce complexity without added benefit.
 
   ##
   ## Uncomment this lines if you have a custome user per environment

data/lib/pg_rls/current/context.rb

@@ -0,0 +1,10 @@
+# frozen_string_literal: true
+
+module PgRls
+  # Current Context
+  module Current
+    class Context < ActiveSupport::CurrentAttributes
+      attribute :tenant
+    end
+  end
+end

data/lib/pg_rls/errors/rake_only_error.rb

@@ -5,7 +5,7 @@ module PgRls
     class RakeOnlyError < StandardError
       def initialize(msg = nil)
         msg ||= 'This method can only be executed through rake tasks'
-        super(msg)
+        super
       end
     end
   end

data/lib/pg_rls/errors/tenant_not_found.rb

@@ -5,9 +5,8 @@ module PgRls
     # Raise Tenant Not found and ensure that the tenant is resetted
     class TenantNotFound < StandardError
       def initialize(msg = nil)
-        PgRls::Tenant.reset_rls!
         msg ||= "Tenant Doesn't exist"
-        super(msg)
+        super
       end
     end
   end

data/lib/pg_rls/logger.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+require 'logger'
+
+module PgRls
+  class Logger
+    def initialize
+      @logger = ::Logger.new($stdout) # You can output to a file if needed
+      @logger.level = ::Logger::DEBUG
+    end
+
+    def log(message, level: :info)
+      case level
+      when :debug
+        @logger.debug(message)
+      when :info
+        @logger.info(message)
+      when :warn
+        @logger.warn(message)
+      when :error
+        @logger.error(message)
+      else
+        @logger.info(message)
+      end
+    end
+
+    def deprecation_warning(message)
+      log("[DEPRECATION WARNING]: #{message}", level: :warn)
+    end
+  end
+end

data/lib/pg_rls/schema/dumper.rb

@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+module PgRls
+  module Schema
+    module Dumper
+      def table(table, stream)
+        temp_stream = StringIO.new
+        super(table, temp_stream)
+        temp_stream_string = temp_stream.string
+        if rls_tenant_table?(table)
+          temp_stream_string.gsub!('create_table', 'create_rls_tenant_table')
+        elsif rls_table?(table)
+          temp_stream_string.gsub!('create_table', 'create_rls_table')
+        end
+
+        stream.print(temp_stream_string)
+      end
+
+      private
+
+      def rls_table?(table_name)
+        # Logic to determine if the table uses RLS
+        # You can check if the table has RLS policies or use a naming convention
+        @connection.execute(<<-SQL.squish).any?
+          SELECT 1 FROM pg_policies WHERE tablename = #{ActiveRecord::Base.connection.quote(table_name)};
+        SQL
+      end
+
+      def rls_tenant_table?(table_name)
+        # Logic to determine if the table is a tenant table
+        # You can check if the table has a specific column or use a naming convention
+        PgRls.table_name.to_s == table_name
+      end
+    end
+  end
+end

data/lib/pg_rls/schema/up_statements.rb

@@ -15,7 +15,6 @@ module PgRls
 
               CREATE USER #{name} WITH PASSWORD '#{password}';
             END IF;
-            GRANT ALL PRIVILEGES ON TABLE schema_migrations TO #{name};
             GRANT USAGE ON SCHEMA #{schema} TO #{name};
             ALTER DEFAULT PRIVILEGES IN SCHEMA #{schema}
               GRANT USAGE, SELECT

data/lib/pg_rls/tenant.rb

@@ -25,7 +25,7 @@ module PgRls
 
           yield(tenant).presence if block_given?
         ensure
-          reset_rls! unless PgRls.test_inline_tenant == true
+          reset_rls! unless PgRls.test_inline_tenant == true || PgRls::Current::Context.tenant.blank?
         end
       end
 
@@ -36,47 +36,92 @@ module PgRls
       end
 
       def fetch!
-        PgRls.main_model.find_by!(
-          tenant_id: PgRls.connection_class.connection.execute(
-            "SELECT current_setting('rls.tenant_id')"
-          ).getvalue(0, 0)
-        )
+        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
+
+      # rubocop:disable Lint/RescueStandardError
+      # rubocop:disable Lint/UselessAssignment
+      def get_tenant_id(connection)
+        connection.execute("SELECT current_setting('rls.tenant_id')").getvalue(0, 0)
+      rescue => e
+        nil
       end
+      # rubocop:enable Lint/RescueStandardError
+      # rubocop:enable Lint/UselessAssignment
 
       def reset_rls!
         PgRls.execute_rls_in_shards do |connection_class|
-          connection_class.transaction do
-            connection_class.connection.execute('RESET rls.tenant_id')
+          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
 
-      def set_rls!(tenant_id)
+      def set_rls!(tenant)
+        tenant_id = tenant.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)))
+          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
+
+      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
 
       private
 
+      def build_on_each_query(ids, scope)
+        return PgRls.main_model.all if ids.empty? && scope.blank?
+
+        return PgRls.main_model.where(id: ids) if scope.blank?
+
+        return scope.where(id: ids) if ids.present?
+
+        scope
+      end
+
       def switch_tenant!(resource)
         tenant = find_tenant(resource)
 
-        set_rls!(tenant.tenant_id)
+        set_rls!(tenant)
 
         tenant
       rescue NoMethodError
         raise PgRls::Errors::TenantNotFound
+      ensure
+        reset_rls! if tenant.blank?
       end
 
       def find_tenant(resource)
-        reset_rls!
-
         tenant = nil
 
         PgRls.search_methods.each do |method|
@@ -85,14 +130,20 @@ module PgRls
           tenant = find_tenant_by_method(resource, method)
         end
 
+        reset_rls! if reset_rls?(tenant)
         raise PgRls::Errors::TenantNotFound if tenant.blank?
 
         tenant
       end
 
+      def reset_rls?(tenant)
+        PgRls::Current::Context.tenant.present? && tenant.present? && PgRls::Current::Context.tenant != tenant
+      end
+
       def find_tenant_by_method(resource, method)
-        look_up_value = resource.is_a?(PgRls.main_model) ? resource.send(method) : resource
-        PgRls.main_model.send("find_by_#{method}!", look_up_value)
+        return resource if resource.is_a?(PgRls.main_model)
+
+        PgRls.main_model.unscoped.send(:"find_by_#{method}!", resource)
       rescue ActiveRecord::RecordNotFound
         nil
       end

data/lib/pg_rls/version.rb

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

data/lib/pg_rls.rb

@@ -5,25 +5,31 @@ require 'forwardable'
 require_relative 'pg_rls/version'
 require_relative 'pg_rls/database/prepared'
 require_relative 'pg_rls/schema/statements'
+require_relative 'pg_rls/schema/dumper'
 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'
+require_relative 'pg_rls/current/context'
+require_relative 'pg_rls/logger'
 
 ActiveRecord::Migrator.prepend PgRls::Admin::ActiveRecord::Migrator
 ActiveRecord::Tasks::DatabaseTasks.prepend PgRls::Admin::ActiveRecord::Tasks::DatabaseTasks
 ActiveRecord::ConnectionAdapters::AbstractAdapter.include PgRls::Schema::Statements
+ActiveRecord::SchemaDumper.prepend PgRls::Schema::Dumper
+
 # PostgreSQL Row Level Security
 module PgRls
   class Error < StandardError; end
   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
+    WRITER_METHODS = %i[table_name class_name search_methods logger excluded_shards].freeze
+    READER_METHODS = %i[connection_class execute table_name class_name search_methods logger excluded_shards].freeze
+    DELEGATORS_METHODS = %i[connection_class execute table_name search_methods class_name main_model logger
+                            excluded_shards].freeze
 
     attr_writer(*WRITER_METHODS)
     attr_reader(*READER_METHODS)
@@ -36,7 +42,7 @@ module PgRls
       yield self
 
       Rails.application.config.to_prepare do
-        PgRls.main_model.ignored_columns = []
+        PgRls.main_model.ignored_columns = [] # rubocop:disable Rails/IgnoredColumnsAssignment
       end
     end
 
@@ -54,7 +60,6 @@ module PgRls
     def establish_new_connection!(admin: false)
       self.as_db_admin = admin
 
-      db_config = PgRls.main_model.connection_db_config.configuration_hash
       execute_rls_in_shards do |connection_class, pool|
         connection_class.connection_pool.disconnect!
         connection_class.remove_connection
@@ -66,20 +71,23 @@ module PgRls
       class_name.to_s.camelize.constantize
     end
 
-    def on_each_tenant(&)
-      with_rls_connection do
-        result = []
-
-        main_model.find_each do |tenant|
-          allowed_search_fields = search_methods.map(&:to_s).intersection(main_model.column_names)
-          Tenant.switch tenant.send(allowed_search_fields.first)
-
-          result << { tenant:, result: ensure_block_execution(tenant, &) }
-        end
-
-        PgRls::Tenant.reset_rls!
+    def on_each_tenant(ids: [], scope: nil, &)
+      logger.deprecation_warning(
+        'PgRls.on_each_tenant is deprecated and will be removed in future versions. ' \
+        'Please use PgRls::Tenant.on_find_each instead.'
+      )
+      Tenant.on_find_each(ids:, scope:, &)
+    end
 
-        result
+    rails_version = Gem.loaded_specs['rails'].version
+    if rails_version >= Gem::Version.new('7.2') && rails_version < Gem::Version.new('7.3')
+      def pool_connection(pool)
+        pool.lease_connection
+      end
+    else
+      def pool_connection(pool)
+        PgRls.logger.deprecation_warning('PgRls.pool_connection is deprecated and will be removed in future PgRls 0.2.0. Please use pool.lease_connection instead.')
+        pool.connection
       end
     end
 
@@ -88,10 +96,13 @@ module PgRls
       result = []
 
       connection_pool_list.each do |pool|
-        pool.connection.transaction do
-          Rails.logger.info("Executing in #{pool.connection.connection_class}")
+        connection = pool_connection(pool)
+        next if excluded_shards.include?(connection.connection_class.connection_db_config.name)
+
+        connection.transaction do
+          Rails.logger.info("Executing in #{connection.connection_class}")
 
-          result << yield(pool.connection.connection_class, pool)
+          result << yield(connection.connection_class, pool)
         end
       end
 
@@ -141,11 +152,9 @@ module PgRls
     end
 
     def execute_query_or_block(query = nil, &)
-      if block_given?
-        ensure_block_execution(&)
-      else
-        execute(query)
-      end
+      return ensure_block_execution(&) if block_given?
+
+      execute(query)
     end
 
     def reset_connection_if_needed(current_tenant, reset_rls_connection)
@@ -173,4 +182,10 @@ module PgRls
 
   mattr_accessor :search_methods
   @@search_methods = %i[subdomain id tenant_id]
+
+  mattr_accessor :logger
+  @@logger = PgRls::Logger.new
+
+  mattr_accessor :excluded_shards
+  @@excluded_shards = []
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: pg_rls
 version: !ruby/object:Gem::Version
-  version: 0.1.10
+  version: 0.2.0
 platform: ruby
 authors:
 - Daniel Laloush
 autorequire: 
 bindir: exe
 cert_chain: []
-date: 2024-07-05 00:00:00.000000000 Z
+date: 2024-09-27 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
@@ -47,6 +47,7 @@ files:
 - lib/generators/templates/pg_rls.rb.tt
 - lib/pg_rls.rb
 - lib/pg_rls/Rakefile
+- lib/pg_rls/current/context.rb
 - lib/pg_rls/database/admin_statements.rb
 - lib/pg_rls/database/configurations.rb
 - lib/pg_rls/database/prepared.rb
@@ -54,6 +55,7 @@ files:
 - lib/pg_rls/errors/index.rb
 - lib/pg_rls/errors/rake_only_error.rb
 - lib/pg_rls/errors/tenant_not_found.rb
+- lib/pg_rls/logger.rb
 - lib/pg_rls/middleware.rb
 - lib/pg_rls/middleware/set_reset_connection.rb
 - lib/pg_rls/middleware/sidekiq.rb
@@ -62,6 +64,7 @@ files:
 - lib/pg_rls/multi_tenancy.rb
 - lib/pg_rls/railtie.rb
 - lib/pg_rls/schema/down_statements.rb
+- lib/pg_rls/schema/dumper.rb
 - lib/pg_rls/schema/statements.rb
 - lib/pg_rls/schema/up_statements.rb
 - lib/pg_rls/tenant.rb
@@ -86,7 +89,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.5.11
+rubygems_version: 3.5.16
 signing_key: 
 specification_version: 4
 summary: Write a short summary, because RubyGems requires one.