Advanced RLS Policies in PostgreSQL for Hierarchical Multi-Tenancy

-- Organizations are the top-level tenants.
CREATE TABLE organizations (
    id BIGINT PRIMARY KEY GENERATED BY DEFAULT AS IDENTITY,
    name TEXT NOT NULL
);

-- Teams belong to an organization.
CREATE TABLE teams (
    id BIGINT PRIMARY KEY GENERATED BY DEFAULT AS IDENTITY,
    organization_id BIGINT NOT NULL REFERENCES organizations(id) ON DELETE CASCADE,
    name TEXT NOT NULL
);

-- Projects belong to a team, and therefore transitively to an organization.
CREATE TABLE projects (
    id BIGINT PRIMARY KEY GENERATED BY DEFAULT AS IDENTITY,
    team_id BIGINT NOT NULL REFERENCES teams(id) ON DELETE CASCADE,
    name TEXT NOT NULL,
    -- We'll denormalize organization_id for easier policy creation.
    organization_id BIGINT NOT NULL REFERENCES organizations(id) ON DELETE CASCADE
);

-- Users are members of a team.
CREATE TABLE users (
    id BIGINT PRIMARY KEY GENERATED BY DEFAULT AS IDENTITY,
    team_id BIGINT NOT NULL REFERENCES teams(id) ON DELETE CASCADE,
    email TEXT NOT NULL UNIQUE
);

-- A simple table to demonstrate RLS on project-related data.
CREATE TABLE tasks (
    id BIGINT PRIMARY KEY GENERATED BY DEFAULT AS IDENTITY,
    project_id BIGINT NOT NULL REFERENCES projects(id) ON DELETE CASCADE,
    description TEXT NOT NULL,
    completed BOOLEAN NOT NULL DEFAULT false
);

-- Create indexes on foreign keys for performance.
CREATE INDEX ON teams (organization_id);
CREATE INDEX ON projects (team_id);
CREATE INDEX ON projects (organization_id);
CREATE INDEX ON users (team_id);
CREATE INDEX ON tasks (project_id);

-- A private schema to hold our security helper functions.
CREATE SCHEMA app_private;

CREATE OR REPLACE FUNCTION app_private.set_current_user_context(user_id BIGINT) 
RETURNS VOID AS $$
DECLARE
    user_team_id BIGINT;
    user_org_id BIGINT;
BEGIN
    -- Find the user's team and organization.
    SELECT u.team_id, t.organization_id 
    INTO user_team_id, user_org_id
    FROM users u
    JOIN teams t ON u.team_id = t.id
    WHERE u.id = user_id;

    IF NOT FOUND THEN
        RAISE EXCEPTION 'User not found: %', user_id;
    END IF;

    -- Set transaction-local settings. These are secure and isolated.
    -- The 'true' flag indicates that an error should not be thrown if the setting doesn't exist.
    PERFORM set_config('app.current_user_id', user_id::TEXT, true);
    PERFORM set_config('app.current_team_id', user_team_id::TEXT, true);
    PERFORM set_config('app.current_org_id', user_org_id::TEXT, true);
END;
$$ LANGUAGE plpgsql SECURITY DEFINER;

-- Revoke execute from public to ensure only trusted roles can call this.
REVOKE EXECUTE ON FUNCTION app_private.set_current_user_context(BIGINT) FROM PUBLIC;

// Example middleware in a Node.js application
async function setUserContext(req, res, next) {
  if (req.user) { // Assuming req.user is populated from a validated JWT
    const client = await pool.connect();
    try {
      // BEGIN transaction implicitly handled by some libraries
      await client.query('SELECT app_private.set_current_user_context($1)', [req.user.id]);
      res.locals.dbClient = client; // Attach client to request for use in handlers
      next();
    } catch (err) {
      client.release();
      next(err);
    }
  } else {
    next();
  }
}

-- A public schema for functions that can be safely used in RLS policies.
CREATE SCHEMA app_public;

