pg_rls 0.0.2.6.11 → 0.1.0

data/lib/pg_rls/schema/up_statements.rb

@@ -5,7 +5,7 @@ module PgRls
     # Up Schema Statements
     module UpStatements
       def create_rls_user(name: PgRls.username, password: PgRls.password, schema: 'public')
-        PgRls.execute <<-SQL
+        ActiveRecord::Migration.execute <<-SQL
           DO
           $do$
           BEGIN
@@ -35,7 +35,7 @@ module PgRls
       end
 
       def create_rls_blocking_function
-        ActiveRecord::Migration.execute <<-SQL
+        ActiveRecord::Migration.execute <<-SQL.squish
           CREATE OR REPLACE FUNCTION id_safe_guard ()
             RETURNS TRIGGER LANGUAGE plpgsql AS $$
               BEGIN
@@ -45,7 +45,7 @@ module PgRls
       end
 
       def create_rls_setter_function
-        ActiveRecord::Migration.execute <<-SQL
+        ActiveRecord::Migration.execute <<-SQL.squish
           CREATE OR REPLACE FUNCTION tenant_id_setter ()
             RETURNS TRIGGER LANGUAGE plpgsql AS $$
               BEGIN
@@ -56,7 +56,7 @@ module PgRls
       end
 
       def append_blocking_function(table_name)
-        ActiveRecord::Migration.execute <<-SQL
+        ActiveRecord::Migration.execute <<-SQL.squish
           CREATE TRIGGER id_safe_guard
             BEFORE UPDATE OF id ON #{table_name}
               FOR EACH ROW EXECUTE PROCEDURE id_safe_guard();
@@ -64,7 +64,7 @@ module PgRls
       end
 
       def append_trigger_function(table_name)
-        ActiveRecord::Migration.execute <<-SQL
+        ActiveRecord::Migration.execute <<-SQL.squish
           CREATE TRIGGER tenant_id_setter
             BEFORE INSERT OR UPDATE ON #{table_name}
               FOR EACH ROW EXECUTE PROCEDURE tenant_id_setter();
@@ -72,7 +72,7 @@ module PgRls
       end
 
       def add_rls_column_to_tenant_table(table_name)
-        ActiveRecord::Migration.execute <<-SQL
+        ActiveRecord::Migration.execute <<-SQL.squish
           ALTER TABLE #{table_name}
             ADD COLUMN IF NOT EXISTS
               tenant_id uuid UNIQUE DEFAULT gen_random_uuid();
@@ -80,7 +80,7 @@ module PgRls
       end
 
       def add_rls_column(table_name)
-        ActiveRecord::Migration.execute <<-SQL
+        ActiveRecord::Migration.execute <<-SQL.squish
           ALTER TABLE #{table_name}
             ADD COLUMN IF NOT EXISTS tenant_id uuid,
             ADD CONSTRAINT fk_#{PgRls.table_name}
@@ -91,7 +91,7 @@ module PgRls
       end
 
       def create_rls_policy(table_name, user = PgRls.username)
-        ActiveRecord::Migration.execute <<-SQL
+        ActiveRecord::Migration.execute <<-SQL.squish
           ALTER TABLE #{table_name} ENABLE ROW LEVEL SECURITY;
           CREATE POLICY #{table_name}_#{user}
             ON #{table_name}

data/lib/pg_rls/tenant.rb

@@ -25,16 +25,10 @@ module PgRls
         raise e
       end
 
-      def find_each(&)
-        PgRls.main_model.find_each do |tenant|
-          with_tenant!(tenant, &)
-        end
-      end
-
-      def with_tenant!(resource, &block)
+      def with_tenant!(resource)
         tenant = switch_tenant!(resource)
 
-        block.call(tenant) if block_given?
+        yield(tenant) if block_given?
       ensure
         reset_rls! unless PgRls.test_inline_tenant == true
       end
@@ -45,47 +39,50 @@ module PgRls
         nil
       end
 
