"""gRPC logging interceptors."""
# Suppress mypy errors related to interceptor base classes.
# https://github.com/grpc/grpc/issues/40550
# https://github.com/python/typeshed/issues/14641
# mypy: allow-any-generics
from __future__ import annotations
import abc
import functools
import logging
import sys
import threading
import time
from collections.abc import Iterator
from types import TracebackType
from typing import TYPE_CHECKING, Any, Callable, Generic, TypeVar
import grpc
from ni_grpc_extensions import _tracelogging
if TYPE_CHECKING:
if sys.version_info >= (3, 11):
from typing import Self
else:
from typing_extensions import Self
_logger = logging.getLogger(__name__)
[docs]
class ClientLogger(
grpc.UnaryUnaryClientInterceptor,
grpc.UnaryStreamClientInterceptor,
grpc.StreamUnaryClientInterceptor,
grpc.StreamStreamClientInterceptor,
):
"""Intercepts gRPC client calls and logs them for debugging."""
[docs]
@classmethod
def is_enabled(cls) -> bool:
"""Indicates whether gRPC client call logging is enabled for the current log level."""
return _ClientCallLogger.is_enabled()
[docs]
def intercept_unary_unary(
self,
continuation: Callable[
[grpc.ClientCallDetails, grpc._TRequest], grpc._CallFuture[grpc._TResponse]
],
client_call_details: grpc.ClientCallDetails,
request: grpc._TRequest,
) -> grpc._CallFuture[grpc._TResponse]:
"""Intercept and log a unary call."""
if _ClientCallLogger.is_enabled():
call_logger = _ClientCallLogger(client_call_details.method)
try:
call = continuation(client_call_details, request)
return _LoggingResponseCallFuture(call_logger, call)
except Exception as e:
call_logger.close(e)
raise
else:
return continuation(client_call_details, request)
[docs]
def intercept_unary_stream(
self,
continuation: Callable[
[grpc.ClientCallDetails, grpc._TRequest], grpc._CallIterator[grpc._TResponse]
],
client_call_details: grpc.ClientCallDetails,
request: grpc._TRequest,
) -> grpc._CallIterator[grpc._TResponse]:
"""Intercept and log a server-streaming call."""
if _ClientCallLogger.is_enabled():
call_logger = _ClientCallLogger(client_call_details.method)
try:
call_iterator = continuation(client_call_details, request)
return _LoggingResponseCallIterator(call_logger, call_iterator)
except Exception as e:
call_logger.close(e)
raise
else:
return continuation(client_call_details, request)
[docs]
def intercept_stream_unary(
self,
continuation: Callable[
[grpc.ClientCallDetails, Iterator[grpc._TRequest]], grpc._CallFuture[grpc._TResponse]
],
client_call_details: grpc.ClientCallDetails,
request_iterator: Iterator[grpc._TRequest],
) -> grpc._CallFuture[grpc._TResponse]:
"""Intercept and log a client-streaming call."""
if _ClientCallLogger.is_enabled():
call_logger = _ClientCallLogger(client_call_details.method)
try:
call = continuation(
client_call_details, _LoggingRequestIterator(call_logger, request_iterator)
)
return _LoggingResponseCallFuture(call_logger, call)
except Exception as e:
call_logger.close(e)
raise
else:
return continuation(client_call_details, request_iterator)
[docs]
def intercept_stream_stream(
self,
continuation: Callable[
[grpc.ClientCallDetails, Iterator[grpc._TRequest]], grpc._CallIterator[grpc._TResponse]
],
client_call_details: grpc.ClientCallDetails,
request_iterator: Iterator[grpc._TRequest],
) -> grpc._CallIterator[grpc._TResponse]:
"""Intercept and log a bidirectional streaming call."""
if _ClientCallLogger.is_enabled():
call_logger = _ClientCallLogger(client_call_details.method)
try:
call_iterator = continuation(
client_call_details, _LoggingRequestIterator(call_logger, request_iterator)
)
return _LoggingResponseCallIterator(call_logger, call_iterator)
except Exception as e:
call_logger.close(e)
raise
else:
return continuation(client_call_details, request_iterator)
[docs]
class ServerLogger(grpc.ServerInterceptor):
"""Intercepts gRPC server calls and logs them for debugging."""
[docs]
@classmethod
def is_enabled(cls) -> bool:
"""Indicates whether gRPC client call logging is enabled for the current log level."""
return _ServerCallLogger.is_enabled()
[docs]
def intercept_service(
self,
continuation: Callable[
[grpc.HandlerCallDetails], grpc.RpcMethodHandler[grpc._TRequest, grpc._TResponse] | None
],
handler_call_details: grpc.HandlerCallDetails,
) -> grpc.RpcMethodHandler[grpc._TRequest, grpc._TResponse] | None:
"""Intercept and log a server call."""
if _ServerCallLogger.is_enabled():
call_logger = _ServerCallLogger(handler_call_details.method)
handler = continuation(handler_call_details)
if handler is None:
return handler
elif handler.unary_unary:
return grpc.unary_unary_rpc_method_handler(
functools.partial(self._log_unary_unary, call_logger, handler.unary_unary),
handler.request_deserializer,
handler.response_serializer,
)
elif handler.unary_stream:
return grpc.unary_stream_rpc_method_handler(
functools.partial(self._log_unary_stream, call_logger, handler.unary_stream),
handler.request_deserializer,
handler.response_serializer,
)
elif handler.stream_unary:
return grpc.stream_unary_rpc_method_handler(
functools.partial(self._log_stream_unary, call_logger, handler.stream_unary),
handler.request_deserializer,
handler.response_serializer,
)
elif handler.stream_stream:
return grpc.stream_stream_rpc_method_handler(
functools.partial(self._log_stream_stream, call_logger, handler.stream_stream),
handler.request_deserializer,
handler.response_serializer,
)
else:
raise RuntimeError("Invalid RpcMethodHandler")
else:
return continuation(handler_call_details)
def _log_unary_unary(
self,
call_logger: _CallLogger,
handler_function: Callable[[grpc._TRequest, grpc.ServicerContext], grpc._TResponse],
request: grpc._TRequest,
context: grpc.ServicerContext,
) -> grpc._TResponse:
with call_logger:
return handler_function(request, context)
def _log_unary_stream(
self,
call_logger: _CallLogger,
handler_function: Callable[
[grpc._TRequest, grpc.ServicerContext], Iterator[grpc._TResponse]
],
request: grpc._TRequest,
context: grpc.ServicerContext,
) -> Iterator[grpc._TResponse]:
try:
return _LoggingResponseIterator(call_logger, handler_function(request, context))
except Exception as e:
call_logger.close(e)
raise
def _log_stream_unary(
self,
call_logger: _CallLogger,
handler_function: Callable[
[Iterator[grpc._TRequest], grpc.ServicerContext], grpc._TResponse
],
request_iterator: Iterator[grpc._TRequest],
context: grpc.ServicerContext,
) -> grpc._TResponse:
with call_logger:
return handler_function(_LoggingRequestIterator(call_logger, request_iterator), context)
def _log_stream_stream(
self,
call_logger: _CallLogger,
handler_function: Callable[
[Iterator[grpc._TRequest], grpc.ServicerContext], Iterator[grpc._TResponse]
],
request_iterator: Iterator[grpc._TRequest],
context: grpc.ServicerContext,
) -> Iterator[grpc._TResponse]:
try:
return _LoggingResponseIterator(
call_logger,
handler_function(_LoggingRequestIterator(call_logger, request_iterator), context),
)
except Exception as e:
call_logger.close(e)
raise
class _CallLogger(abc.ABC):
"""Logs a single call."""
__slots__ = ["_closed"]
# As of 2023, the Python stdlib doesn't support atomic operations, so use a global lock to
# atomically update self._closed.
_lock = threading.Lock()
def __init__(self) -> None:
self._closed = False
def close(self, exception: BaseException | None = None) -> None:
# close() is idempotent to avoid duplicate logs.
with _CallLogger._lock:
if self._closed:
return
self._closed = True
self._close(exception)
@abc.abstractmethod
def _close(self, exception: BaseException | None) -> None:
raise NotImplementedError()
def __enter__(self: Self) -> Self:
return self
def __exit__(
self,
exc_type: type[BaseException] | None,
exc_val: BaseException | None,
traceback: TracebackType | None,
) -> None:
self.close(exc_val)
@abc.abstractmethod
def log_streaming_request(self) -> None:
raise NotImplementedError()
@abc.abstractmethod
def log_streaming_response(self) -> None:
raise NotImplementedError()
class _ClientCallLogger(_CallLogger):
__slots__ = ["_method_name", "_activity_id"]
@classmethod
def is_enabled(cls) -> bool:
return _logger.isEnabledFor(logging.DEBUG) or _tracelogging.is_enabled()
def __init__(self, method_name: str) -> None:
super().__init__()
self._method_name = method_name
_logger.debug("gRPC client call starting: %s", self._method_name)
self._activity_id = _tracelogging.log_grpc_client_call_start(self._method_name)
def _close(self, exception: BaseException | None = None) -> None:
_logger.debug("gRPC client call complete: %s", self._method_name)
_tracelogging.log_grpc_client_call_stop(self._method_name, self._activity_id)
def log_streaming_request(self) -> None:
_logger.debug("gRPC client call streaming request: %s", self._method_name)
_tracelogging.log_grpc_client_call_streaming_request(self._method_name)
def log_streaming_response(self) -> None:
_logger.debug("gRPC client call streaming response: %s", self._method_name)
_tracelogging.log_grpc_client_call_streaming_response(self._method_name)
class _ServerCallLogger(_CallLogger):
__slots__ = ["_method_name", "_start_time", "_activity_id"]
@classmethod
def is_enabled(cls) -> bool:
return _logger.isEnabledFor(logging.INFO) or _tracelogging.is_enabled()
def __init__(self, method_name: str) -> None:
super().__init__()
self._method_name = method_name
self._start_time = time.perf_counter()
_logger.debug("gRPC server call starting: %s", self._method_name)
self._activity_id = _tracelogging.log_grpc_server_call_start(self._method_name)
def _close(self, exception: BaseException | None = None) -> None:
if _logger.isEnabledFor(logging.INFO):
# For production usage, log a concise one-line summary of the call, similar to what
# Serilog provides for ASP.NET Core services. Don't log exception details because
# grpcio's logging_pool already handles this.
elapsed_time = time.perf_counter() - self._start_time
_logger.info(
"gRPC server call %s responded %s in %.4f ms",
self._method_name,
str(_get_status_code(exception)).replace("StatusCode.", ""),
elapsed_time * 1000.0,
)
_logger.debug("gRPC server call complete: %s", self._method_name)
_tracelogging.log_grpc_server_call_stop(self._method_name, self._activity_id)
def log_streaming_request(self) -> None:
_logger.debug("gRPC server call streaming request: %s", self._method_name)
_tracelogging.log_grpc_server_call_streaming_request(self._method_name)
def log_streaming_response(self) -> None:
_logger.debug("gRPC server call streaming response: %s", self._method_name)
_tracelogging.log_grpc_server_call_streaming_response(self._method_name)
def _get_status_code(exception: BaseException | None) -> grpc.StatusCode:
if exception is None:
return grpc.StatusCode.OK
elif isinstance(exception, grpc.RpcError):
return exception.code()
else:
return grpc.StatusCode.UNKNOWN
_T = TypeVar("_T")
class _LoggingRequestIterator(Generic[_T]):
__slots__ = ["_call_logger", "_inner_iterator"]
def __init__(self, call_logger: _CallLogger, inner_iterator: Iterator[_T]) -> None:
self._call_logger = call_logger
self._inner_iterator = inner_iterator
def __iter__(self) -> Iterator[_T]:
return self
def __next__(self) -> _T:
request = next(self._inner_iterator)
self._call_logger.log_streaming_request()
return request
class _LoggingResponseIterator(Generic[_T]):
__slots__ = ["_call_logger", "_inner_iterator"]
def __init__(self, call_logger: _CallLogger, inner_iterator: Iterator[_T]) -> None:
self._call_logger = call_logger
self._inner_iterator = inner_iterator
def __iter__(self) -> Iterator[_T]:
return self
def __next__(self) -> _T:
# For server-streaming and bidirectional RPCs, the call is complete when the response
# stream is closed or it throws an exception.
try:
response = next(self._inner_iterator)
self._call_logger.log_streaming_response()
return response
except StopIteration:
self._call_logger.close()
raise
except Exception as e:
self._call_logger.close(e)
raise
if TYPE_CHECKING:
# These types are marked as @type_check_only and can be used only in type annotations,
# so we redefine them as empty classes at run time. Pyright still warns due to the behavior
# described in https://github.com/microsoft/pyright/issues/9934
_CallFuture = grpc._CallFuture # pyright: ignore[reportGeneralTypeIssues]
_CallIterator = grpc._CallIterator # pyright: ignore[reportGeneralTypeIssues]
else:
class _CallFuture(Generic[_T]):
pass
class _CallIterator(Generic[_T]):
pass
@grpc.Call.register
@grpc.Future.register
class _LoggingResponseCallFuture(_CallFuture[_T]): # pyright: ignore[reportGeneralTypeIssues]
__slots__ = ["_call_logger", "_inner_call_future"]
def __init__(self, call_logger: _CallLogger, inner_call_future: grpc._CallFuture[_T]) -> None:
self._call_logger = call_logger
self._inner_call_future = inner_call_future
def __getattr__(self, name: str) -> Any:
return getattr(self._inner_call_future, name)
def result(self, timeout: float | None = None) -> _T:
# For unary and client-streaming RPCs, the call is complete when it returns a response or
# throws an exception.
with self._call_logger:
return self._inner_call_future.result(timeout)
def exception(self, timeout: float | None = None) -> Exception | None:
with self._call_logger:
return self._inner_call_future.exception(timeout)
def traceback(self, timeout: float | None = None) -> Any:
with self._call_logger:
return self._inner_call_future.traceback(timeout)
# At run time, __getattr__ forwards these to the inner CallFuture.
if TYPE_CHECKING:
def add_callback(self, callback: Callable[[], None]) -> bool: ...
def add_done_callback(
self, fn: Callable[[grpc.Future[grpc._TFutureValue]], None]
) -> None: ...
def cancel(self) -> bool: ...
def cancelled(self) -> bool: ...
def code(self) -> grpc.StatusCode: ...
def details(self) -> str: ...
def done(self) -> bool: ...
def initial_metadata(self) -> grpc._Metadata: ...
def is_active(self) -> bool: ...
def running(self) -> bool: ...
def time_remaining(self) -> float: ...
def trailing_metadata(self) -> grpc._Metadata: ...
@grpc.Call.register
@grpc.Future.register
class _LoggingResponseCallIterator(_CallIterator[_T]): # pyright: ignore
__slots__ = ["_call_logger", "_inner_call_iterator"]
def __init__(
self, call_logger: _CallLogger, inner_call_iterator: grpc._CallIterator[_T]
) -> None:
self._call_logger = call_logger
self._inner_call_iterator = inner_call_iterator
def __getattr__(self, name: str) -> Any:
return getattr(self._inner_call_iterator, name)
def __iter__(self) -> Iterator[_T]:
return self
def __next__(self) -> _T:
# For server-streaming and bidirectional RPCs, the call is complete when the response
# stream is closed or throws an exception.
try:
response = next(self._inner_call_iterator)
self._call_logger.log_streaming_response()
return response
except StopIteration:
self._call_logger.close()
raise
except Exception as e:
self._call_logger.close(e)
raise
# At run time, __getattr__ forwards these to the inner CallIterator.
if TYPE_CHECKING:
def add_callback(self, callback: Callable[[], None]) -> bool: ...
def add_done_callback(
self, fn: Callable[[grpc.Future[grpc._TFutureValue]], None]
) -> None: ...
def cancel(self) -> bool: ...
def cancelled(self) -> bool: ...
def code(self) -> grpc.StatusCode: ...
def details(self) -> str: ...
def done(self) -> bool: ...
def initial_metadata(self) -> grpc._Metadata: ...
def is_active(self) -> bool: ...
def running(self) -> bool: ...
def time_remaining(self) -> float: ...
def trailing_metadata(self) -> grpc._Metadata: ...