Skip to content

Commit

Permalink
Added trace response headers to Flask
Browse files Browse the repository at this point in the history
  • Loading branch information
owais committed Apr 13, 2021
1 parent 370952f commit 1b90220
Show file tree
Hide file tree
Showing 4 changed files with 49 additions and 1 deletion.
2 changes: 1 addition & 1 deletion .github/workflows/test.yml
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ on:
- 'release/*'
pull_request:
env:
CORE_REPO_SHA: cad261e5dae1fe986c87e6965664b45cc9ab73c3
CORE_REPO_SHA: 7b11971c504387341df0c38f5a34d7d1293c7e4f

jobs:
build:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,9 @@ def hello():
from opentelemetry import context, trace
from opentelemetry.instrumentation.flask.version import __version__
from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
from opentelemetry.instrumentation.propagators import (
get_global_back_propagator,
)
from opentelemetry.propagate import extract
from opentelemetry.util._time import _time_ns
from opentelemetry.util.http import get_excluded_urls
Expand Down Expand Up @@ -91,6 +94,13 @@ def _start_response(status, response_headers, *args, **kwargs):
if not _excluded_urls.url_disabled(flask.request.url):
span = flask.request.environ.get(_ENVIRON_SPAN_KEY)

propagator = get_global_back_propagator()
if propagator:
propagator.inject(
response_headers,
setter=otel_wsgi.default_back_propagation_setter,
)

if span:
otel_wsgi.add_response_attributes(
span, status, response_headers
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,11 @@

from opentelemetry import trace
from opentelemetry.instrumentation.flask import FlaskInstrumentor
from opentelemetry.instrumentation.propagators import (
TraceResponsePropagator,
get_global_back_propagator,
set_global_back_propagator,
)
from opentelemetry.test.test_base import TestBase
from opentelemetry.test.wsgitestutil import WsgiTestBase
from opentelemetry.util.http import get_excluded_urls
Expand Down Expand Up @@ -119,6 +124,31 @@ def test_simple(self):
self.assertEqual(span_list[0].kind, trace.SpanKind.SERVER)
self.assertEqual(span_list[0].attributes, expected_attrs)

def test_trace_response(self):
orig = get_global_back_propagator()

set_global_back_propagator(TraceResponsePropagator())
response = self.client.get("/hello/123")
headers = response.headers

span_list = self.memory_exporter.get_finished_spans()
self.assertEqual(len(span_list), 1)
span = span_list[0]

self.assertIn("traceresponse", headers)
self.assertEqual(
headers["access-control-expose-headers"], "traceresponse",
)
self.assertEqual(
headers["traceresponse"],
"00-{0}-{1}-01".format(
trace.format_trace_id(span.get_span_context().trace_id),
trace.format_span_id(span.get_span_context().span_id),
),
)

set_global_back_propagator(orig)

def test_not_recording(self):
mock_tracer = Mock()
mock_span = Mock()
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -253,3 +253,11 @@ def _end_span_after_iterating(iterable, span, tracer, token):
close()
span.end()
context.detach(token)


class BackPropagationSetter:
def set(self, carrier, key, value): # pylint: disable=no-self-use
carrier.append((key, value))


default_back_propagation_setter = BackPropagationSetter()

0 comments on commit 1b90220

Please sign in to comment.