Zero-Leak Multi-TenantRAG Architecture

-- Enable the pgvector extension
CREATE EXTENSION IF NOT EXISTS vector;

-- Create the documents table
CREATE TABLE documents (
  id          UUID DEFAULT gen_random_uuid() PRIMARY KEY,
  tenant_id   TEXT NOT NULL,
  content     TEXT NOT NULL,
  embedding   vector(1536),  -- 1536 dims for text-embedding-3-small
  metadata    JSONB,
  created_at  TIMESTAMPTZ DEFAULT NOW()
);

-- Index for fast vector similarity search
CREATE INDEX ON documents
  USING ivfflat (embedding vector_cosine_ops)
  WITH (lists = 100);

-- Index on tenant_id for the RLS filter
CREATE INDEX ON documents (tenant_id);

-- Enable Row-Level Security
ALTER TABLE documents ENABLE ROW LEVEL SECURITY;

-- Create the isolation policy
-- This reads the session variable we set at connection time
CREATE POLICY tenant_isolation ON documents
  USING (tenant_id = current_setting('app.current_tenant', TRUE));

-- Create a role that is NOT a superuser (RLS applies)
CREATE ROLE app_user WITH LOGIN PASSWORD 'your_strong_password';

-- Grant only the necessary permissions
GRANT CONNECT ON DATABASE your_db TO app_user;
GRANT USAGE ON SCHEMA public TO app_user;
GRANT SELECT, INSERT, UPDATE, DELETE ON documents TO app_user;

-- CRITICAL: Do NOT grant BYPASSRLS privilege to app_user
-- Only your DBA role should have that.

from fastapi import FastAPI, HTTPException, Header
from pydantic import BaseModel
from openai import OpenAI
import asyncpg
import os

app = FastAPI()
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

DATABASE_URL = os.getenv("DATABASE_URL")

class UpsertRequest(BaseModel):
    content: str
    metadata: dict = {}

class QueryRequest(BaseModel):
    query: str
    top_k: int = 5

async def get_embedding(text: str) -> list[float]:
    """Generate an embedding vector for the given text."""
    response = client.embeddings.create(
        input=text,
        model="text-embedding-3-small"
    )
    return response.data[0].embedding

async def get_db_connection(tenant_id: str):
    """
    Get a database connection and immediately set the tenant context.
    This is the critical step that activates RLS.
    """
    conn = await asyncpg.connect(DATABASE_URL)
    
    # THIS IS THE SECURITY-CRITICAL LINE
    # Set the session variable that our RLS policy reads
    await conn.execute(
        "SELECT set_config('app.current_tenant', $1, TRUE)",
        tenant_id
    )
    return conn

@app.post("/documents/upsert")
async def upsert_document(
    req: UpsertRequest,
    x_tenant_id: str = Header(...)  # Extracted from authenticated JWT
):
    if not x_tenant_id:
        raise HTTPException(status_code=401, detail="Tenant ID required")
    
    embedding = await get_embedding(req.content)
    
    conn = await get_db_connection(x_tenant_id)
    try:
        # tenant_id is stored in the row AND enforced by RLS
        # Defense-in-depth: two layers of enforcement
        await conn.execute(
            """
            INSERT INTO documents (tenant_id, content, embedding, metadata)
            VALUES ($1, $2, $3::vector, $4)
            """,
            x_tenant_id,
            req.content,
            embedding,
            req.metadata
        )
    finally:
        await conn.close()
    
    return {"status": "ok"}

@app.post("/documents/query")
async def query_documents(
    req: QueryRequest,
    x_tenant_id: str = Header(...)
):
    if not x_tenant_id:
        raise HTTPException(status_code=401, detail="Tenant ID required")
    
    query_embedding = await get_embedding(req.query)
    
    conn = await get_db_connection(x_tenant_id)
    try:
        # RLS policy automatically filters to current tenant
        # The AND tenant_id = $2 is a redundant application-level check
        rows = await conn.fetch(
            """
            SELECT content, metadata,
                   1 - (embedding <=> $1::vector) AS similarity
            FROM documents
            WHERE tenant_id = $2  -- Redundant but intentional safety check
            ORDER BY embedding <=> $1::vector
            LIMIT $3
            """,
            query_embedding,
            x_tenant_id,  
            req.top_k
        )
    finally:
        await conn.close()
    
    return {
        "results": [
            {
                "content": row["content"],
                "metadata": row["metadata"],
                "similarity": float(row["similarity"])
            }
            for row in rows
        ]
    }

import asyncio
import asyncpg
import pytest

async def test_rls_prevents_cross_tenant_access():
    """
    Proves that RLS blocks cross-tenant data access
    even when the application-level filter is bypassed.
    """
    conn = await asyncpg.connect(DATABASE_URL, user='app_user')
    
    # Insert a document for Tenant A
    await conn.execute(
        "SELECT set_config('app.current_tenant', 'tenant_a', TRUE)"
    )
    await conn.execute(
        "INSERT INTO documents (tenant_id, content, embedding) VALUES ('tenant_a', 'Secret A data', '[0.1,...]'::vector)"
    )
    
    # Now switch context to Tenant B and attempt to read without any filter
    await conn.execute(
        "SELECT set_config('app.current_tenant', 'tenant_b', TRUE)"
    )
    
    # This should return ZERO rows — RLS blocks it at the engine level
    rows = await conn.fetch("SELECT * FROM documents")
    
    assert len(rows) == 0, "CRITICAL SECURITY FAILURE: Cross-tenant data access detected!"
    print("PASS: RLS correctly blocked cross-tenant access.")
    
    await conn.close()

asyncio.run(test_rls_prevents_cross_tenant_access())