-      def fetch!
+      def tenant!
         @tenant ||= PgRls.main_model.find_by!(
           tenant_id: PgRls.connection_class.connection.execute(
             "SELECT current_setting('rls.tenant_id')"
           ).getvalue(0, 0)
         )
       end
-
-      def find_main_model
-        PgRls.main_model.ignored_columns = []
-        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.connection_class.connection.execute('RESET rls.tenant_id')
+        PgRls.execute_rls_in_shards do |connection_class|
+          connection_class.transaction do
+            connection_class.connection.execute('RESET rls.tenant_id')
+          end
+        end
+
+        nil
       end
 
       private
 
       def switch_tenant!(resource)
+        # rubocop: disable Rails/IgnoredColumnsAssignment
         PgRls.main_model.ignored_columns = []
+        # rubocop: enable Rails/IgnoredColumnsAssignment
 
-        connection_adapter = PgRls.connection_class
         find_tenant(resource)
 
-        raise PgRls::Errors::TenantNotFound if tenant.blank?
-
-        connection_adapter.connection.transaction do
-          connection_adapter.connection.execute(format('SET rls.tenant_id = %s',
-                                                       connection_adapter.connection.quote(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.tenant_id)))
+          end
         end
 
         tenant
       end
 
       def find_tenant(resource)
+        raise PgRls::Errors::AdminUsername if admin_username?
+
         reset_rls!
 
         PgRls.search_methods.each do |method|
@@ -95,7 +92,11 @@ module PgRls
           @tenant = find_tenant_by_method(resource, method)
         end
 
-        raise PgRls::Errors::TenantNotFound if tenant.nil?
+        raise PgRls::Errors::TenantNotFound if tenant.blank?
+      end
+
+      def admin_username?
+        PgRls.username != PgRls.current_db_username
       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.0.2.6.11'
+  VERSION = '0.1.0'
 end

data/lib/pg_rls.rb

@@ -5,13 +5,11 @@ require 'forwardable'
 require_relative 'pg_rls/version'
 require_relative 'pg_rls/database/prepared'
 require_relative 'pg_rls/schema/statements'
-require_relative 'pg_rls/schema/solo/statements'
-require_relative 'pg_rls/solo/tenant'
+require_relative 'pg_rls/database/configurations'
 require_relative 'pg_rls/tenant'
-require_relative 'pg_rls/secure_connection'
 require_relative 'pg_rls/multi_tenancy'
 require_relative 'pg_rls/railtie' if defined?(Rails)
-require_relative 'pg_rls/errors/tenant_not_found'
+require_relative 'pg_rls/errors/index'
 
 # PostgreSQL Row Level Security
 module PgRls
@@ -21,13 +19,13 @@ module PgRls
   class << self
     extend Forwardable
 
-    WRITER_METHODS = %i[table_name class_name search_methods establish_default_connection].freeze
+    WRITER_METHODS = %i[table_name class_name search_methods].freeze
     READER_METHODS = %i[
-      connection_class database_configuration execute table_name class_name search_methods establish_default_connection
+      connection_class execute table_name class_name search_methods
     ].freeze
     DELEGATORS_METHODS = %i[
-      connection_class database_configuration execute table_name search_methods
-      class_name all_tenants main_model establish_default_connection
+      connection_class execute table_name search_methods
+      class_name main_model
     ].freeze
 
     attr_writer(*WRITER_METHODS)
@@ -36,100 +34,104 @@ module PgRls
     def_delegators(*DELEGATORS_METHODS)
 
     def setup
+      ActiveRecord::ConnectionAdapters::AbstractAdapter.include PgRls::Schema::Statements
+      ActiveRecord::Base.ignored_columns += %w[tenant_id]
+
       yield self
     end
 
-    def solo_mode?
-      solo_mode
+    def connection_class
+      @connection_class ||= ActiveRecord::Base
     end
 
-    def database_connection_file
-      file = File.read(Rails.root.join('config/database.yml'))
-
-      YAML.safe_load(ERB.new(file).result, aliases: true)
+    def rake_tasks?
+      Rake.application.top_level_tasks.present? || ARGV.any? { |arg| arg =~ /rake|dsl/ }
+    rescue NoMethodError
+      false
     end
 
