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

grpc_server_kit/aio/interceptors/base.py

@@ -0,0 +1,235 @@
+"""Streaming-aware base class for async gRPC server interceptors.
+
+The kit's interceptors wrap *whole* RPCs — including response-streaming ones —
+so timings, error mapping, and cleanup cover the full stream lifetime, not just
+handler creation. Subclasses implement a single async-generator hook,
+:meth:`AsyncServerInterceptor.around_call`, which the base class turns into a
+context manager around each of the four gRPC call kinds (unary-unary,
+unary-stream, stream-unary, stream-stream).
+"""
+
+from __future__ import annotations
+
+import abc
+import contextlib
+import dataclasses
+import functools
+from typing import TYPE_CHECKING, Any
+
+import grpc
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator, Awaitable, Callable, Collection
+
+__all__ = ["AsyncServerInterceptor", "RpcCall", "split_method_name"]
+
+
+@dataclasses.dataclass(frozen=True, slots=True)
+class RpcCall:
+    """One RPC as seen by an interceptor.
+
+    Attributes:
+        method_name: Full RPC name, e.g. ``"/pkg.Service/Method"``.
+        request: The request message (unary request) or an async iterator of
+            request messages (streaming request).
+        context: The gRPC servicer context of this RPC.
+        request_streaming: True for stream-unary / stream-stream calls.
+        response_streaming: True for unary-stream / stream-stream calls.
+    """
+
+    method_name: str
+    request: Any
+    context: grpc.aio.ServicerContext[Any, Any]
+    request_streaming: bool
+    response_streaming: bool
+
+
+def split_method_name(method_name: str) -> tuple[str, str]:
+    """Split ``"/pkg.Service/Method"`` into ``("pkg.Service", "Method")``.
+
+    Returns empty strings for parts that cannot be parsed.
+    """
+    service, _, method = method_name.removeprefix("/").partition("/")
+    return service, method
+
+
+@contextlib.asynccontextmanager
+async def _null_around() -> AsyncIterator[None]:
+    yield
+
+
+def _get_factory_and_behavior(
+    handler: grpc.RpcMethodHandler,
+) -> tuple[Callable[..., grpc.RpcMethodHandler], Callable[..., Any]]:
+    """Return the matching ``*_rpc_method_handler`` factory and the handler behavior."""
+    if handler.unary_unary:
+        return grpc.unary_unary_rpc_method_handler, handler.unary_unary
+    if handler.unary_stream:
+        return grpc.unary_stream_rpc_method_handler, handler.unary_stream
+    if handler.stream_unary:
+        return grpc.stream_unary_rpc_method_handler, handler.stream_unary
+    if handler.stream_stream:
+        return grpc.stream_stream_rpc_method_handler, handler.stream_stream
+    raise NotImplementedError("RPC handler implements no known behavior")
+
+
+class AsyncServerInterceptor(grpc.aio.ServerInterceptor, abc.ABC):
+    """Base class for async server interceptors covering all four RPC kinds.
+
+    Subclasses implement :meth:`around_call` — an async generator that yields
+    exactly once:
+
+    - code before ``yield`` runs before the RPC handler is invoked;
+    - the RPC (including full consumption of a streaming response) happens at
+      the ``yield`` point;
+    - code after ``yield`` (or in a ``finally`` block) runs after the RPC has
+      fully completed.
+
+    An exception raised by the handler — at any point of a streaming response —
+    propagates through the ``yield``, so ordinary ``try/except`` around it
+    covers the whole call. Note that :meth:`grpc.aio.ServicerContext.abort`
+    raises :class:`grpc.aio.AbortError` (not ``grpc.RpcError``); interceptors
+    that must not treat an intentional abort as a failure should re-raise it.
+
+    Args:
+        skip_methods: Full RPC method names (e.g. ``"/grpc.health.v1.Health/Check"``)
+            this interceptor must not wrap at all — handlers for these methods
+            pass through with zero interception overhead.
+
+    Example::
+
+        class TimingInterceptor(AsyncServerInterceptor):
+            async def around_call(self, call: RpcCall) -> AsyncIterator[None]:
+                start = time.perf_counter()
+                try:
+                    yield
+                finally:
+                    print(call.method_name, time.perf_counter() - start)
+    """
+
+    def __init__(self, *, skip_methods: Collection[str] = ()) -> None:
+        super().__init__()
+        self._skip_methods = frozenset(skip_methods)
+
+    @abc.abstractmethod
+    def around_call(self, call: RpcCall) -> AsyncIterator[None]:
+        """Wrap one RPC (async generator: setup before ``yield``, teardown after)."""
+        ...
+
+    @functools.cached_property
+    def _around_factory(self) -> Callable[[RpcCall], contextlib.AbstractAsyncContextManager[None]]:
+        # Built once per interceptor instance; asynccontextmanager itself
+        # produces a fresh context manager per call.
+        return contextlib.asynccontextmanager(self.around_call)
+
+    @functools.cached_property
+    def _wrapped_handlers(self) -> dict[str, tuple[grpc.RpcMethodHandler, grpc.RpcMethodHandler]]:
+        # Per-method cache of (source handler, wrapped handler). Registered
+        # handlers are stable objects, so the identity check on the source
+        # handler keeps the cache correct even for dynamic generic handlers.
+        return {}
+
+    def around(self, call: RpcCall) -> contextlib.AbstractAsyncContextManager[None]:
+        """Return the per-RPC context manager built from :meth:`around_call`.
+
+        Skipped methods get a no-op context manager. Useful for composing
+        interceptors manually and for unit tests.
+        """
+        if call.method_name in self._skip_methods:
+            return _null_around()
+        return self._around_factory(call)
+
+    async def intercept_service(
+        self,
+        continuation: Callable[[grpc.HandlerCallDetails], Awaitable[grpc.RpcMethodHandler | None]],
+        handler_call_details: grpc.HandlerCallDetails,
+    ) -> grpc.RpcMethodHandler | None:
+        """Wrap the resolved RPC handler with :meth:`around_call` (grpc plumbing)."""
+        handler = await continuation(handler_call_details)
+        if handler is None:
+            return None
+
+        raw_method = handler_call_details.method
+        method_name = raw_method.decode("utf-8", errors="replace") if isinstance(raw_method, bytes) else raw_method
+
+        if method_name in self._skip_methods:
+            return handler
+
+        cached = self._wrapped_handlers.get(method_name)
+        if cached is not None and cached[0] is handler:
+            return cached[1]
+
+        factory, behavior = _get_factory_and_behavior(handler)
+
+        wrapped_behavior: Callable[..., Any]
+        if handler.response_streaming:
+            wrapped_behavior = self._wrap_streaming_response(behavior, method_name, bool(handler.request_streaming))
+        else:
+            wrapped_behavior = self._wrap_unary_response(behavior, method_name, bool(handler.request_streaming))
+
+        wrapped = factory(
+            wrapped_behavior,
+            request_deserializer=handler.request_deserializer,
+            response_serializer=handler.response_serializer,
+        )
+        self._wrapped_handlers[method_name] = (handler, wrapped)
+        return wrapped
+
+    def _wrap_unary_response(
+        self,
+        behavior: Callable[..., Any],
+        method_name: str,
+        request_streaming: bool,
+    ) -> Callable[..., Any]:
+        async def invoke(request_or_iterator: Any, context: grpc.aio.ServicerContext[Any, Any]) -> Any:
+            call = RpcCall(
+                method_name=method_name,
+                request=request_or_iterator,
+                context=context,
+                request_streaming=request_streaming,
+                response_streaming=False,
+            )
+            async with self.around(call):
+                response = await behavior(request_or_iterator, context)
+            return response
+
+        return invoke
+
+    def _wrap_streaming_response(
+        self,
+        behavior: Callable[..., Any],
+        method_name: str,
+        request_streaming: bool,
+    ) -> Callable[..., Any]:
+        async def invoke(
+            request_or_iterator: Any,
+            context: grpc.aio.ServicerContext[Any, Any],
+        ) -> AsyncIterator[Any]:
+            call = RpcCall(
+                method_name=method_name,
+                request=request_or_iterator,
+                context=context,
+                request_streaming=request_streaming,
+                response_streaming=True,
+            )
+            async with self.around(call):
+                result = behavior(request_or_iterator, context)
+                if not hasattr(result, "__aiter__"):
+                    # Handlers written against the read/write API are plain
+                    # coroutines (returning None); interceptor-wrapped handlers
+                    # may also be coroutines resolving to an async generator.
+                    result = await result
+                if result is not None:
+                    try:
+                        async for response in result:
+                            yield response
+                    finally:
+                        # Close the inner generator NOW (client gone / early
+                        # break included) so servicer finally-blocks and DI
+                        # teardown run inside the RPC, not later under the
+                        # event loop's async-generator finalizer.
+                        aclose = getattr(result, "aclose", None)
+                        if aclose is not None:
+                            await aclose()
+
+        return invoke