-- Helper to get the current user ID, with casting.
CREATE OR REPLACE FUNCTION app_public.current_user_id() RETURNS BIGINT AS $$
    SELECT nullif(current_setting('app.current_user_id', true), '')::BIGINT;
$$ LANGUAGE sql STABLE;

-- Helper to get the current team ID.
CREATE OR REPLACE FUNCTION app_public.current_team_id() RETURNS BIGINT AS $$
    SELECT nullif(current_setting('app.current_team_id', true), '')::BIGINT;
$$ LANGUAGE sql STABLE;

-- Helper to get the current organization ID.
CREATE OR REPLACE FUNCTION app_public.current_org_id() RETURNS BIGINT AS $$
    SELECT nullif(current_setting('app.current_org_id', true), '')::BIGINT;
$$ LANGUAGE sql STABLE;

-- The core logic for read access: Is the target entity within the user's organization?
CREATE OR REPLACE FUNCTION app_public.is_in_my_organization(target_organization_id BIGINT) RETURNS BOOLEAN AS $$
    SELECT target_organization_id = app_public.current_org_id();
$$ LANGUAGE sql STABLE;

-- The core logic for write access: Is the target entity within the user's immediate team?
CREATE OR REPLACE FUNCTION app_public.is_in_my_team(target_team_id BIGINT) RETURNS BOOLEAN AS $$
    SELECT target_team_id = app_public.current_team_id();
$$ LANGUAGE sql STABLE;

ALTER TABLE projects ENABLE ROW LEVEL SECURITY;
ALTER TABLE tasks ENABLE ROW LEVEL SECURITY;

-- Policy for SELECT access: Users can see all projects in their organization.
CREATE POLICY select_projects ON projects
    FOR SELECT
    USING (app_public.is_in_my_organization(organization_id));

-- Policy for INSERT access: Users can only create projects for their own team.
CREATE POLICY insert_projects ON projects
    FOR INSERT
    WITH CHECK (app_public.is_in_my_team(team_id));

-- Policy for UPDATE access: Users can only update projects within their own team.
CREATE POLICY update_projects ON projects
    FOR UPDATE
    USING (app_public.is_in_my_team(team_id));

-- Policy for DELETE access: Users can only delete projects from their own team.
CREATE POLICY delete_projects ON projects
    FOR DELETE
    USING (app_public.is_in_my_team(team_id));

CREATE POLICY select_tasks ON tasks
    FOR SELECT
    USING (
        EXISTS (
            SELECT 1
            FROM projects p
            WHERE p.id = tasks.project_id
            AND app_public.is_in_my_organization(p.organization_id)
        )
    );

CREATE POLICY insert_tasks ON tasks
    FOR INSERT
    WITH CHECK (
        EXISTS (
            SELECT 1
            FROM projects p
            WHERE p.id = tasks.project_id
            AND app_public.is_in_my_team(p.team_id)
        )
    );

-- For UPDATE and DELETE, we combine the logic.
CREATE POLICY modify_tasks ON tasks
    FOR ALL -- Applies to UPDATE and DELETE, as well as SELECT/INSERT if others not defined
    USING (
        EXISTS (
            SELECT 1
            FROM projects p
            WHERE p.id = tasks.project_id
            AND app_public.is_in_my_team(p.team_id)
        )
    )
    WITH CHECK (
        EXISTS (
            SELECT 1
            FROM projects p
            WHERE p.id = tasks.project_id
            AND app_public.is_in_my_team(p.team_id)
        )
    );

-- A more generic hierarchy model
CREATE TABLE tenants (
    id BIGINT PRIMARY KEY GENERATED BY DEFAULT AS IDENTITY,
    parent_id BIGINT REFERENCES tenants(id),
    name TEXT NOT NULL
);
CREATE INDEX ON tenants (parent_id);

CREATE TABLE user_tenant_memberships (
    user_id BIGINT NOT NULL REFERENCES users(id),
    tenant_id BIGINT NOT NULL REFERENCES tenants(id),
    PRIMARY KEY (user_id, tenant_id)
);

