grpc-server-kit 0.1.0__py3-none-any.whl

grpc_server_kit/aio/health/checkers/utils.py

@@ -0,0 +1,97 @@
+"""Internal utilities for health checkers."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from collections.abc import Awaitable, Callable
+from functools import wraps
+from typing import ParamSpec
+
+from grpc_server_kit.constants import DEFAULT_CHECKER_TIMEOUT
+
+logger = logging.getLogger(__name__)
+
+P = ParamSpec("P")
+
+
+def handle_check_exceptions(
+    check_name: str,
+) -> Callable[[Callable[P, Awaitable[bool]]], Callable[P, Awaitable[bool]]]:
+    """Decorator for consistent exception handling in health check functions.
+
+    This decorator provides uniform error handling for all health check implementations:
+    - Logs timeout errors with warning level
+    - Propagates CancelledError and KeyboardInterrupt for proper shutdown
+    - Catches and logs all other exceptions, returning False (unhealthy)
+
+    Args:
+        check_name: Human-readable name of the check for logging (e.g. "Database", "Redis").
+
+    Returns:
+        Decorator function that wraps health check implementations.
+    """
+
+    def decorator(func: Callable[P, Awaitable[bool]]) -> Callable[P, Awaitable[bool]]:
+        @wraps(func)
+        async def wrapper(*args: P.args, **kwargs: P.kwargs) -> bool:
+            try:
+                return await func(*args, **kwargs)
+            except TimeoutError:
+                timeout = kwargs.get("timeout", "unknown")
+                logger.warning(f"{check_name} health check timed out", extra={"check": check_name, "timeout": timeout})
+                return False
+            except (asyncio.CancelledError, KeyboardInterrupt, SystemExit):
+                raise
+            except (ConnectionError, OSError) as e:
+                # Common networking errors for health checks (e.g. Redis down, Postgres unreachable)
+                logger.warning(
+                    f"{check_name} health check failed (network/OS error)",
+                    extra={"check": check_name, "error": str(e), "error_type": type(e).__name__},
+                )
+                return False
+            except Exception:
+                logger.exception(f"{check_name} health check failed with unexpected error", extra={"check": check_name})
+                return False
+
+        return wrapper
+
+    return decorator
+
+
+class FunctionalHealthChecker:
+    """Generic health checker that wraps a health check function.
+
+    This class provides a common wrapper for any health check function that follows
+    the `check_func(resource, timeout=...)` pattern.
+    """
+
+    def __init__(
+        self,
+        check_func: Callable[..., Awaitable[bool]],
+        resource: object | None = None,
+        timeout: float = DEFAULT_CHECKER_TIMEOUT,
+    ) -> None:
+        """Initialize health checker.
+
+        Args:
+            check_func: Asynchronous function that performs the health check.
+            resource: The resource (e.g. session maker, client) to check.
+            timeout: Maximum time to wait for the check to complete in seconds.
+
+        Raises:
+            ValueError: If timeout is not positive.
+        """
+        if timeout <= 0:
+            raise ValueError(f"timeout must be positive, got {timeout}")
+        self._check_func = check_func
+        self._resource = resource
+        self._timeout = timeout
+
+    async def check(self) -> bool:
+        """Execute the wrapped health check function.
+
+        Returns:
+            True if healthy, False otherwise.
+        """
+        return await self._check_func(self._resource, timeout=self._timeout)

grpc_server_kit/aio/health/orchestrator.py

@@ -0,0 +1,122 @@
+"""Overall health-check orchestration (async)."""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import logging
+from typing import TYPE_CHECKING
+
+from grpc_health.v1 import health_pb2
+
+from grpc_server_kit.constants import DEFAULT_HEALTH_CHECK_TIMEOUT
+
+from .protocols import AsyncHealthCacheProtocol, AsyncHealthChecker
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+logger = logging.getLogger(__name__)
+
+
+async def check_async_overall_health(
+    checkers: list[AsyncHealthChecker] | None = None,
+    cache: AsyncHealthCacheProtocol | None = None,
+    timeout: float = DEFAULT_HEALTH_CHECK_TIMEOUT,
+) -> health_pb2.HealthCheckResponse.ServingStatus:
+    """Check overall service health including all dependencies (async)."""
+    if not checkers:
+        return health_pb2.HealthCheckResponse.SERVING
+
+    if cache is not None:
+        return await cache.get_or_set(
+            check_func=functools.partial(_perform_checks, checkers, timeout=timeout),
+        )
+
+    return await _perform_checks(checkers, timeout=timeout)
+
+
+async def _perform_checks(
+    checkers: Sequence[AsyncHealthChecker], timeout: float = DEFAULT_HEALTH_CHECK_TIMEOUT
+) -> health_pb2.HealthCheckResponse.ServingStatus:
+    """Perform all health checks in parallel and aggregate results.
+
+    Uses asyncio.gather with return_exceptions to collect pass/fail per checker
+    without using exceptions for control flow.
+
+    Args:
+        checkers: Sequence of health checkers to execute. Must not be empty.
+        timeout: Maximum seconds to wait for all health checks to complete.
+
+    Returns:
+        SERVING if all checks pass, NOT_SERVING if any check fails or unexpected error occurs.
+    """
+    loop = asyncio.get_running_loop()
+    start_time = loop.time()
+    checker_count = len(checkers)
+    checker_names = [type(c).__name__ for c in checkers]
+
+    try:
+        async with asyncio.timeout(timeout):
+            results = await asyncio.gather(
+                *(c.check() for c in checkers),
+                return_exceptions=True,
+            )
+    except (asyncio.CancelledError, KeyboardInterrupt):
+        raise
+    except TimeoutError:
+        elapsed = loop.time() - start_time
+        logger.warning(
+            "Health checks timed out",
+            extra={
+                "checker_count": checker_count,
+                "checkers": checker_names,
+                "timeout": timeout,
+                "elapsed_seconds": round(elapsed, 3),
+            },
+        )
+        return health_pb2.HealthCheckResponse.NOT_SERVING
+
+    failed_checkers: list[str] = []
+    unexpected_errors: list[tuple[str, Exception]] = []
+
+    for checker, outcome in zip(checkers, results, strict=True):
+        name = type(checker).__name__
+        if isinstance(outcome, BaseException) and not isinstance(outcome, Exception):
+            raise outcome
+        if outcome is False:
+            failed_checkers.append(name)
+        elif isinstance(outcome, Exception):
+            unexpected_errors.append((name, outcome))
+
+    elapsed = loop.time() - start_time
+    extra = {
+        "checker_count": checker_count,
+        "elapsed_seconds": round(elapsed, 3),
+    }
+
+    if failed_checkers:
+        logger.warning(
+            "Health checks failed",
+            extra={**extra, "failed_count": len(failed_checkers), "failed_checkers": failed_checkers},
+        )
+        return health_pb2.HealthCheckResponse.NOT_SERVING
+
+    if unexpected_errors:
+        # Log each error WITH its own traceback: logger.exception outside an
+        # except block has no active exception and would log "NoneType: None".
+        for name, error in unexpected_errors:
+            logger.error(
+                "Health checker raised unexpected error",
+                extra={**extra, "checker": name, "error_type": type(error).__name__},
+                exc_info=error,
+            )
+        return health_pb2.HealthCheckResponse.NOT_SERVING
+
+    # DEBUG, not INFO: this fires for every Check probe / Watch iteration and
+    # carries no information when everything is healthy.
+    logger.debug(
+        "Health checks passed",
+        extra={**extra, "checkers": checker_names},
+    )
+    return health_pb2.HealthCheckResponse.SERVING

grpc_server_kit/aio/health/protocols.py

@@ -0,0 +1,43 @@
+"""Protocols for health checks."""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from typing import TYPE_CHECKING, Protocol
+
+if TYPE_CHECKING:
+    from grpc_health.v1 import health_pb2
+
+
+class AsyncHealthChecker(Protocol):
+    """Protocol for component health checks (DB, Redis, external service, ...).
+
+    One shape (``async check() -> bool``) serves the kit's gRPC health servicer
+    and any external health registry alike.
+    """
+
+    async def check(self) -> bool:
+        """Check component health. Returns True if healthy."""
+        ...
+
+
+class AsyncHealthCacheProtocol(Protocol):
+    """Protocol for a health status cache with TTL support.
+
+    Implementations must be safe for concurrent coroutine access and must
+    ensure that only one concurrent task performs the underlying check on a
+    cache miss (double-checked locking or equivalent).
+
+    Note: Cache instances should not be shared across different event loops.
+    """
+
+    async def get_or_set(
+        self,
+        check_func: Callable[[], Awaitable[health_pb2.HealthCheckResponse.ServingStatus]],
+        now: float | None = None,
+    ) -> health_pb2.HealthCheckResponse.ServingStatus:
+        """Return the cached health status or perform the check and cache the result."""
+        ...
+
+
+__all__ = ["AsyncHealthCacheProtocol", "AsyncHealthChecker"]

grpc_server_kit/aio/health/servicer.py

@@ -0,0 +1,229 @@
+"""Dynamic health check servicer with real-time dependency checks (async)."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import TYPE_CHECKING, NoReturn
+
+import grpc
+from grpc_health.v1 import health_pb2, health_pb2_grpc
+
+from grpc_server_kit.constants import (
+    DEFAULT_HEALTH_CHECK_INTERVAL,
+    DEFAULT_HEALTH_CHECK_TIMEOUT,
+    DEFAULT_HEALTH_HEARTBEAT_INTERVAL,
+)
+from grpc_server_kit.server import is_valid_service_name
+
+from .orchestrator import check_async_overall_health
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator
+
+    from .protocols import AsyncHealthCacheProtocol, AsyncHealthChecker
+
+logger = logging.getLogger(__name__)
+
+
+class AsyncDynamicHealthServicer(health_pb2_grpc.HealthServicer):
+    """Async health servicer with dynamic dependency checks.
+
+    Implements the standard gRPC Health Checking Protocol v1 with real-time
+    dependency verification. Supports both Check (unary) and Watch (streaming) modes.
+    """
+
+    def __init__(
+        self,
+        checkers: list[AsyncHealthChecker] | None = None,
+        cache: AsyncHealthCacheProtocol | None = None,
+        service_names: list[str] | None = None,
+        check_interval: float = DEFAULT_HEALTH_CHECK_INTERVAL,
+        heartbeat_interval: float = DEFAULT_HEALTH_HEARTBEAT_INTERVAL,
+        check_timeout: float = DEFAULT_HEALTH_CHECK_TIMEOUT,
+    ) -> None:
+        """Initialize health servicer.
+
+        Args:
+            checkers: Optional list of asynchronous health checkers. If None or empty,
+                the service always returns SERVING (no dependencies to check).
+            cache: Optional health status cache following AsyncHealthCacheProtocol.
+                Strongly recommended with multiple Watch clients: watchers share
+                one real dependency check per TTL instead of each running their own.
+            service_names: Optional list of additional service names to support.
+                The empty string "" is always supported and represents overall health.
+            check_interval: Frequency of dependency checks in Watch mode (seconds). Must be positive.
+            heartbeat_interval: Interval to send status even if unchanged in Watch (seconds).
+                Must be >= check_interval.
+            check_timeout: Maximum seconds to wait for each health check run. Must be positive.
+
+        Raises:
+            ValueError: If check_interval, heartbeat_interval or check_timeout are invalid.
+        """
+        if check_interval <= 0:
+            raise ValueError(f"check_interval must be positive, got {check_interval}")
+        if heartbeat_interval < check_interval:
+            raise ValueError(f"heartbeat_interval ({heartbeat_interval}) must be >= check_interval ({check_interval})")
+        if check_timeout <= 0:
+            raise ValueError(f"check_timeout must be positive, got {check_timeout}")
+
+        self._checkers = checkers or []
+        self._cache = cache
+        self._service_names = set(service_names or [])
+        self._service_names.add("")  # Always support overall health check
+        self._check_interval = check_interval
+        self._heartbeat_interval = heartbeat_interval
+        self._check_timeout = check_timeout
+
+    def _validate_service_name(self, service: str) -> bool:
+        """Validate gRPC service name ("" is the overall-health wildcard)."""
+        if not service:
+            return True
+        return is_valid_service_name(service)
+
+    def _is_service_known(self, service: str) -> bool:
+        """Check if the requested service name is registered."""
+        return service in self._service_names
+
+    async def _perform_health_check(self, service: str = "") -> health_pb2.HealthCheckResponse.ServingStatus:
+        """Perform health check with error handling.
+
+        Returns:
+            SERVING if all checks pass, NOT_SERVING on any failure or exception.
+        """
+        try:
+            return await check_async_overall_health(
+                cache=self._cache,
+                checkers=self._checkers,
+                timeout=self._check_timeout,
+            )
+        except (asyncio.CancelledError, KeyboardInterrupt, SystemExit):
+            raise
+        except (RuntimeError, ValueError, TimeoutError) as e:
+            # Expected/recoverable errors during health check execution
+            logger.warning("Health check execution failed (recoverable)", extra={"service": service, "error": str(e)})
+            return health_pb2.HealthCheckResponse.NOT_SERVING
+        except Exception:
+            logger.exception("Health check failed due to unexpected error", extra={"service": service})
+            return health_pb2.HealthCheckResponse.NOT_SERVING
+
+    def _should_send_update(
+        self,
+        status: health_pb2.HealthCheckResponse.ServingStatus,
+        last_status: health_pb2.HealthCheckResponse.ServingStatus | None,
+        last_heartbeat: float,
+        now: float,
+    ) -> bool:
+        """Determine if a status update should be sent in Watch mode.
+
+        Returns:
+            True if update should be sent (status changed or heartbeat interval reached).
+        """
+        is_changed = status != last_status
+        is_heartbeat = (now - last_heartbeat) >= self._heartbeat_interval
+        return is_changed or is_heartbeat
+
+    async def Check(
+        self,
+        request: health_pb2.HealthCheckRequest,
+        _context: grpc.aio.ServicerContext[health_pb2.HealthCheckRequest, health_pb2.HealthCheckResponse],
+    ) -> health_pb2.HealthCheckResponse:
+        """Check service health with real-time dependency checks (unary).
+
+        Implements the Check RPC from gRPC Health Checking Protocol v1.
+
+        Returns:
+            Health check response with status (SERVING, NOT_SERVING, or SERVICE_UNKNOWN).
+        """
+        if not self._validate_service_name(request.service):
+            logger.warning("Health check with invalid service name", extra={"service": request.service})
+            return health_pb2.HealthCheckResponse(status=health_pb2.HealthCheckResponse.SERVICE_UNKNOWN)
+
+        if not self._is_service_known(request.service):
+            logger.debug("Health check for unknown service", extra={"service": request.service})
+            return health_pb2.HealthCheckResponse(status=health_pb2.HealthCheckResponse.SERVICE_UNKNOWN)
+
+        status = await self._perform_health_check(service=request.service)
+        logger.debug(
+            "Health check completed",
+            extra={
+                "service": request.service,
+                "status": health_pb2.HealthCheckResponse.ServingStatus.Name(status),
+            },
+        )
+        return health_pb2.HealthCheckResponse(status=status)
+
+    async def Watch(
+        self,
+        request: health_pb2.HealthCheckRequest,
+        context: grpc.aio.ServicerContext[health_pb2.HealthCheckRequest, health_pb2.HealthCheckResponse],
+    ) -> AsyncGenerator[health_pb2.HealthCheckResponse, None]:
+        """Watch service health status with streaming updates.
+
+        Implements the Watch RPC from gRPC Health Checking Protocol v1.
+        Sends updates when status changes or on heartbeat interval. Per the
+        protocol a Watch stream never completes normally: internal failures
+        abort the stream with INTERNAL so clients can distinguish a broken
+        watch from a deliberate server shutdown.
+
+        Yields:
+            Streaming health check responses as status changes.
+        """
+        if not self._validate_service_name(request.service):
+            logger.warning("Health watch with invalid service name", extra={"service": request.service})
+            yield health_pb2.HealthCheckResponse(status=health_pb2.HealthCheckResponse.SERVICE_UNKNOWN)
+            return
+
+        if not self._is_service_known(request.service):
+            logger.debug("Health watch for unknown service", extra={"service": request.service})
+            yield health_pb2.HealthCheckResponse(status=health_pb2.HealthCheckResponse.SERVICE_UNKNOWN)
+            return
+
+        logger.info(
+            "Health watch started",
+            extra={
+                "service": request.service,
+                "check_interval": self._check_interval,
+                "heartbeat_interval": self._heartbeat_interval,
+            },
+        )
+
+        last_status: health_pb2.HealthCheckResponse.ServingStatus | None = None
+        last_heartbeat = 0.0
+
+        try:
+            while not context.done():
+                status = await self._perform_health_check(service=request.service)
+                now = asyncio.get_running_loop().time()
+
+                if self._should_send_update(status, last_status, last_heartbeat, now):
+                    last_status = status
+                    last_heartbeat = now
+                    yield health_pb2.HealthCheckResponse(status=status)
+
+                if context.done():
+                    break
+
+                await asyncio.sleep(self._check_interval)
+
+        except (asyncio.CancelledError, KeyboardInterrupt, SystemExit):
+            logger.info("Health watch cancelled", extra={"service": request.service})
+            raise
+        except grpc.aio.AbortError:
+            raise
+        except (RuntimeError, ValueError, TimeoutError) as e:
+            logger.warning("Health watch failed (recoverable)", extra={"service": request.service, "error": str(e)})
+            await self._abort_watch(context)
+        except Exception:
+            logger.exception("Health watch failed due to unexpected error", extra={"service": request.service})
+            await self._abort_watch(context)
+        finally:
+            logger.info("Health watch stopped", extra={"service": request.service})
+
+    @staticmethod
+    async def _abort_watch(
+        context: grpc.aio.ServicerContext[health_pb2.HealthCheckRequest, health_pb2.HealthCheckResponse],
+    ) -> NoReturn:
+        """Terminate a broken Watch stream with an error status (never OK EOF)."""
+        await context.abort(grpc.StatusCode.INTERNAL, "Health watch failed")
+        raise RuntimeError("ServicerContext.abort() returned instead of raising")

