Performant RLS Policies in PostgreSQL using JWT Claims for Multi-Tenancy

-- A naive, non-functional example
CREATE POLICY tenant_isolation ON documents
USING (tenant_id = get_current_tenant_id());

// Middleware in a Node.js / Express application
const { Pool } = require('pg');
const jwt = require('jsonwebtoken');

const pool = new Pool(/* connection config */);
const JWT_SECRET = process.env.JWT_SECRET;

// This middleware intercepts requests, verifies the JWT, and sets the DB context.
async function setDatabaseContext(req, res, next) {
  const authHeader = req.headers.authorization;
  if (!authHeader || !authHeader.startsWith('Bearer ')) {
    // Or proceed as an anonymous user
    return next();
  }

  const token = authHeader.split(' ')[1];
  let claims;

  try {
    claims = jwt.verify(token, JWT_SECRET);
  } catch (err) {
    return res.status(401).send('Invalid token');
  }

  // Attach a dedicated client for this request to ensure transaction-level context
  const client = await pool.connect();
  req.dbClient = client;

  try {
    // CRITICAL: Use SET LOCAL. It scopes the setting to the current transaction only.
    // This prevents context from leaking to other requests using the same connection from the pool.
    // We serialize the entire claims object as a JSON string.
    const claimsJson = JSON.stringify(claims);

    // Use pg_escape_literal to prevent SQL injection if claims contained single quotes.
    await client.query(`SET LOCAL myapp.jwt.claims = ${client.escapeLiteral(claimsJson)}`);

    next(); // Proceed to the route handler

  } catch (err) {
    // Ensure client is released on error
    client.release();
    next(err);
  }
}

// Example route handler
app.get('/api/documents', setDatabaseContext, async (req, res, next) => {
  try {
    // The req.dbClient is now context-aware for the duration of this transaction.
    const { rows } = await req.dbClient.query('SELECT id, title, content FROM documents');
    res.json(rows);
  } catch (err) {
    next(err);
  } finally {
    // IMPORTANT: Always release the client back to the pool in a finally block.
    if (req.dbClient) {
      req.dbClient.release();
    }
  }
});

-- DO NOT DO THIS IN PRODUCTION
CREATE OR REPLACE FUNCTION get_current_tenant_from_volatile_json() RETURNS UUID AS $$
BEGIN
  -- This function is VOLATILE by default. It will be re-evaluated for every row.
  -- The JSON parsing and casting happens repeatedly, creating massive overhead.
  RETURN (current_setting('myapp.jwt.claims')::jsonb ->> 'tenant_id')::uuid;
EXCEPTION
  WHEN OTHERS THEN
    RETURN NULL;
END;
$$ LANGUAGE plpgsql;

-- The resulting policy is incredibly slow on large tables.
CREATE POLICY slow_tenant_policy ON documents
USING (tenant_id = get_current_tenant_from_volatile_json());

-- Create a set of robust, performant, STABLE helper functions.

-- Function to get the full claims JSONB object. Caches the result per-query.
CREATE OR REPLACE FUNCTION current_claims() RETURNS jsonb AS $$
BEGIN
  -- The 'true' flag tells current_setting to return NULL if the setting is not found,
  -- instead of throwing an error.
  RETURN current_setting('myapp.jwt.claims', true)::jsonb;
EXCEPTION
  -- Handle cases where the setting is not valid JSON.
  WHEN invalid_text_representation THEN
    RETURN '{}'::jsonb;
END;
$$ LANGUAGE plpgsql STABLE;

-- Function to get the tenant_id. Note the explicit casting and error handling.
CREATE OR REPLACE FUNCTION current_tenant_id() RETURNS UUID AS $$
DECLARE
  tenant_id_text TEXT;
BEGIN
  tenant_id_text := current_claims() ->> 'tenant_id';
  IF tenant_id_text IS NULL THEN
    RETURN NULL;
  END IF;
  -- Defensive check for UUID format before casting
  IF tenant_id_text !~ '^[0-9a-f]{8}-([0-9a-f]{4}-){3}[0-9a-f]{12}$' THEN
      RETURN NULL;
  END IF;
  RETURN tenant_id_text::uuid;
END;
$$ LANGUAGE plpgsql STABLE;

-- Function to get the user_id (subject from JWT).
CREATE OR REPLACE FUNCTION current_user_id() RETURNS UUID AS $$
DECLARE
  user_id_text TEXT;
