Spaces:
Sleeping
Sleeping
File size: 4,597 Bytes
1def50b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | # ============================================================================
# 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()
|