Spaces:
Sleeping
Sleeping
| # ============================================================================ | |
| # IN-MEMORY CACHE & RATE LIMITING | |
| # ============================================================================ | |
| import asyncio | |
| import time | |
| from datetime import datetime, timedelta | |
| from collections import defaultdict | |
| from typing import Dict, Any, Optional, List | |
| from app.config import Config | |
| from app.utils.logger import logger | |
| class TokenCache: | |
| """Thread-safe token cache with TTL""" | |
| def __init__(self): | |
| self._cache: Dict[str, Dict[str, Any]] = {} | |
| self._lock = asyncio.Lock() | |
| async def get(self, email: str) -> Optional[str]: | |
| async with self._lock: | |
| cached = self._cache.get(email) | |
| if cached and cached['expires_at'] > datetime.now(): | |
| logger.debug(f"Token cache hit for {email}") | |
| return cached['token'] | |
| return None | |
| async def set(self, email: str, token: str, ttl: int = Config.TOKEN_CACHE_TTL): | |
| async with self._lock: | |
| self._cache[email] = { | |
| 'token': token, | |
| 'expires_at': datetime.now() + timedelta(seconds=ttl) | |
| } | |
| logger.debug(f"Token cached for {email}") | |
| async def invalidate(self, email: str): | |
| async with self._lock: | |
| if email in self._cache: | |
| del self._cache[email] | |
| logger.debug(f"Token invalidated for {email}") | |
| class RateLimiter: | |
| """Simple in-memory rate limiter""" | |
| def __init__(self, max_requests: int, window: int = 60): | |
| self.max_requests = max_requests | |
| self.window = window # seconds | |
| self.requests: Dict[str, List[float]] = defaultdict(list) | |
| self._lock = asyncio.Lock() | |
| async def check(self, identifier: str) -> bool: | |
| async with self._lock: | |
| now = time.time() | |
| # Clean old requests | |
| self.requests[identifier] = [ | |
| req_time for req_time in self.requests[identifier] | |
| if now - req_time < self.window | |
| ] | |
| if len(self.requests[identifier]) >= self.max_requests: | |
| logger.warning(f"Rate limit exceeded for {identifier}") | |
| return False | |
| self.requests[identifier].append(now) | |
| return True | |
| # Global instances | |
| token_cache = TokenCache() | |
| rate_limiter = RateLimiter(Config.MAX_REQUESTS_PER_MINUTE) | |
| class AttachmentCache: | |
| """Cache for project attachments with TTL""" | |
| def __init__(self, ttl: int = 300): # 5 minutes default | |
| self._cache: Dict[str, Dict[str, Any]] = {} | |
| self._lock = asyncio.Lock() | |
| self.ttl = ttl | |
| async def get_attachments(self, project_name: str) -> Optional[List[Dict]]: | |
| """Get cached attachments for a project""" | |
| async with self._lock: | |
| cached = self._cache.get(project_name) | |
| if cached and cached['expires_at'] > datetime.now(): | |
| logger.debug(f"Attachment cache hit for {project_name}") | |
| return cached['attachments'] | |
| return None | |
| async def set_attachments(self, project_name: str, attachments: List[Dict]): | |
| """Cache attachments for a project""" | |
| async with self._lock: | |
| self._cache[project_name] = { | |
| 'attachments': attachments, | |
| 'expires_at': datetime.now() + timedelta(seconds=self.ttl) | |
| } | |
| logger.debug(f"Cached {len(attachments)} attachments for {project_name}") | |
| async def get_mapping(self) -> Optional[Dict[str, int]]: | |
| """Get cached projectName->projectId mapping""" | |
| async with self._lock: | |
| cached = self._cache.get('_mapping') | |
| if cached and cached['expires_at'] > datetime.now(): | |
| return cached['mapping'] | |
| return None | |
| async def set_mapping(self, mapping: Dict[str, int]): | |
| """Cache projectName->projectId mapping""" | |
| async with self._lock: | |
| self._cache['_mapping'] = { | |
| 'mapping': mapping, | |
| 'expires_at': datetime.now() + timedelta(seconds=self.ttl) | |
| } | |
| logger.debug(f"Cached project mapping for {len(mapping)} projects") | |
| async def invalidate(self, project_name: str = None): | |
| """Invalidate cache""" | |
| async with self._lock: | |
| if project_name: | |
| self._cache.pop(project_name, None) | |
| else: | |
| self._cache.clear() | |
| # Global attachment cache | |
| attachment_cache = AttachmentCache() | |