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

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"

    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()

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

    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

    @property
    def is_connected(self) -> bool:
        """Check if connected to Redis server."""
        res = self._client is not None
        self.logger.debug(res)
        return res

    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 self._client is None or not self.is_connected:
            raise ConnectionError("Not connected to Redis")
        # 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 self._client is None or not self.is_connected:
            raise ConnectionError("Not connected to Redis")
        # 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 self._client is None or not self.is_connected:
            raise ConnectionError("Not connected to Redis")
        # 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 self._client is None or not self.is_connected:
            raise ConnectionError("Not connected to Redis")
        # 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