Noisy Neighbor Prevention.

# noisy_neighbor/rate_limiter.py

"""
Rate limiting implementation using token bucket algorithm.

Token bucket provides:
- Enforces an average rate over time
- Allows short bursts up to the bucket size
- Simple to implement and reason about
"""

from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Optional, Tuple
import asyncio
import time


@dataclass
class RateLimitConfig:
    """Configuration for rate limiting."""
    requests_per_second: float
    burst_size: int  # Maximum tokens (allows short bursts)
    
    @property
    def refill_rate(self) -> float:
        """Tokens added per second."""
        return self.requests_per_second


@dataclass
class RateLimitResult:
    """Result of a rate limit check."""
    allowed: bool
    tokens_remaining: float
    retry_after_seconds: Optional[float] = None
    limit: int = 0
    reset_at: Optional[datetime] = None


class TokenBucket:
    """
    Token bucket rate limiter.
    
    Tokens are added at a constant rate (refill_rate).
    Each request consumes one token.
    Requests are rejected when bucket is empty.
    """
    
    def __init__(self, config: RateLimitConfig):
        self.config = config
        self.tokens = float(config.burst_size)
        self.last_refill = time.monotonic()
        self._lock = asyncio.Lock()
    
    async def try_acquire(self, tokens: int = 1) -> RateLimitResult:
        """
        Try to acquire tokens from the bucket.
        
        Returns RateLimitResult indicating if request is allowed.
        """
        async with self._lock:
            self._refill()
            
            if self.tokens >= tokens:
                self.tokens -= tokens
                return RateLimitResult(
                    allowed=True,
                    tokens_remaining=self.tokens,
                    limit=self.config.burst_size
                )
            else:
                # Calculate when tokens will be available
                tokens_needed = tokens - self.tokens
                wait_seconds = tokens_needed / self.config.refill_rate
                
                return RateLimitResult(
                    allowed=False,
                    tokens_remaining=self.tokens,
                    retry_after_seconds=wait_seconds,
                    limit=self.config.burst_size,
                    reset_at=datetime.now(timezone.utc) + timedelta(seconds=wait_seconds)
                )
    
    def _refill(self):
        """Refill tokens based on elapsed time."""
        now = time.monotonic()
        elapsed = now - self.last_refill
        
        # Add tokens based on elapsed time
        tokens_to_add = elapsed * self.config.refill_rate
        self.tokens = min(
            self.config.burst_size,
            self.tokens + tokens_to_add
        )
        
        self.last_refill = now


class DistributedRateLimiter:
    """
    Distributed rate limiter using Redis.
    
    Uses Redis for coordination across multiple app instances.
    Implements sliding window log algorithm for accuracy:
    at most `limit` requests in any `window_seconds` window.
    """
    
    def __init__(self, redis_client, limit: int, window_seconds: float):
        self.redis = redis_client
        self.limit = limit
        self.window_seconds = window_seconds
    
    async def try_acquire(
        self,
        key: str,
        tokens: int = 1
    ) -> RateLimitResult:
        """
        Try to acquire tokens for a given key.
        
        Key format: "ratelimit:{tenant_id}:{resource}"
        """
        now = time.time()
        window_start = now - self.window_seconds
        
        # Lua script for atomic operation
        script = """
        local key = KEYS[1]
        local now = tonumber(ARGV[1])
        local window_start = tonumber(ARGV[2])
        local limit = tonumber(ARGV[3])
        local tokens = tonumber(ARGV[4])
        local window = tonumber(ARGV[5])
        
        -- Remove old entries outside window
        redis.call('ZREMRANGEBYSCORE', key, '-inf', window_start)
        
        -- Count current entries in window
        local current = redis.call('ZCARD', key)
        
        if current + tokens <= limit then
            -- Add new entries
            for i = 1, tokens do
                redis.call('ZADD', key, now, now .. ':' .. i .. ':' .. math.random())
            end
            redis.call('EXPIRE', key, math.ceil(window) + 1)  -- Outlive the window
            return {1, limit - current - tokens}  -- allowed, remaining
        else
            -- Calculate retry after
            local oldest = redis.call('ZRANGE', key, 0, 0, 'WITHSCORES')
            local retry_after = 0
            if #oldest > 0 then
                retry_after = oldest[2] + window - now
            end
            return {0, retry_after}  -- denied, retry_after
        end
        """
        
        # redis-py signature: eval(script, numkeys, *keys_and_args)
        result = await self.redis.eval(
            script,
            1,  # numkeys: KEYS[1] is the only key
            key,
            now, window_start, self.limit, tokens, self.window_seconds
        )
        
        allowed = result[0] == 1
        
        if allowed:
            return RateLimitResult(
                allowed=True,
                tokens_remaining=result[1],
                limit=self.limit
            )
        else:
            return RateLimitResult(
                allowed=False,
                tokens_remaining=0,
                retry_after_seconds=max(0, result[1]),
                limit=self.limit
            )


