Postgres RLS with Partial Indexes for Performant Multi-Tenancy

-- Create a tenants table for context
CREATE TABLE tenants ( 
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    name TEXT NOT NULL
);

-- The large, monolithic table for our invoices
CREATE TABLE invoices (
    id BIGSERIAL PRIMARY KEY,
    tenant_id UUID NOT NULL, -- Foreign key would normally exist
    invoice_date DATE NOT NULL,
    amount_cents INT NOT NULL,
    status TEXT NOT NULL CHECK (status IN ('draft', 'sent', 'paid', 'void')),
    created_at TIMESTAMPTZ NOT NULL DEFAULT now()
);

-- Populate with 1000 tenants
INSERT INTO tenants (id, name) 
SELECT gen_random_uuid(), 'Tenant ' || i FROM generate_series(1, 1000) i;

-- Populate with 10M invoices, randomly assigned to tenants
INSERT INTO invoices (tenant_id, invoice_date, amount_cents, status)
SELECT
    (SELECT id FROM tenants ORDER BY random() LIMIT 1),
    '2023-01-01'::date + (random() * 365)::int,
    (random() * 100000)::int,
    (ARRAY['draft', 'sent', 'paid', 'void'])[ceil(random()*4)]
FROM generate_series(1, 10000000);

-- Create a standard index that a DBA would add for tenant-specific queries
CREATE INDEX idx_invoices_tenant_id_invoice_date ON invoices(tenant_id, invoice_date);

-- Enable RLS on the table
ALTER TABLE invoices ENABLE ROW LEVEL SECURITY;

-- Create the policy to isolate data by tenant_id
CREATE POLICY tenant_isolation ON invoices
FOR ALL
USING (tenant_id = current_setting('app.tenant_id', true)::uuid);

-- We also need to grant access. For simplicity, we grant to public.
-- In production, this would be a specific application role.
GRANT SELECT, INSERT, UPDATE, DELETE ON invoices TO public;

-- Simulate setting the tenant context in a session
-- Replace with a valid UUID from your tenants table
SET app.tenant_id = 'some-tenant-uuid-from-your-db';

-- Run the query and analyze its performance
EXPLAIN ANALYZE
SELECT id, invoice_date, amount_cents
FROM invoices
WHERE status = 'paid'
  AND invoice_date >= '2023-10-01'
ORDER BY invoice_date DESC
LIMIT 100;

Limit  (cost=189332.61..189332.86 rows=100 width=20) (actual time=1453.315..1453.337 rows=100 loops=1)
  ->  Sort  (cost=189332.61..189498.05 rows=66176 width=20) (actual time=1453.313..1453.323 rows=100 loops=1)
        Sort Key: invoice_date DESC
        Sort Method: top-N heapsort  Memory: 39kB
        ->  Bitmap Heap Scan on invoices  (cost=10173.19..187678.18 rows=66176 width=20) (actual time=146.491..1442.532 rows=2489 loops=1)
              Recheck Cond: ((tenant_id = current_setting('app.tenant_id'::text, true)::uuid) AND (invoice_date >= '2023-10-01'::date) AND (status = 'paid'::text))
              Rows Removed by Index Recheck: 189345
              Heap Blocks: exact=19124
              ->  Bitmap Index Scan on idx_invoices_tenant_id_invoice_date  (cost=0.00..10156.65 rows=66176 width=0) (actual time=143.123..143.123 rows=21613 loops=1)
                    Index Cond: ((tenant_id = current_setting('app.tenant_id'::text, true)::uuid) AND (invoice_date >= '2023-10-01'::date))
Planning Time: 0.315 ms
Execution Time: 1453.489 ms

-- Drop the old monolithic table
DROP TABLE invoices;

-- Create the parent partitioned table
-- Note the PARTITION BY LIST clause
CREATE TABLE invoices (
    id BIGSERIAL,
    tenant_id UUID NOT NULL,
    invoice_date DATE NOT NULL,
    amount_cents INT NOT NULL,
    status TEXT NOT NULL CHECK (status IN ('draft', 'sent', 'paid', 'void')),
    created_at TIMESTAMPTZ NOT NULL DEFAULT now()
) PARTITION BY LIST (tenant_id);

-- The primary key and other indexes must include the partition key (tenant_id)
-- This is a requirement for partitioned tables to allow for local indexes.
ALTER TABLE invoices ADD PRIMARY KEY (id, tenant_id);

-- A global index on the parent table is not what we want. 
-- We will create indexes on each partition instead (local indexes).

CREATE OR REPLACE FUNCTION create_partition_and_insert_invoice()
RETURNS TRIGGER AS $$
DECLARE
    v_partition_name TEXT;
    v_tenant_id_text TEXT;
