Spaces:
Runtime error
Runtime error
File size: 10,819 Bytes
fdc5d7a |
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 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 |
"""Redis client for caching and task queue management."""
import json
from typing import Any, Dict, Optional, TypeVar
from datetime import datetime
import logging
from time import time as _time
import redis.asyncio as redis_async
import redis as redis_sync # Import synchronous Redis client
from pydantic import BaseModel
from tenacity import retry, stop_after_attempt, wait_exponential
from app.core.config import settings
# Type variable for cache
T = TypeVar('T')
# Configure logging
log = logging.getLogger(__name__)
# Redis connection pools for reusing connections
_redis_pool_async = None
_redis_pool_sync = None # Synchronous pool
# Default cache expiration (12 hours)
DEFAULT_CACHE_EXPIRY = 60 * 60 * 12
@retry(stop=stop_after_attempt(5), wait=wait_exponential(multiplier=1, min=1, max=10))
async def get_redis_pool() -> redis_async.Redis:
"""Get or create async Redis connection pool with retry logic."""
global _redis_pool_async
if _redis_pool_async is None:
# Get Redis configuration from settings
redis_url = settings.REDIS_URL or "redis://localhost:6379/0"
try:
# Create connection pool with reasonable defaults
_redis_pool_async = redis_async.ConnectionPool.from_url(
redis_url,
max_connections=10,
decode_responses=True,
health_check_interval=5,
socket_connect_timeout=5,
socket_keepalive=True,
retry_on_timeout=True
)
log.info(f"Created async Redis connection pool with URL: {redis_url}")
except Exception as e:
log.error(f"Error creating async Redis connection pool: {e}")
raise
return redis_async.Redis(connection_pool=_redis_pool_async)
def get_redis_pool_sync() -> redis_sync.Redis:
"""Get or create synchronous Redis connection pool."""
global _redis_pool_sync
if _redis_pool_sync is None:
# Get Redis configuration from settings
redis_url = settings.REDIS_URL or "redis://localhost:6379/0"
try:
# Create connection pool with reasonable defaults
_redis_pool_sync = redis_sync.ConnectionPool.from_url(
redis_url,
max_connections=10,
decode_responses=True,
socket_connect_timeout=5,
socket_keepalive=True,
retry_on_timeout=True
)
log.info(f"Created sync Redis connection pool with URL: {redis_url}")
except Exception as e:
log.error(f"Error creating sync Redis connection pool: {e}")
raise
return redis_sync.Redis(connection_pool=_redis_pool_sync)
@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=1, max=5))
async def get_redis() -> redis_async.Redis:
"""Get Redis client from pool with retry logic."""
try:
redis_client = await get_redis_pool()
return redis_client
except Exception as e:
log.error(f"Error getting Redis client: {e}")
raise
@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=1, max=5))
def get_redis_sync() -> redis_sync.Redis:
"""Get synchronous Redis client from pool with retry logic."""
try:
return get_redis_pool_sync()
except Exception as e:
log.error(f"Error getting synchronous Redis client: {e}")
raise
# Cache key generation
def generate_cache_key(prefix: str, *args: Any) -> str:
"""Generate cache key with prefix and args."""
key_parts = [prefix] + [str(arg) for arg in args if arg]
return ":".join(key_parts)
# JSON serialization helpers
def _json_serialize(obj: Any) -> str:
"""Serialize object to JSON with datetime support."""
def _serialize_datetime(o: Any) -> str:
if isinstance(o, datetime):
return o.isoformat()
if isinstance(o, BaseModel):
return o.dict()
return str(o)
return json.dumps(obj, default=_serialize_datetime)
def _json_deserialize(data: str, model_class: Optional[type] = None) -> Any:
"""Deserialize JSON string to object with datetime support."""
result = json.loads(data)
if model_class and issubclass(model_class, BaseModel):
return model_class.parse_obj(result)
return result
# Async cache operations
async def cache_set(key: str, value: Any, expire: int = DEFAULT_CACHE_EXPIRY) -> bool:
"""Set cache value with expiration (async version)."""
redis_client = await get_redis()
serialized = _json_serialize(value)
try:
await redis_client.set(key, serialized, ex=expire)
log.debug(f"Cached data at key: {key}, expires in {expire}s")
return True
except Exception as e:
log.error(f"Error caching data at key {key}: {e}")
return False
async def cache_get(key: str, model_class: Optional[type] = None) -> Optional[Any]:
"""Get cache value with optional model deserialization (async version)."""
redis_client = await get_redis()
try:
data = await redis_client.get(key)
if not data:
return None
log.debug(f"Cache hit for key: {key}")
return _json_deserialize(data, model_class)
except Exception as e:
log.error(f"Error retrieving cache for key {key}: {e}")
return None
# Synchronous cache operations for Celery tasks
def sync_cache_set(key: str, value: Any, expire: int = DEFAULT_CACHE_EXPIRY) -> bool:
"""Set cache value with expiration (synchronous version for Celery tasks). Logs slow operations."""
redis_client = get_redis_sync()
serialized = _json_serialize(value)
start = _time()
try:
redis_client.set(key, serialized, ex=expire)
elapsed = _time() - start
if elapsed > 2:
log.warning(f"Slow sync_cache_set for key {key}: {elapsed:.2f}s")
log.debug(f"Cached data at key: {key}, expires in {expire}s (sync)")
return True
except Exception as e:
log.error(f"Error caching data at key {key}: {e}")
return False
def sync_cache_get(key: str, model_class: Optional[type] = None) -> Optional[Any]:
"""Get cache value with optional model deserialization (synchronous version for Celery tasks). Logs slow operations."""
redis_client = get_redis_sync()
start = _time()
try:
data = redis_client.get(key)
elapsed = _time() - start
if elapsed > 2:
log.warning(f"Slow sync_cache_get for key {key}: {elapsed:.2f}s")
if not data:
return None
log.debug(f"Cache hit for key: {key} (sync)")
return _json_deserialize(data, model_class)
except Exception as e:
log.error(f"Error retrieving cache for key {key}: {e}")
return None
async def cache_invalidate(key: str) -> bool:
"""Invalidate cache for key."""
redis_client = await get_redis()
try:
await redis_client.delete(key)
log.debug(f"Invalidated cache for key: {key}")
return True
except Exception as e:
log.error(f"Error invalidating cache for key {key}: {e}")
return False
async def cache_invalidate_pattern(pattern: str) -> int:
"""Invalidate all cache keys matching pattern."""
redis_client = await get_redis()
try:
keys = await redis_client.keys(pattern)
if not keys:
return 0
count = await redis_client.delete(*keys)
log.debug(f"Invalidated {count} keys matching pattern: {pattern}")
return count
except Exception as e:
log.error(f"Error invalidating keys with pattern {pattern}: {e}")
return 0
# Task queue operations
async def enqueue_task(queue_name: str, task_id: str, payload: Dict[str, Any]) -> bool:
"""Add task to queue."""
redis_client = await get_redis()
try:
serialized = _json_serialize(payload)
await redis_client.lpush(f"queue:{queue_name}", serialized)
await redis_client.hset(f"tasks:{queue_name}", task_id, "pending")
log.info(f"Enqueued task {task_id} to queue {queue_name}")
return True
except Exception as e:
log.error(f"Error enqueueing task {task_id} to {queue_name}: {e}")
return False
async def mark_task_complete(queue_name: str, task_id: str, result: Optional[Dict[str, Any]] = None) -> bool:
"""Mark task as complete with optional result."""
redis_client = await get_redis()
try:
# Store result if provided
if result:
await redis_client.hset(
f"results:{queue_name}",
task_id,
_json_serialize(result)
)
# Mark task as complete
await redis_client.hset(f"tasks:{queue_name}", task_id, "complete")
await redis_client.expire(f"tasks:{queue_name}", 86400) # Expire after 24 hours
log.info(f"Marked task {task_id} as complete in queue {queue_name}")
return True
except Exception as e:
log.error(f"Error marking task {task_id} as complete: {e}")
return False
async def get_task_status(queue_name: str, task_id: str) -> Optional[str]:
"""Get status of a task."""
redis_client = await get_redis()
try:
status = await redis_client.hget(f"tasks:{queue_name}", task_id)
return status
except Exception as e:
log.error(f"Error getting status for task {task_id}: {e}")
return None
async def get_task_result(queue_name: str, task_id: str) -> Optional[Dict[str, Any]]:
"""Get result of a completed task."""
redis_client = await get_redis()
try:
data = await redis_client.hget(f"results:{queue_name}", task_id)
if not data:
return None
return _json_deserialize(data)
except Exception as e:
log.error(f"Error getting result for task {task_id}: {e}")
return None
# Stream processing for real-time updates
async def add_to_stream(stream: str, data: Dict[str, Any], max_len: int = 1000) -> str:
"""Add event to Redis stream."""
redis_client = await get_redis()
try:
# Convert dict values to strings (Redis streams requirement)
entry = {k: _json_serialize(v) for k, v in data.items()}
# Add to stream with automatic ID generation
event_id = await redis_client.xadd(
stream,
entry,
maxlen=max_len,
approximate=True
)
log.debug(f"Added event {event_id} to stream {stream}")
return event_id
except Exception as e:
log.error(f"Error adding to stream {stream}: {e}")
raise |