ai_v/venv/Lib/site-packages/redis/multidb/healthcheck.py

import logging
from abc import ABC, abstractmethod
from enum import Enum
from time import sleep
from typing import List, Optional, Tuple, Union

from redis import Redis
from redis.backoff import NoBackoff
from redis.http.http_client import DEFAULT_TIMEOUT, HttpClient
from redis.multidb.exception import UnhealthyDatabaseException
from redis.retry import Retry

DEFAULT_HEALTH_CHECK_PROBES = 3
DEFAULT_HEALTH_CHECK_INTERVAL = 5
DEFAULT_HEALTH_CHECK_DELAY = 0.5
DEFAULT_LAG_AWARE_TOLERANCE = 5000

logger = logging.getLogger(__name__)


class HealthCheck(ABC):
    @abstractmethod
    def check_health(self, database) -> bool:
        """Function to determine the health status."""
        pass


class HealthCheckPolicy(ABC):
    """
    Health checks execution policy.
    """

    @property
    @abstractmethod
    def health_check_probes(self) -> int:
        """Number of probes to execute health checks."""
        pass

    @property
    @abstractmethod
    def health_check_delay(self) -> float:
        """Delay between health check probes."""
        pass

    @abstractmethod
    def execute(self, health_checks: List[HealthCheck], database) -> bool:
        """Execute health checks and return database health status."""
        pass


class AbstractHealthCheckPolicy(HealthCheckPolicy):
    def __init__(self, health_check_probes: int, health_check_delay: float):
        if health_check_probes < 1:
            raise ValueError("health_check_probes must be greater than 0")
        self._health_check_probes = health_check_probes
        self._health_check_delay = health_check_delay

    @property
    def health_check_probes(self) -> int:
        return self._health_check_probes

    @property
    def health_check_delay(self) -> float:
        return self._health_check_delay

    @abstractmethod
    def execute(self, health_checks: List[HealthCheck], database) -> bool:
        pass


class HealthyAllPolicy(AbstractHealthCheckPolicy):
    """
    Policy that returns True if all health check probes are successful.
    """

    def __init__(self, health_check_probes: int, health_check_delay: float):
        super().__init__(health_check_probes, health_check_delay)

    def execute(self, health_checks: List[HealthCheck], database) -> bool:
        for health_check in health_checks:
            for attempt in range(self.health_check_probes):
                try:
                    if not health_check.check_health(database):
                        return False
                except Exception as e:
                    raise UnhealthyDatabaseException("Unhealthy database", database, e)

                if attempt < self.health_check_probes - 1:
                    sleep(self._health_check_delay)
        return True


class HealthyMajorityPolicy(AbstractHealthCheckPolicy):
    """
    Policy that returns True if a majority of health check probes are successful.
    """

    def __init__(self, health_check_probes: int, health_check_delay: float):
        super().__init__(health_check_probes, health_check_delay)

    def execute(self, health_checks: List[HealthCheck], database) -> bool:
        for health_check in health_checks:
            if self.health_check_probes % 2 == 0:
                allowed_unsuccessful_probes = self.health_check_probes / 2
            else:
                allowed_unsuccessful_probes = (self.health_check_probes + 1) / 2

            for attempt in range(self.health_check_probes):
                try:
                    if not health_check.check_health(database):
                        allowed_unsuccessful_probes -= 1
                        if allowed_unsuccessful_probes <= 0:
                            return False
                except Exception as e:
                    allowed_unsuccessful_probes -= 1
                    if allowed_unsuccessful_probes <= 0:
                        raise UnhealthyDatabaseException(
                            "Unhealthy database", database, e
                        )

                if attempt < self.health_check_probes - 1:
                    sleep(self._health_check_delay)
        return True


class HealthyAnyPolicy(AbstractHealthCheckPolicy):
    """
    Policy that returns True if at least one health check probe is successful.
    """

    def __init__(self, health_check_probes: int, health_check_delay: float):
        super().__init__(health_check_probes, health_check_delay)

    def execute(self, health_checks: List[HealthCheck], database) -> bool:
        is_healthy = False

        for health_check in health_checks:
            exception = None

            for attempt in range(self.health_check_probes):
                try:
                    if health_check.check_health(database):
                        is_healthy = True
                        break
                    else:
                        is_healthy = False
                except Exception as e:
                    exception = UnhealthyDatabaseException(
                        "Unhealthy database", database, e
                    )

                if attempt < self.health_check_probes - 1:
                    sleep(self._health_check_delay)

            if not is_healthy and not exception:
                return is_healthy
            elif not is_healthy and exception:
                raise exception

        return is_healthy


