def _intercept_call(self, client_call_details, request_iterator,
grpc_type):
metadata = ()
if client_call_details.metadata is not None:
metadata = client_call_details.metadata
# Start a span
current_span = self._start_client_span(client_call_details)
span_context = current_span.context_tracer.span_context
header = self._propagator.to_header(span_context)
grpc_trace_metadata = {
oc_grpc.GRPC_TRACE_KEY: header,
}
if isinstance(metadata, list):
metadata_to_append = list(six.iteritems(grpc_trace_metadata))
else:
metadata_to_append = tuple(six.iteritems(grpc_trace_metadata))
metadata = metadata + metadata_to_append
client_call_details = _ClientCallDetails(
client_call_details.method, client_call_details.timeout, metadata,
client_call_details.credentials)
request_iterator = grpc_utils.wrap_iter_with_message_events(
request_or_response_iter=request_iterator,
span=current_span,
message_event_type=time_event.Type.SENT)
return client_call_details, request_iterator, current_span
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:
_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:
execution_context.get_opencensus_tracer().end_span()
return response_or_iterator
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