Postgres RLS: Advanced Patterns for Hierarchical Multi-Tenancy

-- The 'hello world' of RLS
ALTER TABLE projects ENABLE ROW LEVEL SECURITY;
CREATE POLICY select_project ON projects FOR SELECT
USING (owner_id = auth.uid());

// Conceptual middleware (e.g., Express.js)
async function setDatabaseContext(req, res, next) {
  if (!req.user) { // Assuming previous middleware decodes JWT to req.user
    return next();
  }

  // Assume req.user = { id: 'user-uuid', organizationId: 'org-uuid', roles: ['ADMIN', 'MEMBER'] }
  // Also assume a mechanism to get a DB client for the request, e.g., from a pool.
  const client = await getClientForRequest();
  req.dbClient = client;

  try {
    // A single transaction ensures all settings are applied or none are.
    await client.query('BEGIN');

    // Use SET LOCAL to scope the settings to the current transaction.
    // The 'app.' prefix is a namespace to avoid conflicts with standard Postgres settings.
    await client.query(`SET LOCAL app.user.id = '${req.user.id}';`);
    await client.query(`SET LOCAL app.organization.id = '${req.user.organizationId}';`);

    // It's crucial to properly escape the roles array to prevent SQL injection.
    // A better approach is to use parameterized queries if the driver supports them for SET.
    const rolesString = JSON.stringify(req.user.roles).replace(/'/g, "''");
    await client.query(`SET LOCAL app.user.roles = '${rolesString}';`);

    next(); // Business logic now runs within this transaction with context set.

  } catch (err) {
    // Ensure transaction is rolled back on failure
    await client.query('ROLLBACK');
    client.release();
    return res.status(500).send('Failed to set security context.');
  }
}

// In your route handler, you'd commit or rollback at the end.
app.post('/api/projects', setDatabaseContext, async (req, res) => {
  const client = req.dbClient;
  try {
    // ... your business logic using the client ...
    await client.query('COMMIT');
    res.status(201).send({ message: 'Success' });
  } catch (err) {
    await client.query('ROLLBACK');
    res.status(500).send({ message: 'An error occurred.' });
  } finally {
    client.release();
  }
});

CREATE OR REPLACE FUNCTION get_current_user_id() RETURNS uuid AS $$
BEGIN
  -- The 'true' flag makes it return NULL instead of throwing an error if the setting is not found.
  RETURN current_setting('app.user.id', true)::uuid;
EXCEPTION
  -- Handle invalid UUID format gracefully.
  WHEN invalid_text_representation THEN
    RETURN NULL;
END;
$$ LANGUAGE plpgsql STABLE;

CREATE OR REPLACE FUNCTION get_current_user_roles() RETURNS text[] AS $$
DECLARE
  roles_json text;
BEGIN
  roles_json := current_setting('app.user.roles', true);
  IF roles_json IS NULL OR roles_json = '' THEN
    RETURN ARRAY[]::text[];
  END IF;
  RETURN ARRAY(SELECT json_array_elements_text(roles_json::json));
EXCEPTION
  WHEN others THEN
    -- Catch malformed JSON or other errors.
    RETURN ARRAY[]::text[];
END;
$$ LANGUAGE plpgsql STABLE;

-- Core entities
CREATE TABLE organizations (
  id uuid PRIMARY KEY DEFAULT gen_random_uuid(),
  name text NOT NULL
);

CREATE TABLE teams (
  id uuid PRIMARY KEY DEFAULT gen_random_uuid(),
  organization_id uuid NOT NULL REFERENCES organizations(id) ON DELETE CASCADE,
  name text NOT NULL
);

CREATE TABLE projects (
  id uuid PRIMARY KEY DEFAULT gen_random_uuid(),
  team_id uuid NOT NULL REFERENCES teams(id) ON DELETE CASCADE,
  name text NOT NULL,
  is_public boolean DEFAULT false
);

CREATE TABLE users (
  id uuid PRIMARY KEY DEFAULT gen_random_uuid(),
  email text NOT NULL UNIQUE
);

-- Membership/Junction tables with roles
CREATE TYPE org_role AS ENUM ('ADMIN', 'MEMBER', 'BILLING');
CREATE TYPE team_role AS ENUM ('MAINTAINER', 'DEVELOPER', 'VIEWER');

CREATE TABLE organization_members (
  organization_id uuid NOT NULL REFERENCES organizations(id) ON DELETE CASCADE,
  user_id uuid NOT NULL REFERENCES users(id) ON DELETE CASCADE,
  role org_role NOT NULL,
  PRIMARY KEY (organization_id, user_id)
);

CREATE TABLE team_members (
  team_id uuid NOT NULL REFERENCES teams(id) ON DELETE CASCADE,
  user_id uuid NOT NULL REFERENCES users(id) ON DELETE CASCADE,
  role team_role NOT NULL,
  PRIMARY KEY (team_id, user_id)
);

-- DO NOT USE THIS IN PRODUCTION
CREATE POLICY select_project_naive ON projects FOR SELECT
USING (
  EXISTS (
    -- Check for org-level admin access
    SELECT 1
    FROM organization_members om
    JOIN teams t ON t.organization_id = om.organization_id
    WHERE om.user_id = get_current_user_id()
      AND om.role = 'ADMIN'
      AND t.id = projects.team_id
  )
  OR
  EXISTS (
    -- Check for direct team membership
    SELECT 1
    FROM team_members tm
    WHERE tm.user_id = get_current_user_id()
      AND tm.team_id = projects.team_id
  )
);

CREATE OR REPLACE FUNCTION can_user_view_project(p_user_id uuid, p_project_id uuid) 
RETURNS boolean AS $$
DECLARE
  v_team_id uuid;
  v_organization_id uuid;
  is_org_admin boolean;
  is_team_member boolean;
BEGIN
  -- If user is not logged in, they can't see anything (unless we add public logic later)
  IF p_user_id IS NULL THEN
    RETURN false;
  END IF;

  -- Find the project's hierarchy
  SELECT t.id, t.organization_id
  INTO v_team_id, v_organization_id
  FROM projects p
  JOIN teams t ON p.team_id = t.id
  WHERE p.id = p_project_id;

  -- Short-circuit if project doesn't exist
  IF NOT FOUND THEN
    RETURN false;
  END IF;

  -- Check 1: Is the user an admin of the parent organization?
  SELECT EXISTS (
    SELECT 1 FROM organization_members
    WHERE user_id = p_user_id
      AND organization_id = v_organization_id
      AND role = 'ADMIN'
  ) INTO is_org_admin;

  IF is_org_admin THEN
    RETURN true;
  END IF;

  -- Check 2: Is the user a member of the project's team?
  SELECT EXISTS (
    SELECT 1 FROM team_members
    WHERE user_id = p_user_id
      AND team_id = v_team_id
  ) INTO is_team_member;

  RETURN is_team_member;

END;
$$ LANGUAGE plpgsql STABLE SECURITY DEFINER;

-- The production-ready policy
ALTER TABLE projects ENABLE ROW LEVEL SECURITY;

CREATE POLICY select_project ON projects FOR SELECT
USING ( can_user_view_project(get_current_user_id(), id) );

-- Don't forget policies for other operations!
CREATE POLICY insert_project ON projects FOR INSERT
WITH CHECK ( can_user_create_project_in_team(get_current_user_id(), team_id) ); -- (Function not shown for brevity)

-- Set the context for our test run
BEGIN;
SET LOCAL app.user.id = '...some_user_uuid...';

-- Analyze the query
EXPLAIN ANALYZE SELECT * FROM projects WHERE id = '...some_project_uuid...';

ROLLBACK;

-- Simplified EXPLAIN output snippet
-> Index Scan using projects_pkey on projects
   Filter: can_user_view_project(get_current_user_id(), id)
   -- Inside the function execution:
   -> Seq Scan on organization_members
   -> Seq Scan on team_members

-- Indexes are critical for the function's performance.
CREATE INDEX idx_organization_members_user_id ON organization_members(user_id);
CREATE INDEX idx_team_members_user_id ON team_members(user_id);

-- First, allow access to public projects for everyone, even logged-out users.
CREATE POLICY select_public_projects ON projects FOR SELECT
USING ( is_public = true );

-- The existing policy now only needs to run for non-public projects.
-- We can drop and recreate it.
DROP POLICY select_project ON projects;

CREATE POLICY select_member_projects ON projects FOR SELECT
USING ( is_public = false AND can_user_view_project(get_current_user_id(), id) );

CREATE INDEX idx_projects_public ON projects(id) WHERE is_public = true;

    CREATE ROLE super_admin BYPASSRLS;

    // In the middleware, for a verified super-admin
    await client.query(`SET LOCAL ROLE super_admin;`);

-- Example for UPDATE
CREATE POLICY update_project ON projects FOR UPDATE
USING ( can_user_view_project(get_current_user_id(), id) ) -- Who can even attempt the update
WITH CHECK ( can_user_edit_project(get_current_user_id(), id) ); -- Does the resulting row satisfy edit permissions?

-- A conceptual test case within a single transaction
BEGIN;

-- Set context for a specific user
SET LOCAL app.user.id = 'user_who_is_team_member_uuid';

-- Assert that this user can see the project
-- (Using a testing function that returns true/false or throws error)
SELECT assert_can_select('SELECT * FROM projects WHERE id = \'project_in_their_team_uuid\'');

-- Change context to a different user
SET LOCAL app.user.id = 'user_in_different_org_uuid';

-- Assert that this user CANNOT see the project
SELECT assert_cannot_select('SELECT * FROM projects WHERE id = \'project_in_their_team_uuid\'');

ROLLBACK; -- The test changes nothing in the database.