pg_rls 0.2.5 → 1.0.1

data/app/models/pg_rls/tenant/searchable.rb

@@ -0,0 +1,60 @@
+# frozen_string_literal: true
+
+module PgRls
+  class Tenant
+    # Searchable Class
+    class Searchable
+      SEARCH_METHOD_BY_COLUMN_TYPE = PgRls::Tenant.columns.each_with_object({}) do |column, hash|
+        type = column.sql_type_metadata.type
+        hash[type] ||= []
+        hash[type] << column.name
+      end
+
+      def self.by_rls_object(tenant)
+        case tenant
+        when Tenant then tenant
+        when String, Symbol, Integer then Searchable.by_rls_methods(tenant)
+        when PgRls.main_model then tenant.becomes(Tenant)
+        else raise PgRls::Error::InvalidSearchInput, "Invalid search input: #{tenant}"
+        end
+      end
+
+      def self.by_rls_methods(search_input)
+        new(search_input).search_methods.each do |search_method|
+          tenant = PgRls::Tenant.find_by(search_method => search_input)
+          return tenant if tenant.present?
+        end
+
+        raise PgRls::Error::TenantNotFound, "No tenant found for #{search_input}"
+      end
+
+      def initialize(search_input)
+        @search_input = search_input
+      end
+
+      def search_methods
+        SEARCH_METHOD_BY_COLUMN_TYPE[target_type] || []
+      end
+
+      private
+
+      attr_reader :search_input
+
+      def target_type
+        case search_input
+        when Integer then :integer
+        when String, Symbol then target_type_by_string
+        else raise PgRls::Error::InvalidSearchInput, "Invalid search input: #{search_input}"
+        end
+      end
+
+      def target_type_by_string
+        case search_input
+        when /\A\d+\z/ then :integer
+        when /\A[0-9a-fA-F\-]{36}\z/ then :uuid
+        else :string
+        end
+      end
+    end
+  end
+end

data/app/models/pg_rls/tenant/securable.rb

@@ -0,0 +1,67 @@
+# frozen_string_literal: true
+
+module PgRls
+  class Tenant
+    # Securable Module
+    module Securable
+      extend ::ActiveSupport::Concern
+
+      included do
+        self.table_name = PgRls.table_name
+
+        self.ignored_columns = column_names.reject do |column|
+          PgRls.search_methods.map(&:to_s).include?(column)
+        end
+      end
+
+      class_methods do
+        def reset_rls_used_connections(connection = PgRls::Record.connection)
+          return connection if rls_connection_object_cache_by_thread.nil?
+
+          connection.exec_query("SET rls.tenant_id TO DEFAULT")
+          self.rls_connection_object_cache_by_thread = nil
+          connection
+        end
+
+        def rls_connection_object_cache_by_thread=(value)
+          ::ActiveSupport::IsolatedExecutionState[:active_record_rls_used_connections] = value
+        end
+
+        def rls_connection_object_cache_by_thread
+          ::ActiveSupport::IsolatedExecutionState[:active_record_rls_used_connections]
+        end
+      end
+
+      def set_rls(connection = PgRls::Record.connection)
+        self.class.reset_rls_used_connections if new_tenant?
+        return self if reused_connection?(connection)
+
+        connection.exec_query("SET rls.tenant_id = '#{tenant_id}'")
+
+        self
+      end
+
+      def readonly?
+        true
+      end
+
+      private
+
+      def new_tenant?
+        !rls_used_connections.add?(tenant_id).nil?
+      end
+
+      def reused_connection?(conn)
+        rls_used_connections.add?(conn.object_id).nil?
+      end
+
+      def rls_used_connections
+        if self.class.rls_connection_object_cache_by_thread.nil?
+          self.class.rls_connection_object_cache_by_thread = Set.new([tenant_id])
+        end
+
+        self.class.rls_connection_object_cache_by_thread
+      end
+    end
+  end
+end

data/app/models/pg_rls/tenant/switchable.rb