-    def connection_class
-      @connection_class ||= ActiveRecord::Base
-    end
+    def admin_tasks_execute
+      raise PgRls::Errors::RakeOnlyError unless rake_tasks?
 
-    def establish_new_connection
-      ActiveRecord::Base.connection.disconnect! if ActiveRecord::Base.connection_pool.connected?
+      self.as_db_admin = true
 
-      connection_class.establish_connection(**database_configuration)
+      yield
+    ensure
+      self.as_db_admin = false
     end
 
-    def admin_execute(query = nil)
+    def admin_execute(query = nil, &)
+      current_tenant = PgRls::Tenant.fetch
+
       self.as_db_admin = true
-      establish_new_connection
-      return yield if block_given?
+      establish_new_connection!
+
+      return ensure_block_execution(&) if block_given?
 
       execute(query)
     ensure
       self.as_db_admin = false
-      establish_new_connection
-    end
-
-    def establish_default_connection=(value)
-      ENV['AS_DB_ADMIN'] = value.to_s
-      @default_connection = value
+      establish_new_connection!
+      PgRls::Tenant.switch(current_tenant) if current_tenant.present?
     end
 
-    def default_connection?
-      as_db_admin
+    def establish_new_connection!
+      execute_rls_in_shards do |connection_class, pool|
+        connection_class.remove_connection
+        connection_class.establish_connection(pool.db_config)
+      end
     end
 
     def main_model
       class_name.to_s.camelize.constantize
     end
 
-    def all_tenants
-      main_model.all.each do |tenant|
+    def on_each_tenant(&)
+      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)
 
-        yield(tenant) if block_given?
+        result << { tenant:, result: ensure_block_execution(tenant, &) }
       end
-    end
 
-    def current_connection_username
-      connection_class.connection_db_config.configuration_hash[:username]
-    end
+      PgRls::Tenant.reset_rls!
 
-    def execute(query)
-      ActiveRecord::Migration.execute(query)
+      result
     end
 
-    def database_default_configuration
-      connection_class.connection.pool.db_config.configuration_hash
-    rescue ActiveRecord::NoDatabaseError
-      connection_class.connection_db_config.configuration_hash
-    end
+    def execute_rls_in_shards
+      connection_pool_list = PgRls.connection_class.connection_handler.connection_pool_list
+      result = []
 
-    def database_admin_configuration
-      environment_db_configuration = database_connection_file[Rails.env]
+      connection_pool_list.each do |pool|
+        pool.connection.transaction do
+          Rails.logger.info("Executing in #{pool.connection.connection_class}")
 
-      return environment_db_configuration if environment_db_configuration['username'].present?
+          result << yield(pool.connection.connection_class, pool)
+        end
+      end
 
-      environment_db_configuration.first.last
+      result
     end
 
-    def database_configuration
-      return database_admin_configuration if default_connection?
+    def current_db_username
+      ActiveRecord::Base.connection_db_config.configuration_hash[:username]
+    end
 
-      current_configuration = database_default_configuration.deep_dup
-      current_configuration.tap do |config|
-        config[:username] = PgRls.username
-        config[:password] = PgRls.password
-      end
+    def as_db_admin?
+      @as_db_admin || false
+    end
+
+    private
 
-      current_configuration.freeze
+    attr_writer :as_db_admin
+
+    def ensure_block_execution(*, **)
+      yield(*, **).presence
     end
   end
 
-  mattr_accessor :as_db_admin
-  @@as_db_admin = false
-
   mattr_accessor :table_name
   @@table_name = 'companies'
 
@@ -142,18 +144,9 @@ module PgRls
   mattr_accessor :password
   @@password = 'password'
 
-  mattr_accessor :solo_mode
-  @@solo_mode = false
-
   mattr_accessor :test_inline_tenant
   @@test_inline_tenant = false
 