BEGIN
    -- Sanitize the tenant_id for use in a table name
    v_tenant_id_text := replace(NEW.tenant_id::text, '-', '_');
    v_partition_name := format('invoices_tenant_%s', v_tenant_id_text);

    -- Attempt to insert the row into the partitioned table
    -- This is the "happy path" where the partition already exists
    INSERT INTO invoices VALUES (NEW.*);
    RETURN NEW;

EXCEPTION
    -- If the insert fails because the partition does not exist, create it and retry.
    WHEN OTHERS THEN -- A more specific error code is `23514` (check_violation for partitioning)
                     -- but OTHERS is safer across PG versions.
        BEGIN
            -- Check if another concurrent process created the partition in the meantime
            IF NOT EXISTS (SELECT 1 FROM pg_class WHERE relname = v_partition_name) THEN
                -- Create the partition for the specific tenant_id
                EXECUTE format(
                    'CREATE TABLE %I PARTITION OF invoices FOR VALUES IN (%L)',
                    v_partition_name, NEW.tenant_id
                );

                -- IMPORTANT: Create local indexes on the new partition
                -- This is critical for query performance within the tenant.
                EXECUTE format(
                    'CREATE INDEX %I ON %I (invoice_date DESC)',
                    'idx_' || v_partition_name || '_invoice_date',
                    v_partition_name
                );

                EXECUTE format(
                    'CREATE INDEX %I ON %I (status)',
                    'idx_' || v_partition_name || '_status',
                    v_partition_name
                );
            END IF;
        EXCEPTION
            WHEN duplicate_table THEN
                -- Another process created the partition between our check and our CREATE command. This is fine.
                NULL;
        END;
        
        -- Retry the insert. This time it will succeed.
        INSERT INTO invoices VALUES (NEW.*);
        RETURN NEW;
END;
$$ LANGUAGE plpgsql;

-- This function is intended to be called by application logic, not a standard trigger.
-- A typical application pattern would be:
-- 1. BEGIN transaction
-- 2. Attempt a direct INSERT into `invoices`.
-- 3. If it fails with a partition error, call a function that creates the partition.
-- 4. Retry the INSERT.

-- For simplicity in this example, we'll manually create the first partition.
-- In a real system, a tenant creation hook would call this logic.

-- Manually create a partition for our test tenant
-- Replace with the same UUID used before
CREATE TABLE invoices_tenant_... PARTITION OF invoices FOR VALUES IN ('some-tenant-uuid-from-your-db');
CREATE INDEX idx_invoices_tenant_..._invoice_date ON invoices_tenant_... (invoice_date DESC);
CREATE INDEX idx_invoices_tenant_..._status ON invoices_tenant_... (status);

-- Enable RLS on the parent table
ALTER TABLE invoices ENABLE ROW LEVEL SECURITY;

-- The policy is the same, but it will now work with partition pruning
CREATE POLICY tenant_isolation ON invoices
FOR ALL
USING (tenant_id = current_setting('app.tenant_id', true)::uuid);

-- Grant privileges again
GRANT SELECT, INSERT, UPDATE, DELETE ON invoices TO public;

-- (After populating data into the relevant partitions)

-- Set the same tenant context
SET app.tenant_id = 'some-tenant-uuid-from-your-db';

-- Run the exact same query
EXPLAIN ANALYZE
SELECT id, invoice_date, amount_cents
FROM invoices
WHERE status = 'paid'
  AND invoice_date >= '2023-10-01'
ORDER BY invoice_date DESC
LIMIT 100;

Limit  (cost=0.29..8.40 rows=100 width=20) (actual time=0.041..0.123 rows=100 loops=1)
  ->  Index Scan using idx_invoices_tenant_..._invoice_date on invoices_tenant_... invoices_1  (cost=0.29..20.59 rows=254 width=20) (actual time=0.039..0.101 rows=100 loops=1)
        Index Cond: (invoice_date >= '2023-10-01'::date)
        Filter: (status = 'paid'::text)
Planning Time: 0.451 ms
Execution Time: 0.158 ms

-- Drop the old policy
DROP POLICY tenant_isolation ON invoices;

-- Create a new policy with a role-based bypass
CREATE POLICY tenant_isolation_with_admin_bypass ON invoices
FOR ALL
USING (
    -- Standard tenant isolation
    tenant_id = current_setting('app.tenant_id', true)::uuid
    OR
    -- Bypass for users with the 'app_admin' role
    EXISTS (SELECT 1 FROM pg_roles WHERE rolname = 'app_admin' AND pg_has_role(current_user, 'app_admin', 'MEMBER'))
);

-- Example: Adding a nullable column is fast
ALTER TABLE invoices ADD COLUMN notes TEXT;

-- Example: Adding a non-nullable column with a default requires rewriting all partitions
-- This will acquire a strong lock on each partition in sequence.
ALTER TABLE invoices ADD COLUMN is_archived BOOLEAN NOT NULL DEFAULT false;