BEGIN
  user_id_text := current_claims() ->> 'sub';
  IF user_id_text IS NULL THEN
    RETURN NULL;
  END IF;
  IF user_id_text !~ '^[0-9a-f]{8}-([0-9a-f]{4}-){3}[0-9a-f]{12}$' THEN
      RETURN NULL;
  END IF;
  RETURN user_id_text::uuid;
END;
$$ LANGUAGE plpgsql STABLE;

-- Function to get the user's role.
CREATE OR REPLACE FUNCTION current_user_role() RETURNS TEXT AS $$
BEGIN
  RETURN current_claims() ->> 'role';
END;
$$ LANGUAGE plpgsql STABLE;

-- First, enable RLS on the table.
ALTER TABLE documents ENABLE ROW LEVEL SECURITY;

-- Create a clean, performant policy using our STABLE helper.
CREATE POLICY tenant_isolation ON documents
AS PERMISSIVE FOR ALL
USING (tenant_id = current_tenant_id());

-- Assumes the 'documents' table has an 'owner_id' column.

-- Drop the old policy before creating a new one
DROP POLICY IF EXISTS tenant_and_role_isolation ON documents;

CREATE POLICY tenant_and_role_isolation ON documents
AS PERMISSIVE FOR ALL
USING (
  tenant_id = current_tenant_id() AND (
    -- Admins can see everything within their tenant.
    current_user_role() = 'admin'
    OR
    -- Members can only see their own documents.
    (current_user_role() = 'member' AND owner_id = current_user_id())
  )
);

CREATE TABLE document_shares (
  document_id UUID REFERENCES documents(id) ON DELETE CASCADE,
  shared_with_tenant_id UUID NOT NULL,
  PRIMARY KEY (document_id, shared_with_tenant_id)
);
-- Critical index for policy performance!
CREATE INDEX idx_document_shares_lookup ON document_shares(document_id, shared_with_tenant_id);

-- We might have different policies for SELECT vs. UPDATE/DELETE
DROP POLICY IF EXISTS select_documents ON documents;
DROP POLICY IF EXISTS modify_documents ON documents;

-- Policy for viewing documents: either in your tenant or shared with your tenant.
CREATE POLICY select_documents ON documents
AS PERMISSIVE FOR SELECT
USING (
  tenant_id = current_tenant_id()
  OR
  EXISTS (
    SELECT 1
    FROM document_shares ds
    WHERE ds.document_id = documents.id
      AND ds.shared_with_tenant_id = current_tenant_id()
  )
);

-- Policy for modifying documents: you must be the owner.
CREATE POLICY modify_documents ON documents
AS PERMISSIVE FOR UPDATE, DELETE
USING (tenant_id = current_tenant_id());

    -- Example of a test script using psql or a test runner
    BEGIN;

    -- Impersonate user 1 from tenant A
    SET LOCAL myapp.jwt.claims = '{"sub": "user-id-1", "tenant_id": "tenant-id-a", "role": "member"}';
    -- Run assertions: This SELECT should only return docs for user 1 in tenant A.
    SELECT COUNT(*) FROM documents; 

    -- Impersonate user 2 (admin) from tenant A
    SET LOCAL myapp.jwt.claims = '{"sub": "user-id-2", "tenant_id": "tenant-id-a", "role": "admin"}';
    -- Run assertions: This SELECT should return ALL docs for tenant A.
    SELECT COUNT(*) FROM documents;

    -- Impersonate user 3 from tenant B
    SET LOCAL myapp.jwt.claims = '{"sub": "user-id-3", "tenant_id": "tenant-id-b", "role": "admin"}';
    -- Run assertions: This SELECT should return ALL docs for tenant B.
    SELECT COUNT(*) FROM documents;

    ROLLBACK; -- End the test transaction

CREATE TABLE documents (
  id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
  tenant_id UUID NOT NULL,
  owner_id UUID NOT NULL,
  title TEXT NOT NULL,
  content TEXT,
  created_at TIMESTAMPTZ DEFAULT NOW()
);
CREATE INDEX idx_documents_tenant_id ON documents(tenant_id);
-- Populate with 10M rows

    Aggregate  (cost=150000.00..150000.01 rows=1 width=8)
      ->  Seq Scan on documents  (cost=0.00..145000.00 rows=10000000 width=0)
            Filter: (tenant_id = get_current_tenant_from_volatile_json()) -- HIGH COST

    Aggregate  (cost=1500.00..1500.01 rows=1 width=8)
      ->  Index Only Scan using idx_documents_tenant_id on documents (cost=0.43..1450.00 rows=1000 width=0)
            Index Cond: (tenant_id = current_tenant_id()) -- LOW COST