-  mattr_accessor :session_key
-  @@session_key = '_hub_sessions'
-
-  mattr_accessor :session_prefix
-  @@session_prefix = '_session_id:2::'
-
   mattr_accessor :search_methods
   @@search_methods = %i[subdomain id tenant_id]
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: pg_rls
 version: !ruby/object:Gem::Version
-  version: 0.0.2.6.11
+  version: 0.1.0
 platform: ruby
 authors:
 - Daniel Laloush
 autorequire: 
 bindir: exe
 cert_chain: []
-date: 2023-02-20 00:00:00.000000000 Z
+date: 2023-10-14 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: bundler
@@ -61,8 +61,12 @@ files:
 - lib/generators/templates/pg_rls.rb.tt
 - lib/pg_rls.rb
 - lib/pg_rls/Rakefile
+- 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
 - lib/pg_rls/middleware.rb
 - lib/pg_rls/middleware/set_reset_connection.rb
@@ -72,12 +76,8 @@ files:
 - lib/pg_rls/multi_tenancy.rb
 - lib/pg_rls/railtie.rb
 - lib/pg_rls/schema/down_statements.rb
-- lib/pg_rls/schema/solo/statements.rb
-- lib/pg_rls/schema/solo/up_statements.rb
 - lib/pg_rls/schema/statements.rb
 - lib/pg_rls/schema/up_statements.rb
-- lib/pg_rls/secure_connection.rb
-- lib/pg_rls/solo/tenant.rb
 - lib/pg_rls/tenant.rb
 - lib/pg_rls/version.rb
 homepage: https://github.com/Dandush03/pg_rls
@@ -93,14 +93,14 @@ required_ruby_version: !ruby/object:Gem::Requirement
   requirements:
   - - "~>"
     - !ruby/object:Gem::Version
-      version: '3.1'
+      version: '3.2'
 required_rubygems_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.3.26
+rubygems_version: 3.4.19
 signing_key: 
 specification_version: 4
 summary: Write a short summary, because RubyGems requires one.

data/lib/pg_rls/schema/solo/statements.rb

@@ -1,24 +0,0 @@
-# frozen_string_literal: true
-
-require_relative '../statements'
-require_relative './up_statements'
-
-module PgRls
-  module Schema
-    module Solo
-      # Schema Solo Statements
-      module Statements
-        include PgRls::Schema::Statements
-        include PgRls::Schema::Solo::UpStatements
-
-        def create_rls_table(table_name, **options, &)
-          setup_rls_tenant_table
-          create_table(table_name, **options, &)
-          add_rls_column(table_name)
-          create_rls_policy(table_name)
-          append_trigger_function(table_name)
-        end
-      end
-    end
-  end
-end

data/lib/pg_rls/schema/solo/up_statements.rb

