""" Tenant Scope Enforcement for AegisLM Multi-Tenant System Provides tenant isolation utilities and context management. """ import uuid from contextlib import asynccontextmanager from typing import AsyncGenerator, Optional from dataclasses import dataclass, field @dataclass class TenantScope: """ Represents the current tenant scope for request processing. This ensures all database queries are properly scoped to the current tenant to prevent cross-tenant data leakage. """ tenant_id: uuid.UUID active: bool = True def filter_query(self, table_name: str) -> str: """ Get the tenant filter SQL for a given table. This should be used in all raw SQL queries to ensure tenant isolation. """ return f"tenant_id = '{self.tenant_id}'" def get_tenant_filter(self) -> dict: """Get the tenant filter as a dictionary for ORM queries.""" return {"tenant_id": self.tenant_id} # Thread-local storage for tenant context _tenant_context: Optional[TenantScope] = None def get_tenant_scope() -> Optional[TenantScope]: """ Get the current tenant scope from context. Returns None if no tenant context is set (e.g., for system tasks). """ return _tenant_context def set_tenant_scope(scope: TenantScope) -> None: """ Set the current tenant scope in context. This should be called after authentication to set the tenant for the current request processing. """ global _tenant_context _tenant_context = scope def clear_tenant_scope() -> None: """Clear the current tenant scope.""" global _tenant_context _tenant_context = None @asynccontextmanager async def tenant_context(tenant_id: uuid.UUID) -> AsyncGenerator[TenantScope, None]: """ Context manager for tenant scope. Usage: async with tenant_context(request.tenant_id): # All queries will be scoped to this tenant results = await get_evaluation_runs() """ scope = TenantScope(tenant_id=tenant_id) old_context = _tenant_context try: set_tenant_scope(scope) yield scope finally: clear_tenant_scope() if old_context is not None: set_tenant_scope(old_context) def require_tenant(tenant_id: uuid.UUID) -> TenantScope: """ Require a tenant scope to be set. Raises an exception if no tenant context is active. """ scope = get_tenant_scope() if scope is None: raise RuntimeError("No tenant context set") if scope.tenant_id != tenant_id: raise PermissionError(f"Access denied: tenant {tenant_id} does not match context {scope.tenant_id}") return scope class TenantFilter: """ Utility class for adding tenant filters to queries. """ @staticmethod def for_evaluation_runs(tenant_id: uuid.UUID) -> dict: """Get filter for evaluation_runs table.""" return {"tenant_id": tenant_id} @staticmethod def for_evaluation_results(tenant_id: uuid.UUID) -> dict: """Get filter for evaluation_results table - requires join through runs.""" return {"run": {"tenant_id": tenant_id}} @staticmethod def for_monitoring_metrics(tenant_id: uuid.UUID) -> dict: """Get filter for monitoring_metrics table.""" return {"tenant_id": tenant_id} @staticmethod def for_alerts(tenant_id: uuid.UUID) -> dict: """Get filter for alerts table.""" return {"tenant_id": tenant_id} @staticmethod def for_release_validations(tenant_id: uuid.UUID) -> dict: """Get filter for release_validations table.""" return {"tenant_id": tenant_id} @staticmethod def for_api_keys(tenant_id: uuid.UUID) -> dict: """Get filter for api_keys table.""" return {"tenant_id": tenant_id} @staticmethod def for_users(tenant_id: uuid.UUID) -> dict: """Get filter for users table.""" return {"tenant_id": tenant_id} def validate_tenant_access(resource_tenant_id: uuid.UUID, user_tenant_id: uuid.UUID) -> bool: """ Validate that a user can access a resource based on tenant. This is a safety check to prevent cross-tenant access. Returns True if access is allowed, False otherwise. """ return resource_tenant_id == user_tenant_id def get_safe_tenant_filter(tenant_id: Optional[uuid.UUID]) -> dict: """ Get a tenant filter that handles None gracefully. If tenant_id is None, returns an empty filter (for system tasks). """ if tenant_id is None: return {} return {"tenant_id": tenant_id}