/
__init__.py
127 lines (107 loc) · 4.63 KB
/
__init__.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
# Copyright The OpenTelemetry Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import typing
from typing import Collection
from starlette import applications
from starlette.routing import Match
from opentelemetry.instrumentation.asgi import OpenTelemetryMiddleware
from opentelemetry.instrumentation.asgi.package import _instruments
from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
from opentelemetry.semconv.trace import SpanAttributes
from opentelemetry.trace import Span
from opentelemetry.util.http import get_excluded_urls
_excluded_urls = get_excluded_urls("STARLETTE")
_ServerRequestHookT = typing.Optional[typing.Callable[[Span, dict], None]]
_ClientRequestHookT = typing.Optional[typing.Callable[[Span, dict], None]]
_ClientResponseHookT = typing.Optional[typing.Callable[[Span, dict], None]]
class StarletteInstrumentor(BaseInstrumentor):
"""An instrumentor for starlette
See `BaseInstrumentor`
"""
_original_starlette = None
@staticmethod
def instrument_app(
app: applications.Starlette,
server_request_hook: _ServerRequestHookT = None,
client_request_hook: _ClientRequestHookT = None,
client_response_hook: _ClientResponseHookT = None,
tracer_provider=None,
):
"""Instrument an uninstrumented Starlette application."""
if not getattr(app, "is_instrumented_by_opentelemetry", False):
app.add_middleware(
OpenTelemetryMiddleware,
excluded_urls=_excluded_urls,
default_span_details=_get_route_details,
server_request_hook=server_request_hook,
client_request_hook=client_request_hook,
client_response_hook=client_response_hook,
tracer_provider=tracer_provider,
)
app.is_instrumented_by_opentelemetry = True
def instrumentation_dependencies(self) -> Collection[str]:
return _instruments
def _instrument(self, **kwargs):
self._original_starlette = applications.Starlette
_InstrumentedStarlette._tracer_provider = kwargs.get("tracer_provider")
_InstrumentedStarlette._server_request_hook = kwargs.get(
"server_request_hook"
)
_InstrumentedStarlette._client_request_hook = kwargs.get(
"client_request_hook"
)
_InstrumentedStarlette._client_response_hook = kwargs.get(
"client_response_hook"
)
applications.Starlette = _InstrumentedStarlette
def _uninstrument(self, **kwargs):
applications.Starlette = self._original_starlette
class _InstrumentedStarlette(applications.Starlette):
_tracer_provider = None
_server_request_hook: _ServerRequestHookT = None
_client_request_hook: _ClientRequestHookT = None
_client_response_hook: _ClientResponseHookT = None
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.add_middleware(
OpenTelemetryMiddleware,
excluded_urls=_excluded_urls,
default_span_details=_get_route_details,
server_request_hook=_InstrumentedStarlette._server_request_hook,
client_request_hook=_InstrumentedStarlette._client_request_hook,
client_response_hook=_InstrumentedStarlette._client_response_hook,
tracer_provider=_InstrumentedStarlette._tracer_provider,
)
def _get_route_details(scope):
"""Callback to retrieve the starlette route being served.
TODO: there is currently no way to retrieve http.route from
a starlette application from scope.
See: https://github.com/encode/starlette/pull/804
"""
app = scope["app"]
route = None
for starlette_route in app.routes:
match, _ = starlette_route.matches(scope)
if match == Match.FULL:
route = starlette_route.path
break
if match == Match.PARTIAL:
route = starlette_route.path
# method only exists for http, if websocket
# leave it blank.
span_name = route or scope.get("method", "")
attributes = {}
if route:
attributes[SpanAttributes.HTTP_ROUTE] = route
return span_name, attributes