pg_rls 1.0.0 → 1.0.1

data/lib/pg_rls/errors/index.rb

@@ -1,4 +0,0 @@
-# frozen_string_literal: true
-
-require_relative 'rake_only_error'
-require_relative 'tenant_not_found'

data/lib/pg_rls/errors/rake_only_error.rb

@@ -1,12 +0,0 @@
-# frozen_string_literal: true
-
-module PgRls
-  module Errors
-    class RakeOnlyError < StandardError
-      def initialize(msg = nil)
-        msg ||= 'This method can only be executed through rake tasks'
-        super
-      end
-    end
-  end
-end

data/lib/pg_rls/errors/tenant_not_found.rb

@@ -1,13 +0,0 @@
-# frozen_string_literal: true
-
-module PgRls
-  module Errors
-    # Raise Tenant Not found and ensure that the tenant is resetted
-    class TenantNotFound < StandardError
-      def initialize(msg = nil)
-        msg ||= "Tenant Doesn't exist"
-        super
-      end
-    end
-  end
-end

data/lib/pg_rls/logger.rb

@@ -1,31 +0,0 @@
-# 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/middleware/set_reset_connection.rb

@@ -1,93 +0,0 @@
-# frozen_string_literal: true
-
-module PgRls
-  class FrozenConfiguration < StandardError; end
-
-  def self.sessions
-    @sessions ||= Redis.new(@session_store_server)
-  end
-
-  def self.session_prefix
-    @session_prefix ||= begin
-      store_default_warden_key = @session_store_default_warden_key || '2'
-
-      "#{session_key_prefix}:#{store_default_warden_key}::"
-    end
-  end
-
-  def self.session_store_server=(opts = {})
-    raise Errors::FrozenConfiguration unless @sessions.nil?
-
-    @session_store_server = opts.deep_symbolize_keys
-  end
-
-  def self.session_store_default_warden_key=(val)
-    raise Errors::FrozenConfiguration unless @sessions.nil?
-
-    @session_store_default_warden_key = val
-  end
-
-  def self.session_key_prefix
-    @session_key_prefix ||= @session_key_prefix || @session_store_server[:key_prefix]
-  end
-
-  def self.session_key_prefix=(val)
-    raise Errors::FrozenConfiguration unless @sessions.nil?
-
-    @session_key_prefix = val
-  end
-end
-
-module PgRls
-  module Middleware
-    # Set RLS if sessions present.
-    class SetResetConnection
-      def initialize(app)
-        @app = app
-      end
-
-      def call(env)
-        tenant = load_tenant_thought_session(env)
-
-        return @app.call(env) if tenant.blank?
-
-        PgRls::Tenant.with_tenant!(tenant) { @app.call(env) }
-      rescue PgRls::Errors::TenantNotFound
-        @app.call(env)
-      rescue ActiveRecord::RecordNotFound => e
-        raise e unless rails_active_storage_request?(env)
-
-        [404, { 'Content-Type' => 'text/plain' }, ['Could not find asset']]
-      end
-
-      def load_session_cookie_value(env)
-        cookie_string = env['HTTP_COOKIE']
-        return if cookie_string.nil?
-
-        cookie_regex = /#{PgRls.session_key_prefix}=([^;]+)/
-        match = cookie_regex.match(cookie_string)
-        match[1] if match
-      end
-
-      def load_tenant_thought_session(env)
-        cookie = load_session_cookie_value(env)
-
-        return if cookie.blank?
-
-        redis_session_key = "#{PgRls.session_prefix}#{Digest::SHA256.hexdigest(cookie)}"
-        tenant_session = Marshal.load(PgRls.sessions.get(redis_session_key))
-
-        return if tenant_session.blank?
-        return if tenant_session['_tenant'].blank?
-
-        tenant_session['_tenant'] if tenant_session.present?
-      rescue TypeError
-        nil
-      end
-
-      def rails_active_storage_request?(env)
-        env['PATH_INFO'].start_with?('/rails/active_storage/')
-      end
-    end
-  end
-end

data/lib/pg_rls/middleware/sidekiq/client.rb

@@ -1,22 +0,0 @@
-# frozen_string_literal: true
-
-# :nocov:
-module PgRls
-  module Middleware
-    module Sidekiq
-      # Set PgRls Policies
-      class Client
-        def call(_job_class, msg, _queue, _redis_pool)
-          load_tenant_attribute!(msg)
-          yield
-        end
-
-        def load_tenant_attribute!(msg)
-          return msg['admin'] = true if PgRls.admin_connection?
-
-          msg['pg_rls'] ||= PgRls::Tenant.fetch&.id
-        end
-      end
-    end
-  end
-end

data/lib/pg_rls/middleware/sidekiq/server.rb