-- Our core helper function, now using recursion
CREATE OR REPLACE FUNCTION app_public.is_descendant_or_self(ancestor_id BIGINT, descendant_id BIGINT) 
RETURNS BOOLEAN AS $$
BEGIN
    -- Fast path for direct equality
    IF ancestor_id = descendant_id THEN
        RETURN TRUE;
    END IF;

    RETURN EXISTS (
        WITH RECURSIVE hierarchy AS (
            SELECT id FROM tenants WHERE id = descendant_id
            UNION ALL
            SELECT t.parent_id FROM tenants t JOIN hierarchy h ON t.id = h.id WHERE t.parent_id IS NOT NULL
        )
        SELECT 1 FROM hierarchy WHERE id = ancestor_id
    );
END;
$$ LANGUAGE plpgsql STABLE;

-- Assuming 'resources' table has a 'tenant_id'
CREATE POLICY select_resources ON resources
FOR SELECT USING (
    EXISTS (
        SELECT 1
        FROM user_tenant_memberships utm
        WHERE utm.user_id = app_public.current_user_id()
        AND app_public.is_descendant_or_self(utm.tenant_id, resources.tenant_id)
    )
);

CREATE EXTENSION IF NOT EXISTS ltree;

ALTER TABLE tenants ADD COLUMN path LTREE;

-- We'd need a trigger to maintain this path on insert/update
CREATE OR REPLACE FUNCTION update_tenant_path() RETURNS TRIGGER AS $$
BEGIN
    IF NEW.parent_id IS NULL THEN
        NEW.path = NEW.id::text::ltree;
    ELSE
        SELECT path || NEW.id::text FROM tenants WHERE id = NEW.parent_id INTO NEW.path;
        IF NOT FOUND THEN
            RAISE EXCEPTION 'Invalid parent_id %', NEW.parent_id;
        END IF;
    END IF;
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

CREATE TRIGGER tenant_path_trigger
BEFORE INSERT OR UPDATE ON tenants
FOR EACH ROW EXECUTE FUNCTION update_tenant_path();

-- Create a GiST index for fast path lookups
CREATE INDEX tenant_path_idx ON tenants USING gist (path);

CREATE OR REPLACE FUNCTION app_public.is_descendant_or_self_ltree(ancestor_id BIGINT, descendant_id BIGINT)
RETURNS BOOLEAN AS $$
DECLARE
    ancestor_path LTREE;
    descendant_path LTREE;
BEGIN
    SELECT path INTO ancestor_path FROM tenants WHERE id = ancestor_id;
    SELECT path INTO descendant_path FROM tenants WHERE id = descendant_id;

    -- The '<@' operator means 'is ancestor of or equal to'
    RETURN ancestor_path <@ descendant_path;
END;
$$ LANGUAGE plpgsql STABLE;

CREATE TABLE user_accessible_tenants (
    user_id BIGINT NOT NULL,
    tenant_id BIGINT NOT NULL,
    PRIMARY KEY (user_id, tenant_id)
);

-- This table would be populated by a trigger on 'tenants' and 'user_tenant_memberships'
-- The trigger would calculate all descendants of a user's assigned tenants and insert them.

CREATE POLICY select_resources_fast ON resources
FOR SELECT USING (
    EXISTS (
        SELECT 1
        FROM user_accessible_tenants uat
        WHERE uat.user_id = app_public.current_user_id()
        AND uat.tenant_id = resources.tenant_id
    )
);

    -- In set_current_user_context, check if the user is an admin and set the flag
    -- PERFORM set_config('app.is_admin', 'true', true);

    -- Modify the policy
    CREATE POLICY select_projects ON projects FOR SELECT
    USING (
        (current_setting('app.is_admin', true) = 'true')
        OR
        (app_public.is_in_my_organization(organization_id))
    );

    BEGIN;
    -- Set the context for the user you want to test
    SELECT app_private.set_current_user_context(123);
    -- Impersonate the role your application uses
    SET ROLE my_application_role;
    
    -- Run your test query
    EXPLAIN ANALYZE SELECT * FROM tasks WHERE id = 456;
    -- This will show you the query plan with the RLS policy applied.

    ROLLBACK;