python-repositories/python_repositories/adapters/redis_adapter.py

"""Definition of RedisAdapter class."""

from __future__ import annotations
from typing import cast

from python_repositories.adapters.connection_aware_adapter import (
    ConnectionAwareAdapter,
)
from python_repositories.config import RedisConfig
from python_repositories.interfaces import 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, ConnectionAwareAdapter):
    """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"
    connection_name: str = "Redis"

    def __init__(
        self,
        *,
        config: RedisConfig | None = None,
        client: redis.Redis | None = None,
    ) -> None:
        super().__init__()
        if client is not None and config is None:
            raise ValueError("config is required when client is provided")
        if config is None and client is None:
            config = RedisConfig.from_env(self.uri_env_var_name)
        assert config is not None
        self._config = config
        self._client_injected = client is not None
        self._client: redis.Redis | None = client
        self.path: str = RedisPath.root_path()
        if self._client_injected:
            self._invalidate_health_cache()

    def _is_client_ready(self) -> bool:
        return self._client is not None

    def connect(self) -> None:
        """Connect to the Redis server."""
        if self._client_injected:
            if self._client is not None:
                try:
                    if not self._client.ping():
                        raise ConnectionError(
                            f"Could not connect to Redis at {self._config.uri}"
                        )
                except (redis.ConnectionError, redis.TimeoutError) as exc:
                    raise ConnectionError(
                        f"Could not connect to Redis at {self._config.uri}"
                    ) from exc
            self._invalidate_health_cache()
            return
        if self._client is not None:
            self._client.close()
            self._client = None
            self._invalidate_health_cache()
        uri = self._config.uri
        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
        self._client = client
        self._invalidate_health_cache()

    def disconnect(self) -> None:
        """Disconnect from the Redis server."""
        if self._client is not None and not self._client_injected:
            self._client.close()
        self._client = None
        self._invalidate_health_cache()

    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 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
        self._require_connected()
        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
        self._require_connected()
        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
        self._require_connected()
        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
        self._require_connected()
        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