pg_rls 0.2.4 → 1.0.0

data/README.md

@@ -1,24 +1,26 @@
-<!--
-  Title: PgRls Rails
-  Description: rails multitenancy with pg rls
-  Author: dandush03
-  -->
-<meta name="google-site-verification" content="Mc1vBv8PRYPw_cdd3EiKhF2vlOeIEIk3VYhAg75ertI" />
+# PgRls Rails
+
+> PostgreSQL Row Level Security: The Rails right way to do multitenancy
 
 [![Contributors][contributors-shield]][contributors-url]
 [![Forks][forks-shield]][forks-url]
 [![Stargazers][stars-shield]][stars-url]
 [![Issues][issues-shield]][issues-url]
-[![LinkedIn][linkedin-shield2]][linkedin-url2]
-[![Hireable][hireable]][hireable-url]
+[![MIT License][license-shield]][license-url]
+[![LinkedIn][linkedin-shield]][linkedin-url]
 [![Donate][donate]][paypal-donate-code]
+[![Hireable][hireable]][hireable-url]
 
-<!-- PROJECT LOGO -->
-<br />
 <p align="center">
- <h1 align="center">PgRls<h2 align="center">PostgreSQL Row Level Security<br />The Rails right way to do multitenancy</h2></h1>
+
+  <h3 align="center">
+    <a href="https://github.com/Dandush03/pg_rls">
+        <img src="./assets/logo.svg" alt="Logo" width="80" height="80">
+    </a>
+  </h3>
 
   <p align="center">
+    PostgreSQL Row Level Security: The Rails right way to do multitenancy
     <br />
     <a href="https://github.com/Dandush03/pg_rls/wiki"><strong>Explore the docs »</strong></a>
     <br />
@@ -26,125 +28,288 @@
     <a href="https://github.com/Dandush03/pg_rls/issues">Report Bug</a>
     ·
     <a href="https://github.com/Dandush03/pg_rls/issues">Request Feature</a>
-    ·
-    <a href="https://github.com/Dandush03/pg_rls">API Repo</a>
   </p>
-
 </p>
 