class TenantRateLimiter:
    """
    Rate limiter with per-tenant configuration.
    
    Different tenants can have different limits based on their plan.
    """
    
    def __init__(self, redis_client, quota_service):
        self.redis = redis_client
        self.quota_service = quota_service
        self._limiters: dict = {}
    
    async def check_rate_limit(
        self,
        tenant_id: str,
        resource: str = "api",
        tokens: int = 1
    ) -> RateLimitResult:
        """
        Check rate limit for a tenant and resource.
        """
        # Quotas are named "<resource>_rate" (e.g. "api_rate", "export_rate");
        # fall back to the general API quota if the resource has none
        quota = (
            await self.quota_service.get_quota(tenant_id, f"{resource}_rate")
            or await self.quota_service.get_quota(tenant_id, "api_rate")
        )
        
        if not quota:
            # No quota configured = unlimited (careful!)
            return RateLimitResult(allowed=True, tokens_remaining=float('inf'))
        
        # Same key the QuotaService reads for rate usage
        limiter_key = f"rate:{tenant_id}:{quota.name}"
        
        # e.g. api_rate: limit=60, window_seconds=60 -> 60 requests/minute
        limiter = DistributedRateLimiter(
            self.redis,
            limit=int(quota.limit),
            window_seconds=quota.window_seconds or 1
        )
        
        result = await limiter.try_acquire(limiter_key, tokens)
        
        # Track usage for analytics
        await self._record_usage(tenant_id, resource, tokens, result.allowed)
        
        return result
    
    async def _record_usage(
        self,
        tenant_id: str,
        resource: str,
        tokens: int,
        allowed: bool
    ):
        """Record usage for monitoring and billing."""
        await self.redis.hincrby(
            f"usage:{tenant_id}:{resource}",
            "total_requests",
            tokens
        )
        
        if not allowed:
            await self.redis.hincrby(
                f"usage:{tenant_id}:{resource}",
                "throttled_requests",
                tokens
            )

# noisy_neighbor/middleware.py

"""
Middleware for enforcing rate limits on incoming requests.
"""

from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import JSONResponse
import logging
import math

logger = logging.getLogger(__name__)


class RateLimitMiddleware(BaseHTTPMiddleware):
    """
    Middleware that enforces rate limits per tenant.
    """
    
    def __init__(self, app, rate_limiter: TenantRateLimiter):
        super().__init__(app)
        self.rate_limiter = rate_limiter
    
    async def dispatch(self, request: Request, call_next):
        # Skip rate limiting for health checks
        if request.url.path in ["/health", "/metrics"]:
            return await call_next(request)
        
        # Get tenant from context (set by auth middleware)
        tenant_id = getattr(request.state, "tenant_id", None)
        
        if not tenant_id:
            # No tenant = no rate limit (or reject, depending on policy)
            return await call_next(request)
        
        # Determine resource type based on endpoint
        resource = self._get_resource_type(request)
        
        # Check rate limit
        result = await self.rate_limiter.check_rate_limit(
            tenant_id=tenant_id,
            resource=resource
        )
        
        if not result.allowed:
            logger.warning(
                "Rate limit exceeded",
                extra={
                    "tenant_id": tenant_id,
                    "resource": resource,
                    "path": request.url.path,
                    "retry_after": result.retry_after_seconds
                }
            )
            
            return JSONResponse(
                status_code=429,
                content={
                    "error": "rate_limit_exceeded",
                    "message": "Too many requests. Please slow down.",
                    "retry_after_seconds": result.retry_after_seconds,
                    "limit": result.limit
                },
                headers={
                    "Retry-After": str(max(1, math.ceil(result.retry_after_seconds or 1))),
                    "X-RateLimit-Limit": str(result.limit),
                    "X-RateLimit-Remaining": str(int(result.tokens_remaining)),
                }
            )
        
        # Process request
        response = await call_next(request)
        
        # Add rate limit headers (skip when unlimited: int(inf) would raise)
        if result.limit:
            response.headers["X-RateLimit-Limit"] = str(result.limit)
            response.headers["X-RateLimit-Remaining"] = str(int(result.tokens_remaining))
        
        return response
    
    def _get_resource_type(self, request: Request) -> str:
        """
        Determine resource type for rate limiting.
        
        Different endpoints might have different limits.
        """
        path = request.url.path
        
        if path.startswith("/api/search"):
            return "search"  # Search might have lower limits
        elif path.startswith("/api/export"):
            return "export"  # Export might have much lower limits
        elif path.startswith("/api/bulk"):
            return "bulk"    # Bulk operations limited
        else:
            return "api"     # Default API limit

# noisy_neighbor/quota_service.py

"""
Service for managing and enforcing tenant resource quotas.
"""

from dataclasses import dataclass, field
from typing import Dict, Optional, List, Tuple
from datetime import datetime, timedelta
from enum import Enum
import logging
import time

logger = logging.getLogger(__name__)