grpc_server_kit/aio/interceptors/context.py

@@ -0,0 +1,179 @@
+"""Context interceptor: extract gRPC metadata into context variables / structlog."""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, NoReturn
+
+import grpc
+
+from .base import AsyncServerInterceptor, RpcCall
+from .metadata import get_metadata_dict
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator, Callable
+
+    import structlog
+else:
+    try:
+        import structlog
+    except ImportError:
+        structlog = None
+
+__all__ = ["AsyncContextInterceptor", "HeaderConfig"]
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class HeaderConfig:
+    """Configuration for a single header-to-context-variable mapping."""
+
+    header_name: str
+    context_var_name: str
+    context_setter: Callable[[str], Callable[[], None] | Any] | None = None
+    default_factory: Callable[[], str] | None = None
+    validator: Callable[[str], bool] | None = None
+    required: bool = False
+    allow_empty: bool = False
+
+    def __post_init__(self) -> None:
+        """Validate configuration."""
+        if not self.header_name:
+            raise ValueError("header_name cannot be empty")
+        if not self.context_var_name:
+            raise ValueError("context_var_name cannot be empty")
+
+
+class AsyncContextInterceptor(AsyncServerInterceptor):
+    """Interceptor that extracts gRPC metadata and binds it to context variables.
+
+    This interceptor should be early in the chain so subsequent interceptors and
+    the handler see the bound context variables. It supports extracting standard
+    headers like ``x-request-id`` as well as custom header mappings. Bindings
+    stay in place for the whole RPC — including the full response stream of
+    streaming calls — and are removed afterwards.
+    """
+
+    def __init__(
+        self,
+        header_configs: list[HeaderConfig],
+        bind_method_name: bool = True,
+        bind_structlog: bool = True,
+        method_key: str = "grpc_method",
+    ) -> None:
+        """Initialize interceptor with header configurations.
+
+        Args:
+            header_configs: List of header configurations to extract and bind.
+            bind_method_name: Whether to bind the gRPC method name to context.
+            bind_structlog: Whether to bind extracted variables to structlog context
+                (no-op if structlog is not installed).
+            method_key: Key for the gRPC method name in context.
+        """
+        super().__init__()
+        self._header_configs = header_configs
+        self._bind_method_name = bind_method_name
+        self._bind_structlog = bind_structlog
+        self._method_key = method_key
+
+    async def around_call(self, call: RpcCall) -> AsyncIterator[None]:
+        """Bind context variables from metadata for the duration of the RPC."""
+        metadata = get_metadata_dict(call.context)
+        bound_vars, removers = await self._process_headers(metadata, call.context, call.method_name)
+
+        if self._bind_method_name:
+            bound_vars[self._method_key] = call.method_name
+
+        self._bind_context(bound_vars)
+        try:
+            yield
+        finally:
+            self._unbind_context(bound_vars, removers)
+
+    async def _process_headers(
+        self,
+        metadata: dict[str, str],
+        context: grpc.aio.ServicerContext[Any, Any],
+        method_name: str,
+    ) -> tuple[dict[str, str], list[Callable[[], None]]]:
+        """Process configured headers and extract values."""
+        bound_vars: dict[str, str] = {}
+        removers: list[Callable[[], None]] = []
+
+        for config in self._header_configs:
+            value = self._get_header_value(config, metadata)
+
+            if value is None:
+                if config.required:
+                    await self._abort_request(
+                        context, method_name, config.header_name, f"Required header '{config.header_name}' not found"
+                    )
+                continue
+
+            if not value and not config.allow_empty:
+                if config.required:
+                    await self._abort_request(
+                        context, method_name, config.header_name, f"Header '{config.header_name}' cannot be empty"
+                    )
+                continue
+
+            if config.validator and not config.validator(value):
+                if config.required:
+                    await self._abort_request(
+                        context, method_name, config.header_name, f"Invalid value for header '{config.header_name}'"
+                    )
+                continue
+
+            if config.context_setter:
+                # The setter may return a cleanup callable (remover)
+                remover = config.context_setter(value)
+                if callable(remover):
+                    removers.append(remover)
+
+            bound_vars[config.context_var_name] = value
+
+        logger.debug("Extracted context variables", extra=bound_vars)
+        return bound_vars, removers
+
+    async def _abort_request(
+        self,
+        context: grpc.aio.ServicerContext[Any, Any],
+        method_name: str,
+        header_name: str,
+        details: str,
+    ) -> NoReturn:
+        """Log and abort a gRPC request (abort raises ``grpc.aio.AbortError``)."""
+        logger.warning(
+            "gRPC request context extraction failed",
+            extra={"method": method_name, "header": header_name, "details": details},
+        )
+        await context.abort(grpc.StatusCode.INVALID_ARGUMENT, details)
+        # grpc.aio contexts never return from abort(); guard against non-conforming ones.
+        raise RuntimeError("ServicerContext.abort() returned instead of raising")
+
+    def _get_header_value(self, config: HeaderConfig, metadata: dict[str, str]) -> str | None:
+        """Get header value from metadata or default factory."""
+        value = metadata.get(config.header_name)
+        if value is None and config.default_factory:
+            return config.default_factory()
+        return value
+
+    def _bind_context(self, bound_vars: dict[str, str]) -> None:
+        """Bind variables to structlog context."""
+        if self._bind_structlog and structlog is not None:
+            structlog.contextvars.bind_contextvars(**bound_vars)
+
+    def _unbind_context(self, bound_vars: dict[str, str], removers: list[Callable[[], None]]) -> None:
+        """Unbind variables from structlog context and call cleanup removers."""
+        if self._bind_structlog and structlog is not None and bound_vars:
+            structlog.contextvars.unbind_contextvars(*bound_vars.keys())
+
+        # Call cleanup removers in reverse order
+        while removers:
+            remover = removers.pop()
+            try:
+                remover()
+            except Exception:
+                logger.exception("Failed to call context cleanup remover")

