PostgreSQL RLS: Dynamic Policies for Complex SaaS Architectures

-- The basic, well-known RLS pattern
ALTER TABLE projects ENABLE ROW LEVEL SECURITY;

CREATE POLICY tenant_isolation ON projects
USING (tenant_id = current_setting('app.tenant_id')::uuid);

-- Simplified Schema for Demonstration
CREATE EXTENSION IF NOT EXISTS "uuid-ossp";

CREATE TABLE organizations (
    id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
    name TEXT NOT NULL
);

CREATE TABLE users (
    id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
    email TEXT NOT NULL UNIQUE
);

CREATE TABLE projects (
    id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
    org_id UUID NOT NULL REFERENCES organizations(id),
    name TEXT NOT NULL
);

CREATE TABLE tasks (
    id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
    project_id UUID NOT NULL REFERENCES projects(id),
    title TEXT NOT NULL,
    completed BOOLEAN DEFAULT FALSE
);

-- Junction tables for permissions
CREATE TYPE org_role AS ENUM ('ADMIN', 'MEMBER');
CREATE TABLE organization_members (
    user_id UUID NOT NULL REFERENCES users(id),
    org_id UUID NOT NULL REFERENCES organizations(id),
    role org_role NOT NULL,
    PRIMARY KEY (user_id, org_id)
);

CREATE TYPE project_role AS ENUM ('EDITOR', 'VIEWER');
CREATE TABLE project_members (
    user_id UUID NOT NULL REFERENCES users(id),
    project_id UUID NOT NULL REFERENCES projects(id),
    role project_role NOT NULL,
    PRIMARY KEY (user_id, project_id)
);

-- DO NOT USE THIS IN PRODUCTION
CREATE POLICY select_tasks_naive ON tasks FOR SELECT
USING (
    EXISTS (
        SELECT 1
        FROM projects p
        JOIN organization_members om ON p.org_id = om.org_id
        WHERE p.id = tasks.project_id
          AND om.user_id = current_setting('app.user_id')::uuid
          AND om.role = 'ADMIN'
    )
    OR
    EXISTS (
        SELECT 1
        FROM project_members pm
        WHERE pm.project_id = tasks.project_id
          AND pm.user_id = current_setting('app.user_id')::uuid
    )
);

ALTER TABLE tasks ENABLE ROW LEVEL SECURITY;

-- Benchmarking the naive approach
-- (Imagine millions of tasks and thousands of users)
EXPLAIN ANALYZE
SELECT count(*) FROM tasks;

{
  "sub": "user-uuid-c4a1...",
  "exp": 1678886400,
  "iat": 1678882800,
  "https://myapp.com/claims": {
    "org_id": "org-uuid-b21a...",
    "org_role": "MEMBER",
    "projects": [
      {"id": "project-uuid-1111...", "role": "EDITOR"},
      {"id": "project-uuid-2222...", "role": "VIEWER"}
    ]
  }
}

{
  "sub": "user-uuid-d5b2...",
  "exp": 1678886400,
  "iat": 1678882800,
  "https://myapp.com/claims": {
    "org_id": "org-uuid-b21a...",
    "org_role": "ADMIN"
  }
}

// Node.js (Express.js) middleware example
const { Pool } = require('pg');
const pool = new Pool();

// Middleware to be used on authenticated routes
async function setDatabaseContext(req, res, next) {
  // Assume req.user is populated by a JWT validation middleware
  if (!req.user || !req.user['https://myapp.com/claims']) {
    return next(); // Or throw an error
  }

  const client = await pool.connect();
  req.dbClient = client; // Attach client to request for use in route handlers

  try {
    // The JWT claims object is stringified and sent to PostgreSQL
    const claims = JSON.stringify(req.user['https://myapp.com/claims']);

    // IMPORTANT: Use SET LOCAL to scope the setting to the transaction
    await client.query('BEGIN');
    await client.query(`SET LOCAL app.jwt.claims = '${claims.replace(/'/g, "''")}';`);
    
    // Also set user_id for auditing or simple checks
    const userId = req.user.sub;
    await client.query(`SET LOCAL app.user_id = '${userId.replace(/'/g, "''")}';`);
    
    next();
  } catch (err) {
    await client.query('ROLLBACK');
    client.release();
    next(err);
  }
}

// In your route handler, you would use req.dbClient
app.get('/tasks', setDatabaseContext, async (req, res, next) => {
  try {
    const result = await req.dbClient.query('SELECT * FROM tasks LIMIT 100;');
    await req.dbClient.query('COMMIT');
    res.json(result.rows);
  } catch (err) {
    await req.dbClient.query('ROLLBACK');
    next(err);
  } finally {
    if (req.dbClient) {
      req.dbClient.release();
    }
  }
});

-- Helper functions to safely access JWT claims

-- Returns the entire claims object as JSONB
CREATE OR REPLACE FUNCTION get_jwt_claims() 
RETURNS jsonb AS $$
BEGIN
    -- Use coalesce to return an empty JSON object if the setting is not present
    -- This prevents errors and allows for unauthenticated access patterns if needed
    RETURN coalesce(current_setting('app.jwt.claims', true), '{}')::jsonb;
EXCEPTION
    -- In case of invalid JSON, return an empty object
    WHEN invalid_text_representation THEN
        RETURN '{}'::jsonb;
END;
$$ LANGUAGE plpgsql STABLE;

-- Returns the current user's organization role
CREATE OR REPLACE FUNCTION get_org_role() 
RETURNS text AS $$
    SELECT get_jwt_claims() ->> 'org_role';
