flask: allow custom resource names (#410)

Fixes #353

Author: clutchski

ddtrace/contrib/flask/middleware.py

@@ -21,6 +21,9 @@
 log = logging.getLogger(__name__)
 
 
+SPAN_NAME = 'flask.request'
+
+
 class TraceMiddleware(object):
 
     def __init__(self, app, tracer, service="flask", use_signals=True, distributed_tracing=False):
@@ -112,7 +115,7 @@ def _start_span(self):
                 self._tracer.context_provider.activate(context)
         try:
             g.flask_datadog_span = self._tracer.trace(
-                "flask.request",
+                SPAN_NAME,
                 service=self._service,
                 span_type=http.TYPE,
             )
@@ -161,8 +164,12 @@ def _finish_span(self, span, exception=None):
             endpoint = request.endpoint or code
             url = request.base_url or ''
 
-        resource = endpoint or code
-        span.resource = compat.to_unicode(resource).lower()
+        # Let users specify their own resource in middleware if they so desire.
+        # See case https://github.com/DataDog/dd-trace-py/issues/353
+        if span.resource == SPAN_NAME:
+            resource = endpoint or code
+            span.resource = compat.to_unicode(resource).lower()
+
         span.set_tag(http.URL, compat.to_unicode(url))
         span.set_tag(http.STATUS_CODE, code)
         span.set_tag(http.METHOD, method)

tests/contrib/flask/test_flask.py

@@ -74,6 +74,13 @@ def child():
         span.set_tag('a', 'b')
         return 'child'
 
+@app.route("/custom_span")
+def custom_span():
+    span = tracer.current_span()
+    assert span
+    span.resource = "overridden"
+    return 'hiya'
+
 
 def unicode_view():
     return u'üŋïĉóđē'
@@ -395,3 +402,20 @@ def test_propagation(self):
         eq_(s.trace_id, 1234)
         eq_(s.parent_id, 4567)
         eq_(s.get_metric(SAMPLING_PRIORITY_KEY), 2)
+
+    def test_custom_span(self):
+        rv = app.get('/custom_span')
+        eq_(rv.status_code, 200)
+        # ensure trace worked
+        assert not tracer.current_span(), tracer.current_span().pprint()
+        spans = writer.pop()
+        eq_(len(spans), 1)
+        s = spans[0]
+        eq_(s.service, service)
+        eq_(s.resource, "overridden")
+        eq_(s.error, 0)
+        eq_(s.meta.get(http.STATUS_CODE), '200')
+        eq_(s.meta.get(http.METHOD), 'GET')
+
+
+