class QuotaType(Enum):
    RATE = "rate"           # Per time window
    CONCURRENT = "concurrent"  # Simultaneous
    VOLUME = "volume"       # Total amount
    COMPUTE = "compute"     # Processing time


@dataclass
class QuotaDefinition:
    """Definition of a quota."""
    name: str
    quota_type: QuotaType
    limit: float
    window_seconds: Optional[int] = None  # For rate quotas
    burst_size: Optional[int] = None      # For rate quotas
    
    # Soft limit for warnings (percentage of limit)
    warning_threshold: float = 0.8
    
    # Action when exceeded
    action: str = "reject"  # reject, throttle, warn


@dataclass
class QuotaUsage:
    """Current usage against a quota."""
    quota_name: str
    current: float
    limit: float
    percentage: float
    is_exceeded: bool
    is_warning: bool
    reset_at: Optional[datetime] = None


@dataclass
class TenantQuotas:
    """All quotas for a tenant."""
    tenant_id: str
    plan: str
    quotas: Dict[str, QuotaDefinition] = field(default_factory=dict)


# Default quotas by plan
DEFAULT_QUOTAS = {
    "free": {
        "api_rate": QuotaDefinition(
            name="api_rate",
            quota_type=QuotaType.RATE,
            limit=60,
            window_seconds=60,
            burst_size=10
        ),
        "storage_gb": QuotaDefinition(
            name="storage_gb",
            quota_type=QuotaType.VOLUME,
            limit=1
        ),
        "concurrent_requests": QuotaDefinition(
            name="concurrent_requests",
            quota_type=QuotaType.CONCURRENT,
            limit=5
        ),
        "query_timeout_seconds": QuotaDefinition(
            name="query_timeout_seconds",
            quota_type=QuotaType.COMPUTE,
            limit=30
        ),
        "users": QuotaDefinition(
            name="users",
            quota_type=QuotaType.VOLUME,
            limit=5
        ),
    },
    "pro": {
        "api_rate": QuotaDefinition(
            name="api_rate",
            quota_type=QuotaType.RATE,
            limit=1000,
            window_seconds=60,
            burst_size=100
        ),
        "storage_gb": QuotaDefinition(
            name="storage_gb",
            quota_type=QuotaType.VOLUME,
            limit=100
        ),
        "concurrent_requests": QuotaDefinition(
            name="concurrent_requests",
            quota_type=QuotaType.CONCURRENT,
            limit=50
        ),
        "query_timeout_seconds": QuotaDefinition(
            name="query_timeout_seconds",
            quota_type=QuotaType.COMPUTE,
            limit=120
        ),
        "users": QuotaDefinition(
            name="users",
            quota_type=QuotaType.VOLUME,
            limit=50
        ),
    },
    "enterprise": {
        "api_rate": QuotaDefinition(
            name="api_rate",
            quota_type=QuotaType.RATE,
            limit=10000,
            window_seconds=60,
            burst_size=1000
        ),
        "storage_gb": QuotaDefinition(
            name="storage_gb",
            quota_type=QuotaType.VOLUME,
            limit=10000  # 10TB, essentially unlimited
        ),
        "concurrent_requests": QuotaDefinition(
            name="concurrent_requests",
            quota_type=QuotaType.CONCURRENT,
            limit=200
        ),
        "query_timeout_seconds": QuotaDefinition(
            name="query_timeout_seconds",
            quota_type=QuotaType.COMPUTE,
            limit=300
        ),
        "users": QuotaDefinition(
            name="users",
            quota_type=QuotaType.VOLUME,
            limit=100000  # Essentially unlimited
        ),
    }
}