@@ -1,19 +0,0 @@
-# frozen_string_literal: true
-
-# :nocov:
-module PgRls
-  module Middleware
-    module Sidekiq
-      # Set PgRls Policies
-      class Server
-        def call(_job_instance, msg, _queue, &)
-          if msg['admin']
-            PgRls.admin_execute(&)
-          else
-            PgRls::Tenant.with_tenant!(msg['pg_rls'], &)
-          end
-        end
-      end
-    end
-  end
-end

data/lib/pg_rls/middleware/sidekiq.rb

@@ -1,11 +0,0 @@
-# frozen_string_literal: true
-
-require_relative 'sidekiq/client'
-require_relative 'sidekiq/server'
-
-module PgRls
-  module Middleware
-    module Sidekiq
-    end
-  end
-end

data/lib/pg_rls/middleware.rb

@@ -1,8 +0,0 @@
-# frozen_string_literal: true
-
-require_relative 'middleware/sidekiq'
-
-module PgRls
-  module Middleware
-  end
-end

data/lib/pg_rls/multi_tenancy.rb

@@ -1,32 +0,0 @@
-# frozen_string_literal: true
-
-module PgRls
-  # Ensure Connection is with App_use
-  module MultiTenancy
-    def self.included(base)
-      base.class_eval do
-        around_action :switch_tenant!
-
-        def current_tenant
-          @current_tenant ||= request.subdomain
-        end
-        helper_method :current_tenant
-      end
-    end
-
-    private
-
-    def switch_tenant!
-      fetched_tenant = session[:_tenant] || current_tenant
-      return yield if PgRls::Tenant.fetch.present?
-
-      Tenant.with_tenant!(fetched_tenant) do |tenant|
-        session[:_tenant] = tenant.id
-        yield(tenant)
-      end
-    rescue NoMethodError
-      session[:_tenant] = nil
-      raise PgRls::Errors::TenantNotFound, 'No tenant was found'
-    end
-  end
-end

data/lib/pg_rls/schema/down_statements.rb

@@ -1,54 +0,0 @@
-# frozen_string_literal: true
-
-module PgRls
-  module Schema
-    # Down Schema Statements
-    module DownStatements
-      def drop_rls_user
-        ActiveRecord::Migration.execute <<~SQL.squish
-          DROP OWNED BY #{PgRls.username};
-          REVOKE ALL PRIVILEGES ON ALL TABLES IN SCHEMA public FROM #{PgRls.username};
-          REVOKE ALL PRIVILEGES ON ALL SEQUENCES IN SCHEMA public FROM #{PgRls.username};
-          REVOKE ALL PRIVILEGES ON ALL FUNCTIONS IN SCHEMA public FROM #{PgRls.username};
-          DROP USER #{PgRls.username};
-        SQL
-      end
-
-      def drop_rls_blocking_function
-        ActiveRecord::Migration.execute 'DROP FUNCTION IF EXISTS id_safe_guard ();'
-      end
-
-      def drop_rls_setter_function
-        ActiveRecord::Migration.execute 'DROP FUNCTION IF EXISTS tenant_id_setter ();'
-      end
-
-      def detach_blocking_function(table_name)
-        ActiveRecord::Migration.execute <<-SQL.squish
-          DROP TRIGGER IF EXISTS id_safe_guard
-            ON #{table_name};
-        SQL
-      end
-
-      def detach_trigger_function(table_name)
-        ActiveRecord::Migration.execute <<-SQL.squish
-          DROP TRIGGER IF EXISTS tenant_id_setter
-            ON #{table_name};
-        SQL
-      end
-
-      def drop_rls_column(table_name)
-        ActiveRecord::Migration.execute <<-SQL.squish
-          ALTER TABLE #{table_name}
-            DROP COLUMN IF EXISTS tenant_id;
-        SQL
-      end
-
-      def drop_rls_policy(table_name)
-        ActiveRecord::Migration.execute <<-SQL.squish
-          DROP POLICY #{table_name}_#{PgRls.username} ON #{table_name};
-          ALTER TABLE #{table_name} DISABLE ROW LEVEL SECURITY;
-        SQL
-      end
-    end
-  end
-end

data/lib/pg_rls/schema/dumper.rb

@@ -1,36 +0,0 @@
-# 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/statements.rb