grpc_server_kit/aio/interceptors/__init__.py

@@ -0,0 +1,62 @@
+"""Async gRPC server interceptors.
+
+All interceptors are streaming-aware: they wrap the whole RPC (including full
+response streams), so timings, error mapping, and cleanup cover the entire
+call. Build custom interceptors by subclassing
+:class:`~grpc_server_kit.aio.interceptors.base.AsyncServerInterceptor`.
+
+Canonical execution order (outermost first):
+1. ``AsyncMetricsInterceptor``            (metrics — measures everything below)
+2. ``AsyncContextInterceptor``            (extract metadata into context vars)
+3. ``AsyncRequestLoggerInterceptor``      (request/response logging)
+4. ``AsyncTracingInterceptor``            (tracing span)
+5. ``AsyncSentryInterceptor``             (error tracking)
+6. ``AsyncExceptionHandlerInterceptor``   (exception → gRPC status mapping)
+"""
+
+from .base import AsyncServerInterceptor, RpcCall, split_method_name
+from .context import AsyncContextInterceptor, HeaderConfig
+from .exception_handler import (
+    GRPC_DEFAULT_ERROR_STATUS_MAP,
+    GRPC_SAFE_ERROR_MESSAGES,
+    AsyncExceptionHandlerInterceptor,
+    ErrorDetailFactory,
+    default_error_detail,
+    find_mapped_status,
+)
+from .metadata import get_metadata_dict
+from .metrics import AsyncMetricsInterceptor
+from .protocols import (
+    ErrorReporterProtocol,
+    GrpcServerMetricsProtocol,
+    SpanAttributeValue,
+    SpanProtocol,
+    TracerProtocol,
+)
+from .request_logger import AsyncRequestLoggerInterceptor
+from .sentry import AsyncSentryInterceptor
+from .tracing import AsyncTracingInterceptor
+
+__all__ = [
+    "GRPC_DEFAULT_ERROR_STATUS_MAP",
+    "GRPC_SAFE_ERROR_MESSAGES",
+    "AsyncContextInterceptor",
+    "AsyncExceptionHandlerInterceptor",
+    "AsyncMetricsInterceptor",
+    "AsyncRequestLoggerInterceptor",
+    "AsyncSentryInterceptor",
+    "AsyncServerInterceptor",
+    "AsyncTracingInterceptor",
+    "ErrorDetailFactory",
+    "ErrorReporterProtocol",
+    "GrpcServerMetricsProtocol",
+    "HeaderConfig",
+    "RpcCall",
+    "SpanAttributeValue",
+    "SpanProtocol",
+    "TracerProtocol",
+    "default_error_detail",
+    "find_mapped_status",
+    "get_metadata_dict",
+    "split_method_name",
+]