class QuotaService:
    """
    Service for managing tenant quotas.
    """
    
    def __init__(self, db, cache, metrics_client):
        self.db = db
        self.cache = cache
        self.metrics = metrics_client
    
    async def get_tenant_quotas(self, tenant_id: str) -> TenantQuotas:
        """
        Get all quotas for a tenant.
        """
        # Check cache
        cache_key = f"quotas:{tenant_id}"
        cached = await self.cache.get(cache_key)
        
        if cached:
            return TenantQuotas(**cached)
        
        # Load tenant plan
        tenant = await self.db.fetchrow(
            "SELECT plan, custom_quotas FROM tenants WHERE id = $1",
            tenant_id
        )
        
        if not tenant:
            raise ValueError(f"Tenant not found: {tenant_id}")
        
        # Start with default quotas for plan
        plan = tenant["plan"]
        quotas = DEFAULT_QUOTAS.get(plan, DEFAULT_QUOTAS["free"]).copy()
        
        # Apply any custom overrides
        if tenant["custom_quotas"]:
            for name, override in tenant["custom_quotas"].items():
                if name in quotas:
                    quotas[name] = QuotaDefinition(**{**quotas[name].__dict__, **override})
        
        tenant_quotas = TenantQuotas(
            tenant_id=tenant_id,
            plan=plan,
            quotas=quotas
        )
        
        # Cache for 5 minutes
        await self.cache.set(cache_key, tenant_quotas.__dict__, ttl=300)
        
        return tenant_quotas
    
    async def get_quota(
        self,
        tenant_id: str,
        quota_name: str
    ) -> Optional[QuotaDefinition]:
        """Get a specific quota for a tenant."""
        tenant_quotas = await self.get_tenant_quotas(tenant_id)
        return tenant_quotas.quotas.get(quota_name)
    
    async def check_quota(
        self,
        tenant_id: str,
        quota_name: str,
        requested_amount: float = 1
    ) -> QuotaUsage:
        """
        Check if a quota allows the requested amount.
        """
        quota = await self.get_quota(tenant_id, quota_name)
        
        if not quota:
            # No quota defined = allowed
            return QuotaUsage(
                quota_name=quota_name,
                current=0,
                limit=float('inf'),
                percentage=0,
                is_exceeded=False,
                is_warning=False
            )
        
        # Get current usage based on quota type
        if quota.quota_type == QuotaType.RATE:
            current = await self._get_rate_usage(tenant_id, quota_name, quota)
        elif quota.quota_type == QuotaType.CONCURRENT:
            current = await self._get_concurrent_usage(tenant_id, quota_name)
        elif quota.quota_type == QuotaType.VOLUME:
            current = await self._get_volume_usage(tenant_id, quota_name)
        elif quota.quota_type == QuotaType.COMPUTE:
            current = await self._get_compute_usage(tenant_id, quota_name, quota)
        else:
            current = 0
        
        percentage = (current + requested_amount) / quota.limit
        is_exceeded = (current + requested_amount) > quota.limit
        is_warning = percentage >= quota.warning_threshold
        
        usage = QuotaUsage(
            quota_name=quota_name,
            current=current,
            limit=quota.limit,
            percentage=percentage,
            is_exceeded=is_exceeded,
            is_warning=is_warning
        )
        
        # Record metrics
        await self._record_quota_check(tenant_id, usage)
        
        return usage
    
    async def _get_rate_usage(
        self,
        tenant_id: str,
        quota_name: str,
        quota: QuotaDefinition
    ) -> float:
        """Get rate usage from sliding window (same key as the rate limiter)."""
        key = f"rate:{tenant_id}:{quota_name}"
        window_start = time.time() - quota.window_seconds
        
        # Count requests in window
        count = await self.cache.zcount(
            key,
            window_start,
            '+inf'
        )
        
        return count
    
    async def _get_concurrent_usage(
        self,
        tenant_id: str,
        quota_name: str
    ) -> float:
        """Get current concurrent operations (ConcurrencyLimiter's sorted set)."""
        key = f"concurrent:{tenant_id}"
        # Members are scored by expiry time; count the unexpired ones
        return await self.cache.zcount(key, time.time(), '+inf')
    
    async def _get_volume_usage(
        self,
        tenant_id: str,
        quota_name: str
    ) -> float:
        """Get total volume usage from database."""
        # Different queries based on quota type
        if quota_name == "storage_gb":
            result = await self.db.fetchrow(
                """
                SELECT COALESCE(SUM(size_bytes), 0) / 1073741824.0 as usage_gb
                FROM files WHERE tenant_id = $1
                """,
                tenant_id
            )
            return result["usage_gb"]
        
        elif quota_name == "users":
            result = await self.db.fetchrow(
                "SELECT COUNT(*) as count FROM users WHERE tenant_id = $1",
                tenant_id
            )
            return result["count"]
        
        return 0
    
    async def _get_compute_usage(
        self,
        tenant_id: str,
        quota_name: str,
        quota: QuotaDefinition
    ) -> float:
        """Get compute usage in current window."""
        # Compute usage tracked per query, not cumulative
        # For query timeout, we return the timeout limit itself
        return 0  # Single query can use up to the limit
    
    async def _record_quota_check(self, tenant_id: str, usage: QuotaUsage):
        """Record quota check for monitoring."""
        self.metrics.gauge(
            "quota_usage_percentage",
            usage.percentage * 100,
            tags={
                "tenant_id": tenant_id,
                "quota_name": usage.quota_name
            }
        )
        
        if usage.is_warning:
            logger.warning(
                "Quota warning threshold reached",
                extra={
                    "tenant_id": tenant_id,
                    "quota_name": usage.quota_name,
                    "percentage": usage.percentage
                }
            )


class QuotaEnforcer:
    """
    Enforces quotas by rejecting or throttling operations.
    """
    
    def __init__(self, quota_service: QuotaService):
        self.quota_service = quota_service
    
    async def enforce(
        self,
        tenant_id: str,
        quota_name: str,
        requested_amount: float = 1
    ) -> Tuple[bool, Optional[str]]:
        """
        Enforce a quota.
        
        Returns (allowed, error_message).
        """
        usage = await self.quota_service.check_quota(
            tenant_id,
            quota_name,
            requested_amount
        )
        
        if usage.is_exceeded:
            quota = await self.quota_service.get_quota(tenant_id, quota_name)
            
            error_message = (
                f"Quota exceeded for {quota_name}. "
                f"Current: {usage.current:.1f}, Limit: {usage.limit:.1f}. "
                f"Please upgrade your plan or wait for quota reset."
            )
            
            return False, error_message
        
        return True, None