@@ -0,0 +1,40 @@
+# frozen_string_literal: true
+
+module PgRls
+  class Tenant
+    # Switchable Module
+    module Switchable
+      extend ::ActiveSupport::Concern
+
+      def switch(tenant)
+        set_rls!(tenant)
+      rescue PgRls::Error::TenantNotFound, PgRls::Error::InvalidSearchInput
+        nil
+      end
+
+      def run_within(tenant)
+        set_rls!(tenant)
+        yield(PgRls::Current.tenant).presence
+      ensure
+        PgRls::Tenant.reset_rls
+      end
+
+      def with_tenant!(...)
+        PgRls::Deprecation.warn("This method is deprecated and will be removed in future versions. " \
+                                "please use PgRls::Tenant.run_within instead.")
+        run_within(...)
+      end
+
+      def set_rls!(tenant)
+        PgRls::Current.tenant = Searchable.by_rls_object(tenant)
+        raise PgRls::Error::TenantNotFound, "No tenant found for #{tenant}" unless PgRls::Current.tenant.present?
+
+        PgRls::Current.tenant
+      end
+
+      def reset_rls
+        PgRls::Current.reset
+      end
+    end
+  end
+end

data/app/models/pg_rls/tenant.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+module PgRls
+  # Tenant model
+  class Tenant < Record
+    include Securable
+    extend Switchable
+  end
+end

data/assets/logo.svg

@@ -0,0 +1,8 @@
+<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 200 200">
+  <rect width="200" height="200" fill="#f8f8f8"/>
+  <g transform="translate(20, 20)">
+    <path d="M80 0 L160 40 L160 120 L80 160 L0 120 L0 40 Z" fill="#ff4444"/>
+    <path d="M80 20 L140 52 L140 108 L80 140 L20 108 L20 52 Z" fill="#ffffff"/>
+    <text x="80" y="95" font-family="Arial, sans-serif" font-size="36" font-weight="bold" text-anchor="middle" fill="#ff4444">PgRls</text>
+  </g>
+</svg>

data/docker-compose.yml

@@ -0,0 +1,14 @@
+services:
+  postgres:
+    image: postgres:latest
+    environment:
+      POSTGRES_PASSWORD: password
+    healthcheck:
+      test: pg_isready -U postgres
+    ports:
+      - "5432:5432"
+    volumes:
+      - postgres:/var/lib/postgresql/data
+
+volumes:
+  postgres:

data/lib/generators/pg_rls/active_record/active_record_generator.rb

@@ -1,115 +1,112 @@
 # frozen_string_literal: true
 
-require 'rails/generators/active_record/model/model_generator'
-require File.join(File.dirname(__FILE__), '../base')
+require "rails/generators/active_record/model/model_generator"
+require_relative "../install/install_generator"
 
 module PgRls
   module Generators
     # Active Record Generator
     class ActiveRecordGenerator < ::ActiveRecord::Generators::ModelGenerator
-      include ::PgRls::Base
+      source_root File.expand_path("../templates", __dir__.to_s)
 
-      source_root File.expand_path('./templates', __dir__)
+      class_option :parent, type: :string, default: "ApplicationRecord",
+                            desc: "The parent class for the generated model"
 
-      def check_class_collision; end
+      class_option :rls_parent, type: :string, default: "PgRls::Record",
+                                desc: "The parent class for the rls generated model"
 
-      def create_migration_file; end
-
-      def migration_exist?
-        @migration_exist ||= Dir.glob("#{migration_path}/*create_#{table_name}.rb").present?
+      # Need to override so it will not check for class collision
+      def check_class_collision
+        super
+        @class_coalescing = false
+      rescue Rails::Generators::Error
+        @class_coalescing = true
       end
 
-      def create_tenant_migration_file
-        return unless creating?
+      def create_model_file
+        return if class_coalescing?
 
-        migration_template(create_migration_template_path,
-                           "#{migration_path}/#{create_file_sub_name}_#{table_name}.rb",
-                           migration_version:)
+        generate_abstract_class if database && !custom_parent?
+        template("app/models/model.rb", model_path, parent_class_name: parent_class_name)
       end
 
