def intercept_stream_stream(self, continuation, client_call_details,
request_iterator):
if CLOUD_TRACE in client_call_details.method:
response = continuation(client_call_details, request_iterator)
return response
new_details, new_request_iterator, current_span = self._intercept_call(
client_call_details=client_call_details,
request_iterator=request_iterator,
grpc_type=oc_grpc.STREAM_STREAM)
response_it = continuation(new_details, new_request_iterator)
response_it = grpc_utils.wrap_iter_with_message_events(
request_or_response_iter=response_it,
span=current_span,
message_event_type=time_event.Type.RECEIVED)
response_it = grpc_utils.wrap_iter_with_end_span(response_it)
return response_it
def new_behavior(request_or_iterator, servicer_context):
span = self._start_server_span(servicer_context)
try:
if request_streaming:
request_or_iterator = grpc_utils.wrap_iter_with_message_events( # noqa: E501
request_or_response_iter=request_or_iterator,
span=span,
message_event_type=time_event.Type.RECEIVED)
else:
grpc_utils.add_message_event(
proto_message=request_or_iterator,
span=span,
message_event_type=time_event.Type.RECEIVED,
)
# invoke the original rpc behavior
response_or_iterator = behavior(request_or_iterator,
servicer_context)
if response_streaming:
response_or_iterator = grpc_utils.wrap_iter_with_message_events( # noqa: E501
request_or_response_iter=response_or_iterator,
span=span,
message_event_type=time_event.Type.SENT)
response_or_iterator = grpc_utils.wrap_iter_with_end_span( # noqa: E501
response_or_iterator)
else:
grpc_utils.add_message_event(
proto_message=response_or_iterator,
span=span,
message_event_type=time_event.Type.SENT,
)
except Exception as exc:
logging.exception(exc)
_add_exc_info(span)
raise
finally:
# if the response is unary, end the span here. Otherwise
# it will be closed when the response iter completes
if not response_streaming:
self.tracer.end_span()
return response_or_iterator