# noisy_neighbor/concurrency_limiter.py

"""
Concurrency limiting to prevent too many simultaneous operations.
"""

from typing import Optional
import asyncio


class ConcurrencyLimiter:
    """
    Limits concurrent operations per tenant.
    
    Prevents a tenant from overwhelming the system with
    many parallel requests.
    """
    
    def __init__(self, redis_client, default_limit: int = 10):
        self.redis = redis_client
        self.default_limit = default_limit
    
    async def try_acquire(
        self,
        tenant_id: str,
        operation_id: str,
        limit: Optional[int] = None,
        ttl_seconds: int = 300
    ) -> bool:
        """
        Try to acquire a concurrency slot.
        
        Args:
            tenant_id: Tenant identifier
            operation_id: Unique ID for this operation
            limit: Max concurrent operations (None = use default)
            ttl_seconds: Auto-release after this time (prevents leaks)
        
        Returns:
            True if slot acquired, False if limit reached
        """
        limit = limit or self.default_limit
        key = f"concurrent:{tenant_id}"
        
        # Lua script for atomic check-and-set
        script = """
        local key = KEYS[1]
        local operation_id = ARGV[1]
        local limit = tonumber(ARGV[2])
        local ttl = tonumber(ARGV[3])
        local now = tonumber(ARGV[4])
        
        -- Remove expired entries
        redis.call('ZREMRANGEBYSCORE', key, '-inf', now)
        
        -- Check current count
        local current = redis.call('ZCARD', key)
        
        if current < limit then
            -- Add this operation with expiry timestamp
            redis.call('ZADD', key, now + ttl, operation_id)
            return 1
        else
            return 0
        end
        """
        
        import time
        now = time.time()
        
        # redis-py signature: eval(script, numkeys, *keys_and_args)
        result = await self.redis.eval(
            script,
            1,  # numkeys: KEYS[1] is the only key
            key,
            operation_id, limit, ttl_seconds, now
        )
        
        return result == 1
    
    async def release(self, tenant_id: str, operation_id: str):
        """
        Release a concurrency slot.
        
        Call this when operation completes.
        """
        key = f"concurrent:{tenant_id}"
        await self.redis.zrem(key, operation_id)
    
    async def get_current_count(self, tenant_id: str) -> int:
        """Get current concurrent operation count."""
        key = f"concurrent:{tenant_id}"
        
        import time
        now = time.time()
        
        # Clean up expired and count
        await self.redis.zremrangebyscore(key, '-inf', now)
        return await self.redis.zcard(key)


class ConcurrencyContext:
    """
    Context manager for concurrent operations.
    
    Usage:
        async with ConcurrencyContext(limiter, tenant_id, op_id) as acquired:
            if acquired:
                # Do operation
            else:
                # Handle limit exceeded
    """
    
    def __init__(
        self,
        limiter: ConcurrencyLimiter,
        tenant_id: str,
        operation_id: str,
        limit: Optional[int] = None
    ):
        self.limiter = limiter
        self.tenant_id = tenant_id
        self.operation_id = operation_id
        self.limit = limit
        self.acquired = False
    
    async def __aenter__(self) -> bool:
        self.acquired = await self.limiter.try_acquire(
            self.tenant_id,
            self.operation_id,
            self.limit
        )
        return self.acquired
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if self.acquired:
            await self.limiter.release(self.tenant_id, self.operation_id)
        return False

# noisy_neighbor/query_analyzer.py

"""
Analyze query complexity to prevent expensive operations.
"""

from dataclasses import dataclass
from typing import List, Optional
from enum import Enum
import json
import re


class QueryComplexity(Enum):
    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"
    VERY_HIGH = "very_high"
    BLOCKED = "blocked"


@dataclass
class QueryAnalysis:
    """Result of query analysis."""
    complexity: QueryComplexity
    estimated_rows: int
    estimated_time_ms: int
    warnings: List[str]
    blocked_reason: Optional[str] = None


