pg_rls 0.2.5 → 1.0.0

data/lib/pg_rls/active_record/connection_adapters/postgre_sql/rls_user_statements.rb

@@ -0,0 +1,127 @@
+# frozen_string_literal: true
+
+module PgRls
+  module ActiveRecord
+    module ConnectionAdapters
+      module PostgreSQL
+        # This module contains the logic to grant user privileges
+        module RlsUserStatements
+          include SqlHelperMethod
+
+          def create_rls_role(name, password)
+            create_rls_user(name, password)
+            assign_user_to_group(name)
+          end
+
+          def drop_rls_role(name)
+            remove_user_from_group(name)
+            drop_rls_user(name)
+          end
+
+          def user_exists?(name)
+            execute_sql!(user_exists_sql(name)).first.present?
+          end
+
+          def drop_rls_user(name)
+            execute_sql!(drop_rls_user_sql(name))
+          end
+
+          def create_rls_user(name, password)
+            execute_sql!(create_rls_user_sql(name, password))
+          end
+
+          def create_rls_group(name = PgRls.rls_role_group)
+            execute_sql!(create_rls_group_sql(name))
+          end
+
+          def drop_rls_group(name = PgRls.rls_role_group)
+            execute_sql!(drop_rls_group_sql(name))
+          end
+
+          def assign_user_to_group(name)
+            execute_sql!(assign_user_to_group_sql(name))
+          end
+
+          def remove_user_from_group(name)
+            execute_sql!(remove_user_from_group_sql(name))
+          end
+
+          private
+
+          def user_exists_sql(name)
+            <<~SQL
+              SELECT 1
+              FROM pg_catalog.pg_roles
+              WHERE rolname = '#{name}';
+            SQL
+          end
+
+          def drop_rls_user_sql(name)
+            <<~SQL
+              DO $do$ BEGIN
+                IF EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = '#{name}') THEN
+                  DROP ROLE #{name};
+                END IF; END $do$;
+            SQL
+          end
+
+          def create_rls_user_sql(name, password)
+            <<~SQL
+              DO $do$ BEGIN
+                IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = '#{name}') THEN
+                  CREATE ROLE #{name} LOGIN PASSWORD '#{password}';
+                END IF;
+              END $do$;
+            SQL
+          end
+
+          def create_rls_group_sql(role_name)
+            <<~SQL
+              DO $do$ BEGIN
+                IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = '#{role_name}') THEN
+                  CREATE ROLE #{role_name} NOLOGIN;
+                END IF;
+              END $do$;
+            SQL
+          end
+
+          def drop_rls_group_sql(role_name)
+            <<~SQL
+              DO $do$ BEGIN
+                IF EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = '#{role_name}') THEN
+                  DROP ROLE #{role_name};
+                END IF;
+              END $do$;
+            SQL
+          end
+
+          def assign_user_to_group_sql(name)
+            <<~SQL
+              DO $$ BEGIN IF NOT EXISTS (
+                  SELECT FROM pg_catalog.pg_auth_members
+                    JOIN pg_roles AS rls_group ON rls_group.oid = pg_auth_members.roleid
+                    JOIN pg_roles AS rls_member ON rls_member.oid = pg_auth_members.member
+                    WHERE rls_group.rolname = 'rls_group' AND rls_member.rolname = '#{name}'
+                ) THEN
+                  GRANT rls_group TO #{name};
+                END IF; END $$;
+            SQL
+          end
+
+          def remove_user_from_group_sql(name)
+            <<~SQL
+              DO $$ BEGIN IF EXISTS (
+                  SELECT FROM pg_catalog.pg_auth_members
+                    JOIN pg_roles AS rls_group ON rls_group.oid = pg_auth_members.roleid
+                    JOIN pg_roles AS rls_member ON rls_member.oid = pg_auth_members.member
+                    WHERE rls_group.rolname = 'rls_group' AND rls_member.rolname = '#{name}'
+                ) THEN
+                  REVOKE rls_group FROM #{name};
+                END IF; END $$;
+            SQL
+          end
+        end
+      end
+    end
+  end
+end

data/lib/pg_rls/active_record/connection_adapters/postgre_sql/schema_dumper.rb

