feat: add client debug logging support for unary-stream gRPC calls

Author: parthea

google/api_core/grpc_helpers.py

@@ -17,9 +17,13 @@
 
 import collections
 import functools
+import logging
+import pickle
 import warnings
 
+import google.protobuf.json_format
 import grpc
+import proto
 
 from google.api_core import exceptions
 import google.auth
@@ -48,6 +52,7 @@
 else:
     HAS_GRPC_GCP = False
 
+_LOGGER = logging.getLogger(__name__)
 
 # The list of gRPC Callable interfaces that return iterators.
 _STREAM_WRAP_CLASSES = (grpc.UnaryStreamMultiCallable, grpc.StreamStreamMultiCallable)
@@ -113,7 +118,29 @@ def __next__(self) -> P:
                 result = self._stored_first_result
                 del self._stored_first_result
                 return result
-            return next(self._wrapped)
+            result = next(self._wrapped)
+
+            logging_enabled = _LOGGER.isEnabledFor(logging.DEBUG)
+            if logging_enabled:  # pragma: NO COVER
+                if isinstance(result, proto.Message):
+                    response_payload = type(result).to_json(result)
+                elif isinstance(result, google.protobuf.message.Message):
+                    response_payload = google.protobuf.json_format.MessageToJson(result)
+                else:
+                    response_payload = (
+                        f"{type(result).__name__}: {pickle.dumps(result)}"
+                    )
+                grpc_response = {
+                    "payload": response_payload,
+                    "status": "OK",
+                }
+                _LOGGER.debug(
+                    f"Received response of type {type(result)} via gRPC stream",
+                    extra={
+                        "response": grpc_response,
+                    },
+                )
+            return result
         except grpc.RpcError as exc:
             # If the stream has already returned data, we cannot recover here.
             raise exceptions.from_grpc_error(exc) from exc

google/api_core/grpc_helpers_async.py

@@ -20,17 +20,23 @@
 
 import asyncio
 import functools
+import logging
+import pickle
 
 from typing import AsyncGenerator, Generic, Iterator, Optional, TypeVar
 
+import google.protobuf.json_format
 import grpc
 from grpc import aio
+import proto
 
 from google.api_core import exceptions, grpc_helpers
 
 # denotes the proto response type for grpc calls
 P = TypeVar("P")
 
+_LOGGER = logging.getLogger(__name__)
+
 # NOTE(lidiz) Alternatively, we can hack "__getattribute__" to perform
 # automatic patching for us. But that means the overhead of creating an
 # extra Python function spreads to every single send and receive.
@@ -94,7 +100,28 @@ def __init__(self):
 
     async def read(self) -> P:
         try:
-            return await self._call.read()
+            result = await self._call.read()
+            logging_enabled = _LOGGER.isEnabledFor(logging.DEBUG)
+            if logging_enabled:  # pragma: NO COVER
+                if isinstance(result, proto.Message):
+                    response_payload = type(result).to_json(result)
+                elif isinstance(result, google.protobuf.message.Message):
+                    response_payload = google.protobuf.json_format.MessageToJson(result)
+                else:
+                    response_payload = (
+                        f"{type(result).__name__}: {pickle.dumps(result)}"
+                    )
+                grpc_response = {
+                    "payload": response_payload,
+                    "status": "OK",
+                }
+                _LOGGER.debug(
+                    f"Received response of type {type(result)} via gRPC stream",
+                    extra={
+                        "response": grpc_response,
+                    },
+                )
+            return result
         except grpc.RpcError as rpc_error:
             raise exceptions.from_grpc_error(rpc_error) from rpc_error
 
@@ -219,7 +246,7 @@ def create_channel(
     default_host=None,
     compression=None,
     attempt_direct_path: Optional[bool] = False,
-    **kwargs
+    **kwargs,
 ):
     """Create an AsyncIO secure channel with credentials.