class QueryComplexityAnalyzer:
    """
    Analyzes SQL queries for complexity.
    
    Used to prevent queries that would consume too many resources.
    """
    
    # Patterns that indicate expensive queries
    EXPENSIVE_PATTERNS = [
        (r'SELECT\s+\*', 'SELECT * can be expensive, specify columns'),
        (r'(?i)CROSS\s+JOIN', 'CROSS JOIN can explode row count'),
        (r'(?i)(?<!NOT\s)LIKE\s+[\'"]%', 'Leading wildcard LIKE prevents index use'),
        (r'(?i)ORDER\s+BY.*RANDOM', 'ORDER BY RANDOM is expensive'),
        (r'(?i)GROUP\s+BY.*HAVING', 'Complex GROUP BY with HAVING'),
    ]
    
    # Patterns that should be blocked entirely
    BLOCKED_PATTERNS = [
        (r'(?i)SELECT.*FROM.*,.*,.*,.*,', 'Too many table joins'),
        (r'(?is)^\s*(DELETE|UPDATE)\b(?!.*\bWHERE\b)', 'DELETE/UPDATE without WHERE'),
    ]
    
    def __init__(self, db):
        self.db = db
    
    async def analyze(
        self,
        query: str,
        tenant_id: str
    ) -> QueryAnalysis:
        """
        Analyze a query for complexity.
        """
        warnings = []
        
        # Check for blocked patterns
        for pattern, reason in self.BLOCKED_PATTERNS:
            if re.search(pattern, query):
                return QueryAnalysis(
                    complexity=QueryComplexity.BLOCKED,
                    estimated_rows=0,
                    estimated_time_ms=0,
                    warnings=[],
                    blocked_reason=reason
                )
        
        # Check for expensive patterns
        for pattern, warning in self.EXPENSIVE_PATTERNS:
            if re.search(pattern, query):
                warnings.append(warning)
        
        # Use EXPLAIN to estimate cost
        explain_result = await self._explain_query(query, tenant_id)
        
        estimated_rows = explain_result.get("rows", 0)
        estimated_cost = explain_result.get("cost", 0)
        
        # Determine complexity based on estimates
        if estimated_rows > 10_000_000 or estimated_cost > 100000:
            complexity = QueryComplexity.VERY_HIGH
        elif estimated_rows > 1_000_000 or estimated_cost > 10000:
            complexity = QueryComplexity.HIGH
        elif estimated_rows > 100_000 or estimated_cost > 1000:
            complexity = QueryComplexity.MEDIUM
        else:
            complexity = QueryComplexity.LOW
        
        # Estimate time (very rough)
        estimated_time_ms = int(estimated_cost * 0.1)  # Rough estimate
        
        return QueryAnalysis(
            complexity=complexity,
            estimated_rows=estimated_rows,
            estimated_time_ms=estimated_time_ms,
            warnings=warnings
        )
    
    async def _explain_query(self, query: str, tenant_id: str) -> dict:
        """Run EXPLAIN on the query."""
        try:
            # Plain EXPLAIN plans without executing (never EXPLAIN ANALYZE
            # here); run it on a tenant-scoped (RLS) connection
            explain_query = f"EXPLAIN (FORMAT JSON) {query}"
            
            result = await self.db.fetchrow(explain_query)
            
            if result:
                # asyncpg returns the json column as text: parse it first
                plan = json.loads(result[0])[0]["Plan"]
                return {
                    "rows": plan.get("Plan Rows", 0),
                    "cost": plan.get("Total Cost", 0)
                }
        except Exception as e:
            # If EXPLAIN fails, assume worst case (classifies as VERY_HIGH)
            return {"rows": 100_000_000, "cost": 1_000_000}
        
        return {"rows": 0, "cost": 0}


class QueryGuard:
    """
    Guards query execution with complexity and timeout limits.
    """
    
    def __init__(
        self,
        analyzer: QueryComplexityAnalyzer,
        quota_service: QuotaService
    ):
        self.analyzer = analyzer
        self.quota_service = quota_service
    
    async def guard_query(
        self,
        query: str,
        tenant_id: str
    ) -> tuple[bool, Optional[str], Optional[int]]:
        """
        Check if query should be allowed.
        
        Returns:
            (allowed, error_message, timeout_seconds)
        """
        # Get tenant's query timeout quota
        timeout_quota = await self.quota_service.get_quota(
            tenant_id,
            "query_timeout_seconds"
        )
        timeout = int(timeout_quota.limit) if timeout_quota else 30
        
        # Analyze query complexity
        analysis = await self.analyzer.analyze(query, tenant_id)
        
        if analysis.blocked_reason:
            return False, f"Query blocked: {analysis.blocked_reason}", None
        
        if analysis.complexity == QueryComplexity.VERY_HIGH:
            # Check if tenant has permission for very expensive queries
            tenant_quotas = await self.quota_service.get_tenant_quotas(tenant_id)
            
            if tenant_quotas.plan != "enterprise":
                return (
                    False,
                    "Query too complex for your plan. Estimated rows: "
                    f"{analysis.estimated_rows:,}. Consider adding filters or "
                    "upgrading to Enterprise.",
                    None
                )
        
        # Add warnings to response headers
        for warning in analysis.warnings:
            # These would be added to response headers
            pass
        
        return True, None, timeout

# noisy_neighbor/connection_pool.py

"""
Per-tenant connection pool management.

Prevents one tenant from exhausting all database connections.
"""

from dataclasses import dataclass
from typing import Dict, Optional
import asyncio
import asyncpg
import logging

logger = logging.getLogger(__name__)


