Files
Agentswarm/orchestrator/redis_client.py
T
2026-06-08 17:32:34 +08:00

144 lines
5.1 KiB
Python

"""Redis client for orchestrator state management.
Redis is the canonical store. An in-memory fallback (backed by the in-process
``fakeredis`` emulator, which faithfully implements the Redis list/hash API) is
available ONLY for local development and CI, and ONLY when explicitly enabled via
``REDIS_FAKE=1`` or ``ALLOW_MEMORY_STORE=1``. In production neither flag is set, so a
Redis outage fails fast on startup instead of silently dropping durability/Manager
state.
"""
import os
from typing import Optional
import logging
try: # redis-py is required in production; guarded so dev/CI can run on fakeredis only.
import redis.asyncio as redis
except Exception: # pragma: no cover - exercised only when redis-py is absent
redis = None
logger = logging.getLogger(__name__)
def _truthy(value: Optional[str]) -> bool:
return (value or "").strip().lower() in {"1", "true", "yes", "on"}
class RedisClient:
"""Async Redis client wrapper for orchestrator operations."""
def __init__(self):
self.client = None
self.host = os.getenv("REDIS_HOST", "redis-service")
self.port = int(os.getenv("REDIS_PORT", "6379"))
self.db = int(os.getenv("REDIS_DB", "0"))
def _fallback_allowed(self) -> bool:
return _truthy(os.getenv("REDIS_FAKE")) or _truthy(os.getenv("ALLOW_MEMORY_STORE"))
def _make_fake_client(self):
"""Return an in-process fakeredis client (dev/CI fallback only)."""
import fakeredis.aioredis as fakeredis # imported lazily; dev/CI dependency
return fakeredis.FakeRedis(decode_responses=True)
async def connect(self):
"""Establish Redis connection, or a gated in-memory fallback for dev/CI."""
# Explicit fake mode (used by local runs and tests) short-circuits real Redis.
if _truthy(os.getenv("REDIS_FAKE")):
self.client = self._make_fake_client()
logger.warning("REDIS_FAKE enabled; using in-memory fakeredis (NOT for production)")
return
try:
if redis is None:
raise RuntimeError("redis-py is not installed")
self.client = redis.Redis(
host=self.host,
port=self.port,
db=self.db,
decode_responses=True,
socket_connect_timeout=5,
socket_keepalive=True,
)
await self.client.ping()
logger.info(f"Connected to Redis at {self.host}:{self.port}")
except Exception as e:
if self._fallback_allowed():
logger.warning(
f"Redis unavailable ({e}); ALLOW_MEMORY_STORE set, using in-memory fakeredis "
"fallback (NOT for production)"
)
self.client = self._make_fake_client()
return
logger.error(f"Failed to connect to Redis: {e}")
raise
async def disconnect(self):
"""Close Redis connection."""
if self.client:
await self.client.close()
logger.info("Disconnected from Redis")
async def set(self, key: str, value: str, ex: Optional[int] = None):
"""Set key-value pair with optional expiration."""
await self.client.set(key, value, ex=ex)
async def get(self, key: str) -> Optional[str]:
"""Get value by key."""
return await self.client.get(key)
async def delete(self, key: str):
"""Delete key."""
await self.client.delete(key)
async def exists(self, key: str) -> bool:
"""Check if key exists."""
return await self.client.exists(key) > 0
async def hset(self, name: str, key: str, value: str):
"""Set hash field."""
await self.client.hset(name, key, value)
async def hget(self, name: str, key: str) -> Optional[str]:
"""Get hash field."""
return await self.client.hget(name, key)
async def hgetall(self, name: str) -> dict:
"""Get all hash fields."""
return await self.client.hgetall(name)
async def hdel(self, name: str, *keys: str):
"""Delete hash fields."""
await self.client.hdel(name, *keys)
async def keys(self, pattern: str) -> list:
"""Get keys matching pattern."""
return await self.client.keys(pattern)
async def lpush(self, key: str, *values: str):
"""Push values to list head."""
await self.client.lpush(key, *values)
async def rpush(self, key: str, *values: str):
"""Push values to list tail."""
await self.client.rpush(key, *values)
async def rpop(self, key: str) -> Optional[str]:
"""Pop value from list tail."""
return await self.client.rpop(key)
async def llen(self, key: str) -> int:
"""Get list length."""
return await self.client.llen(key)
async def lrange(self, key: str, start: int, end: int) -> list:
"""Get list range."""
return await self.client.lrange(key, start, end)
async def lrem(self, key: str, count: int, value: str) -> int:
"""Remove values from list."""
return await self.client.lrem(key, count, value)
# Global Redis client instance
redis_client = RedisClient()