-      def convert_tenant_migration_file
-        unless creating?
-          migration_template(convert_migration_template_path,
-                             "#{migration_path}/#{convert_file_sub_name}_#{table_name}.rb",
-                             migration_version:)
-        end
-
-        return unless installation_in_progress?
+      def upgrade_model_file
+        return unless class_coalescing? && model_file_exists?
 
-        migration_template('convert_migration_backport.rb.tt',
-                           "#{migration_path}/pg_rls_backport_#{table_name}.rb",
-                           migration_version:)
+        gsub_file(model_path, /< ApplicationRecord/, "< #{parent_class_name}")
       end
 
-      def create_model_file
-        return if migration_exist?
-
-        generate_abstract_class if database && !parent
+      def create_migration_file
+        return if skip_migration_creation? || class_coalescing?
 
-        template model_template_path, model_file
+        clean_indexes_attributes!
+        migration_template "db/migrate/create_#{migration_template_prefix}_table.rb",
+                           "db/migrate/create_#{migration_template_prefix}_#{table_name}.rb"
       end
 
-      def inject_method_to_model
-        return unless installation_in_progress?
-
-        gsub_file(model_file, /Class #{class_name} < #{parent_class_name.classify}/mi) do |match|
-          "#{match}\n  def self.current\n    PgRls::Tenant.fetch\n  end\n"
-        end
-      end
+      def upgrade_migration_file
+        return if skip_migration_creation? || !class_coalescing?
 
-      def model_file
-        File.join('app/models', class_path, "#{file_name}.rb")
+        migration_template "db/migrate/convert_to_#{migration_template_prefix}_table.rb",
+                           "db/migrate/convert_to_#{migration_template_prefix}_#{table_name}.rb"
       end
 
-      def create_migration_template_path
-        return 'init_migration.rb.tt' if installation_in_progress?
+      def backport_migration_file
+        return if skip_migration_creation? || installing? || !class_coalescing?
 
-        'migration.rb.tt'
+        migration_template "db/migrate/backport_pg_rls_table.rb",
+                           "db/migrate/backport_pg_rls_to_#{table_name}.rb"
       end
 
-      def convert_migration_template_path
-        return 'init_convert_migration.rb.tt' if installation_in_progress?
+      private
 
-        'convert_migration.rb.tt'
+      def installing?
+        Kernel.const_defined?("PgRls::InstallGenerator")
       end
 
-      def model_template_path
-        return 'init_model.rb.tt' if installation_in_progress?
+      def parent_class_name
+        return rls_parent unless installing?
 
-        'model.rb.tt'
+        super
       end
 
-      def create_file_sub_name
-        return 'pg_rls_create_tenant' if installation_in_progress?
+      def rls_parent
+        options[:rls_parent]
+      end
 
-        'pg_rls_create'
+      def migration_template_prefix
+        installing? ? "pg_rls_tenant" : "pg_rls"
       end
 
-      def convert_file_sub_name
-        return 'pg_rls_convert_tenant' if installation_in_progress?
+      def generate_abstract_class
+        return if File.exist?(generate_abstract_class_path)
 
-        'pg_rls_convert'
+        template "app/models/abstract_base_class.rb", generate_abstract_class_path
       end
 
-      def installation_in_progress?
-        shell.base.class.name.include?('Install')
+      def model_path
+        File.join("app/models", class_path, "#{file_name}.rb")
       end
 
-      def migration_version
-        "[#{Rails::VERSION::MAJOR}.#{Rails::VERSION::MINOR}]"
+      def generate_abstract_class_path
+        File.join("app/models", "#{database.underscore}_record.rb")
       end
 
-      def migration_path
-        db_migrate_path
+      def class_coalescing?
+        @class_coalescing
       end
 
-      def creating?
-        @creating ||= !migration_exist?
-      end
+      def clean_indexes_attributes!
+        return unless options[:indexes] == false
 
-      protected
+        attributes.each do |a|
+          a.attr_options.delete(:index) if a.reference? && !a.has_index?
+        end
+      end
 
-      def migration_action = 'add'
+      def model_file_exists?
+        File.exist?(File.join(destination_root, model_path))
+      end
     end
   end
 end

data/lib/generators/pg_rls/install/install_generator.rb

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+module PgRls
+  # Generator to install the gem
+  class InstallGenerator < Rails::Generators::Base
+    source_root File.expand_path("../templates", __dir__.to_s)
+
+    attr_reader :attributes
+
+    def initialize(args, *options)
+      pg_rls_config(args.first)
+      super
+    end
+
+    def create_install_config
+      template "config/initializers/pg_rls.rb"
+    end
+
+    hook_for :orm, required: true
+
+    def show_readme
+      readme "USAGE" if invoke?
+    end
+
+    private
+
+    def pg_rls_config(tenant_model_or_table)
+      raise "Tenant model or table name is required" if tenant_model_or_table.blank?
+
+      PgRls.class_name = tenant_model_or_table.capitalize.singularize.to_sym
+      PgRls.table_name = tenant_model_or_table.underscore.pluralize.to_sym
+    end
+
+    def invoke?
+      behavior == :invoke
+    end
+  end
+end

data/lib/generators/pg_rls/pg_rls_generator.rb

@@ -1,9 +1,10 @@
 # frozen_string_literal: true
 
-require_relative File.join(File.dirname(__FILE__), 'active_record/active_record_generator')
+require_relative File.join(File.dirname(__FILE__), "active_record/active_record_generator")
 
 module PgRls
   module Generators
+    # PgRls Generator
     class PgRlsGenerator < ::Rails::Generators::NamedBase
       # override ModelGenerator
       hook_for :orm, required: true

data/lib/generators/pg_rls/templates/USAGE

@@ -0,0 +1,28 @@
+===============================================================================
+
+🎉 Congratulations! The PgRls gem has been successfully installed.
+
+To generate a model with secure Row-Level Security (RLS), run the following command:
+
+   rails g pg_rls ModelName
+
+or
+
+   rails generate pg_rls ModelName
+
+📌 Remember:
+
+- Replace `ModelName` with the actual name of your model (e.g., `User`, `Post`).
+- This will set up the necessary RLS policies and table configurations for multi-tenant support.
+
+🛠  Next Steps:
+
+1. After generating the model, review the generated migration and model files.
+2. Ensure your database is configured to support Row-Level Security.
+3. Run `rails db:migrate` to apply the generated migration.
+4. Customize the generated policies to fit your application's security needs.
+
+💡 Tip: Check the documentation at <https://github.com/Dandush03/pg_rls>
+   for more details on how to fine-tune RLS policies.
+
+===============================================================================

data/lib/generators/pg_rls/{active_record/templates → templates/app/models}/abstract_base_class.rb.tt

@@ -1,7 +1,5 @@
-# frozen_string_literal: true
-
 <% module_namespacing do -%>
-class <%= abstract_class_name %> < ApplicationRecord
+class <%= abstract_class_name %> < <%= parent_class_name.classify %>
   self.abstract_class = true
 
   connects_to database: { <%= ActiveRecord.writing_role %>: :<%= database -%> }

data/lib/generators/pg_rls/{active_record/templates → templates/app/models}/model.rb.tt

@@ -1,5 +1,3 @@
-# frozen_string_literal: true
-
 <% module_namespacing do -%>
 class <%= class_name %> < <%= parent_class_name.classify %>
 <% attributes.select(&:reference?).each do |attribute| -%>

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

@@ -0,0 +1,58 @@
+# frozen_string_literal: true
+
+# PgRls Configuration
+# ----------------------
+# Use this file to configure PgRls to fit your application's needs.
+# Below are the default settings and examples of how to customize them.
+#
+# You can modify these settings as needed and restart your application.
+
+PgRls.setup do |config|
+  # The main class name used for the RLS tenant model.
+  # This should be the base model representing your tenant (e.g., :organization, :account, etc.).
+  # Default: Derived from your tenant model specified during installation.
+  config.class_name = :<%= PgRls.class_name %>
+
+  # The name of the database table where tenant information is stored.
+  # Ensure this table is properly configured with RLS enabled.
+  # Default: Derived from your tenant model.
+  config.table_name = :<%= PgRls.table_name %>
+
+  # The methods used to search for the tenant. These methods should exist on the tenant model.
+  # Default: Search methods specified during installation.
+  config.search_methods = <%= "%i[#{PgRls.search_methods.join(' ')}]" %>
+
+  # The attributes that should be tracked in the request context.
+  # This allows you to specify which attributes PgRls::Current should manage dynamically.
+  # By default, this is an empty array, as only the current tenant is tracked for RLS purposes.
+  # You can add additional attributes when needed, such as theme, branch, etc.
+  config.current_attributes = []
+
+  # If you need to track additional attributes dynamically, use the `__` convention.
+  # This allows calling nested models with a double underscore, inspired by Stimulus controllers.
+  # Example: If `organization__branch` is added, it will resolve to Organization::Branch.first
+  # config.current_attributes = %i[organization__branch]
+
+  # Database credentials for the RLS user (if needed).
+  # If using Rails credentials, ensure they are correctly stored in config/credentials.yml.enc.
+  # Uncomment and update these lines if you need to provide explicit credentials.
+  # config.username = Rails.application.credentials.dig(:database, :username)
+  # config.password = Rails.application.credentials.dig(:database, :password)
+
+  # The base class from which PgRls::Record should inherit.
+  # This allows your PgRls models to inherit from the same base class as your application models.
+  # For Rails 5+ applications that use ApplicationRecord, this is particularly useful.
+  # Default: 'ActiveRecord::Base'
+  # config.abstract_base_record_class = 'ApplicationRecord'
+
+  # The database role group used for RLS operations.
+  # This role acts as a base role, allowing multiple users to be assigned this role for consistent access control.
+  # Default: "rls_group"
+  # config.rls_role_group = "rls_group"
+
+  # The database schema where your tenant tables are located.
+  # Default: "public"
+  # config.schema = "public"
+
+  # You can add additional custom configurations here if needed.
+end

data/lib/generators/pg_rls/templates/db/migrate/backport_pg_rls_table.rb.tt

@@ -0,0 +1,14 @@
+class BackportPgRlsTo<%= table_name.camelize %> < ActiveRecord::Migration[<%= ActiveRecord::Migration.current_version %>]
+  def up
+    # Suggested Code:
+    # PgRls.on_each_tenant do |tenant|
+    #   <%= class_name.camelize %>.where(identifier_reference_for_tenant: tenant.id)
+    #      .in_batches.update_all(tenant_id: tenant.tenant_id)
+    # end
+  end
+
+  def down
+    # Suggested Code:
+    # raise ActiveRecord::IrreversibleMigration, 'This migration is irreversible, please restore from backup.'
+  end
+end

data/lib/generators/pg_rls/templates/db/migrate/convert_to_pg_rls_table.rb.tt

@@ -0,0 +1,5 @@
+class ConvertToPgRls<%= table_name.camelize %> < ActiveRecord::Migration[<%= ActiveRecord::Migration.current_version %>]
+  def change
+    convert_to_rls_table :<%= table_name %>
+  end
+end

data/lib/generators/pg_rls/templates/db/migrate/convert_to_pg_rls_tenant_table.rb.tt

@@ -0,0 +1,5 @@
+class ConvertToPgRlsTenant<%= table_name.camelize %> < ActiveRecord::Migration[<%= ActiveRecord::Migration.current_version %>]
+  def change
+    convert_to_rls_tenant_table :<%= table_name %>
+  end
+end

data/lib/generators/pg_rls/templates/db/migrate/create_pg_rls_table.rb.tt

@@ -0,0 +1,29 @@
+class CreatePgRls<%= table_name.camelize %> < ActiveRecord::Migration[<%= ActiveRecord::Migration.current_version %>]
+  def change
+    create_rls_table :<%= table_name %><%= primary_key_type %> do |t|
+<% attributes.each do |attribute| -%>
+<% if attribute.password_digest? -%>
+      t.string :password_digest<%= attribute.inject_options %>
+<% elsif attribute.token? -%>
+      t.string :<%= attribute.name %><%= attribute.inject_options %>
+<% elsif attribute.reference? -%>
+      t.<%= attribute.type %> :<%= attribute.name %><%= attribute.inject_options %><%= foreign_key_type %>
+<% elsif !attribute.virtual? -%>
+      t.<%= attribute.type %> :<%= attribute.name %><%= attribute.inject_options %>
+<% end -%>
+<% end -%>
+<% unless attributes.empty? -%>
+
+<% end -%>
+<% if options[:timestamps] -%>
+      t.timestamps
+<% end -%>
+    end
+<% attributes.select(&:token?).each do |attribute| -%>
+    add_index :<%= table_name %>, :<%= attribute.index_name %><%= attribute.inject_index_options %>, unique: true
+<% end -%>
+<% attributes_with_index.each do |attribute| -%>
+    add_index :<%= table_name %>, :<%= attribute.index_name %><%= attribute.inject_index_options %>
+<% end -%>
+  end
+end

data/lib/generators/pg_rls/templates/db/migrate/create_pg_rls_tenant_table.rb.tt

@@ -0,0 +1,29 @@
+class CreatePgRlsTenant<%= table_name.camelize %> < ActiveRecord::Migration[<%= ActiveRecord::Migration.current_version %>]
+  def change
+    create_rls_tenant_table :<%= table_name %><%= primary_key_type %> do |t|
+<% attributes.each do |attribute| -%>
+<% if attribute.password_digest? -%>
+      t.string :password_digest<%= attribute.inject_options %>
+<% elsif attribute.token? -%>
+      t.string :<%= attribute.name %><%= attribute.inject_options %>
+<% elsif attribute.reference? -%>
+      t.<%= attribute.type %> :<%= attribute.name %><%= attribute.inject_options %><%= foreign_key_type %>
+<% elsif !attribute.virtual? -%>
+      t.<%= attribute.type %> :<%= attribute.name %><%= attribute.inject_options %>
+<% end -%>
+<% end -%>
+<% unless attributes.empty? -%>
+
+<% end -%>
+<% if options[:timestamps] -%>
+      t.timestamps
+<% end -%>
+    end
+<% attributes.select(&:token?).each do |attribute| -%>
+    add_index :<%= table_name %>, :<%= attribute.index_name %><%= attribute.inject_index_options %>, unique: true
+<% end -%>
+<% attributes_with_index.each do |attribute| -%>
+    add_index :<%= table_name %>, :<%= attribute.index_name %><%= attribute.inject_index_options %>
+<% end -%>
+  end
+end

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

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+module PgRls
+  module ActiveRecord
+    module ConnectionAdapters
+      # ActiveRecord ConnectionPool Connection Adapter Extension
+      module ConnectionPool
+        def checkout(checkout_timeout = @checkout_timeout)
+          conn = super
+          return conn unless rls_connection?
+          return reset_rls_used_connections(conn) if PgRls::Current.tenant.nil?
+
+          PgRls::Current.tenant.set_rls(conn)
+          conn
+        end
+
+        def rls_connection?
+          pool_config.db_config.configuration_hash[:rls] == true
+        end
+
+        def reset_rls_used_connections(connection)
+          PgRls::Tenant.reset_rls_used_connections(connection)
+        end
+      end
+    end
+  end
+end
+
+ActiveRecord::ConnectionAdapters::ConnectionPool.prepend(
+  PgRls::ActiveRecord::ConnectionAdapters::ConnectionPool
+)