@@ -0,0 +1,71 @@
+# frozen_string_literal: true
+
+module PgRls
+  module ActiveRecord
+    module ConnectionAdapters
+      module PostgreSQL
+        # This module is responsible for changing the `create_table` method to `create_rls_table`
+        # when the table is a RLS table.
+        module SchemaDumper
+          def tables(stream)
+            tmp_stream = StringIO.new
+            super(tmp_stream)
+
+            stream_content = tmp_stream.string
+
+            return stream.puts stream_content if @rls_tenant_table.nil?
+
+            stream.puts @rls_tenant_table
+            stream.puts unless stream_content.nil?
+            stream.puts stream_content.chomp
+          end
+
+          def table(table_name, stream)
+            original_table_method = method(:table).super_method
+            stream_content = dump_table_to_string(original_table_method, table_name)
+
+            if rls_tenant_table?(table_name)
+              @rls_tenant_table = stream_content.gsub!("create_table", "create_rls_tenant_table").to_s
+              return
+            elsif rls_table?(table_name)
+              stream_content.gsub!("create_table", "create_rls_table")
+            end
+
+            stream.print stream_content
+          end
+
+          private
+
+          def dump_table_to_string(original_table_method, table_name)
+            temp_stream = StringIO.new
+            original_table_method.call(table_name, temp_stream)
+            temp_stream.string
+          end
+
+          def rls_table?(table_name)
+            rls_table_array.include?(table_name)
+          end
+
+          def rls_tenant_table?(table_name)
+            PgRls.table_name.to_s == table_name.to_s
+          end
+
+          def rls_table_array
+            @rls_table_array ||= fetch_rls_tables
+          end
+
+          def fetch_rls_tables
+            statement = <<-SQL
+              SELECT tablename FROM pg_policies WHERE schemaname = '#{PgRls.schema}';
+            SQL
+            @connection.execute(statement).to_a.map { |table| table["tablename"] }
+          end
+        end
+      end
+    end
+  end
+end
+
+ActiveRecord::ConnectionAdapters::PostgreSQL::SchemaDumper.prepend(
+  PgRls::ActiveRecord::ConnectionAdapters::PostgreSQL::SchemaDumper
+)

data/lib/pg_rls/active_record/connection_adapters/postgre_sql/schema_statements.rb

