Lesson 16.16: Metadata Filtering & Permission-Aware Retrieval

from enum import Enum
from typing import List, Optional
from pydantic import BaseModel, Field

class DocumentVisibility(str, Enum):
    PUBLIC = "public"          # Accessible across the entire tenant
    RESTRICTED = "restricted"  # Limited to specific roles or users
    CONFIDENTIAL = "confidential" # Strict C-Level or executive only

class ChunkSecurityPayload(BaseModel):
    chunk_id: str
    document_id: str
    tenant_id: str = Field(..., description="Mandatory tenant boundary identifier")
    
    # Access Control Matrix
    visibility: DocumentVisibility = DocumentVisibility.RESTRICTED
    allowed_roles: List[str] = Field(default_factory=list)
    allowed_user_ids: List[str] = Field(default_factory=list)
    
    # Classification Metadata
    department: str
    category: str
    security_level: int = Field(default=1, ge=1, le=5) # 1: Low, 5: Secret
    
    # Provenance and Versioning
    source_filename: str
    page_number: int
    acl_version: int = 1

from qdrant_client import QdrantClient
from qdrant_client.http import models

client = QdrantClient(url="http://localhost:6333")
COLLECTION_NAME = "enterprise_knowledge"

# 1. Critical: Tenant ID Index (Must be KEYWORD)
client.create_payload_index(
    collection_name=COLLECTION_NAME,
    field_name="tenant_id",
    field_schema=models.PayloadSchemaType.KEYWORD,
    is_tenant=True  # Enables Qdrant multi-tenant memory optimizations!
)

# 2. Roles Index (For fast array membership lookup)
client.create_payload_index(
    collection_name=COLLECTION_NAME,
    field_name="allowed_roles",
    field_schema=models.PayloadSchemaType.KEYWORD
)

# 3. User IDs Index
client.create_payload_index(
    collection_name=COLLECTION_NAME,
    field_name="allowed_user_ids",
    field_schema=models.PayloadSchemaType.KEYWORD
)

# 4. Document Visibility Index
client.create_payload_index(
    collection_name=COLLECTION_NAME,
    field_name="visibility",
    field_schema=models.PayloadSchemaType.KEYWORD
)

from typing import Set
from qdrant_client.http import models

def build_qdrant_permission_filter(
    tenant_id: str,
    user_id: str,
    roles: Set[str]
) -> models.Filter:
    """
    Constructs an airtight boolean filter ensuring:
    1. Zero cross-tenant leakage (must tenant_id)
    2. Precise RBAC & ACL enforcement (should clauses)
    """
    # Clause 1: Hard Tenant Isolation (Non-negotiable)
    tenant_constraint = models.FieldCondition(
        key="tenant_id",
        match=models.MatchValue(value=tenant_id)
    )

    # Clause 2: Permission Matrix (Public OR Explicit User OR Matching Role)
    permission_clauses = [
        # Accessible if marked tenant-public
        models.FieldCondition(
            key="visibility",
            match=models.MatchValue(value="tenant")
        ),
        # Accessible if user is explicitly named in ACL
        models.FieldCondition(
            key="allowed_user_ids",
            match=models.MatchValue(value=user_id)
        )
    ]

    # Accessible if any of the user's active roles match
    for role in roles:
        permission_clauses.append(
            models.FieldCondition(
                key="allowed_roles",
                match=models.MatchValue(value=role)
            )
        )

    return models.Filter(
        must=[tenant_constraint],
        should=permission_clauses
    )

from typing import List, Dict, Any, Set
from qdrant_client import QdrantClient

def search_authorized_knowledge(
    qdrant: QdrantClient,
    query_vector: List[float],
    tenant_id: str,
    user_id: str,
    roles: Set[str],
    limit: int = 5
) -> List[Dict[str, Any]]:
    # 1. Build the permission filter
    security_filter = build_qdrant_permission_filter(
        tenant_id=tenant_id,
        user_id=user_id,
        roles=roles
    )

    # 2. Execute Pre-Filtered Vector Search
    search_results = qdrant.search(
        collection_name="enterprise_knowledge",
        query_vector=query_vector,
        query_filter=security_filter,
        limit=limit,
        with_payload=True
    )

    # 3. Format safe payload results
    authorized_chunks = []
    for hit in search_results:
        payload = hit.payload
        authorized_chunks.append({
            "chunk_id": payload.get("chunk_id"),
            "document_id": payload.get("document_id"),
            "tenant_id": payload.get("tenant_id"),
            "score": hit.score,
            "text": payload.get("text"),
            "source": payload.get("source_filename"),
            "category": payload.get("category")
        })

    return authorized_chunks