$$ LANGUAGE sql STABLE;

-- Checks if the user is an organization admin
CREATE OR REPLACE FUNCTION is_org_admin() 
RETURNS boolean AS $$
    SELECT get_org_role() = 'ADMIN';
$$ LANGUAGE sql STABLE;

-- Returns an array of project UUIDs the user has explicit access to
CREATE OR REPLACE FUNCTION get_accessible_project_ids() 
RETURNS uuid[] AS $$
    SELECT ARRAY(
        SELECT (value ->> 'id')::uuid
        FROM jsonb_array_elements(get_jwt_claims() -> 'projects')
    );
$$ LANGUAGE sql STABLE;

-- Returns the current user's ID from the session
CREATE OR REPLACE FUNCTION get_current_user_id()
RETURNS uuid AS $$
    SELECT current_setting('app.user_id', true)::uuid;
$$ LANGUAGE sql STABLE;

-- First, drop the old, inefficient policy
DROP POLICY IF EXISTS select_tasks_naive ON tasks;

-- The new, dynamic policy for SELECT
CREATE POLICY select_tasks_dynamic ON tasks FOR SELECT
USING (
    EXISTS (
        SELECT 1
        FROM projects p
        WHERE p.id = tasks.project_id
        AND (
            -- Condition 1: User is an org admin in the project's organization
            (is_org_admin() AND p.org_id = (get_jwt_claims() ->> 'org_id')::uuid)
            OR
            -- Condition 2: User has explicit access to this project
            p.id = ANY(get_accessible_project_ids())
        )
    )
);

-- Always remember to enable RLS on the table
ALTER TABLE tasks ENABLE ROW LEVEL SECURITY;

-- Policy to allow task creation
CREATE POLICY insert_tasks_dynamic ON tasks FOR INSERT
WITH CHECK (
    EXISTS (
        SELECT 1
        FROM projects p
        WHERE p.id = tasks.project_id
        AND (
            (is_org_admin() AND p.org_id = (get_jwt_claims() ->> 'org_id')::uuid)
            OR
            -- More granular: only allow EDITORs to create tasks
            (p.id, 'EDITOR') IN (SELECT (value ->> 'id')::uuid, value ->> 'role' FROM jsonb_array_elements(get_jwt_claims() -> 'projects'))
        )
    )
);

-- Policy for updates and deletes
-- Let's say only Org Admins or Project Editors can modify/delete tasks
CREATE POLICY update_delete_tasks_dynamic ON tasks FOR ALL -- Covers UPDATE and DELETE
USING (
    -- The USING clause is the same as our SELECT policy for visibility
    EXISTS (SELECT 1 FROM projects p WHERE p.id = tasks.project_id AND ((is_org_admin() AND p.org_id = (get_jwt_claims() ->> 'org_id')::uuid) OR p.id = ANY(get_accessible_project_ids())))
)
WITH CHECK (
    -- The WITH CHECK clause ensures they can't move a task to a project they don't have edit access to
    EXISTS (
        SELECT 1
        FROM projects p
        WHERE p.id = tasks.project_id
        AND (
            (is_org_admin() AND p.org_id = (get_jwt_claims() ->> 'org_id')::uuid)
            OR
            (p.id, 'EDITOR') IN (SELECT (value ->> 'id')::uuid, value ->> 'role' FROM jsonb_array_elements(get_jwt_claims() -> 'projects'))
        )
    )
);

CREATE OR REPLACE FUNCTION cache_user_permissions() 
RETURNS void AS $$
BEGIN
    -- Check if the cache already exists for this transaction to make the function idempotent
    IF EXISTS (SELECT 1 FROM pg_class WHERE relname = 'session_permissions_cache' AND relpersistence = 't') THEN
        RETURN;
    END IF;

    -- Create a temporary table that is dropped at the end of the session
    CREATE TEMP TABLE session_permissions_cache (
        project_id UUID PRIMARY KEY,
        role project_role
    ) ON COMMIT DROP;

    -- Populate the cache with a single, complex query
    INSERT INTO session_permissions_cache (project_id, role)
    SELECT pm.project_id, pm.role
    FROM project_members pm
    WHERE pm.user_id = get_current_user_id();

    -- Add any other complex permission logic here, e.g., permissions inherited from teams

END;
$$ LANGUAGE plpgsql VOLATILE; -- VOLATILE because it has side effects (creates a table)

// In the Node.js middleware, after setting the context:
await client.query('SELECT cache_user_permissions();');

-- RLS policy using the temporary cache table
CREATE POLICY select_tasks_cached ON tasks FOR SELECT
USING (
    EXISTS (
        SELECT 1
        FROM projects p
        WHERE p.id = tasks.project_id
        AND (
            (is_org_admin() AND p.org_id = (get_jwt_claims() ->> 'org_id')::uuid)
            OR
            EXISTS (SELECT 1 FROM session_permissions_cache spc WHERE spc.project_id = p.id)
        )
    )
);

    BEGIN;
    -- Set the context as the application would
    SET LOCAL app.jwt.claims = '{"org_id": "...", "org_role": "MEMBER", "projects": [{"id": "..."}]}';
    SET LOCAL app.user_id = '...';

    -- To test a specific user's view, you can set the role if you're not using JWT-based roles
    -- SET ROLE some_user_role;

    -- Now run your query and see what the user would see
    SELECT * FROM tasks;

    -- Check the query plan
    EXPLAIN ANALYZE SELECT * FROM tasks;

    ROLLBACK; -- Clean up