@dataclass
class PoolConfig:
    """Configuration for a connection pool."""
    min_size: int
    max_size: int
    statement_timeout_ms: int


class TenantConnectionPoolManager:
    """
    Manages database connection pools per tenant.
    
    Each tenant gets their own pool with limits, preventing
    one tenant from exhausting shared connections.
    """
    
    # Pool configurations by plan
    POOL_CONFIGS = {
        "free": PoolConfig(min_size=1, max_size=5, statement_timeout_ms=30000),
        "pro": PoolConfig(min_size=2, max_size=20, statement_timeout_ms=120000),
        "enterprise": PoolConfig(min_size=5, max_size=50, statement_timeout_ms=300000),
    }
    
    def __init__(self, dsn: str, tenant_service):
        self.dsn = dsn
        self.tenant_service = tenant_service
        self._pools: Dict[str, asyncpg.Pool] = {}
        self._lock = asyncio.Lock()
    
    async def get_pool(self, tenant_id: str) -> asyncpg.Pool:
        """
        Get connection pool for a tenant.
        
        Creates pool on first access.
        """
        if tenant_id in self._pools:
            return self._pools[tenant_id]
        
        async with self._lock:
            # Double-check after acquiring lock
            if tenant_id in self._pools:
                return self._pools[tenant_id]
            
            # Get tenant's plan
            tenant = await self.tenant_service.get_tenant(tenant_id)
            plan = tenant.plan if tenant else "free"
            config = self.POOL_CONFIGS.get(plan, self.POOL_CONFIGS["free"])
            
            # Create pool with tenant-specific limits
            pool = await asyncpg.create_pool(
                self.dsn,
                min_size=config.min_size,
                max_size=config.max_size,
                # Server-side: PostgreSQL cancels the statement itself
                server_settings={
                    "statement_timeout": str(config.statement_timeout_ms)
                },
                # Client-side backstop (asyncpg stops waiting)
                command_timeout=config.statement_timeout_ms / 1000,
                setup=self._setup_connection
            )
            
            self._pools[tenant_id] = pool
            
            logger.info(
                f"Created connection pool for tenant {tenant_id}",
                extra={
                    "tenant_id": tenant_id,
                    "plan": plan,
                    "max_connections": config.max_size
                }
            )
            
            return pool
    
    async def _setup_connection(self, conn):
        """Per-connection setup hook."""
        # Note: statement_timeout is set at pool creation via server_settings
        pass
    
    async def get_pool_stats(self, tenant_id: str) -> dict:
        """Get statistics for a tenant's pool."""
        pool = self._pools.get(tenant_id)
        
        if not pool:
            return {"exists": False}
        
        return {
            "exists": True,
            "size": pool.get_size(),
            "free_size": pool.get_idle_size(),
            "used_size": pool.get_size() - pool.get_idle_size(),
            "min_size": pool.get_min_size(),
            "max_size": pool.get_max_size(),
        }
    
    async def close_pool(self, tenant_id: str):
        """Close a tenant's connection pool."""
        pool = self._pools.pop(tenant_id, None)
        
        if pool:
            await pool.close()
            logger.info(f"Closed connection pool for tenant {tenant_id}")
    
    async def close_all(self):
        """Close all connection pools."""
        for tenant_id, pool in list(self._pools.items()):
            await pool.close()
        
        self._pools.clear()

# noisy_neighbor/fair_scheduler.py

"""
Fair scheduling to ensure all tenants get reasonable service.
"""

from dataclasses import dataclass, field
from typing import Dict, List, Optional, Any, Callable
from datetime import datetime, timezone
import asyncio
import heapq
import logging

logger = logging.getLogger(__name__)


@dataclass
class ScheduledTask:
    """A task scheduled for execution."""
    task_id: str
    tenant_id: str
    priority: int
    enqueued_at: datetime
    payload: Any
    callback: Callable
    
    def __lt__(self, other):
        # Lower priority number = higher priority
        # If same priority, earlier enqueue time wins
        if self.priority != other.priority:
            return self.priority < other.priority
        return self.enqueued_at < other.enqueued_at


