python-repositories/python_repositories/adapters/redis_adapter.py

"""Definition of RedisAdapter class."""

from __future__ import annotations
from collections.abc import Iterator
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:
            config = RedisConfig.from_env(self.uri_env_var_name)
        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 _validate_injected_client(self) -> None:
        if self._client is None:
            return
        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

    def _establish_connection(self) -> None:
        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

    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 _validate_pattern(self, pattern: str) -> None:
        if not isinstance(pattern, str) or len(pattern) == 0:
            raise ValueError("Pattern must be a non-empty string")

    def list_keys(self, pattern: str) -> list[str]:
        """List keys in Redis using KEYS; may block on large datasets."""
        self._validate_pattern(pattern)
        # 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

    def scan_keys(
        self,
        pattern: str,
        *,
        count: int | None = None,
    ) -> Iterator[str]:
        """Yield keys in Redis using SCAN to avoid blocking large datasets."""
        self._validate_pattern(pattern)
        self._require_connected()
        assert self._client is not None
        client = self._client

        def _decode(key: bytes | str) -> str:
            return key if isinstance(key, str) else key.decode(self.encoding)

        def _iter() -> Iterator[str]:
            scan_iter = (
                client.scan_iter(match=pattern, count=count)
                if count is not None
                else client.scan_iter(match=pattern)
            )
            for key_raw in scan_iter:
                yield _decode(key_raw)

        return _iter()