Skip to content

Commit

Permalink
refactor: extract set_send_span and set_server_span logic to helpers
Browse files Browse the repository at this point in the history
  • Loading branch information
omgitsaheadcrab committed Aug 21, 2024
1 parent ddfb5e6 commit 3fc263e
Showing 1 changed file with 70 additions and 43 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -829,6 +829,66 @@ async def otel_receive():

return otel_receive

def _set_send_span(
self,
server_span_name,
scope,
send,
message,
status_code,
expecting_trailers,
):
"""Set send span attributes and status code."""
with self.tracer.start_as_current_span(
" ".join((server_span_name, scope["type"], "send"))
) as send_span:
if callable(self.client_response_hook):
self.client_response_hook(send_span, scope, message)

if send_span.is_recording():
if message["type"] == "http.response.start":
expecting_trailers = message.get("trailers", False)
send_span.set_attribute("asgi.event.type", message["type"])

if status_code:
set_status_code(
send_span,
status_code,
None,
self._sem_conv_opt_in_mode,
)
return expecting_trailers

def _set_server_span(
self, server_span, message, status_code, duration_attrs
):
"""Set server span attributes and status code."""
if (
server_span.is_recording()
and server_span.kind == trace.SpanKind.SERVER
and "headers" in message
):
custom_response_attributes = (
collect_custom_headers_attributes(
message,
self.http_capture_headers_sanitize_fields,
self.http_capture_headers_server_response,
normalise_response_header_name,
)
if self.http_capture_headers_server_response
else {}
)
if len(custom_response_attributes) > 0:
server_span.set_attributes(custom_response_attributes)

if status_code:
set_status_code(
server_span,
status_code,
duration_attrs,
self._sem_conv_opt_in_mode,
)

def _get_otel_send(
self,
server_span,
Expand All @@ -850,52 +910,19 @@ async def otel_send(message: dict[str, Any]):
status_code = 200

if not self.exclude_send_span:
with self.tracer.start_as_current_span(
" ".join((server_span_name, scope["type"], "send"))
) as send_span:
if callable(self.client_response_hook):
self.client_response_hook(send_span, scope, message)

if send_span.is_recording():
if message["type"] == "http.response.start":
expecting_trailers = message.get("trailers", False)
send_span.set_attribute(
"asgi.event.type", message["type"]
)
if status_code:
set_status_code(
send_span,
status_code,
None,
self._sem_conv_opt_in_mode,
)

if (
server_span.is_recording()
and server_span.kind == trace.SpanKind.SERVER
and "headers" in message
):
custom_response_attributes = (
collect_custom_headers_attributes(
message,
self.http_capture_headers_sanitize_fields,
self.http_capture_headers_server_response,
normalise_response_header_name,
)
if self.http_capture_headers_server_response
else {}
)
if len(custom_response_attributes) > 0:
server_span.set_attributes(custom_response_attributes)

if status_code:
set_status_code(
server_span,
expecting_trailers = self._set_send_span(
server_span_name,
scope,
send,
message,
status_code,
duration_attrs,
self._sem_conv_opt_in_mode,
expecting_trailers,
)

self._set_server_span(
server_span, message, status_code, duration_attrs
)

propagator = get_global_response_propagator()
if propagator:
propagator.inject(
Expand Down

0 comments on commit 3fc263e

Please sign in to comment.