Advanced PostgreSQL RLS Patterns for Multi-Tenant SaaS Architectures

-- tenants table to store organization info
CREATE TABLE tenants (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    name TEXT NOT NULL
);

-- projects table, where each project belongs to a tenant
CREATE TABLE projects (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    tenant_id UUID NOT NULL REFERENCES tenants(id) ON DELETE CASCADE,
    name TEXT NOT NULL,
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

-- Enable RLS on the projects table
ALTER TABLE projects ENABLE ROW LEVEL SECURITY;
-- The owner can still do anything, useful for migrations
ALTER TABLE projects FORCE ROW LEVEL SECURITY;

-- The core RLS policy for tenant isolation
CREATE POLICY tenant_isolation_policy ON projects
AS PERMISSIVE FOR ALL
TO authenticated_user -- Apply this policy only to a specific role
USING (tenant_id = current_setting('app.current_tenant_id')::uuid)
WITH CHECK (tenant_id = current_setting('app.current_tenant_id')::uuid);

// assume 'jwt' is a decoded JWT payload: { sub: 'user-123', tenant_id: 'tenant-abc', ... }
const { Pool } = require('pg');

const pool = new Pool({ /* connection config */ });

async function getTenantDbClient(jwt) {
    if (!jwt || !jwt.tenant_id) {
        throw new Error('Invalid JWT or missing tenant_id');
    }

    const client = await pool.connect();

    try {
        // Set the session variable for the current tenant. THIS IS THE CRITICAL STEP.
        // We use pg_typeof to ensure the parameter is treated as text.
        await client.query(`SET LOCAL app.current_tenant_id = $1`, [jwt.tenant_id]);

        // You might also set user ID for audit trails or more granular permissions
        await client.query(`SET LOCAL app.current_user_id = $1`, [jwt.sub]);

        // The client is now "primed" for this tenant and user.
        // Return it to be used by the application logic.
        return client;
    } catch (err) {
        // If setup fails, release the client to prevent leaks.
        client.release();
        throw err;
    }
}

// Example usage in an Express.js route handler
app.get('/api/projects', async (req, res) => {
    let client;
    try {
        // req.user is populated by your authentication middleware
        client = await getTenantDbClient(req.user);
        
        // Notice the query has NO `WHERE tenant_id = ...` clause.
        // RLS handles it automatically and securely.
        const { rows } = await client.query('SELECT id, name FROM projects');
        
        res.json(rows);
    } catch (error) {
        console.error(error);
        res.status(500).send('Internal Server Error');
    } finally {
        if (client) {
            client.release();
        }
    }
});

-- Drop the old policy first
DROP POLICY tenant_isolation_policy ON projects;

-- Create a new, role-aware policy
CREATE POLICY tenant_and_admin_policy ON projects
AS PERMISSIVE FOR ALL
TO authenticated_user
USING (
    -- Condition 1: The user is a super_admin
    current_setting('app.current_user_role', true) = 'super_admin' 
    OR
    -- Condition 2: The user is a regular user matching the tenant_id
    tenant_id = current_setting('app.current_tenant_id')::uuid
)
WITH CHECK (
    -- The WITH CHECK clause must also be updated
    current_setting('app.current_user_role', true) = 'super_admin' 
    OR
    tenant_id = current_setting('app.current_tenant_id')::uuid
);

async function getTenantDbClient(jwt) {
    // ... (previous connection and tenant_id logic)

    try {
        await client.query(`SET LOCAL app.current_tenant_id = $1`, [jwt.tenant_id || null]);
        await client.query(`SET LOCAL app.current_user_id = $1`, [jwt.sub]);
        
        // New logic for handling roles
        if (jwt.role) {
            await client.query(`SET LOCAL app.current_user_role = $1`, [jwt.role]);
        } else {
             // Explicitly set to a non-privileged role if not present
            await client.query(`SET LOCAL app.current_user_role = 'user'`);
        }

        return client;
    } catch (err) {
        // ... (error handling)
    }
}

-- 1. Create the role. NOINHERIT is a good security practice.
-- LOGIN allows it to connect. BYPASSRLS is the magic attribute.
CREATE ROLE billing_service_role WITH NOINHERIT LOGIN BYPASSRLS PASSWORD 'a-very-strong-secret-password';

-- 2. Grant only the necessary permissions.
-- The webapp role (authenticated_user) should NOT have these privileges.
GRANT SELECT ON projects TO billing_service_role;
GRANT SELECT, USAGE ON SEQUENCE projects_id_seq TO billing_service_role; -- If needed

-- Revoke any unwanted default privileges
REVOKE ALL ON SCHEMA public FROM billing_service_role;

CREATE TABLE tasks (
    id BIGSERIAL PRIMARY KEY,
    project_id UUID REFERENCES projects(id),
    tenant_id UUID NOT NULL, -- Denormalized for direct RLS
    title TEXT NOT NULL,
    status TEXT NOT NULL
);

-- Add RLS policy similar to projects
ALTER TABLE tasks ENABLE ROW LEVEL SECURITY;
ALTER TABLE tasks FORCE ROW LEVEL SECURITY;
CREATE POLICY tenant_isolation_policy ON tasks ...

EXPLAIN ANALYZE SELECT * FROM tasks WHERE status = 'pending';

-- RESULT (Simplified)
Seq Scan on tasks  (cost=0.00..250000.00 rows=50000 width=128) (actual time=0.015..3500.123 rows=50000 loops=1)
  Filter: ((status = 'pending'::text) AND (tenant_id = current_setting('app.current_tenant_id')::uuid))
  Rows Removed by Filter: 9950000
Planning Time: 0.150 ms
Execution Time: 3501.234 ms

CREATE INDEX idx_tasks_tenant_id ON tasks(tenant_id);

EXPLAIN ANALYZE SELECT * FROM tasks WHERE status = 'pending';

-- RESULT (Simplified)
Bitmap Heap Scan on tasks  (cost=500.00..25000.00 rows=50 width=128) (actual time=0.500..1.234 rows=50 loops=1)
  Recheck Cond: (tenant_id = current_setting('app.current_tenant_id')::uuid)
  Filter: (status = 'pending'::text)
  Rows Removed by Filter: 950
  ->  Bitmap Index Scan on idx_tasks_tenant_id  (cost=0.00..499.00 rows=1000) (actual time=0.450..0.450 rows=1000 loops=1)
        Index Cond: (tenant_id = current_setting('app.current_tenant_id')::uuid)
Planning Time: 0.250 ms
Execution Time: 1.456 ms

CREATE INDEX idx_tasks_tenant_id_status ON tasks(tenant_id, status);

CREATE TABLE project_shares (
    project_id UUID NOT NULL REFERENCES projects(id) ON DELETE CASCADE,
    shared_with_user_id UUID NOT NULL, -- References a global users table
    -- Storing the tenant_id of the user it's shared with can optimize lookups
    shared_with_tenant_id UUID NOT NULL REFERENCES tenants(id) ON DELETE CASCADE,
    permission_level TEXT NOT NULL DEFAULT 'read' -- e.g., 'read', 'write'
);

CREATE INDEX idx_project_shares_user_project ON project_shares(shared_with_user_id, project_id);

DROP POLICY tenant_and_admin_policy ON projects;

CREATE POLICY complex_access_policy ON projects
AS PERMISSIVE FOR SELECT
TO authenticated_user
USING (
    -- Admins can see everything
    current_setting('app.current_user_role', true) = 'super_admin'
    OR
    -- Users can see projects in their own tenant
    tenant_id = current_setting('app.current_tenant_id')::uuid
    OR
    -- Users can see projects shared with them
    EXISTS (
        SELECT 1
        FROM project_shares
        WHERE project_shares.project_id = projects.id
          AND project_shares.shared_with_user_id = current_setting('app.current_user_id')::uuid
    )
);

-- A separate, simpler policy for write operations (INSERT, UPDATE, DELETE)
-- Typically, only project owners can modify them.
CREATE POLICY write_access_policy ON projects
AS PERMISSIVE FOR ALL
TO authenticated_user
USING (tenant_id = current_setting('app.current_tenant_id')::uuid)
WITH CHECK (tenant_id = current_setting('app.current_tenant_id')::uuid);

-- A helper function to apply a standard RLS policy to a table
CREATE OR REPLACE FUNCTION apply_tenant_rls(table_name_param TEXT)
RETURNS void AS $$
BEGIN
    -- First, ensure RLS is enabled on the target table
    EXECUTE format('ALTER TABLE %I ENABLE ROW LEVEL SECURITY', table_name_param);
    EXECUTE format('ALTER TABLE %I FORCE ROW LEVEL SECURITY', table_name_param);

    -- Drop existing policy to make this function idempotent
    EXECUTE format('DROP POLICY IF EXISTS tenant_isolation_policy ON %I', table_name_param);

    -- Create the standard policy
    EXECUTE format('CREATE POLICY tenant_isolation_policy ON %I ' ||
                   'AS PERMISSIVE FOR ALL TO authenticated_user ' ||
                   'USING (tenant_id = current_setting(''app.current_tenant_id'')::uuid) ' ||
                   'WITH CHECK (tenant_id = current_setting(''app.current_tenant_id'')::uuid)', 
                   table_name_param);
END;
$$ LANGUAGE plpgsql;

-- In a new migration file (e.g., V3__add_invoices_table.sql)

CREATE TABLE invoices (
    id BIGSERIAL PRIMARY KEY,
    tenant_id UUID NOT NULL REFERENCES tenants(id),
    amount_cents INTEGER NOT NULL
);

-- Apply our standardized RLS policy
SELECT apply_tenant_rls('invoices');