class FairScheduler:
    """
    Fair scheduler that ensures no tenant monopolizes resources.
    
    Uses weighted fair queuing:
    - Each tenant gets a weight based on their plan
    - Tasks are ordered by (running tasks + 1) / weight, then enqueue time
    - Caveat: the priority is fixed at submit time, so sustained high-weight
      traffic can starve free tenants. Add aging (improve priority with
      wait time) or use round-robin (TenantPriorityQueue) to bound waits.
    """
    
    # Weight by plan (higher = more priority)
    PLAN_WEIGHTS = {
        "free": 1,
        "pro": 5,
        "enterprise": 20,
    }
    
    def __init__(self, max_concurrent: int = 10):
        self.max_concurrent = max_concurrent
        self._queue: List[ScheduledTask] = []
        self._tenant_usage: Dict[str, int] = {}  # Running tasks per tenant
        self._tenant_weights: Dict[str, int] = {}
        self._running = 0
        self._lock = asyncio.Lock()
        self._condition = asyncio.Condition()
    
    async def submit(
        self,
        tenant_id: str,
        task_id: str,
        payload: Any,
        callback: Callable,
        plan: str = "free"
    ):
        """
        Submit a task for fair scheduling.
        """
        weight = self.PLAN_WEIGHTS.get(plan, 1)
        self._tenant_weights[tenant_id] = weight
        
        # Calculate priority based on weight and current usage
        # Lower number = higher priority
        current_usage = self._tenant_usage.get(tenant_id, 0)
        priority = (current_usage + 1) / weight  # More usage = lower priority
        
        task = ScheduledTask(
            task_id=task_id,
            tenant_id=tenant_id,
            priority=int(priority * 1000),  # Scale for heap
            enqueued_at=datetime.now(timezone.utc),
            payload=payload,
            callback=callback
        )
        
        async with self._lock:
            heapq.heappush(self._queue, task)
        
        # Signal that new work is available
        async with self._condition:
            self._condition.notify()
    
    async def run(self):
        """
        Main scheduler loop.
        """
        while True:
            task = await self._get_next_task()
            
            if task:
                # Execute task in background
                asyncio.create_task(self._execute_task(task))
    
    async def _get_next_task(self) -> Optional[ScheduledTask]:
        """Get the next task to execute."""
        async with self._condition:
            while True:
                # Wait if at max concurrency
                while self._running >= self.max_concurrent:
                    await self._condition.wait()
                
                # Try to get a task
                async with self._lock:
                    if self._queue:
                        task = heapq.heappop(self._queue)
                        self._running += 1
                        self._tenant_usage[task.tenant_id] = \
                            self._tenant_usage.get(task.tenant_id, 0) + 1
                        return task
                
                # No tasks, wait for signal
                await self._condition.wait()
    
    async def _execute_task(self, task: ScheduledTask):
        """Execute a task and update counts."""
        try:
            await task.callback(task.payload)
        except Exception as e:
            logger.error(
                f"Task failed: {e}",
                extra={
                    "task_id": task.task_id,
                    "tenant_id": task.tenant_id
                }
            )
        finally:
            async with self._condition:
                self._running -= 1
                self._tenant_usage[task.tenant_id] = \
                    max(0, self._tenant_usage.get(task.tenant_id, 0) - 1)
                self._condition.notify()
    
    def get_stats(self) -> dict:
        """Get scheduler statistics."""
        return {
            "queue_length": len(self._queue),
            "running": self._running,
            "max_concurrent": self.max_concurrent,
            "tenant_usage": dict(self._tenant_usage)
        }


class TenantPriorityQueue:
    """
    Priority queue with per-tenant fairness.
    
    Ensures that a tenant with many queued tasks doesn't
    starve tenants with fewer tasks.
    """
    
    def __init__(self, max_per_tenant: int = 100):
        self.max_per_tenant = max_per_tenant
        self._queues: Dict[str, List[Any]] = {}  # Per-tenant queues
        self._round_robin_index = 0
        self._tenant_order: List[str] = []
        self._lock = asyncio.Lock()
    
    async def enqueue(
        self,
        tenant_id: str,
        item: Any
    ) -> bool:
        """
        Enqueue an item for a tenant.
        
        Returns False if tenant's queue is full.
        """
        async with self._lock:
            if tenant_id not in self._queues:
                self._queues[tenant_id] = []
                self._tenant_order.append(tenant_id)
            
            if len(self._queues[tenant_id]) >= self.max_per_tenant:
                return False
            
            self._queues[tenant_id].append(item)
            return True
    
    async def dequeue(self) -> Optional[tuple[str, Any]]:
        """
        Dequeue an item using round-robin across tenants.
        
        Returns (tenant_id, item) or None if empty.
        """
        async with self._lock:
            if not self._tenant_order:
                return None
            
            # Try each tenant in round-robin order
            attempts = len(self._tenant_order)
            
            for _ in range(attempts):
                tenant_id = self._tenant_order[self._round_robin_index]
                self._round_robin_index = \
                    (self._round_robin_index + 1) % len(self._tenant_order)
                
                if self._queues[tenant_id]:
                    item = self._queues[tenant_id].pop(0)
                    
                    # Remove tenant if queue empty
                    if not self._queues[tenant_id]:
                        del self._queues[tenant_id]
                        self._tenant_order.remove(tenant_id)
                        if self._round_robin_index >= len(self._tenant_order):
                            self._round_robin_index = 0
                    
                    return tenant_id, item
            
            return None
    
    def get_queue_depth(self, tenant_id: str) -> int:
        """Get queue depth for a tenant."""
        return len(self._queues.get(tenant_id, []))
    
    def get_total_depth(self) -> int:
        """Get total queue depth across all tenants."""
        return sum(len(q) for q in self._queues.values())