@@ -1,72 +0,0 @@
-# frozen_string_literal: true
-
-require_relative 'down_statements'
-require_relative 'up_statements'
-
-module PgRls
-  module Schema
-    # Schema Statements
-    module Statements
-      include UpStatements
-      include DownStatements
-
-      def create_rls_tenant_table(table_name, **, &)
-        create_rls_user
-        create_rls_setter_function
-        create_rls_blocking_function
-        create_table(table_name, **, &)
-        add_rls_column_to_tenant_table(table_name)
-        append_blocking_function(table_name)
-      end
-
-      def create_rls_table(table_name, **, &)
-        create_table(table_name, **, &)
-        add_rls_column(table_name)
-        create_rls_policy(table_name)
-        append_trigger_function(table_name)
-      end
-
-      def drop_rls_tenant_table(table_name)
-        drop_rls_setter_function
-        detach_blocking_function(table_name)
-        drop_table(table_name)
-        drop_rls_blocking_function
-        drop_rls_user
-      end
-
-      def drop_rls_table(table_name)
-        detach_trigger_function(table_name)
-        drop_rls_policy(table_name)
-        drop_table(table_name)
-      end
-
-      def convert_to_rls_tenant_table(table_name, **_options)
-        create_rls_user
-        create_rls_setter_function
-        create_rls_blocking_function
-        add_rls_column_to_tenant_table(table_name)
-        append_blocking_function(table_name)
-      end
-
-      def revert_rls_tenant_table(table_name)
-        drop_rls_setter_function
-        detach_blocking_function(table_name)
-        drop_rls_blocking_function
-        drop_rls_user
-        drop_rls_column(table_name)
-      end
-
-      def convert_to_rls_table(table_name)
-        add_rls_column(table_name)
-        create_rls_policy(table_name)
-        append_trigger_function(table_name)
-      end
-
-      def revert_rls_table(table_name)
-        detach_trigger_function(table_name)
-        drop_rls_policy(table_name)
-        drop_rls_column(table_name)
-      end
-    end
-  end
-end

data/lib/pg_rls/schema/up_statements.rb

@@ -1,104 +0,0 @@
-# frozen_string_literal: true
-
-module PgRls
-  module Schema
-    # Up Schema Statements
-    module UpStatements
-      def create_rls_user(name: PgRls.username, password: PgRls.password, schema: 'public')
-        ActiveRecord::Migration.execute <<-SQL
-          DO
-          $do$
-          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;
-          $do$;
-        SQL
-      end
-
-      def create_rls_blocking_function
-        ActiveRecord::Migration.execute <<-SQL.squish
-          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_setter_function
-        ActiveRecord::Migration.execute <<-SQL.squish
-          CREATE OR REPLACE FUNCTION tenant_id_setter ()
-            RETURNS TRIGGER LANGUAGE plpgsql AS $$
-              BEGIN
-                new.tenant_id:= (current_setting('rls.tenant_id'));
-                RETURN new;
-              END $$;
-        SQL
-      end
-
-      def append_blocking_function(table_name)
-        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();
-        SQL
-      end
-
-      def append_trigger_function(table_name)
-        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();
-        SQL
-      end
-
-      def add_rls_column_to_tenant_table(table_name)
-        ActiveRecord::Migration.execute <<-SQL.squish
-          ALTER TABLE #{table_name}
-            ADD COLUMN IF NOT EXISTS
-              tenant_id uuid UNIQUE DEFAULT gen_random_uuid();
-        SQL
-      end
-
-      def add_rls_column(table_name)
-        ActiveRecord::Migration.execute <<-SQL.squish
-          ALTER TABLE #{table_name}
-            ADD COLUMN IF NOT EXISTS tenant_id uuid,
-            ADD CONSTRAINT fk_#{PgRls.table_name}
-              FOREIGN KEY (tenant_id)
-              REFERENCES #{PgRls.table_name}(tenant_id)
-              ON DELETE CASCADE;
-        SQL
-      end
-
-      def create_rls_policy(table_name, user = PgRls.username)
-        ActiveRecord::Migration.execute <<-SQL.squish
-          ALTER TABLE #{table_name} ENABLE ROW LEVEL SECURITY;
-          CREATE POLICY #{table_name}_#{user}
-            ON #{table_name}
-            TO #{user}
-            USING (tenant_id = NULLIF(current_setting('rls.tenant_id', TRUE), '')::uuid);
-        SQL
-      end
-    end
-  end
-end

data/lib/pg_rls/tenant.rb

@@ -1,153 +0,0 @@
-# frozen_string_literal: true
-
-module PgRls
-  # Tenant Controller
-  module Tenant
-    class << self
-      def switch(resource)
-        switch!(resource)
-      rescue PgRls::Errors::TenantNotFound
-        nil
-      end
-
-      def switch!(resource)
-        tenant = switch_tenant!(resource)
-
-        "RLS changed to '#{tenant.id}'"
-      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)
-
-          yield(tenant).presence if block_given?
-        ensure
-          reset_rls! unless PgRls.test_inline_tenant == true || PgRls::Current::Context.tenant.blank?
-        end
-      end
-
-      def fetch
-        fetch!
-      rescue ActiveRecord::StatementInvalid, ActiveRecord::RecordNotFound
-        nil
-      end
-
-      def fetch!
-        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.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)
-        tenant_id = tenant.tenant_id
-        PgRls.execute_rls_in_shards do |connection_class|
-          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)
-
-        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)
-        tenant = nil
-
-        PgRls.search_methods.each do |method|
-          break if tenant.present?
-
-          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)
-        return resource if resource.is_a?(PgRls.main_model)
-
-        PgRls.main_model.unscoped.send(:"find_by_#{method}!", resource)
-      rescue ActiveRecord::RecordNotFound
-        nil
-      end
-    end
-  end
-end