python-repositories/python_repositories/adapters/redis_adapter.py

"""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