class HealthCheckPolicies(Enum):
    HEALTHY_ALL = HealthyAllPolicy
    HEALTHY_MAJORITY = HealthyMajorityPolicy
    HEALTHY_ANY = HealthyAnyPolicy


DEFAULT_HEALTH_CHECK_POLICY: HealthCheckPolicies = HealthCheckPolicies.HEALTHY_ALL


class PingHealthCheck(HealthCheck):
    """
    Health check based on PING command.
    """

    def check_health(self, database) -> bool:
        if isinstance(database.client, Redis):
            return database.client.execute_command("PING")
        else:
            # For a cluster checks if all nodes are healthy.
            all_nodes = database.client.get_nodes()
            for node in all_nodes:
                if not node.redis_connection.execute_command("PING"):
                    return False

            return True


class LagAwareHealthCheck(HealthCheck):
    """
    Health check available for Redis Enterprise deployments.
    Verify via REST API that the database is healthy based on different lags.
    """

    def __init__(
        self,
        rest_api_port: int = 9443,
        lag_aware_tolerance: int = DEFAULT_LAG_AWARE_TOLERANCE,
        timeout: float = DEFAULT_TIMEOUT,
        auth_basic: Optional[Tuple[str, str]] = None,
        verify_tls: bool = True,
        # TLS verification (server) options
        ca_file: Optional[str] = None,
        ca_path: Optional[str] = None,
        ca_data: Optional[Union[str, bytes]] = None,
        # Mutual TLS (client cert) options
        client_cert_file: Optional[str] = None,
        client_key_file: Optional[str] = None,
        client_key_password: Optional[str] = None,
    ):
        """
        Initialize LagAwareHealthCheck with the specified parameters.

        Args:
            rest_api_port: Port number for Redis Enterprise REST API (default: 9443)
            lag_aware_tolerance: Tolerance in lag between databases in MS (default: 100)
            timeout: Request timeout in seconds (default: DEFAULT_TIMEOUT)
            auth_basic: Tuple of (username, password) for basic authentication
            verify_tls: Whether to verify TLS certificates (default: True)
            ca_file: Path to CA certificate file for TLS verification
            ca_path: Path to CA certificates directory for TLS verification
            ca_data: CA certificate data as string or bytes
            client_cert_file: Path to client certificate file for mutual TLS
            client_key_file: Path to client private key file for mutual TLS
            client_key_password: Password for encrypted client private key
        """
        self._http_client = HttpClient(
            timeout=timeout,
            auth_basic=auth_basic,
            retry=Retry(NoBackoff(), retries=0),
            verify_tls=verify_tls,
            ca_file=ca_file,
            ca_path=ca_path,
            ca_data=ca_data,
            client_cert_file=client_cert_file,
            client_key_file=client_key_file,
            client_key_password=client_key_password,
        )
        self._rest_api_port = rest_api_port
        self._lag_aware_tolerance = lag_aware_tolerance

    def check_health(self, database) -> bool:
        if database.health_check_url is None:
            raise ValueError(
                "Database health check url is not set. Please check DatabaseConfig for the current database."
            )

        if isinstance(database.client, Redis):
            db_host = database.client.get_connection_kwargs()["host"]
        else:
            db_host = database.client.startup_nodes[0].host

        base_url = f"{database.health_check_url}:{self._rest_api_port}"
        self._http_client.base_url = base_url

        # Find bdb matching to the current database host
        matching_bdb = None
        for bdb in self._http_client.get("/v1/bdbs"):
            for endpoint in bdb["endpoints"]:
                if endpoint["dns_name"] == db_host:
                    matching_bdb = bdb
                    break

                # In case if the host was set as public IP
                for addr in endpoint["addr"]:
                    if addr == db_host:
                        matching_bdb = bdb
                        break

        if matching_bdb is None:
            logger.warning("LagAwareHealthCheck failed: Couldn't find a matching bdb")
            raise ValueError("Could not find a matching bdb")

        url = (
            f"/v1/bdbs/{matching_bdb['uid']}/availability"
            f"?extend_check=lag&availability_lag_tolerance_ms={self._lag_aware_tolerance}"
        )
        self._http_client.get(url, expect_json=False)

        # Status checked in an http client, otherwise HttpError will be raised
        return True