Add custom span name hook for Flask

instrumentation/opentelemetry-instrumentation-flask/src/opentelemetry/instrumentation/flask/__init__.py

@@ -20,7 +20,7 @@ This library builds on the OpenTelemetry WSGI middleware to track web requests
 in Flask applications. In addition to opentelemetry-instrumentation-wsgi, it supports
 flask-specific features such as:
 
-* The Flask endpoint name is used as the Span name.
+* The Flask url rule pattern is used as the Span name.
 * The ``http.route`` Span attribute is set so that one can see which URL rule
   matched a request.
 
@@ -74,6 +74,13 @@ def get_excluded_urls():
 
 _excluded_urls = get_excluded_urls()
 
+def get_default_span_name():
+    span_name = None
+    try:
+        span_name = flask.request.url_rule.rule
+    except AttributeError:
+        span_name = otel_wsgi.get_default_span_name(flask.request.environ)
+    return span_name
 
 def _rewrapped_app(wsgi_app):
     def _wrapped_app(environ, start_response):
@@ -104,45 +111,39 @@ def _rewrapped_app(wsgi_app):
 
     return _wrapped_app
 
+def _wrapped_before_request(name_callback):
+    def _before_request():
+        if _excluded_urls.url_disabled(flask.request.url):
+            return
 
-def _before_request():
-    if _excluded_urls.url_disabled(flask.request.url):
-        return
+        environ = flask.request.environ
+        span_name = name_callback()
+        token = context.attach(
+            propagators.extract(otel_wsgi.carrier_getter, environ)
+        )
 
-    environ = flask.request.environ
-    span_name = None
-    try:
-        span_name = flask.request.url_rule.rule
-    except AttributeError:
-        pass
-    if span_name is None:
-        span_name = otel_wsgi.get_default_span_name(environ)
-    token = context.attach(
-        propagators.extract(otel_wsgi.carrier_getter, environ)
-    )
+        tracer = trace.get_tracer(__name__, __version__)
 
-    tracer = trace.get_tracer(__name__, __version__)
-
-    span = tracer.start_span(
-        span_name,
-        kind=trace.SpanKind.SERVER,
-        start_time=environ.get(_ENVIRON_STARTTIME_KEY),
-    )
-    if span.is_recording():
-        attributes = otel_wsgi.collect_request_attributes(environ)
-        if flask.request.url_rule:
-            # For 404 that result from no route found, etc, we
-            # don't have a url_rule.
-            attributes["http.route"] = flask.request.url_rule.rule
-        for key, value in attributes.items():
-            span.set_attribute(key, value)
-
-    activation = tracer.use_span(span, end_on_exit=True)
-    activation.__enter__()
-    environ[_ENVIRON_ACTIVATION_KEY] = activation
-    environ[_ENVIRON_SPAN_KEY] = span
-    environ[_ENVIRON_TOKEN] = token
+        span = tracer.start_span(
+            span_name,
+            kind=trace.SpanKind.SERVER,
+            start_time=environ.get(_ENVIRON_STARTTIME_KEY),
+        )
+        if span.is_recording():
+            attributes = otel_wsgi.collect_request_attributes(environ)
+            if flask.request.url_rule:
+                # For 404 that result from no route found, etc, we
+                # don't have a url_rule.
+                attributes["http.route"] = flask.request.url_rule.rule
+            for key, value in attributes.items():
+                span.set_attribute(key, value)
 
+        activation = tracer.use_span(span, end_on_exit=True)
+        activation.__enter__()
+        environ[_ENVIRON_ACTIVATION_KEY] = activation
+        environ[_ENVIRON_SPAN_KEY] = span
+        environ[_ENVIRON_TOKEN] = token
+    return _before_request
 
 def _teardown_request(exc):
     if _excluded_urls.url_disabled(flask.request.url):
@@ -173,6 +174,8 @@ class _InstrumentedFlask(flask.Flask):
         self._original_wsgi_ = self.wsgi_app
         self.wsgi_app = _rewrapped_app(self.wsgi_app)
 
+        _before_request = _wrapped_before_request(get_default_span_name)
+        self._before_request = _before_request
         self.before_request(_before_request)
         self.teardown_request(_teardown_request)
 
@@ -188,7 +191,7 @@ class FlaskInstrumentor(BaseInstrumentor):
         self._original_flask = flask.Flask
         flask.Flask = _InstrumentedFlask
 
-    def instrument_app(self, app):  # pylint: disable=no-self-use
+    def instrument_app(self, app, name_callback=get_default_span_name):  # pylint: disable=no-self-use
         if not hasattr(app, "_is_instrumented"):
             app._is_instrumented = False
 
@@ -196,6 +199,8 @@ class FlaskInstrumentor(BaseInstrumentor):
             app._original_wsgi_app = app.wsgi_app
             app.wsgi_app = _rewrapped_app(app.wsgi_app)
 
+            _before_request = _wrapped_before_request(name_callback)
+            app._before_request = _before_request
             app.before_request(_before_request)
             app.teardown_request(_teardown_request)
             app._is_instrumented = True
@@ -215,7 +220,7 @@ class FlaskInstrumentor(BaseInstrumentor):
             app.wsgi_app = app._original_wsgi_app
 
             # FIXME add support for other Flask blueprints that are not None
-            app.before_request_funcs[None].remove(_before_request)
+            app.before_request_funcs[None].remove(app._before_request)
             app.teardown_request_funcs[None].remove(_teardown_request)
             del app._original_wsgi_app
 

instrumentation/opentelemetry-instrumentation-flask/tests/test_programmatic.py

@@ -178,3 +178,28 @@ class TestProgrammatic(InstrumentationTest, TestBase, WsgiTestBase):
         self.client.get("/excluded_noarg2")
         span_list = self.memory_exporter.get_finished_spans()
         self.assertEqual(len(span_list), 1)
+
+class TestProgrammaticCustomSpanName(InstrumentationTest, TestBase, WsgiTestBase):
+    def setUp(self):
+        super().setUp()
+
+        def custom_span_name():
+            return "flask-custom-span-name"
+
+        self.app = Flask(__name__)
+
+        FlaskInstrumentor().instrument_app(self.app, name_callback=custom_span_name)
+
+        self._common_initialization()
+
+    def tearDown(self):
+        super().tearDown()
+        with self.disable_logging():
+            FlaskInstrumentor().uninstrument_app(self.app)
+
+    def test_custom_span_name(self):
+        self.client.get("/hello/123")
+
+        span_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(span_list), 1)
+        self.assertEqual(span_list[0].name, "flask-custom-span-name")