@@ -0,0 +1,120 @@
+# frozen_string_literal: true
+
+module PgRls
+  module ActiveRecord
+    module ConnectionAdapters
+      module PostgreSQL
+        # This module contains the logic to create rls policies tables and functions
+        module SchemaStatements
+          def create_rls_tenant_table(table_name, ...)
+            create_rls_initialize_setup
+            create_table(table_name, ...)
+            create_rls_tenant_table_setup(table_name)
+          end
+
+          def convert_to_rls_tenant_table(table_name)
+            create_rls_initialize_setup
+            create_rls_tenant_table_setup(table_name)
+          end
+
+          def create_rls_table(table_name, ...)
+            create_table(table_name, ...)
+            create_rls_table_setup(table_name) if check_rls_user_privileges!(PgRls.username, PgRls.schema)
+          end
+
+          def convert_to_rls_table(table_name)
+            create_rls_table_setup(table_name) if check_rls_user_privileges!(PgRls.username, PgRls.schema)
+          end
+
+          def drop_rls_tenant_table(table_name, ...)
+            drop_rls_initialize_setup(table_name)
+            drop_table(table_name, ...)
+          end
+
+          def revert_from_rls_tenant_table(table_name)
+            drop_rls_initialize_setup(table_name)
+            remove_column(table_name, :tenant_id)
+          end
+
+          def drop_rls_table(table_name, ...)
+            drop_rls_table_setup(table_name)
+            drop_table(table_name, ...)
+          end
+
+          def revert_from_rls_table(table_name)
+            drop_rls_table_setup(table_name)
+          end
+
+          def create_rls_index(table_name, columns, **options)
+            add_index(table_name, rls_index_columns(columns), **options)
+          end
+
+          def drop_rls_index(table_name, columns, **options)
+            remove_index(table_name, rls_index_columns(columns), **options)
+          end
+
+          {
+            create_rls_tenant_table: :drop_rls_tenant_table,
+            convert_to_rls_tenant_table: :revert_from_rls_tenant_table,
+            create_rls_table: :drop_rls_table,
+            convert_to_rls_table: :revert_from_rls_table,
+            create_rls_index: :drop_rls_index
+          }.each do |cmd, inv|
+            [[inv, cmd], [cmd, inv]].uniq.each do |method, inverse|
+              class_eval <<-RUBY, __FILE__, __LINE__ + 1
+                def invert_#{method}(args, &block)          # def invert_create_table(args, &block)
+                  [:#{inverse}, args, block]                #   [:drop_table, args, block]
+                end                                         # end
+              RUBY
+            end
+          end
+
+          private
+
+          def create_rls_table_setup(table_name)
+            add_column(table_name, :tenant_id, :uuid, null: false) unless column_exists?(table_name, :tenant_id)
+            enable_table_rls(table_name, PgRls.username)
+            append_rls_table_triggers(table_name)
+          end
+
+          def create_rls_tenant_table_setup(table_name)
+            unless column_exists?(table_name, :tenant_id)
+              add_column(
+                table_name, :tenant_id, :uuid, default: "gen_random_uuid()", null: false
+              )
+            end
+            add_index(table_name, :tenant_id, unique: true) unless index_exists?(table_name, :tenant_id, unique: true)
+            append_tenant_table_triggers(table_name)
+          end
+
+          def create_rls_initialize_setup
+            create_rls_group
+            create_rls_role(PgRls.username, PgRls.password)
+            grant_rls_user_privileges(PgRls.schema)
+            create_rls_functions
+          end
+
+          def drop_rls_table_setup(table_name)
+            drop_rls_table_triggers(table_name)
+            disable_table_rls(table_name, PgRls.username)
+            remove_column(table_name, :tenant_id, if_exists: true) if table_exists?(table_name)
+          end
+
+          def drop_rls_initialize_setup(table_name)
+            drop_tenant_table_triggers(table_name)
+            drop_rls_functions
+            revoke_rls_user_privileges(PgRls.schema)
+            drop_rls_role(PgRls.username)
+            drop_rls_group
+          end
+
+          def rls_index_columns(columns)
+            cols = Array(columns).map(&:to_sym)
+            cols << :tenant_id unless cols.include?(:tenant_id)
+            cols
+          end
+        end
+      end
+    end
+  end
+end

data/lib/pg_rls/active_record/connection_adapters/postgre_sql/sql_helper_method.rb

@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+module PgRls
+  module ActiveRecord
+    module ConnectionAdapters
+      module PostgreSQL
+        # This module contains the logic to validate user privileges
+        module SqlHelperMethod
+          private
+
+          def execute_sql!(statement)
+            ::ActiveRecord::Base.transaction(requires_new: true) do
+              execute(statement.sanitize_sql)
+            rescue StandardError => e
+              raise e unless rescue_sql_error?(e)
+
+              ::ActiveRecord::Base.connection.rollback_db_transaction
+              retry
+            end
+          end
+
+          def rescue_sql_error?(error)
+            error.message.include?("PG::InFailedSqlTransaction") ||
+              error.message.include?("PG::TRDeadlockDetected")
+          end
+        end
+      end
+    end
+  end
+end

data/lib/pg_rls/active_record/connection_adapters/postgre_sql.rb

@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+require "active_record/connection_adapters/postgresql_adapter"
+require_relative "postgre_sql/errors"
+require_relative "postgre_sql/sql_helper_method"
+require_relative "postgre_sql/rls_functions"
+require_relative "postgre_sql/rls_triggers"
+require_relative "postgre_sql/rls_user_statements"
+require_relative "postgre_sql/check_rls_user_privileges"
+require_relative "postgre_sql/grant_rls_user_privileges"
+require_relative "postgre_sql/rls_policies"
+require_relative "postgre_sql/schema_statements"
+require_relative "postgre_sql/schema_dumper"
+
+module PgRls
+  module ActiveRecord
+    module ConnectionAdapters
+      # ActiveRecord PostgreSQL Connection Adapter Extension
+      module PostgreSQL
+        def self.included(base)
+          # Dynamically include all modules into the adapter
+          constants.each do |const_name|
+            next if const_name == :SchemaDumper
+
+            mod = const_get(const_name)
+            base.include(mod) if mod.is_a?(Module) && !mod.is_a?(Class)
+          end
+        end
+      end
+    end
+  end
+end
+
+ActiveRecord::ConnectionAdapters::PostgreSQL::SchemaStatements.include(
+  PgRls::ActiveRecord::ConnectionAdapters::PostgreSQL
+)

data/lib/pg_rls/active_record/connection_adapters.rb

@@ -0,0 +1,12 @@
+# frozen_string_literal: true
+
+require_relative "connection_adapters/postgre_sql"
+require_relative "connection_adapters/connection_pool"
+
+module PgRls
+  module ActiveRecord
+    # ActiveRecord Connection Adapter Extension
+    module ConnectionAdapters
+    end
+  end
+end

data/lib/pg_rls/active_record/database_shards.rb

@@ -0,0 +1,74 @@
+# frozen_string_literal: true
+
+module PgRls
+  module ActiveRecord
+    # Overwrite the configurations method to add the RLS configurations
+    module DatabaseShards
+      REQUIRED_CONFIGURATION_KEYS = %w[adapter host database username password rls_mode].freeze
+
+      def add_rls_configurations(config, new_config = {})
+        current_config = config[Rails.env]
+
+        if rls_shard_config?(current_config)
+          add_primary_and_rls_config(current_config, new_config)
+        else
+          current_config.each do |key, value|
+            add_primary_and_rls_config(value, new_config, key)
+          end
+        end
+
+        { Rails.env => new_config }
+      end
+
+      def configurations=(config)
+        new_config = add_rls_configurations(config)
+        super(new_config)
+      end
+
+      private
+
+      def add_primary_and_rls_config(config, new_config, key = "primary")
+        new_config[key] = config
+
+        return new_config unless rls_shard_config?(config)
+
+        configuration = adapter_configurations(config, new_config, key)
+
+        if configuration.nil?
+          raise ArgumentError,
+                "Invalid RLS mode: #{config["rls_mode"]}. valid options are: dual, single, none"
+        end
+
+        new_config
+      end
+
+      def adapter_configurations(config, new_config, key)
+        case config["rls_mode"]
+        when "dual"
+          new_config["rls_#{key}"] = config.merge(rls_configuration)
+        when "single"
+          new_config[key] = config.merge(rls_configuration)
+        when "none"
+          new_config[key] = config
+        end
+      end
+
+      def rls_configuration
+        {
+          "username" => PgRls.username.to_s,
+          "password" => PgRls.password.to_s,
+          "database_tasks" => false,
+          "rls" => true
+        }
+      end
+
+      def rls_shard_config?(config)
+        return false unless config.is_a?(Hash)
+
+        REQUIRED_CONFIGURATION_KEYS.all? { |key| config.key?(key) }
+      end
+    end
+  end
+end
+
+ActiveRecord::Base.singleton_class.prepend(PgRls::ActiveRecord::DatabaseShards)

data/lib/pg_rls/active_record/migration/command_recorder.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+module PgRls
+  module ActiveRecord
+    module Migration
+      # ActiveRecord Migration Command Recorder Extension
+      module CommandRecorder
+        REVERSIBLE_AND_IRREVERSIBLE_METHODS = %i[
+          create_rls_tenant_table convert_to_rls_tenant_table drop_rls_tenant_table
+          create_rls_table convert_to_rls_table drop_rls_table
+        ].freeze
+
+        def self.included(base)
+          REVERSIBLE_AND_IRREVERSIBLE_METHODS.each do |method|
+            base.class_eval <<-RUBY, __FILE__, __LINE__ + 1
+              def #{method}(*args, &block)          # def create_table(*args, &block)
+                record(:"#{method}", args, &block)  #   record(:create_table, args, &block)
+              end                                   # end
+            RUBY
+            base.send(:ruby2_keywords, method)
+          end
+        end
+      end
+    end
+  end
+end
+
+ActiveRecord::Migration::CommandRecorder.include(PgRls::ActiveRecord::Migration::CommandRecorder)

data/lib/pg_rls/active_record/migration.rb

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+require_relative "migration/command_recorder"
+
+module PgRls
+  module ActiveRecord
+    # ActiveRecord Migration Extension
+    module Migration
+    end
+  end
+end

data/lib/pg_rls/active_record/test_databases.rb

@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+module PgRls
+  module ActiveRecord
+    module TestDatabases # :nodoc:
+      def create_and_load_schema(i, env_name:)
+        super
+
+        PgRls::Record.configurations.configs_for(env_name: env_name, include_hidden: true).each do |db_config|
+          next if db_config.name == "primary"
+
+          db_config._database = "#{db_config.database}-#{i}"
+        end
+      end
+    end
+  end
+end
+
+ActiveRecord::TestDatabases.singleton_class.prepend(PgRls::ActiveRecord::TestDatabases)

data/lib/pg_rls/active_record.rb

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+require_relative "active_record/migration"
+require_relative "active_record/connection_adapters"
+require_relative "active_record/database_shards"
+
+module PgRls
+  # ActiveRecord Extension
+  module ActiveRecord
+  end
+end

data/lib/pg_rls/active_support/string_ext.rb

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+module PgRls
+  module ActiveSupport
+    # Extensions to the String class
+    module StringExt
+      def sanitize_sql
+        str = dup
+        str.gsub!(/[[:space:]]+/, " ")
+        str.strip!
+        str.to_s
+      end
+    end
+  end
+end
+
+String.include(PgRls::ActiveSupport::StringExt)

data/lib/pg_rls/active_support.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+require_relative "active_support/string_ext"
+
+module PgRls
+  # Active Support Extension
+  module ActiveSupport
+  end
+end

data/lib/pg_rls/connection_config.rb

@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+module PgRls
+  # The ConnectionConfig class provides methods to configure and manage
+  # the connection settings for Row Level Security (RLS) in PostgreSQL.
+  # It includes methods to look up and validate the connection configuration,
+  # as well as helper methods to build the configuration hash based on different modes.
+  class ConnectionConfig
+    def initialize(db_config = ::ActiveRecord::Base.connection_db_config)
+      @db_config = db_config
+      @connection_name = db_config.name
+    end
+
+    def look_up_connection_config
+      config_hash = build_config_hash(@db_config, @connection_name)
+
+      return invalid_connection_config unless config_hash
+
+      PgRls.connects_to = config_hash.deep_transform_values(&:to_sym)
+    end
+
+    def connection_config?
+      PgRls.connects_to&.key?(:database) || false
+    end
+
+    def invalid_connection_config
+      raise PgRls::Error::InvalidConnectionConfig,
+            "you must edit your database.yml file to include the RLS configuration, " \
+            "or set the RLS configuration manually in the PgRls initializer"
+    end
+
+    private
+
+    def build_config_hash(db_config, connection_name)
+      case db_config.configuration_hash[:rls_mode]
+      when "dual"
+        build_dual_mode_config(connection_name)
+      when "single", "none"
+        build_single_mode_config(connection_name)
+      end
+    end
+
+    def build_dual_mode_config(connection_name)
+      {
+        shards: {
+          rls: { writing: "rls_#{connection_name}", reading: "rls_#{connection_name}" },
+          admin: { writing: connection_name, reading: connection_name }
+        }
+      }
+    end
+
+    def build_single_mode_config(connection_name)
+      {
+        database: {
+          writing: connection_name,
+          reading: connection_name
+        }
+      }
+    end
+  end
+end

data/lib/pg_rls/deprecation.rb

@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+module PgRls
+  # Deprecator Module
+  module Deprecation
+    def self.warn(message)
+      logger.warn(message)
+    end
+
+    def self.logger
+      @logger ||= ::ActiveSupport::Deprecation.new(PgRls::VERSION, "PgRls")
+    end
+  end
+end

data/lib/pg_rls/engine.rb

@@ -0,0 +1,8 @@
+# frozen_string_literal: true
+
+module PgRls
+  # Engine
+  class Engine < ::Rails::Engine
+    isolate_namespace PgRls
+  end
+end

data/lib/pg_rls/error.rb

@@ -0,0 +1,10 @@
+# frozen_string_literal: true
+
+module PgRls
+  # Main PgRls Error Class
+  class Error < StandardError
+    class InvalidConnectionConfig < Error; end
+    class InvalidSearchInput < Error; end
+    class TenantNotFound < Error; end
+  end
+end