-### Table of Contents
-* [Required Installations](#required-Installations)
-  * [Installing](#installing)
-  * [Instructions](#instructions)
-  * [Testing](#Testing)
-* [Development](#testing)
-* [Contact](#contact)
-* [Contributing](#contributing)
-* [License](#license)
-* [Code of Conduct](#Code-of-Conduct)
-* [Show your support](#Show-your-support)
+## Table of Contents
+
+- [About The Project](#about-the-project)
+- [Getting Started](#getting-started)
+  - [Prerequisites](#prerequisites)
+  - [Installation](#installation)
+  - [Configuration](#configuration)
+  - [How It Works](#how-it-works)
+- [Usage](#usage)
+- [RLS Index Management Methods](#rls-index-management-methods)
+- [Testing](#testing)
+- [Development](#development)
+  - [Development Workflow](#development-workflow)
+  - [Releasing a New Version](#releasing-a-new-version)
+- [Contributing](#contributing)
+- [License](#license)
+- [Contact](#contact)
+- [Acknowledgements](#acknowledgements)
+
+## About The Project
+
+It's time we start doing multitenancy right! You can avoid creating a separate Postgres schema/databases for each customer or trying to ensure the WHERE clause of every single query includes the particular company. Just integrate PgRls seamlessly to your application.
+
+This gem will integrate PostgreSQL RLS to help you develop a great multitenancy application.
+
+## Getting Started
+
+### Prerequisites
 
-### It's time we start doing multitenancy right! You can avoid creating a separate Postgres schema/databases for each customer or trying to ensure the WHERE clause of every single query includes the particular company. Just integrate PgRls seamlessly to your application.
+- Ruby (~> 3.0)
+- ActiveRecord (~> 7.0)
+- PostgreSQL (> 9.0)
+- Warden
+- pg (~> 1.2)
 
-### This gem will integrate PostgreSQL RLS to help you develop a great multitenancy application.
+### Installation
 
-## Required Installation
-### Installing
+1. Add this line to your application's Gemfile:
 
-Add this line to your application's Gemfile:
+   ```ruby
+   gem 'pg_rls'
+   ```
+
+2. Execute:
+
+   ```bash
+   bundle install
+   ```
+
+   Or install it yourself with:
+
+   ```bash
+   gem install pg_rls
+   ```
+
+### Configuration
+
+You must configure the `rls_mode` in your `database.yml` file. This setting controls how RLS (Row-Level Security) connections are handled for your app. It supports three modes:
+
+- none: No RLS connections.
+- single: Only RLS connections, which is ideal for production environments.
+- dual: Both RLS and non-RLS connections, mainly for development and testing.
+
+Example configuration in database.yml for development:
 
 ```ruby
-gem 'pg_rls'
+development:
+  <<: *default
+  database: dev_db
+  # Use 'dual' for development to switch between RLS and non-RLS connections.
+  rls_mode: <%= ENV.fetch('RLS_MODE', 'dual') %>
 ```
 
-And then execute:
+#### Using the `dual` mode is not recommended in high-demand environments, as it will duplicate the connection pool for each RLS shard, leading to unnecessary overhead. Instead, configure the RLS mode to `single` or `none` and balance your requests accordingly. The `single` mode ensures only RLS connections are used, while `none` disables RLS for this environment
 
-    $ bundle install
+#### For flexible production configurations, you can use an environment variable to set the rls_mode
 
-Or install it yourself with:
+```ruby
+production:
+  <<: *default
+  database: prod_db
+  rls_mode: <%= ENV.fetch('RLS_MODE', 'dual') %>
+```
 
-    $ gem install pg_rls
+### How It Works
 
-### Instructions
+The `rls_mode` setting in your `database.yml` controls how your application handles database connections with Row-Level Security (RLS). Here's a breakdown of how each mode functions:
 
-```bash
-rails generate pg_rls:install company #=> where company eq tenant model name
+1. **Single Mode:** In this mode, the application will modify the database connection’s `username` to `PgRls.username`, ensuring that all queries are executed with RLS rules enabled.
+   - **Effect:** Only RLS connections are used, and all operations are securely performed within the specified tenant context. This is recommended for production environments where strict RLS enforcement is needed, or when the application does not need to execute queries as an "admin" user.
+   - **Use Case:** When RLS is required for all operations to prevent unauthorized data access.
+
+2. **None Mode**: This mode does not modify any shards or connections. The username and connection behavior remain as defined in your database.yml, without applying the RLS username.
+   - **Effect:** No RLS rules are applied, and the application operates in a traditional mode without tenant-based restrictions.
+   - **Use Case:** Useful when you do not need to enforce tenant isolation through RLS, such as for administrative tasks or environments that do not require multi-tenancy.
+
+3. **Dual Mode:** In this mode, the application will duplicate each shard that has RLS enabled, adding a prefix of `rls_` to the shard name. Both RLS and non-RLS connections will be available. For example, if your shard is named animals, the RLS version will be named `rls_animals`.
+   - **Effect:** This maintains two connection pools per shard, one with RLS (`rls_` prefixed) and one without. While it provides flexibility, it also increases resource consumption by duplicating the connection pool.
+   - **Use Case:** This mode can be used in production environments for applications that require both RLS and non-RLS connections but is not recommended for extremely high-demand environments due to the overhead caused by duplicating the connection pools. In less demanding production settings, it offers useful flexibility.
+
+### Configuring `PgRls::Current`
+
+You can configure `PgRls::Current` dynamically using an initializer. This allows you to specify the attributes that should be tracked in the request context. By default, `PgRls::Current` stores tenant-related information, but you can extend it as needed.
+
+#### **Defining Custom Current Attributes**
+To configure the attributes, update your Rails initializer (`config/initializers/pg_rls.rb`):
+
+```ruby
+PgRls.setup do |config|
+  current_attributes = %i[organization__branch]
+end
 ```
-You can change company to anything you'd like, for example, `tenant`
-This will generate the model and inject all the required code
 
-For any new model that needs to be under rls, you can generate it by writing
+This ensures that your custom attributes are loaded and available across requests.
 
-```bash
-rails generate pg_rls user #=> where user eq model name
+#### **Using `__` Convention for Subclasses**
+Inspired by Stimulus controllers, `PgRls::Current` supports a **double underscore (`__`) convention** to allow easy reference to subclasses of your models.
+
+For example, if you have the following models:
+
+```ruby
+class Organization < ApplicationRecord; end
+class Organization::Branch < ApplicationRecord; end
 ```
-and it will generate all the necesary information for you.
 
-You can swtich to another tenant by using
+You can dynamically access `Organization::Branch.first` using:
+
 ```ruby
-PgRls::Tenant.switch :app #=> where app eq tenant name
+PgRls::Current.organization__branch  # Resolves to Organization::Branch.first
 ```
-Don't forget to update how you want `PgRls` to find your tenant, you can set multiple options by modifying `api/config/initializers/pg_rls.rb` `search_methods`
 
-```yml
-# app/config/database.yml
-<% def db_username
-    return PgRls.username unless ENV['AS_DB_ADMIN']
+This works because the `PgRls::Current` implementation automatically transforms attribute names with `__` into proper class names, making it easy to extend without manual configurations.
 
-    Rails.application.credentials.dig(:database, :server_1, :username)
-   end %>
+This approach provides a flexible way to structure your tenant-based logic without requiring manual mappings for every subclass.
 
-...
+#### **Ensuring Proper RLS Configuration**
+Since `PgRls::Current` uses `.first` to retrieve the record, you should ensure that the table is under **Row-Level Security (RLS)** and that the attribute used is **unique** within the tenant's scope. If the record is not unique, it is recommended to **manually set the attribute** to avoid unintended results from querying the first available record.
 
-development:
-  <<: *default
-  database: example_development
-  username: <%= db_username %> # Apply this to production and all env including tests
+Example of setting the attribute manually:
+
+```ruby
+PgRls::Current.organization__branch = Organization::Branch.find_by(name: 'Main Branch')
+```
+
+## Usage
 
-...
+1. Generate the necessary files:
 
+   ```bash
+   rails generate pg_rls:install company  # where 'company' is your tenant model name
+   ```
+
+   You can change 'company' to anything you'd like, for example, 'tenant'.
+
+2. For any new model that needs to be under RLS:
+
+   ```bash
+   rails generate pg_rls user  # where 'user' is your model name
+   ```
+
+3. Switch to another tenant:
+
+   ```ruby
+   PgRls::Tenant.switch :app  # where 'app' is your tenant name
+   ```
+
+## RLS Index Management Methods
+
+These functions help manage indexes on tables protected by Row Level Security (RLS), ensuring that the `tenant_id` field is always included in the indexes to maintain integrity and multitenant isolation.
+
+### `create_rls_index`
+Creates an index on an RLS-enabled table, automatically adding the `tenant_id` field if it is not present in the list of columns.
+
+**Usage:**
+```ruby
+create_rls_index(:users, [:email])
+# This will create an index on [:email, :tenant_id] for the users table
 ```
-### Testing
 
-If you are getting `PG::InsufficientPrivilege: ERROR:  permission denied ` you can override does permistion by running `RAILS_ENV=test rake db:grant_usage`
+You can also pass additional options compatible with `add_index`:
+```ruby
+create_rls_index(:users, [:email], unique: true, name: 'index_users_on_email_and_tenant_id')
+```
 
-Many application uses some sort of database cleaner before running thair spec so on each test that we run we'll have an empty state. Usually, those gems clear our user configuration for the database. To solve this issue, we must implement the following:
+### `drop_rls_index`
+Removes an index created with `create_rls_index`, ensuring that the same columns (including `tenant_id`) are used.
 
+**Usage:**
 ```ruby
-# spec/rails_helper.rb
+drop_rls_index(:users, [:email])
+# Removes the index on [:email, :tenant_id] for the users table
+```
+
+These functions are defined in `lib/pg_rls/active_record/connection_adapters/postgre_sql/schema_statements.rb` and are useful for maintaining index consistency in multitenant environments with RLS.
 
-...
-# some database cleaning strategy
+## Testing
 
+If you encounter `PG::InsufficientPrivilege: ERROR: permission denied`, override permissions by running:
+
+```bash
+RAILS_ENV=test rake db:grant_usage
+```
+
+For database cleaning strategies, implement the following in your `spec/rails_helper.rb`:
+
+```ruby
 config.before(:suite) do
-  # Create the tenant which in this example is company and we are using FactoryBot
   FactoryBot.create(:company, subdomain: 'app')
-  # In this default case our initializer is set to search by subdomain so will use it
   PgRls::Tenant.switch :app
 end
-...
 ```
+
+### Running tests in parallel
+
+If you want to run your tests using `parallelize`, make sure to include the following in your test helper file (for example, `test_helper.rb` or `rails_helper.rb`):
+
+```ruby
+require "pg_rls/active_record/test_databases"
+```
+
+This is required for proper test database setup when running tests in parallel. You can see an example in the `test/test_helper.rb` file in this repository.
+
 ## Development
 
-After checking out the repo, run `bin/setup` to install dependencies. Then, run `rake spec` to run the tests. You can also run `bin/console` for an interactive prompt that will allow you to experiment.
+After checking out the repo, run `bin/setup` to install dependencies. Then, run `rake spec` to run the tests. You can also run `bin/console` for an interactive prompt.
+
+### Development Workflow
+
+Before each push, follow this workflow:
+
+1. Run quality checks:
+
+   ```bash
+   ./review_code.sh
+   ```
+
+   This script performs:
+   - Rubocop
+   - RSpec (100% code coverage required)
+   - Steep (type checking)
+
+2. Ensure 100% documentation coverage.
 
-To install this gem onto your local machine, run `bundle exec rake install`. To release a new version, update the version number in `version.rb`, and then run `bundle exec rake release`, which will create a git tag for the version, push git commits and the created tag, and push the `.gem` file to [rubygems.org](https://rubygems.org).
+3. Run tests:
+
+   ```bash
+   bin/test
+   ```
+
+### Releasing a New Version
+
+1. Update the version number in `version.rb`
+2. Run `bundle exec rake release`
 
 ## Contributing
 
-Bug reports and pull requests are welcome on GitHub at https://github.com/dandush03/pg_rls. This project is intended to be a safe, welcoming space for collaboration, and contributors are expected to adhere to the [code of conduct](https://github.com/dandush03/pg_rls/blob/master/CODE_OF_CONDUCT.md).
+1. Fork the Project
+2. Create your Feature Branch (`git checkout -b feature/AmazingFeature`)
+3. Commit your Changes (`git commit -m 'Add some AmazingFeature'`)
+4. Push to the Branch (`git push origin feature/AmazingFeature`)
+5. Open a Pull Request
 
 ## License
 
-The gem is available as open source under the terms of the [MIT License](https://opensource.org/licenses/MIT).
+Distributed under the MIT License. See `LICENSE` for more information.
+
+## Contact
+
+If you need help, feel free to reach out through the repository [issues](https://github.com/dandush03/pg_rls/issues) page or contact me via [LinkedIn](https://www.linkedin.com/in/daniel-laloush/).
+
+Project Link: [https://github.com/Dandush03/pg_rls](https://github.com/Dandush03/pg_rls)
 
-## Code of Conduct
+## Acknowledgements
 
-Everyone interacting in the PgRls project's codebases, issue trackers, chat rooms and mailing lists is expected to follow the [code of conduct](https://github.com/dandush03/pg_rls/blob/master/CODE_OF_CONDUCT.md).
+- [GitHub Emoji Cheat Sheet](https://www.webpagefx.com/tools/emoji-cheat-sheet)
+
+- [Choose an Open Source License](https://choosealicense.com)
+- [GitHub Pages](https://pages.github.com)
 
-## Note
-Currently we only support subdomain as a searcher but will soon integrate slug/domain and cookies support
-we recommed the use of ``
 ## Show your support
 
 Give a ⭐️ if you like this project!
@@ -153,8 +318,7 @@ If this project help you reduce time to develop, you can give me a cup of coffee
 
 [![paypal][paypal-url]][paypal-donate-code]
 
-<!-- MARKDOWN LINKS & IMAGES -->
-[contributors-shield]: https://img.shields.io/github/contributors/Dandush03/React-Calculator.svg?style=flat-square
+[contributors-shield]: https://img.shields.io/github/contributors/Dandush03/pg_rls.svg?style=flat-square
 [contributors-url]: https://github.com/Dandush03/pg_rls/graphs/contributors
 [forks-shield]: https://img.shields.io/github/forks/Dandush03/pg_rls.svg?style=flat-square
 [forks-url]: https://github.com/Dandush03/pg_rls/network/members
@@ -164,10 +328,10 @@ If this project help you reduce time to develop, you can give me a cup of coffee
 [issues-url]: https://github.com/Dandush03/pg_rls/issues
 [license-shield]: https://img.shields.io/github/license/Dandush03/pg_rls.svg?style=flat-square
 [license-url]: https://github.com/Dandush03/pg_rls/blob/master/LICENSE.txt
-[linkedin-shield2]: https://img.shields.io/badge/-LinkedIn-black.svg?style=flat-square&logo=linkedin&colorB=555
-[linkedin-url2]: https://www.linkedin.com/in/daniel-laloush/
-[hireable]: https://cdn.rawgit.com/hiendv/hireable/master/styles/flat/yes.svg
+[linkedin-shield]: https://img.shields.io/badge/-LinkedIn-black.svg?style=flat-square&logo=linkedin&colorB=555
+[linkedin-url]: https://www.linkedin.com/in/daniel-laloush/
+[hireable-url]: https://www.linkedin.com/in/daniel-laloush/
 [paypal-url]: https://www.paypalobjects.com/en_US/i/btn/btn_donateCC_LG.gif
 [paypal-donate-code]: https://www.paypal.com/donate?hosted_button_id=QKZFZAMQNC8JL
-[hireable-url]: https://www.linkedin.com/in/daniel-laloush/
 [donate]: https://img.shields.io/badge/Donate-PayPal-blue.svg
+[hireable]: https://cdn.rawgit.com/hiendv/hireable/master/styles/flat/yes.svg

data/Rakefile

@@ -1,20 +1,13 @@
 # frozen_string_literal: true
 
-require 'bundler/gem_tasks'
-require 'rspec/core/rake_task'
-require 'rubocop/rake_task'
+require "bundler/gem_tasks"
 
-RSpec::Core::RakeTask.new(:spec)
+require "rubocop/rake_task"
 
 RuboCop::RakeTask.new
 
 task default: %i[spec rubocop]
 
-desc 'Generate documentation for PgRls.'
-Rake::RDocTask.new(:rdoc) do |rdoc|
-  rdoc.rdoc_dir = 'rdoc'
-  rdoc.title    = 'PgRls'
-  rdoc.options << '--line-numbers' << '--inline-source'
-  rdoc.rdoc_files.include('README.md')
-  rdoc.rdoc_files.include('lib/**/*.rb')
-end
+require_relative "test/dummy/config/application"
+
+Rails.application.load_tasks

data/Steepfile

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+D = Steep::Diagnostic
+
+target :pg_rls do
+  signature "sig"
+
+  check "lib"
+  ignore "lib/pg_rls/active_record/test_databases.rb"
+
+  # library "pathname"              # Standard libraries
+  # library "strong_json"           # Gems
+
+  # configure_code_diagnostics(D::Ruby.default)      # `default` diagnostics setting (applies by default)
+  configure_code_diagnostics(D::Ruby.strict) # `strict` diagnostics setting
+  # configure_code_diagnostics(D::Ruby.lenient)      # `lenient` diagnostics setting
+  # configure_code_diagnostics(D::Ruby.silent)       # `silent` diagnostics setting
+  configure_code_diagnostics do |hash| # You can setup everything yourself
+    hash[D::Ruby::UnexpectedSuper] = :hint
+  end
+end
+
+# target :test do
+#   signature "sig", "sig-private"
+#
+#   check "test"
+#
+#   # library "pathname"              # Standard libraries
+# end

data/UPGRADE.md

@@ -0,0 +1,106 @@
+
+# Upgrade Guide
+
+## Upgrading from 0.2 to 1.0.1
+
+### 1. Switch from `admin_execute` to Native Sharding
+
+We are no longer using the `admin_execute` methods to handle admin connections. Instead, Rails' native sharding system is used for handling admin connections.
+
+#### Steps
+
+- **Remove `admin_execute` methods** from your codebase, as they are no longer needed.
+- **Modify your `database.yml`** to leverage Rails' native sharding system.
+
+### 2. `database.yml` Configuration Changes
+
+You must now configure sharding in `database.yml` using the following pattern:
+
+- The first shard should be named `primary` and use the `PgRls.username` and password for regular operations.
+- The second shard should be named `admin` to handle administrative connections.
+
+#### Example `database.yml` for **Development**
+
+NOTE: Both shards connects to the same database
+
+```yaml
+development:
+  primary:
+    database: some_database_primary
+    adapter: postgresql
+    username: <%= ENV['PG_RLS_USERNAME'] %>
+    password: <%= ENV['PG_RLS_PASSWORD'] %>
+  admin:
+    database: some_database_primary
+    adapter: postgresql
+    username: <%= ENV['USERNAME'] %>
+    password: <%= ENV['PASSWORD'] %>
+```
+
+**Recommendation**: In production, it is recommended to configure only one user per server to optimize query performance.
+
+### 3. **Assign Users to the `pg_rls` Group**
+
+It is **mandatory** that all users of the `pg_rls` gem be assigned to the `pg_rls` group. This ensures the gem functions correctly and enforces security and access control.
+
+#### Steps
+
+- Assign each relevant user to the `pg_rls` group by running:
+
+  ```sql
+  GRANT pg_rls TO your_user;
+  ```
+
+- You can check if the user is already part of the group with the following query:
+
+  ```sql
+  SELECT rolname 
+  FROM pg_auth_members 
+  JOIN pg_roles ON pg_roles.oid = pg_auth_members.member 
+  WHERE rolname = 'your_user' 
+  AND roleid = (SELECT oid FROM pg_roles WHERE rolname = 'pg_rls');
+  ```
+
+### 4. Boot Protection with Pre-Checks
+
+There is now a pre-check method that will **not allow your server to boot** if the following conditions are not met:
+
+- The user is assigned to the `pg_rls` group.
+- The `database.yml` is properly configured with the shards (`primary` and `admin`).
+
+Ensure these conditions are satisfied before attempting to start your server.
+
+### 5. Optimizing Query Performance
+
+For optimal performance in production:
+
+- It is **recommended to use only one configuration** (i.e., one user) per server.
+- If necessary, use two machines for processes that require deep analysis or sudo powers. However, note that multitenancy may not provide significant benefits if an admin has access to everything.
+
+### 6. Role and Privilege Checks
+
+Ensure that proper role and privilege checks are added to your codebase using the `pg_class`, `pg_roles`, and `pg_namespace` tables for table and schema-specific privileges. You can ensure users have proper access by querying and validating against these PostgreSQL system tables.
+
+#### Example Check for Role Existence
+
+```ruby
+check_rls_user_privileges!(role_name, schema)
+```
+
+#### Example Grant Privileges for Role
+
+```ruby
+# Make sure rls_group role is created before running this command
+grant_rls_user_privileges(schema)
+```
+
+#### Example of Creating A RLS Group and User
+
+```ruby
+create_rls_role(name, password)
+# This creates a role named rls_group. Additionally, it creates a user with the
+# name and password you've passed to this method and assigns that new user
+# to the rls_group.
+```
+
+Make sure to follow the updated security practices for checking user access and permissions within the `pg_rls` framework.

data/app/models/pg_rls/admin.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+# This file of the PgRls module provides methods to interact with the admin shard
+module PgRls
+  # The Admin class provides the admin_execute method to execute
+  class Admin
+    def self.execute(sql = nil)
+      PgRls::Record.connected_to(shard: :admin) do
+        return yield.presence if block_given?
+
+        PgRls::Record.connection.execute(sql).presence
+      end
+    end
+  end
+
+  # Alias for the Admin.execute method
+  def self.admin_execute(sql = nil, &block)
+    Deprecation.warn(
+      "This method is deprecated and will be removed in future versions. " \
+      "please use PgRls::Admin.execute instead."
+    )
+    Admin.execute(sql, &block)
+  end
+end

data/app/models/pg_rls/current.rb

@@ -0,0 +1,48 @@
+# frozen_string_literal: true
+
+module PgRls
+  # Current Tenant State
+  class Current < ::ActiveSupport::CurrentAttributes
+    attribute(*PgRls.current_attributes.dup.push(:tenant, :tenant_history))
+
+    PgRls.current_attributes.each do |attribute|
+      define_method(attribute) do
+        @attributes[attribute] ||= fetch_attribute(attribute)
+      end
+    end
+
+    def fetch_attribute(attribute)
+      klass_name = attribute.to_s.gsub("__", "/").classify
+
+      send(:"#{attribute}=", klass_name.constantize.first)
+    end
+
+    def tenant=(tenant)
+      add_tenant_to_history
+      super
+      tenant&.set_rls
+    end
+
+    def reset
+      history = Array @attributes[:tenant_history]
+      super
+      @attributes[:tenant_history] = history
+      restore_most_recent_tenant
+      @attributes
+    end
+
+    private
+
+    def add_tenant_to_history
+      @attributes[:tenant_history] ||= []
+      @attributes[:tenant_history] << @attributes[:tenant] unless @attributes[:tenant].nil?
+    end
+
+    def restore_most_recent_tenant
+      @attributes[:tenant] = @attributes[:tenant_history].pop
+      return PgRls::Tenant.reset_rls_used_connections if @attributes[:tenant].nil?
+
+      @attributes[:tenant].set_rls
+    end
+  end
+end

data/app/models/pg_rls/record.rb

@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+# Row Level Security for PostgreSQL
+module PgRls
+  # Base class for all models that should be protected by RLS
+  class Record < PgRls.abstract_base_record_class.constantize
+    self.abstract_class = true
+
+    self.ignored_columns += %w[tenant_id]
+
+    connects_to(**PgRls.connects_to)
+  end
+end