Skip to content

Commit ab3f832

Browse files
committed
add debug logs for rest streaming
1 parent db32877 commit ab3f832

File tree

1 file changed

+24
-2
lines changed

1 file changed

+24
-2
lines changed

google/api_core/_rest_streaming_base.py

Lines changed: 24 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
"""Helpers for server-side streaming in REST."""
1616

1717
from collections import deque
18+
import logging
1819
import string
1920
from typing import Deque, Union
2021
import types
@@ -23,6 +24,8 @@
2324
import google.protobuf.message
2425
from google.protobuf.json_format import Parse
2526

27+
_LOGGER = logging.getLogger(__name__)
28+
2629

2730
class BaseResponseIterator:
2831
"""Base Iterator over REST API responses. This class should not be used directly.
@@ -97,19 +100,38 @@ def _process_chunk(self, chunk: str):
97100
self._obj += char
98101
self._escape_next = not self._escape_next if char == "\\" else False
99102

103+
def _log_response_payload(self, response_payload: str):
104+
rest_response = {
105+
"payload": response_payload,
106+
"status": "OK",
107+
}
108+
_LOGGER.debug(
109+
"Received response via REST stream",
110+
extra={
111+
"response": rest_response,
112+
},
113+
)
114+
100115
def _create_grab(self):
116+
logging_enabled = _LOGGER.isEnabledFor(logging.DEBUG)
101117
if issubclass(self._response_message_cls, proto.Message):
102118

103119
def grab(this):
120+
result = this._ready_objs.popleft()
121+
if logging_enabled: # pragma: NO COVER
122+
self._log_result(result)
104123
return this._response_message_cls.from_json(
105-
this._ready_objs.popleft(), ignore_unknown_fields=True
124+
result, ignore_unknown_fields=True
106125
)
107126

108127
return grab
109128
elif issubclass(self._response_message_cls, google.protobuf.message.Message):
110129

111130
def grab(this):
112-
return Parse(this._ready_objs.popleft(), this._response_message_cls())
131+
result = this._ready_objs.popleft()
132+
if logging_enabled: # pragma: NO COVER
133+
self._log_result(result)
134+
return Parse(result, this._response_message_cls())
113135

114136
return grab
115137
else:

0 commit comments

Comments
 (0)