from dataclasses import dataclass
from fastapi import FastAPI, Depends, Header
from pydantic import BaseModel, Field

app = FastAPI()

@dataclass(frozen=True)
class AuthenticatedPrincipal:
    user_id: str
    tenant_id: str
    roles: frozenset[str]

class KnowledgeSearchRequest(BaseModel):
    query: str = Field(..., min_length=2, max_length=1000)
    limit: int = Field(default=5, ge=1, le=20)
    # Notice: Zero tenant_id or role fields in request body!
    # Clients cannot forge or elevate their authorization context.

async def authenticate_request(
    authorization: str = Header(..., description="Bearer JWT Token")
) -> AuthenticatedPrincipal:
    """Verifies cryptographic signature and decodes trusted principal identity."""
    token = authorization.replace("Bearer ", "").strip()
    payload = verify_and_decode_jwt(token) # Internal JWT validation function
    
    return AuthenticatedPrincipal(
        user_id=payload["sub"],
        tenant_id=payload["tenant_id"],
        roles=frozenset(payload.get("roles", []))
    )

@app.post("/api/v1/knowledge/search")
async def secure_knowledge_search(
    req: KnowledgeSearchRequest,
    principal: AuthenticatedPrincipal = Depends(authenticate_request)
):
    # Only trusted principal credentials gate the search
    results = search_authorized_knowledge(
        qdrant=qdrant_client,
        query_vector=generate_embedding(req.query),
        tenant_id=principal.tenant_id,
        user_id=principal.user_id,
        roles=set(principal.roles),
        limit=req.limit
    )
    return {"query": req.query, "results": results}

-- 1. Enable Row-Level Security on document chunks table
ALTER TABLE document_chunks ENABLE ROW LEVEL SECURITY;

-- 2. Define Tenant Isolation Policy
-- Reads are permitted ONLY if the row's tenant_id matches current session variable
CREATE POLICY tenant_isolation_policy ON document_chunks
    FOR ALL
    USING (tenant_id = current_setting('app.current_tenant_id', true));

-- 3. Application Connection Setup in Python/SQLAlchemy
-- When acquiring a connection from the pool, set session context:
-- SET LOCAL app.current_tenant_id = 'tenant_alpha';
-- Now any SELECT * automatically isolates Tenant Alpha data!

import pytest

def test_cross_tenant_isolation(qdrant_mock):
    """Tenant A's user must NEVER receive chunks belonging to Tenant B."""
    results = search_authorized_knowledge(
        qdrant=qdrant_mock,
        query_vector=[0.1] * 1536,
        tenant_id="tenant_alpha",
        user_id="alice",
        roles={"employee"}
    )
    for r in results:
        assert r["tenant_id"] == "tenant_alpha"
        assert r["tenant_id"] != "tenant_beta"

def test_confidential_document_blocked_for_normal_employee(qdrant_mock):
    """Ordinary employees must never retrieve confidential finance audit chunks."""
    results = search_authorized_knowledge(
        qdrant=qdrant_mock,
        query_vector=[0.1] * 1536,
        tenant_id="tenant_alpha",
        user_id="alice",
        roles={"employee"}
    )
    categories = [r["category"] for r in results]
    assert "finance_audit" not in categories

def test_finance_manager_can_access_audit_report(qdrant_mock):
    """Authorized finance manager must successfully retrieve audit chunks."""
    results = search_authorized_knowledge(
        qdrant=qdrant_mock,
        query_vector=[0.1] * 1536,
        tenant_id="tenant_alpha",
        user_id="charlie",
        roles={"employee", "finance_manager"}
    )
    categories = [r["category"] for r in results]
    assert "finance_audit" in categories