|
15 | 15 | """Helpers for server-side streaming in REST."""
|
16 | 16 |
|
17 | 17 | from collections import deque
|
| 18 | +import logging |
18 | 19 | import string
|
19 | 20 | from typing import Deque, Union
|
20 | 21 | import types
|
|
23 | 24 | import google.protobuf.message
|
24 | 25 | from google.protobuf.json_format import Parse
|
25 | 26 |
|
| 27 | +_LOGGER = logging.getLogger(__name__) |
| 28 | + |
26 | 29 |
|
27 | 30 | class BaseResponseIterator:
|
28 | 31 | """Base Iterator over REST API responses. This class should not be used directly.
|
@@ -97,19 +100,38 @@ def _process_chunk(self, chunk: str):
|
97 | 100 | self._obj += char
|
98 | 101 | self._escape_next = not self._escape_next if char == "\\" else False
|
99 | 102 |
|
| 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 | + |
100 | 115 | def _create_grab(self):
|
| 116 | + logging_enabled = _LOGGER.isEnabledFor(logging.DEBUG) |
101 | 117 | if issubclass(self._response_message_cls, proto.Message):
|
102 | 118 |
|
103 | 119 | def grab(this):
|
| 120 | + result = this._ready_objs.popleft() |
| 121 | + if logging_enabled: # pragma: NO COVER |
| 122 | + self._log_result(result) |
104 | 123 | return this._response_message_cls.from_json(
|
105 |
| - this._ready_objs.popleft(), ignore_unknown_fields=True |
| 124 | + result, ignore_unknown_fields=True |
106 | 125 | )
|
107 | 126 |
|
108 | 127 | return grab
|
109 | 128 | elif issubclass(self._response_message_cls, google.protobuf.message.Message):
|
110 | 129 |
|
111 | 130 | 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()) |
113 | 135 |
|
114 | 136 | return grab
|
115 | 137 | else:
|
|
0 commit comments