|
3 | 3 |
|
4 | 4 | import logging
|
5 | 5 |
|
6 |
| -from flask import request as flask_request |
7 |
| - |
8 | 6 | import opentelemetry.ext.wsgi as otel_wsgi
|
9 | 7 | from opentelemetry import propagators, trace
|
10 |
| -from opentelemetry.ext.flask.version import __version__ |
11 | 8 | from opentelemetry.util import time_ns
|
| 9 | +import flask |
12 | 10 |
|
13 | 11 | logger = logging.getLogger(__name__)
|
14 | 12 |
|
|
17 | 15 | _ENVIRON_ACTIVATION_KEY = object()
|
18 | 16 |
|
19 | 17 |
|
20 |
| -def instrument_app(flask): |
21 |
| - """Makes the passed-in Flask object traced by OpenTelemetry. |
22 |
| -
|
23 |
| - You must not call this function multiple times on the same Flask object. |
24 |
| - """ |
25 |
| - |
26 |
| - wsgi = flask.wsgi_app |
27 |
| - |
28 |
| - def wrapped_app(environ, start_response): |
29 |
| - # We want to measure the time for route matching, etc. |
30 |
| - # In theory, we could start the span here and use update_name later |
31 |
| - # but that API is "highly discouraged" so we better avoid it. |
32 |
| - environ[_ENVIRON_STARTTIME_KEY] = time_ns() |
33 |
| - |
34 |
| - def _start_response(status, response_headers, *args, **kwargs): |
35 |
| - span = flask_request.environ.get(_ENVIRON_SPAN_KEY) |
36 |
| - if span: |
37 |
| - otel_wsgi.add_response_attributes( |
38 |
| - span, status, response_headers |
| 18 | +def patch(): |
| 19 | + |
| 20 | + class PatchedFlask(flask.Flask): |
| 21 | + |
| 22 | + def __init__(self, *args, **kwargs): |
| 23 | + |
| 24 | + super().__init__(*args, **kwargs) |
| 25 | + |
| 26 | + # Single use variable here to avoid recursion issues. |
| 27 | + wsgi = self.wsgi_app |
| 28 | + |
| 29 | + def wrapped_app(environ, start_response): |
| 30 | + # We want to measure the time for route matching, etc. |
| 31 | + # In theory, we could start the span here and use |
| 32 | + # update_name later but that API is "highly discouraged" so |
| 33 | + # we better avoid it. |
| 34 | + environ[_ENVIRON_STARTTIME_KEY] = time_ns() |
| 35 | + |
| 36 | + def _start_response( |
| 37 | + status, response_headers, *args, **kwargs |
| 38 | + ): |
| 39 | + span = flask.request.environ.get(_ENVIRON_SPAN_KEY) |
| 40 | + if span: |
| 41 | + otel_wsgi.add_response_attributes( |
| 42 | + span, status, response_headers |
| 43 | + ) |
| 44 | + else: |
| 45 | + logger.warning( |
| 46 | + "Flask environ's OpenTelemetry span " |
| 47 | + "missing at _start_response(%s)", |
| 48 | + status, |
| 49 | + ) |
| 50 | + |
| 51 | + return start_response( |
| 52 | + status, response_headers, *args, **kwargs |
| 53 | + ) |
| 54 | + return wsgi(environ, _start_response) |
| 55 | + |
| 56 | + self.wsgi_app = wrapped_app |
| 57 | + |
| 58 | + @self.before_request |
| 59 | + def _before_flask_request(): |
| 60 | + environ = flask.request.environ |
| 61 | + span_name = ( |
| 62 | + flask.request.endpoint |
| 63 | + or otel_wsgi.get_default_span_name(environ) |
| 64 | + ) |
| 65 | + parent_span = propagators.extract( |
| 66 | + otel_wsgi.get_header_from_environ, environ |
39 | 67 | )
|
40 |
| - else: |
41 |
| - logger.warning( |
42 |
| - "Flask environ's OpenTelemetry span missing at _start_response(%s)", |
43 |
| - status, |
| 68 | + |
| 69 | + tracer = trace.tracer() |
| 70 | + |
| 71 | + attributes = otel_wsgi.collect_request_attributes(environ) |
| 72 | + if flask.request.url_rule: |
| 73 | + # For 404 that result from no route found, etc, we don't |
| 74 | + # have a url_rule. |
| 75 | + attributes["http.route"] = flask.request.url_rule.rule |
| 76 | + span = tracer.start_span( |
| 77 | + span_name, |
| 78 | + parent_span, |
| 79 | + kind=trace.SpanKind.SERVER, |
| 80 | + attributes=attributes, |
| 81 | + start_time=environ.get(_ENVIRON_STARTTIME_KEY), |
44 | 82 | )
|
45 |
| - return start_response(status, response_headers, *args, **kwargs) |
46 |
| - |
47 |
| - return wsgi(environ, _start_response) |
48 |
| - |
49 |
| - flask.wsgi_app = wrapped_app |
50 |
| - |
51 |
| - flask.before_request(_before_flask_request) |
52 |
| - flask.teardown_request(_teardown_flask_request) |
53 |
| - |
54 |
| - |
55 |
| -def _before_flask_request(): |
56 |
| - environ = flask_request.environ |
57 |
| - span_name = flask_request.endpoint or otel_wsgi.get_default_span_name( |
58 |
| - environ |
59 |
| - ) |
60 |
| - parent_span = propagators.extract( |
61 |
| - otel_wsgi.get_header_from_environ, environ |
62 |
| - ) |
63 |
| - |
64 |
| - tracer = trace.tracer_source().get_tracer(__name__, __version__) |
65 |
| - |
66 |
| - attributes = otel_wsgi.collect_request_attributes(environ) |
67 |
| - if flask_request.url_rule: |
68 |
| - # For 404 that result from no route found, etc, we don't have a url_rule. |
69 |
| - attributes["http.route"] = flask_request.url_rule.rule |
70 |
| - span = tracer.start_span( |
71 |
| - span_name, |
72 |
| - parent_span, |
73 |
| - kind=trace.SpanKind.SERVER, |
74 |
| - attributes=attributes, |
75 |
| - start_time=environ.get(_ENVIRON_STARTTIME_KEY), |
76 |
| - ) |
77 |
| - activation = tracer.use_span(span, end_on_exit=True) |
78 |
| - activation.__enter__() |
79 |
| - environ[_ENVIRON_ACTIVATION_KEY] = activation |
80 |
| - environ[_ENVIRON_SPAN_KEY] = span |
81 |
| - |
82 |
| - |
83 |
| -def _teardown_flask_request(exc): |
84 |
| - activation = flask_request.environ.get(_ENVIRON_ACTIVATION_KEY) |
85 |
| - if not activation: |
86 |
| - logger.warning( |
87 |
| - "Flask environ's OpenTelemetry activation missing at _teardown_flask_request(%s)", |
88 |
| - exc, |
89 |
| - ) |
90 |
| - return |
91 |
| - |
92 |
| - if exc is None: |
93 |
| - activation.__exit__(None, None, None) |
94 |
| - else: |
95 |
| - activation.__exit__( |
96 |
| - type(exc), exc, getattr(exc, "__traceback__", None) |
97 |
| - ) |
| 83 | + activation = tracer.use_span(span, end_on_exit=True) |
| 84 | + activation.__enter__() |
| 85 | + environ[_ENVIRON_ACTIVATION_KEY] = activation |
| 86 | + environ[_ENVIRON_SPAN_KEY] = span |
| 87 | + |
| 88 | + @self.teardown_request |
| 89 | + def _teardown_flask_request(exc): |
| 90 | + activation = flask.request.environ.get(_ENVIRON_ACTIVATION_KEY) |
| 91 | + if not activation: |
| 92 | + logger.warning( |
| 93 | + "Flask environ's OpenTelemetry activation missing at " |
| 94 | + "_teardown_flask_request(%s)", |
| 95 | + exc, |
| 96 | + ) |
| 97 | + return |
| 98 | + |
| 99 | + if exc is None: |
| 100 | + activation.__exit__(None, None, None) |
| 101 | + else: |
| 102 | + activation.__exit__( |
| 103 | + type(exc), exc, getattr(exc, "__traceback__", None) |
| 104 | + ) |
| 105 | + |
| 106 | + def tato(self): |
| 107 | + pass |
| 108 | + |
| 109 | + flask.Flask = PatchedFlask |
0 commit comments