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

parthea · parthea · commit ae07484aa631 · 2025-02-14T22:34:19.000Z
diff --git a/google/api_core/grpc_helpers.py b/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,31 @@ 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
diff --git a/google/api_core/grpc_helpers_async.py b/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,30 @@ 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
 