@@ -1,106 +0,0 @@
-# frozen_string_literal: true
-
-module PgRls
-  module Schema
-    module Solo
-      # Up Schema Solo Statements
-      module UpStatements
-        def setup_rls_tenant_table
-          ActiveRecord::Migration.execute <<-SQL
-            DO
-            $do$
-              BEGIN
-                IF NOT EXISTS (
-                  SELECT FROM pg_tables
-                  WHERE schemaname = 'public' AND tablename = '#{PgRls.table_name}') THEN
-                    #{create_rls_user}
-                    #{create_rls_setter_function}
-                    #{create_rls_blocking_function}
-                    #{create_rls_solo_tenant_table}
-                    #{append_blocking_function}
-                END IF;
-              END;
-            $do$;
-          SQL
-        end
-
-        def create_rls_user(name: PgRls.username, password: PgRls.password, schema: 'public')
-          <<~SQL
-            -- Grant Role Permissions
-            BEGIN
-              IF NOT EXISTS (
-                SELECT FROM pg_catalog.pg_roles  -- SELECT list can be empty for this
-                WHERE  rolname = '#{name}') THEN
-
-                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
-                ON SEQUENCES TO #{name};
-              ALTER DEFAULT PRIVILEGES IN SCHEMA #{schema}
-                GRANT SELECT, INSERT, UPDATE, DELETE
-                ON TABLES TO #{name};
-              GRANT SELECT, INSERT, UPDATE, DELETE
-                ON ALL TABLES IN SCHEMA #{schema}
-                TO #{name};
-              GRANT USAGE, SELECT
-                ON ALL SEQUENCES IN SCHEMA #{schema}
-                TO #{name};
-            END;
-          SQL
-        end
-
-        def create_rls_setter_function
-          <<~SQL
-            -- Create RLS Setter Function
-            CREATE OR REPLACE FUNCTION tenant_id_setter ()
-              RETURNS TRIGGER LANGUAGE plpgsql AS $$
-                BEGIN
-                  IF NOT EXISTS (
-                    SELECT FROM #{PgRls.table_name}
-                      WHERE tenant_id = (current_setting('rls.tenant_id'))::uuid
-                  ) THEN
-                    INSERT INTO #{PgRls.table_name} (tenant_id)
-                      VALUES ((current_setting('rls.tenant_id'))::uuid);
-                  END IF;
-
-                  NEW.tenant_id:= (current_setting('rls.tenant_id'));
-                  RETURN NEW;
-                END $$;
-          SQL
-        end
-
-        def create_rls_blocking_function
-          <<~SQL
-            -- Create RLS Blocking Function
-            CREATE OR REPLACE FUNCTION id_safe_guard ()
-              RETURNS TRIGGER LANGUAGE plpgsql AS $$
-                BEGIN
-                  RAISE EXCEPTION 'This column is guarded due to tenancy dependency';
-                END $$;
-          SQL
-        end
-
-        def create_rls_solo_tenant_table
-          <<~SQL
-            -- Create Tenant Table
-            CREATE TABLE #{PgRls.table_name} (
-              tenant_id uuid PRIMARY KEY
-            );
-          SQL
-        end
-
-        def append_blocking_function
-          <<~SQL
-            -- Append Blocking Function
-            CREATE TRIGGER id_safe_guard
-              BEFORE UPDATE OF tenant_id ON #{PgRls.table_name}
-                FOR EACH ROW EXECUTE PROCEDURE id_safe_guard();
-          SQL
-        end
-      end
-    end
-  end
-end

data/lib/pg_rls/secure_connection.rb

@@ -1,23 +0,0 @@
-# frozen_string_literal: true
-
-module PgRls
-  # Ensure Connection is with App_use
-  module SecureConnection
-    def self.establish_secure_connection
-      return if secure_connection_established?
-
-      return if PgRls.default_connection?
-
-      PgRls.establish_new_connection
-    end
-
-    def self.secure_connection_established?
-      PgRls.current_connection_username == PgRls.username
-    end
-
-    def self.included(base)
-      establish_secure_connection
-      base.ignored_columns = %w[tenant_id]
-    end
-  end
-end

data/lib/pg_rls/solo/tenant.rb

@@ -1,50 +0,0 @@
-# frozen_string_literal: true
-
-module PgRls
-  module Solo
-    # Set and Fetch Tenant without loading a model
-    class Tenant
-      class << self
-        attr_reader :tenant
-
-        def switch!(resource)
-          switch_tenant!(resource)
-        rescue StandardError => e
-          Rails.logger.info('connection was not made')
-          raise e
-        end
-
-        def fetch
-          @fetch ||= PgRls.connection_class.connection.execute(
-            "SELECT current_setting('rls.tenant_id')"
-          ).getvalue(0, 0)
-        end
-
-        def around(resource)
-          switch_tenant!(resource)
-          yield
-        ensure
-          reset_rls!
-        end
-
-        private
-
-        def reset_rls!
-          @fetch = nil
-          @tenant = nil
-          PgRls.connection_class.connection.execute('RESET rls.tenant_id')
-        end
-
-        def switch_tenant!(resource)
-          connection_adapter = PgRls.connection_class
-
-          raise PgRls::Errors::TenantNotFound if resource.blank?
-
-          connection_adapter.connection.execute(format('SET rls.tenant_id = %s',
-                                                       connection_adapter.connection.quote(resource)))
-          "RLS changed to '#{resource}'"
-        end
-      end
-    end
-  end
-end