"""Definition of RedisAdapter class.""" from __future__ import annotations from typing import Self, cast import os import structlog import time from python_utils import check_env from python_repositories.interfaces import ( ConnectionAwareInterface, ContextAwareInterface, JsonRepositoryInterface, ) try: import redis from redis.commands.json.path import Path as RedisPath except ImportError as exc: raise ImportError( "Redis support requires the redis extra. " "Install with: pip install python-repositories[redis]" ) from exc class RedisAdapter( JsonRepositoryInterface, ContextAwareInterface, ConnectionAwareInterface, ): """Redis adapter exposing basic CRUD functionality.""" uri_env_var_name: str = "REDIS_URI" path: str = "." # JSON root path, updated in __init__ encoding: str = "UTF-8" health_check_ttl_seconds: float = 1.0 def __init__(self) -> None: # Setup logger self.logger = structlog.get_logger( self.__class__.__name__, ) # Check environment variables check_env(self.uri_env_var_name) # Prepare internal variables self._client: redis.Redis | None = None self.path: str = RedisPath.root_path() self._health_check_at: float | None = None self._health_check_ok: bool = False def __enter__(self) -> Self: """Enter the context.""" self.connect() return self def __exit__( self, exc_type: type | None, exc_val: object | None, exc_tb: object | None ) -> None: """Exit the context.""" ctx_info = {"exc_type": exc_type, "exc_val": exc_val, "exc_tb": exc_tb} if any( ( exc_type is not None, exc_val is not None, exc_tb is not None, ), ): self.logger.error("Error while exiting context", **ctx_info) self.disconnect() def connect(self) -> None: """Connect to the Redis server.""" if self._client is not None: self._client.close() self._client = None self._invalidate_health_cache() # Prepare arguments uri = str(os.getenv(self.uri_env_var_name)) # Connect client try: client = redis.Redis.from_url( url=uri, socket_connect_timeout=10, ) if not client.ping(): raise ConnectionError(f"Could not connect to Redis at {uri}") except (redis.ConnectionError, redis.TimeoutError) as exc: raise ConnectionError(f"Could not connect to Redis at {uri}") from exc # Persist client self._client = client self._invalidate_health_cache() def disconnect(self) -> None: """Disconnect from the Redis server.""" # Close connection if self._client is not None: self._client.close() # Reset client self._client = None self._invalidate_health_cache() def _invalidate_health_cache(self) -> None: self._health_check_at = None self._health_check_ok = False def _probe_connection(self) -> bool: assert self._client is not None try: return bool(self._client.ping()) except (redis.ConnectionError, redis.TimeoutError): return False def is_connected(self) -> bool: """Check if connected to the Redis server.""" if self._client is None: return False now = time.monotonic() if self._health_check_at is not None: seconds_since_last_health_check = now - self._health_check_at cache_is_fresh = ( seconds_since_last_health_check < self.health_check_ttl_seconds ) if cache_is_fresh: return self._health_check_ok result = self._probe_connection() self._health_check_at = now self._health_check_ok = result self.logger.debug("Connection status", connected=result) return result def set(self, key: str, data: dict) -> None: """Set a JSON object in Redis.""" # Check input if not isinstance(key, str) or len(key) == 0: raise ValueError("Key must be a non-empty string") if not isinstance(data, dict) or len(data) == 0: raise ValueError("Data must be a non-empty dictionary") # Check connection if not self.is_connected(): raise ConnectionError("Not connected to Redis") assert self._client is not None # Set data self._client.json().set(key, self.path, data) self.logger.debug(f"Set {key} to {data}") def get(self, key: str) -> dict | None: """Get a JSON object from Redis.""" # Check input if not isinstance(key, str) or len(key) == 0: raise ValueError("Key must be a non-empty string") # Check connection if not self.is_connected(): raise ConnectionError("Not connected to Redis") assert self._client is not None # Get data data = cast( dict | None, self._client.json().get(key), ) self.logger.debug(f"Got {data} from {key}") return data def delete(self, key: str) -> None: """Delete data from Redis.""" # Check input if not isinstance(key, str) or len(key) == 0: raise ValueError("Key must be a non-empty string") # Check connection if not self.is_connected(): raise ConnectionError("Not connected to Redis") assert self._client is not None # Delete data self._client.json().delete(key) self.logger.debug(f"Deleted {key}") def list_keys(self, pattern: str) -> list[str]: """List keys in Redis matching a pattern.""" # Check input if not isinstance(pattern, str) or len(pattern) == 0: raise ValueError("Pattern must be a non-empty string") # Check connection if not self.is_connected(): raise ConnectionError("Not connected to Redis") assert self._client is not None # List keys keys_raw = cast( list[bytes], self._client.keys(pattern), ) keys: list[str] = [key.decode(self.encoding) for key in keys_raw] self.logger.debug(f"Got {keys} matching {pattern}") return keys