pg_rls 0.1.11 → 0.2.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 0e815f8e1e99ddf19dcd53be6876034dbeb535985c46ab7111555ac0fa02c02a
-  data.tar.gz: 4ce207b4766242cdc97940b313bb9a8ea864250ebd953bf4c36168d6ceced848
+  metadata.gz: 6bd9ecf91e6966f35c1bdc24a9f312209bc505dcaa220957437eaf4b1b9c6d42
+  data.tar.gz: 605a1950d3f4c338f808bb49129564f3ae2d7fe20d40a00698e22bc7f4603bfd
 SHA512:
-  metadata.gz: 9369bf467e43421b42093783428b3fdb544806ee378ad67e43500fecd530e971bc50b3671faf7cb54d126281bf21c5f88a7f9bb3eeffa52bdce0068578cfae04
-  data.tar.gz: d6fef41724951488a13c21fad439cc3f1f0c781e918d022f426adce50d81e8133f98961695042143dd8c6f96a229d8e6d42baed8cadf563b14b5ece4f8960611
+  metadata.gz: e6ab0a689e37a49cf95bf5126f7897902abb15bfb061a750486806e0b2040d30d92a71f4bc0747a33e5121c0a94a9dee4e90c690511716a93eab2118a04d9486
+  data.tar.gz: 05fd9172dc9b97ca4cda73d6a5dd8accea22d022786b638a5268f47fd1f6d99c7a598594eaf914fed47541690e985de0fb27c6fcce76bcb54f5757ee6eba80d1

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/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,17 +36,32 @@ module PgRls
       end
 
       def fetch!
-        PgRls::Current::Context.tenant ||= 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
 
@@ -54,30 +69,60 @@ module PgRls
         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)
+        PgRls.establish_new_connection! if PgRls.admin_connection?
+        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|
@@ -86,14 +131,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.11'
+  VERSION = '0.2.1'
 end

data/lib/pg_rls.rb

@@ -5,6 +5,7 @@ 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'
@@ -12,19 +13,23 @@ 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)
@@ -37,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
 
@@ -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
 
@@ -171,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.11
+  version: 0.2.1
 platform: ruby
 authors:
 - Daniel Laloush
 autorequire: 
 bindir: exe
 cert_chain: []
-date: 2024-08-28 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
@@ -55,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
@@ -63,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
@@ -87,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.