Initial ScopeManager integration.

flask_opentracing/tracing.py

@@ -24,7 +24,7 @@ def __init__(self, tracer=None, trace_all_requests=False, app=None,
 
         self._trace_all_requests = trace_all_requests
         self._start_span_cb = start_span_cb
-        self._current_spans = {}
+        self._current_scopes = {}
 
         # tracing all requests requires that app != None
         if self._trace_all_requests:
@@ -85,26 +85,29 @@ def get_span(self, request=None):
         """
         if request is None and stack.top:
             request = stack.top.request
-        return self._current_spans.get(request, None)
+
+        scope = self._current_scopes.get(request, None)
+        return None if scope is None else scope.span
 
     def _before_request_fn(self, attributes):
         request = stack.top.request
         operation_name = request.endpoint
         headers = {}
         for k, v in request.headers:
             headers[k.lower()] = v
-        span = None
+
         try:
             span_ctx = self.tracer.extract(opentracing.Format.HTTP_HEADERS,
                                            headers)
-            span = self.tracer.start_span(operation_name=operation_name,
-                                          child_of=span_ctx)
+            scope = self.tracer.start_active_span(operation_name,
+                                                  child_of=span_ctx)
         except (opentracing.InvalidCarrierException,
                 opentracing.SpanContextCorruptedException):
-            span = self.tracer.start_span(operation_name=operation_name)
+            scope = self.tracer.start_active_span(operation_name)
 
-        self._current_spans[request] = span
+        self._current_scopes[request] = scope
 
+        span = scope.span
         span.set_tag(tags.COMPONENT, 'Flask')
         span.set_tag(tags.HTTP_METHOD, request.method)
         span.set_tag(tags.HTTP_URL, request.base_url)
@@ -123,12 +126,12 @@ def _after_request_fn(self, response=None):
 
         # the pop call can fail if the request is interrupted by a
         # `before_request` method so we need a default
-        span = self._current_spans.pop(request, None)
-        if span is not None:
+        scope = self._current_scopes.pop(request, None)
+        if scope is not None:
             if response is not None:
-                span.set_tag(tags.HTTP_STATUS_CODE, response.status_code)
+                scope.span.set_tag(tags.HTTP_STATUS_CODE, response.status_code)
 
-            span.finish()
+            scope.close()
 
     def _call_start_span_cb(self, span, request):
         if self._start_span_cb is None:

tests/test_flask_tracing.py

@@ -18,9 +18,9 @@
 
 
 def flush_spans(tcr):
-    for req in tcr._current_spans:
-        tcr._current_spans[req].finish()
-    tcr._current_spans = {}
+    for req in tcr._current_scopes:
+        tcr._current_scopes[req].close()
+    tcr._current_scopes = {}
 
 
 @app.route('/test')
@@ -40,6 +40,13 @@ def decorated_fn_simple():
     return 'Success again'
 
 
+@app.route('/decorated_child_span_test')
+@tracing.trace()
+def decorated_fn_with_child_span():
+    with tracing.tracer.start_active_span('child'):
+        return 'Success'
+
+
 @app.route('/wire')
 def send_request():
     span = tracing_all.get_span()
@@ -63,15 +70,19 @@ def test_span_creation(self):
             assert tracing_all.get_span(request)
             assert not tracing.get_span(request)
             assert tracing_deferred.get_span(request)
+
+            active_span = tracing_all.tracer.active_span
+            assert tracing_all.get_span(request) is active_span
+
             flush_spans(tracing_all)
             flush_spans(tracing_deferred)
 
     def test_span_deletion(self):
-        assert not tracing_all._current_spans
-        assert not tracing_deferred._current_spans
+        assert not tracing_all._current_scopes
+        assert not tracing_deferred._current_scopes
         test_app.get('/test')
-        assert not tracing_all._current_spans
-        assert not tracing_deferred._current_spans
+        assert not tracing_all._current_scopes
+        assert not tracing_deferred._current_scopes
 
     def test_span_tags(self):
         test_app.get('/another_test_simple')
@@ -90,9 +101,9 @@ def test_requests_distinct(self):
             app.preprocess_request()
         with app.test_request_context('/test'):
             app.preprocess_request()
-            second_span = tracing_all._current_spans.pop(request)
-            assert second_span
-            second_span.finish()
+            second_scope = tracing_all._current_scopes.pop(request)
+            assert second_scope
+            second_scope.close()
             assert not tracing_all.get_span(request)
         # clear current spans
         flush_spans(tracing_all)
@@ -102,16 +113,20 @@ def test_decorator(self):
         with app.test_request_context('/another_test'):
             app.preprocess_request()
             assert not tracing.get_span(request)
-            assert len(tracing_deferred._current_spans) == 1
-            assert len(tracing_all._current_spans) == 1
+            assert len(tracing_deferred._current_scopes) == 1
+            assert len(tracing_all._current_scopes) == 1
+
+            active_span = tracing_all.tracer.active_span
+            assert tracing_all.get_span(request) is active_span
+
         flush_spans(tracing)
         flush_spans(tracing_all)
         flush_spans(tracing_deferred)
 
         test_app.get('/another_test')
-        assert not tracing_all._current_spans
-        assert not tracing._current_spans
-        assert not tracing_deferred._current_spans
+        assert not tracing_all._current_scopes
+        assert not tracing._current_scopes
+        assert not tracing_deferred._current_scopes
 
     def test_over_wire(self):
         rv = test_app.get('/wire')
@@ -122,6 +137,15 @@ def test_over_wire(self):
         assert spans[0].context.trace_id == spans[1].context.trace_id
         assert spans[0].parent_id == spans[1].context.span_id
 
+    def test_child_span(self):
+        rv = test_app.get('/decorated_child_span_test')
+        assert '200' in str(rv.status_code)
+
+        spans = tracing.tracer.finished_spans()
+        assert len(spans) == 2
+        assert spans[0].context.trace_id == spans[1].context.trace_id
+        assert spans[0].parent_id == spans[1].context.span_id
+
 
 class TestTracingStartSpanCallback(unittest.TestCase):
     def test_simple(self):