grpc_server_kit/aio/interceptors/exception_handler.py

@@ -0,0 +1,167 @@
+"""Exception handler interceptor: map exceptions to gRPC errors (framework-agnostic).
+
+This interceptor is intentionally decoupled from any error/domain model. It maps an
+exception type to a ``grpc.StatusCode`` (walking the MRO) and aborts the request with
+a details string produced by a pluggable ``detail_factory``. The default factory emits
+a safe, non-leaking message per status code.
+
+To expose a structured error model (e.g. a domain ``DomainError.to_json()``), inject a
+custom ``error_status_map`` and ``detail_factory`` — no subclassing required.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import TYPE_CHECKING, NoReturn
+
+import grpc
+
+from grpc_server_kit.interceptors.utils import is_server_error
+
+from .base import AsyncServerInterceptor, RpcCall
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator, Callable
+
+__all__ = [
+    "GRPC_DEFAULT_ERROR_STATUS_MAP",
+    "GRPC_SAFE_ERROR_MESSAGES",
+    "AsyncExceptionHandlerInterceptor",
+    "ErrorDetailFactory",
+    "default_error_detail",
+    "find_mapped_status",
+]
+
+logger = logging.getLogger(__name__)
+
+# Pluggable factory: given the raised exception and the resolved status code, return the
+# gRPC ``details`` string to send to the client (e.g. a plain message or serialized JSON).
+type ErrorDetailFactory = Callable[[Exception, grpc.StatusCode], str]
+
+# Default mapping for common built-in exceptions (no domain/framework types).
+# Deliberately excludes TypeError: it almost always signals a server-side
+# programming bug, which must surface as INTERNAL (logged with a stack trace),
+# not be blamed on the client as INVALID_ARGUMENT.
+GRPC_DEFAULT_ERROR_STATUS_MAP: dict[type[Exception], grpc.StatusCode] = {
+    ValueError: grpc.StatusCode.INVALID_ARGUMENT,
+    PermissionError: grpc.StatusCode.PERMISSION_DENIED,
+    NotImplementedError: grpc.StatusCode.UNIMPLEMENTED,
+    TimeoutError: grpc.StatusCode.DEADLINE_EXCEEDED,
+    FileNotFoundError: grpc.StatusCode.NOT_FOUND,
+}
+
+# Safe, non-leaking messages per status code (used by the default detail factory).
+GRPC_SAFE_ERROR_MESSAGES: dict[grpc.StatusCode, str] = {
+    grpc.StatusCode.INVALID_ARGUMENT: "Invalid request data",
+    grpc.StatusCode.PERMISSION_DENIED: "Permission denied",
+    grpc.StatusCode.UNIMPLEMENTED: "Method is not implemented",
+    grpc.StatusCode.NOT_FOUND: "Resource not found",
+    grpc.StatusCode.ALREADY_EXISTS: "Resource already exists",
+    grpc.StatusCode.FAILED_PRECONDITION: "Failed precondition",
+    grpc.StatusCode.UNAUTHENTICATED: "Unauthenticated",
+    grpc.StatusCode.RESOURCE_EXHAUSTED: "Resource exhausted",
+    grpc.StatusCode.DEADLINE_EXCEEDED: "Deadline exceeded",
+    grpc.StatusCode.UNAVAILABLE: "Service unavailable",
+    grpc.StatusCode.ABORTED: "Operation aborted",
+    grpc.StatusCode.INTERNAL: "Internal server error",
+}
+
+_DEFAULT_FALLBACK_MESSAGE = "Request processing failed"
+
+
+def default_error_detail(exc: Exception, status: grpc.StatusCode) -> str:
+    """Default detail factory: a safe, non-leaking message for the resolved status code.
+
+    Never exposes the exception's own message/internals to the client.
+    """
+    return GRPC_SAFE_ERROR_MESSAGES.get(status, _DEFAULT_FALLBACK_MESSAGE)
+
+
+def find_mapped_status(
+    exc_type: type[BaseException],
+    error_status_map: dict[type[Exception], grpc.StatusCode],
+) -> grpc.StatusCode:
+    """Resolve the gRPC status an exception type maps to (walking the MRO).
+
+    Unmapped types resolve to ``INTERNAL``. Shared by the exception handler and
+    by capture filters that must classify errors with the SAME mapping.
+    """
+    for base in exc_type.__mro__:
+        if base in error_status_map:
+            return error_status_map[base]
+    return grpc.StatusCode.INTERNAL
+
+
+class AsyncExceptionHandlerInterceptor(AsyncServerInterceptor):
+    """Interceptor that catches exceptions and converts them to gRPC errors.
+
+    Rules:
+    - ``grpc.aio.AbortError`` (the handler already aborted with its own status) and
+      ``grpc.RpcError`` are re-raised as-is.
+    - ``CancelledError`` / ``KeyboardInterrupt`` / ``SystemExit`` are re-raised.
+    - Any other exception is mapped to a ``grpc.StatusCode`` (via the MRO) and the request
+      is aborted with details from the ``detail_factory``. Unmapped exceptions default to
+      ``INTERNAL`` with a safe message (no internal details leaked).
+
+    Covers streaming RPCs: exceptions raised while producing a response stream are
+    mapped and aborted the same way as unary failures.
+    """
+
+    def __init__(
+        self,
+        error_status_map: dict[type[Exception], grpc.StatusCode] | None = None,
+        *,
+        detail_factory: ErrorDetailFactory | None = None,
+        merge_defaults: bool = True,
+    ) -> None:
+        """Initialize interceptor.
+
+        Args:
+            error_status_map: Mapping of exception types to gRPC status codes. User entries
+                take precedence over the built-in defaults.
+            detail_factory: Callable producing the gRPC ``details`` string from
+                ``(exception, status_code)``. Defaults to :func:`default_error_detail`.
+            merge_defaults: When True (default), merge ``error_status_map`` over the built-in
+                defaults; when False, use only the provided map.
+        """
+        super().__init__()
+        base = dict(GRPC_DEFAULT_ERROR_STATUS_MAP) if merge_defaults else {}
+        if error_status_map:
+            base.update(error_status_map)
+        self._error_status_map = base
+        self._detail_factory = detail_factory or default_error_detail
+
+    async def around_call(self, call: RpcCall) -> AsyncIterator[None]:
+        """Run the RPC; map any unexpected exception to a gRPC abort."""
+        try:
+            yield
+        except (asyncio.CancelledError, KeyboardInterrupt, SystemExit):
+            raise
+        except grpc.aio.AbortError:
+            # The handler (or an inner interceptor) already aborted this RPC with
+            # its own status — pass it through untouched.
+            raise
+        except grpc.RpcError:
+            raise
+        except Exception as exc:
+            await self._abort_with_mapped_status(exc, call)
+
+    async def _abort_with_mapped_status(self, exc: Exception, call: RpcCall) -> NoReturn:
+        """Map the exception to a gRPC status and abort the request."""
+        status_code = find_mapped_status(type(exc), self._error_status_map)
+        details = self._detail_factory(exc, status_code)
+        log_extra = {
+            "error_type": type(exc).__name__,
+            "method": call.method_name,
+            "status_code": status_code.name,
+        }
+
+        if is_server_error(status_code):
+            logger.exception("gRPC request failed (server error)", extra=log_extra)
+        else:
+            logger.info("gRPC request failed (client error)", extra=log_extra)
+
+        await call.context.abort(status_code, details)
+        # grpc.aio contexts never return from abort(); guard against non-conforming ones.
+        raise RuntimeError("ServicerContext.abort() returned instead of raising")

grpc_server_kit/aio/interceptors/metadata.py

@@ -0,0 +1,57 @@
+"""Per-request invocation-metadata parsing with a context-local cache."""
+
+from __future__ import annotations
+
+import contextvars
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    import grpc
+
+__all__ = ["get_metadata_dict"]
+
+# Max bytes to represent as hex before truncating (avoids memory spike on large binary metadata).
+_METADATA_HEX_TRUNCATE_BYTES = 512
+
+# Context-local cache: each RPC runs in its own asyncio task (and therefore its
+# own execution context), so a ContextVar isolates entries per request without a
+# global registry or explicit cleanup. The cached context object is kept to
+# guard against reuse across nested calls within one task.
+_metadata_cache: contextvars.ContextVar[tuple[Any, dict[str, str]] | None] = contextvars.ContextVar(
+    "grpc_server_kit_metadata_cache",
+    default=None,
+)
+
+
+def get_metadata_dict(context: grpc.aio.ServicerContext[Any, Any]) -> dict[str, str]:
+    """Return invocation metadata as a dictionary, cached for the current RPC.
+
+    Avoids repeated ``dict()`` conversions of the metadata across interceptors.
+    All values are converted to strings (bytes are decoded as UTF-8; undecodable
+    bytes are hex-encoded and truncated).
+    """
+    cached = _metadata_cache.get()
+    if cached is not None and cached[0] is context:
+        return cached[1]
+
+    raw_metadata = context.invocation_metadata()
+    metadata: dict[str, str] = {}
+    if raw_metadata is not None:
+        for key, value in raw_metadata:
+            if isinstance(value, bytes):
+                try:
+                    metadata[key] = value.decode("utf-8")
+                except UnicodeDecodeError:
+                    metadata[key] = _bytes_to_safe_str(value)
+            else:
+                metadata[key] = value
+
+    _metadata_cache.set((context, metadata))
+    return metadata
+
+
+def _bytes_to_safe_str(value: bytes) -> str:
+    """Convert bytes to string for metadata; avoid unbounded .hex() on large payloads."""
+    if len(value) <= _METADATA_HEX_TRUNCATE_BYTES:
+        return value.hex()
+    return value[:_METADATA_HEX_TRUNCATE_BYTES].hex() + "...(truncated)"