Using Row-Level Security In SQLAlchemy

Row-level security (RLS) in PostgreSQL enables tables to implement policies that limit access or modification of rows according to the user's role, enhancing the basic SQL-standard privileges provided by GRANT.

Once activated, all access to the table has to adhere to these policies. If no policies are defined on the table, it defaults to a deny-all rule, meaning no rows can be seen or mutated. These policies can be tailored to specific commands, roles, or both, allowing for detailed management of who can access or change data.

This guide explains how to attach RLS Policies to your SQLAlchemy models definition and configure Atlas to manage both the RLS and the SQLAlchemy models as a single migration unit.

Atlas Pro Feature

Atlas support for Row-Level Security Policies used in this guide is available exclusively to Pro users. To use this feature, run:

atlas login

Getting started with Atlas and SQLAlchemy

Before we continue, ensure you have installed the Atlas SQLAlchemy Provider on your SQLAlchemy project.

To set it up, follow along the getting started guide for SQLAlchemy and Atlas.

Use a Composite Schema to leverage Atlas' RLS feature

In this guide, we will configure Atlas to read the state of the schema from a Composite Schema data source. This will allow us to use table modeling capabilities SQLAlchemy, and to combine them with Atlas features to fill the gap.

1. Let's define a simple schema with two models - users and tenants:

models.py
from sqlalchemy import String, ForeignKey
from sqlalchemy.orm import (
    DeclarativeBase,
    Mapped,
    mapped_column,
    relationship
)


class Base(DeclarativeBase):
    pass


class User(Base):
    __tablename__ = "users"
    id: Mapped[int] = mapped_column(primary_key=True)
    name: Mapped[str] = mapped_column(String(30))
    tenant_id: Mapped[int] = mapped_column(ForeignKey("tenants.id"))
    tenant: Mapped["Tenant"] = relationship(back_populates="users")


class Tenant(Base):
    __tablename__ = "tenants"
    id: Mapped[int] = mapped_column(primary_key=True)
    name: Mapped[str] = mapped_column(String(30))
    users: Mapped[list["User"]] = relationship(
        back_populates="tenant", cascade="all, delete-orphan"
    )

2. Now, suppose we want to limit access to the users table based on the tenant_id field. We can achieve this by defining a Row-Level Security (RLS) policy on the users table. Below is the SQL code that defines the RLS policy:

schema.sql
--- Enable row-level security on the users table.
ALTER TABLE "users" ENABLE ROW LEVEL SECURITY;

-- Create a policy that restricts access to rows in the users table based on the current tenant.
CREATE POLICY tenant_isolation ON "users"
    USING ("tenant_id" = current_setting('app.current_tenant')::integer);

3. In your atlas.hcl config file, add a composite_schema that includes both your custom security policies in schema.sql and your SQLAlchemy models:

atlas.hcl
data "composite_schema" "example" {
  // First, load the schema with the SQLAlchemy provider
  schema "public" {
    url = data.external_schema.sqlalchemy.url
  }
  // Next, load the RLS policies
  schema "public" {
    url = "file://schema.sql"
  }
}

env "local" {
  src = data.composite_schema.example.url
  dev = "docker://postgres/16/dev?search_path=public"
}

Usage

After setting up our composite schema, we can get its representation using the atlas schema inspect command, generate schema migrations for it, apply them to a database, and more. Below are a few commands to get you started with Atlas:

Inspect the Schema

The atlas schema inspect command is commonly used to inspect databases. However, we can also use it to inspect our composite_schema and print the SQL representation of it:

atlas schema inspect \
  --env local \
  --url env://src \
  --format '{{ sql . }}'

The command above prints the following SQL. Note, Atlas knows to create the tenant_isolation policy after the users table:

-- Create "tenants" table
CREATE TABLE "tenants" ("id" serial NOT NULL, "name" character varying(30) NOT NULL, PRIMARY KEY ("id"));
-- Create "users" table
CREATE TABLE "users" ("id" serial NOT NULL, "name" character varying(30) NOT NULL, "tenant_id" integer NOT NULL, PRIMARY KEY ("id"), CONSTRAINT "users_tenant_id_fkey" FOREIGN KEY ("tenant_id") REFERENCES "tenants" ("id") ON UPDATE NO ACTION ON DELETE NO ACTION);
-- Enable row-level security for "users" table
ALTER TABLE "users" ENABLE ROW LEVEL SECURITY;
-- Create policy "tenant_isolation"
CREATE POLICY "tenant_isolation" ON "users" AS PERMISSIVE FOR ALL TO PUBLIC USING (tenant_id = (current_setting('app.current_tenant'::text))::integer);

Generate Migrations

To generate a migration for the schema, run the following command:

atlas migrate diff --env local

Note that a new migration file is created with the following content:

migrations/20240712090543.sql
-- Create "tenants" table
CREATE TABLE "public"."tenants" (
  "id" serial NOT NULL,
  "name" character varying(30) NOT NULL,
  PRIMARY KEY ("id")
);
-- Create "users" table
CREATE TABLE "public"."users" (
  "id" serial NOT NULL,
  "name" character varying(30) NOT NULL,
  "tenant_id" integer NOT NULL,
  PRIMARY KEY ("id"),
  CONSTRAINT "users_tenant_id_fkey" FOREIGN KEY ("tenant_id") REFERENCES "public"."tenants" ("id") ON UPDATE NO ACTION ON DELETE NO ACTION
);
-- Enable row-level security for "users" table
ALTER TABLE "public"."users" ENABLE ROW LEVEL SECURITY;
-- Create policy "tenant_isolation"
CREATE POLICY "tenant_isolation" ON "public"."users" AS PERMISSIVE FOR ALL TO PUBLIC USING (tenant_id = (current_setting('app.current_tenant'::text))::integer);

Apply the Migrations

To apply the migration generated above to a database, run the following command:

atlas migrate apply \
  --env local \
  --url "postgres://postgres:pass@localhost:5432/database?search_path=public&sslmode=disable"

Apply the Schema Directly on the Database

Sometimes, there is a need to apply the schema directly to the database without generating a migration file. For example, when experimenting with schema changes, spinning up a database for testing, etc. In such cases, you can use the command below to apply the schema directly to the database:

atlas schema apply \
  --env local \
  --url "postgres://postgres:pass@localhost:5432/database?search_path=public&sslmode=disable"

Code Example

The following example assumes that the current state of the users table is as follows:

id	name	tenant_id
1	John	1
2	Sarah	2

main.py
from db.models import User
from sqlalchemy import create_engine, select, text
from sqlalchemy.orm import Session


def main():
    engine = create_engine("postgresql://client:pass@localhost/example?sslmode=disable")

    with Session(engine) as session:
        session.execute(text("select set_config('app.current_tenant', '1', true)"))
        query = select(User)
        for user1 in session.scalars(query):
            print(f"{user1.name=}")

    with Session(engine) as session:
        session.execute(text("select set_config('app.current_tenant', '2', true)"))
        query = select(User)
        for user2 in session.scalars(query):
            print(f"{user2.name=}")


if __name__ == '__main__':
    main()


# Output:
# user1.name='John'
# user2.name='Sarah'

Getting started with Atlas and SQLAlchemy​

Use a Composite Schema to leverage Atlas' RLS feature​

Usage​

Inspect the Schema​

Generate Migrations​

Apply the Migrations​

Code Example​