Spec: PostgreSQL Row-Level Security (Gap #5)¶

CREATE POLICY tenant_isolation ON messages
  USING (conversation_id IN (
    SELECT id FROM conversations 
    WHERE tenant_id = current_setting('app.tenant_id', true)
  ));

"""Enable Row-Level Security on patient-data tables.

Defense-in-depth: DAO layer (Wave 0.1) provides application-level
tenant isolation. RLS adds database-level enforcement.
"""

from alembic import op


def upgrade():
    # ── Step 1: Create application role (non-superuser) ──
    # Railway PostgreSQL default user is superuser.
    # Create a restricted role for the application.
    op.execute("""
        DO $$
        BEGIN
            IF NOT EXISTS (SELECT FROM pg_roles WHERE rolname = 'curaway_app') THEN
                CREATE ROLE curaway_app WITH NOSUPERUSER NOINHERIT LOGIN PASSWORD 'SET_VIA_ENV';
            END IF;
        END $$;
    """)

    # ── Step 2: Create indexes on tenant_id for performance ─��
    # RLS evaluates tenant_id per row. Index prevents full table scans.
    tables_with_tenant_id = [
        'patients', 'cases', 'fhir_resources', 'document_references',
        'consent_records', 'conversations', 'match_results', 'notifications',
        'consultations', 'feedback_records', 'device_registrations',
        'data_forwarding_audits',
    ]
    for table in tables_with_tenant_id:
        op.execute(f"""
            CREATE INDEX IF NOT EXISTS idx_{table}_tenant_id 
            ON {table} (tenant_id);
        """)

    # ── Step 3: Enable RLS on all patient-data tables ──
    all_rls_tables = tables_with_tenant_id + ['messages']
    for table in all_rls_tables:
        op.execute(f"ALTER TABLE {table} ENABLE ROW LEVEL SECURITY;")

    # ── Step 4: Create tenant isolation policies ──
    # Standard policy for tables with direct tenant_id
    for table in tables_with_tenant_id:
        op.execute(f"""
            CREATE POLICY tenant_isolation ON {table}
                USING (tenant_id = current_setting('app.tenant_id', true))
                WITH CHECK (tenant_id = current_setting('app.tenant_id', true));
        """)

    # Special policy for messages (no direct tenant_id — join through conversations)
    op.execute("""
        CREATE POLICY tenant_isolation ON messages
            USING (conversation_id IN (
                SELECT id FROM conversations 
                WHERE tenant_id = current_setting('app.tenant_id', true)
            ))
            WITH CHECK (conversation_id IN (
                SELECT id FROM conversations 
                WHERE tenant_id = current_setting('app.tenant_id', true)
            ));
    """)

    # ── Step 5: Grant permissions to application role ──
    for table in all_rls_tables:
        op.execute(f"GRANT SELECT, INSERT, UPDATE, DELETE ON {table} TO curaway_app;")

    # Grant on system tables too (no RLS, but role needs access)
    system_tables = [
        'audit_logs', 'events', 'providers', 'provider_procedures',
        'provider_facilities', 'doctors', 'doctor_procedures',
        'treatment_categories', 'procedure_requirements', 'consent_purposes',
        'notification_templates', 'notification_preferences',
        'idempotency_keys', 'legal_agreements', 'user_agreement_acceptances',
        'tenants', 'tenant_settings',
    ]
    for table in system_tables:
        op.execute(f"GRANT SELECT, INSERT, UPDATE, DELETE ON {table} TO curaway_app;")

    # Grant sequence access for auto-increment
    op.execute("GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA public TO curaway_app;")


def downgrade():
    """Reverse all RLS changes. Safe — drops policies and disables RLS."""
    all_rls_tables = [
        'patients', 'cases', 'fhir_resources', 'document_references',
        'consent_records', 'conversations', 'messages', 'match_results',
        'notifications', 'consultations', 'feedback_records',
        'device_registrations', 'data_forwarding_audits',
    ]
    for table in all_rls_tables:
        op.execute(f"DROP POLICY IF EXISTS tenant_isolation ON {table};")
        op.execute(f"ALTER TABLE {table} DISABLE ROW LEVEL SECURITY;")

    # Drop indexes (optional — they help performance regardless of RLS)
    # Keep them: they benefit regular queries too.

    # Drop role
    op.execute("DROP ROLE IF EXISTS curaway_app;")

async def get_db(request: Request) -> AsyncGenerator[AsyncSession, None]:
    tenant_id = request.headers.get("X-Tenant-ID")
    async with async_session() as session:
        if tenant_id:
            await session.execute(
                text("SET app.tenant_id = :tid"),
                {"tid": tenant_id},
            )
        yield session
        await session.commit()

async def qstash_callback(document_id: str, db: AsyncSession):
    # Step 1: Lookup tenant_id using superuser (no RLS context yet)
    result = await db.execute(
        text("SELECT tenant_id FROM document_references WHERE id = :did"),
        {"did": document_id},
    )
    tenant_id = result.scalar()

    # Step 2: SET tenant context for subsequent queries
    await db.execute(text("SET app.tenant_id = :tid"), {"tid": tenant_id})

    # Step 3: Proceed with tenant-scoped operations
    ...

# In GDPR erasure handler — use superuser connection (default)
# Do NOT set app.tenant_id — superuser bypasses RLS
async def gdpr_erase_patient(patient_id: str, db: AsyncSession):
    # This runs as superuser — no RLS filtering
    # Delete across all tenants
    await db.execute(
        text("DELETE FROM fhir_resources WHERE patient_id = :pid"),
        {"pid": patient_id},
    )
    # ... cascade through all tables

# Via Railway CLI or psql
CREATE ROLE curaway_app WITH NOSUPERUSER NOINHERIT LOGIN PASSWORD 'xxx';

def test_rls_enabled_on_all_required_tables():
    """Verify RLS is enabled on all 13 patient-data tables."""
    for table in REQUIRED_RLS_TABLES:
        result = db.execute(text(f"""
            SELECT rowsecurity FROM pg_tables 
            WHERE tablename = '{table}'
        """))
        assert result.scalar() == True, f"RLS not enabled on {table}"

def test_rls_policy_exists_on_all_tables():
    """Verify tenant_isolation policy exists on each table."""
    for table in REQUIRED_RLS_TABLES:
        result = db.execute(text(f"""
            SELECT COUNT(*) FROM pg_policies 
            WHERE tablename = '{table}' AND policyname = 'tenant_isolation'
        """))
        assert result.scalar() == 1, f"Policy missing on {table}"

async def test_cross_tenant_select_blocked():
    """Tenant A cannot SELECT tenant B's patients."""
    # Create patients in two tenants
    patient_a = await create_patient(tenant="tenant-a")
    patient_b = await create_patient(tenant="tenant-b")

    # Set context to tenant A
    await db.execute(text("SET app.tenant_id = 'tenant-a'"))
    result = await db.execute(select(Patient))
    patients = result.scalars().all()

    assert patient_a in patients
    assert patient_b not in patients

async def test_cross_tenant_update_blocked():
    """Tenant A cannot UPDATE tenant B's patients."""
    patient_b = await create_patient(tenant="tenant-b")

    await db.execute(text("SET app.tenant_id = 'tenant-a'"))
    rows = await db.execute(
        text("UPDATE patients SET intake_status = 'hacked' WHERE id = :id"),
        {"id": patient_b.id},
    )
    assert rows.rowcount == 0  # No rows updated

async def test_cross_tenant_delete_blocked():
    """Tenant A cannot DELETE tenant B's patients."""
    patient_b = await create_patient(tenant="tenant-b")

    await db.execute(text("SET app.tenant_id = 'tenant-a'"))
    rows = await db.execute(
        text("DELETE FROM patients WHERE id = :id"),
        {"id": patient_b.id},
    )
    assert rows.rowcount == 0  # No rows deleted

async def test_null_tenant_id_returns_nothing():
    """If SET is not called, no rows are visible."""
    patient = await create_patient(tenant="tenant-a")
    # Do NOT set app.tenant_id
    result = await db.execute(select(Patient))
    assert len(result.scalars().all()) == 0

async def test_messages_rls_via_conversation_join():
    """Messages are filtered through conversation's tenant_id."""
    conv_a = await create_conversation(tenant="tenant-a")
    msg_a = await create_message(conversation_id=conv_a.id)

    await db.execute(text("SET app.tenant_id = 'tenant-b'"))
    result = await db.execute(select(Message))
    assert msg_a not in result.scalars().all()

async def test_rls_performance_overhead():
    """RLS adds <5% overhead on 1000-row SELECT."""
    # Insert 1000 patients in tenant-a
    await bulk_insert_patients(1000, tenant="tenant-a")

    # Measure without RLS context (superuser, bypasses RLS)
    start = time.monotonic()
    await db.execute(text("SELECT COUNT(*) FROM patients"))
    baseline_ms = (time.monotonic() - start) * 1000

    # Measure with RLS context
    await db.execute(text("SET app.tenant_id = 'tenant-a'"))
    start = time.monotonic()
    await db.execute(text("SELECT COUNT(*) FROM patients"))
    rls_ms = (time.monotonic() - start) * 1000

    overhead = (rls_ms - baseline_ms) / baseline_ms
    assert overhead < 0.05, f"RLS overhead {overhead:.1%} exceeds 5%"