Skip to content

Commit

Permalink
Add hooks for aiohttp, asgi, starlette, fastAPI, urllib, urllib3 (#576)
Browse files Browse the repository at this point in the history
  • Loading branch information
Ryo Kather authored Jul 26, 2021
1 parent 1157eb2 commit c5c6977
Show file tree
Hide file tree
Showing 15 changed files with 556 additions and 222 deletions.
5 changes: 5 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,13 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).

## [Unreleased](https://github.com/open-telemetry/opentelemetry-python/compare/v1.4.0-0.23b0...HEAD)

### Added
- `opentelemetry-sdk-extension-aws` Add AWS resource detectors to extension package
([#586](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/586))
- `opentelemetry-instrumentation-asgi`, `opentelemetry-instrumentation-aiohttp-client`, `openetelemetry-instrumentation-fastapi`,
`opentelemetry-instrumentation-starlette`, `opentelemetry-instrumentation-urllib`, `opentelemetry-instrumentation-urllib3` Added `request_hook` and `response_hook` callbacks
([#576](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/576))

## [1.4.0-0.23b0](https://github.com/open-telemetry/opentelemetry-python/releases/tag/v1.4.0-0.23b0) - 2021-07-21

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -81,13 +81,25 @@ def strip_query_params(url: yarl.URL) -> str:
)
from opentelemetry.propagate import inject
from opentelemetry.semconv.trace import SpanAttributes
from opentelemetry.trace import SpanKind, TracerProvider, get_tracer
from opentelemetry.trace import Span, SpanKind, TracerProvider, get_tracer
from opentelemetry.trace.status import Status, StatusCode
from opentelemetry.util.http import remove_url_credentials

_UrlFilterT = typing.Optional[typing.Callable[[str], str]]
_SpanNameT = typing.Optional[
typing.Union[typing.Callable[[aiohttp.TraceRequestStartParams], str], str]
_RequestHookT = typing.Optional[
typing.Callable[[Span, aiohttp.TraceRequestStartParams], None]
]
_ResponseHookT = typing.Optional[
typing.Callable[
[
Span,
typing.Union[
aiohttp.TraceRequestEndParams,
aiohttp.TraceRequestExceptionParams,
],
],
None,
]
]


Expand All @@ -108,7 +120,8 @@ def url_path_span_name(params: aiohttp.TraceRequestStartParams) -> str:

def create_trace_config(
url_filter: _UrlFilterT = None,
span_name: _SpanNameT = None,
request_hook: _RequestHookT = None,
response_hook: _ResponseHookT = None,
tracer_provider: TracerProvider = None,
) -> aiohttp.TraceConfig:
"""Create an aiohttp-compatible trace configuration.
Expand All @@ -134,15 +147,16 @@ def create_trace_config(
it as a span attribute. This can be useful to remove sensitive data
such as API keys or user personal information.
:param str span_name: Override the default span name.
:param Callable request_hook: Optional callback that can modify span name and request params.
:param Callable response_hook: Optional callback that can modify span name and response params.
:param tracer_provider: optional TracerProvider from which to get a Tracer
:return: An object suitable for use with :py:class:`aiohttp.ClientSession`.
:rtype: :py:class:`aiohttp.TraceConfig`
"""
# `aiohttp.TraceRequestStartParams` resolves to `aiohttp.tracing.TraceRequestStartParams`
# which doesn't exist in the aiottp intersphinx inventory.
# Explicitly specify the type for the `span_name` param and rtype to work
# which doesn't exist in the aiohttp intersphinx inventory.
# Explicitly specify the type for the `request_hook` and `response_hook` param and rtype to work
# around this issue.

tracer = get_tracer(__name__, __version__, tracer_provider)
Expand All @@ -161,17 +175,15 @@ async def on_request_start(
return

http_method = params.method.upper()
if trace_config_ctx.span_name is None:
request_span_name = "HTTP {}".format(http_method)
elif callable(trace_config_ctx.span_name):
request_span_name = str(trace_config_ctx.span_name(params))
else:
request_span_name = str(trace_config_ctx.span_name)
request_span_name = "HTTP {}".format(http_method)

trace_config_ctx.span = trace_config_ctx.tracer.start_span(
request_span_name, kind=SpanKind.CLIENT,
)

if callable(request_hook):
request_hook(trace_config_ctx.span, params)

if trace_config_ctx.span.is_recording():
attributes = {
SpanAttributes.HTTP_METHOD: http_method,
Expand All @@ -198,6 +210,9 @@ async def on_request_end(
if trace_config_ctx.span is None:
return

if callable(response_hook):
response_hook(trace_config_ctx.span, params)

if trace_config_ctx.span.is_recording():
trace_config_ctx.span.set_status(
Status(http_status_to_status_code(int(params.response.status)))
Expand All @@ -215,6 +230,9 @@ async def on_request_exception(
if trace_config_ctx.span is None:
return

if callable(response_hook):
response_hook(trace_config_ctx.span, params)

if trace_config_ctx.span.is_recording() and params.exception:
trace_config_ctx.span.set_status(Status(StatusCode.ERROR))
trace_config_ctx.span.record_exception(params.exception)
Expand All @@ -223,7 +241,7 @@ async def on_request_exception(
def _trace_config_ctx_factory(**kwargs):
kwargs.setdefault("trace_request_ctx", {})
return types.SimpleNamespace(
span_name=span_name, tracer=tracer, url_filter=url_filter, **kwargs
tracer=tracer, url_filter=url_filter, **kwargs
)

trace_config = aiohttp.TraceConfig(
Expand All @@ -240,7 +258,8 @@ def _trace_config_ctx_factory(**kwargs):
def _instrument(
tracer_provider: TracerProvider = None,
url_filter: _UrlFilterT = None,
span_name: _SpanNameT = None,
request_hook: _RequestHookT = None,
response_hook: _ResponseHookT = None,
):
"""Enables tracing of all ClientSessions
Expand All @@ -256,7 +275,8 @@ def instrumented_init(wrapped, instance, args, kwargs):

trace_config = create_trace_config(
url_filter=url_filter,
span_name=span_name,
request_hook=request_hook,
response_hook=response_hook,
tracer_provider=tracer_provider,
)
trace_config._is_instrumented_by_opentelemetry = True
Expand Down Expand Up @@ -304,12 +324,14 @@ def _instrument(self, **kwargs):
``url_filter``: A callback to process the requested URL prior to adding
it as a span attribute. This can be useful to remove sensitive data
such as API keys or user personal information.
``span_name``: Override the default span name.
``request_hook``: An optional callback that is invoked right after a span is created.
``response_hook``: An optional callback which is invoked right before the span is finished processing a response.
"""
_instrument(
tracer_provider=kwargs.get("tracer_provider"),
url_filter=kwargs.get("url_filter"),
span_name=kwargs.get("span_name"),
request_hook=kwargs.get("request_hook"),
response_hook=kwargs.get("response_hook"),
)

def _uninstrument(self, **kwargs):
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@
from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
from opentelemetry.semconv.trace import SpanAttributes
from opentelemetry.test.test_base import TestBase
from opentelemetry.trace import StatusCode
from opentelemetry.trace import Span, StatusCode


def run_with_test_server(
Expand Down Expand Up @@ -161,46 +161,51 @@ def test_not_recording(self):
self.assertFalse(mock_span.set_attribute.called)
self.assertFalse(mock_span.set_status.called)

def test_span_name_option(self):
for span_name, method, path, expected in (
("static", "POST", "/static-span-name", "static"),
(
lambda params: "{} - {}".format(
params.method, params.url.path
),
"PATCH",
"/some/path",
"PATCH - /some/path",
),
def test_hooks(self):
method = "PATCH"
path = "/some/path"
expected = "PATCH - /some/path"

def request_hook(span: Span, params: aiohttp.TraceRequestStartParams):
span.update_name("{} - {}".format(params.method, params.url.path))

def response_hook(
span: Span,
params: typing.Union[
aiohttp.TraceRequestEndParams,
aiohttp.TraceRequestExceptionParams,
],
):
with self.subTest(span_name=span_name, method=method, path=path):
host, port = self._http_request(
trace_config=aiohttp_client.create_trace_config(
span_name=span_name
),
method=method,
url=path,
status_code=HTTPStatus.OK,
)
span.set_attribute("response_hook_attr", "value")

self.assert_spans(
[
(
expected,
(StatusCode.UNSET, None),
{
SpanAttributes.HTTP_METHOD: method,
SpanAttributes.HTTP_URL: "http://{}:{}{}".format(
host, port, path
),
SpanAttributes.HTTP_STATUS_CODE: int(
HTTPStatus.OK
),
},
)
]
)
self.memory_exporter.clear()
host, port = self._http_request(
trace_config=aiohttp_client.create_trace_config(
request_hook=request_hook, response_hook=response_hook,
),
method=method,
url=path,
status_code=HTTPStatus.OK,
)

for span in self.memory_exporter.get_finished_spans():
self.assertEqual(span.name, expected)
self.assertEqual(
(span.status.status_code, span.status.description),
(StatusCode.UNSET, None),
)
self.assertEqual(
span.attributes[SpanAttributes.HTTP_METHOD], method
)
self.assertEqual(
span.attributes[SpanAttributes.HTTP_URL],
"http://{}:{}{}".format(host, port, path),
)
self.assertEqual(
span.attributes[SpanAttributes.HTTP_STATUS_CODE], HTTPStatus.OK
)
self.assertIn("response_hook_attr", span.attributes)
self.assertEqual(span.attributes["response_hook_attr"], "value")
self.memory_exporter.clear()

def test_url_filter_option(self):
# Strips all query params from URL before adding as a span attribute.
Expand Down Expand Up @@ -501,19 +506,32 @@ def strip_query_params(url: yarl.URL) -> str:
span.attributes[SpanAttributes.HTTP_URL],
)

def test_span_name(self):
def span_name_callback(params: aiohttp.TraceRequestStartParams) -> str:
return "{} - {}".format(params.method, params.url.path)
def test_hooks(self):
def request_hook(span: Span, params: aiohttp.TraceRequestStartParams):
span.update_name("{} - {}".format(params.method, params.url.path))

def response_hook(
span: Span,
params: typing.Union[
aiohttp.TraceRequestEndParams,
aiohttp.TraceRequestExceptionParams,
],
):
span.set_attribute("response_hook_attr", "value")

AioHttpClientInstrumentor().uninstrument()
AioHttpClientInstrumentor().instrument(span_name=span_name_callback)
AioHttpClientInstrumentor().instrument(
request_hook=request_hook, response_hook=response_hook
)

url = "/test-path"
run_with_test_server(
self.get_default_request(url), url, self.default_handler
)
span = self.assert_spans(1)
self.assertEqual("GET - /test-path", span.name)
self.assertIn("response_hook_attr", span.attributes)
self.assertEqual(span.attributes["response_hook_attr"], "value")


class TestLoadingAioHttpInstrumentor(unittest.TestCase):
Expand Down
Loading

0 comments on commit c5c6977

Please sign in to comment.