Repository: scoutapp/scout_apm_python (python)

Issue:
Add span context manager
As per c46c79523358f88180caac77be8c5761e6482bba  - would make several pieces of integrations simpler, removing unnecessary `try/finally`

Pull request description:
Add TrackedRequest.span() context manager
Fixes #557. Reduces duplication.

Patch:
diff --git a/src/scout_apm/async_/instruments/jinja2.py b/src/scout_apm/async_/instruments/jinja2.py
index 4e3408f..41820e7 100644
--- a/src/scout_apm/async_/instruments/jinja2.py
+++ b/src/scout_apm/async_/instruments/jinja2.py
@@ -9,9 +9,6 @@ from scout_apm.core.tracked_request import TrackedRequest
 @wrapt.decorator
 async def wrapped_render_async(wrapped, instance, args, kwargs):
     tracked_request = TrackedRequest.instance()
-    span = tracked_request.start_span(operation="Template/Render")
-    span.tag("name", instance.name)
-    try:
+    with tracked_request.span(operation="Template/Render") as span:
+        span.tag("name", instance.name)
         return await wrapped(*args, **kwargs)
-    finally:
-        tracked_request.stop_span()
diff --git a/src/scout_apm/async_/starlette.py b/src/scout_apm/async_/starlette.py
index 0d6c7c3..e4c4292 100644
--- a/src/scout_apm/async_/starlette.py
+++ b/src/scout_apm/async_/starlette.py
@@ -81,14 +81,12 @@ def install_background_instrumentation():
     async def wrapped_background_call(wrapped, instance, args, kwargs):
         tracked_request = TrackedRequest.instance()
         tracked_request.is_real_request = True
-        tracked_request.start_span(
+
+        with tracked_request.span(
             operation="Job/{}.{}".format(
                 instance.func.__module__, instance.func.__qualname__
             )
-        )
-        try:
+        ):
             return await wrapped(*args, **kwargs)
-        finally:
-            tracked_request.stop_span()
 
     BackgroundTask.__call__ = wrapped_background_call(BackgroundTask.__call__)
diff --git a/src/scout_apm/bottle.py b/src/scout_apm/bottle.py
index f7c7d58..d4ea8fb 100644
--- a/src/scout_apm/bottle.py
+++ b/src/scout_apm/bottle.py
@@ -60,9 +60,6 @@ def wrap_callback(wrapped, instance, args, kwargs):
         controller_name = "/home"
     if not controller_name.startswith("/"):
         controller_name = "/" + controller_name
-    tracked_request.start_span(
-        operation="Controller{}".format(controller_name), should_capture_backtrace=False
-    )
 
     if scout_config.value("collect_remote_ip"):
         # Determine a remote IP to associate with the request. The
@@ -84,14 +81,15 @@ def wrap_callback(wrapped, instance, args, kwargs):
         amazon_queue_time = request.headers.get("x-amzn-trace-id", "")
         track_amazon_request_queue_time(amazon_queue_time, tracked_request)
 
-    try:
-        value = wrapped(*args, **kwargs)
-    except Exception:
-        tracked_request.tag("error", "true")
-        raise
-    else:
-        if 500 <= response.status_code <= 599:
+    with tracked_request.span(
+        operation="Controller{}".format(controller_name), should_capture_backtrace=False
+    ):
+        try:
+            value = wrapped(*args, **kwargs)
+        except Exception:
             tracked_request.tag("error", "true")
-        return value
-    finally:
-        tracked_request.stop_span()
+            raise
+        else:
+            if 500 <= response.status_code <= 599:
+                tracked_request.tag("error", "true")
+            return value
diff --git a/src/scout_apm/core/stacktracer.py b/src/scout_apm/core/stacktracer.py
index a262d72..120c451 100644
--- a/src/scout_apm/core/stacktracer.py
+++ b/src/scout_apm/core/stacktracer.py
@@ -18,14 +18,11 @@ def trace_method(cls, method_name=None):
                 operation = operation + "/" + detail["name"]
 
             tracked_request = TrackedRequest.instance()
-            span = tracked_request.start_span(operation=operation)
-            for key, value in detail.items():
-                span.tag(key, value)
+            with tracked_request.span(operation=operation) as span:
+                for key, value in detail.items():
+                    span.tag(key, value)
 
-            try:
                 return wrapped(*args, **kwargs)
-            finally:
-                tracked_request.stop_span()
 
         setattr(cls, method_to_patch, wrapper(getattr(cls, method_to_patch)))
 
diff --git a/src/scout_apm/core/tracked_request.py b/src/scout_apm/core/tracked_request.py
index 9be78c0..da79cab 100644
--- a/src/scout_apm/core/tracked_request.py
+++ b/src/scout_apm/core/tracked_request.py
@@ -3,6 +3,7 @@ from __future__ import absolute_import, division, print_function, unicode_litera
 
 import datetime as dt
 import logging
+from contextlib import contextmanager
 from uuid import uuid4
 
 from scout_apm.core import backtrace, objtrace
@@ -122,6 +123,14 @@ class TrackedRequest(object):
         if len(self.active_spans) == 0:
             self.finish()
 
+    @contextmanager
+    def span(self, *args, **kwargs):
+        span = self.start_span(*args, **kwargs)
+        try:
+            yield span
+        finally:
+            self.stop_span()
+
     def current_span(self):
         if self.active_spans:
             return self.active_spans[-1]
diff --git a/src/scout_apm/django/middleware.py b/src/scout_apm/django/middleware.py
index 004365a..2968910 100644
--- a/src/scout_apm/django/middleware.py
+++ b/src/scout_apm/django/middleware.py
@@ -165,9 +165,6 @@ class MiddlewareTimingMiddleware(object):
 
         tracked_request = TrackedRequest.instance()
 
-        tracked_request.start_span(
-            operation="Middleware", should_capture_backtrace=False
-        )
         queue_time = request.META.get("HTTP_X_QUEUE_START") or request.META.get(
             "HTTP_X_REQUEST_START", ""
         )
@@ -177,10 +174,11 @@ class MiddlewareTimingMiddleware(object):
                 request.META.get("HTTP_X_AMZN_TRACE_ID", ""), tracked_request
             )
 
-        try:
+        with tracked_request.span(
+            operation="Middleware",
+            should_capture_backtrace=False,
+        ):
             return self.get_response(request)
-        finally:
-            tracked_request.stop_span()
 
 
 class ViewTimingMiddleware(object):
@@ -210,15 +208,12 @@ class ViewTimingMiddleware(object):
 
         # This operation name won't be recorded unless changed later in
         # process_view
-        tracked_request.start_span(operation="Unknown", should_capture_backtrace=False)
-        try:
+        with tracked_request.span(operation="Unknown", should_capture_backtrace=False):
             response = self.get_response(request)
             track_request_view_data(request, tracked_request)
             if 500 <= response.status_code <= 599:
                 tracked_request.tag("error", "true")
             return response
-        finally:
-            tracked_request.stop_span()
 
     def process_view(self, request, view_func, view_args, view_kwargs):
         """
diff --git a/src/scout_apm/flask/__init__.py b/src/scout_apm/flask/__init__.py
index abd2cb8..ad4dc0e 100644
--- a/src/scout_apm/flask/__init__.py
+++ b/src/scout_apm/flask/__init__.py
@@ -44,24 +44,23 @@ class ScoutApm(object):
         tracked_request = TrackedRequest.instance()
         tracked_request.is_real_request = True
         request._scout_tracked_request = tracked_request
-        span = tracked_request.start_span(
-            operation=operation, should_capture_backtrace=False
-        )
-        request._scout_view_span = span
 
         werkzeug_track_request_data(request, tracked_request)
 
-        try:
-            response = wrapped(*args, **kwargs)
-        except Exception as exc:
-            tracked_request.tag("error", "true")
-            raise exc
-        else:
-            if 500 <= response.status_code <= 599:
+        with tracked_request.span(
+            operation=operation, should_capture_backtrace=False
+        ) as span:
+            request._scout_view_span = span
+
+            try:
+                response = wrapped(*args, **kwargs)
+            except Exception as exc:
                 tracked_request.tag("error", "true")
-            return response
-        finally:
-            tracked_request.stop_span()
+                raise exc
+            else:
+                if 500 <= response.status_code <= 599:
+                    tracked_request.tag("error", "true")
+                return response
 
     @wrapt.decorator
     def wrapped_preprocess_request(self, wrapped, instance, args, kwargs):
@@ -78,11 +77,8 @@ class ScoutApm(object):
         if not have_before_request_funcs:
             return wrapped(*args, **kwargs)
 
-        tracked_request.start_span("PreprocessRequest", should_capture_backtrace=False)
-        try:
+        with tracked_request.span("PreprocessRequest", should_capture_backtrace=False):
             return wrapped(*args, **kwargs)
-        finally:
-            tracked_request.stop_span()
 
     def extract_flask_settings(self):
         """
diff --git a/src/scout_apm/instruments/elasticsearch.py b/src/scout_apm/instruments/elasticsearch.py
index 9747443..5b5eb60 100644
--- a/src/scout_apm/instruments/elasticsearch.py
+++ b/src/scout_apm/instruments/elasticsearch.py
@@ -132,12 +132,8 @@ def wrap_client_index_method(wrapped, instance, args, kwargs):
     camel_name = "".join(c.title() for c in wrapped.__name__.split("_"))
     operation = "Elasticsearch/{}/{}".format(index, camel_name)
     tracked_request = TrackedRequest.instance()
-    tracked_request.start_span(operation=operation, ignore_children=True)
-
-    try:
+    with tracked_request.span(operation=operation, ignore_children=True):
         return wrapped(*args, **kwargs)
-    finally:
-        tracked_request.stop_span()
 
 
 @wrapt.decorator
@@ -145,12 +141,8 @@ def wrap_client_method(wrapped, instance, args, kwargs):
     camel_name = "".join(c.title() for c in wrapped.__name__.split("_"))
     operation = "Elasticsearch/{}".format(camel_name)
     tracked_request = TrackedRequest.instance()
-    tracked_request.start_span(operation=operation, ignore_children=True)
-
-    try:
+    with tracked_request.span(operation=operation, ignore_children=True):
         return wrapped(*args, **kwargs)
-    finally:
-        tracked_request.stop_span()
 
 
 have_patched_transport = False
@@ -217,11 +209,8 @@ def wrapped_perform_request(wrapped, instance, args, kwargs):
         op = "Unknown"
 
     tracked_request = TrackedRequest.instance()
-    tracked_request.start_span(
-        operation="Elasticsearch/{}".format(op), ignore_children=True
-    )
-
-    try:
+    with tracked_request.span(
+        operation="Elasticsearch/{}".format(op),
+        ignore_children=True,
+    ):
         return wrapped(*args, **kwargs)
-    finally:
-        tracked_request.stop_span()
diff --git a/src/scout_apm/instruments/jinja2.py b/src/scout_apm/instruments/jinja2.py
index 08bf54a..31970d5 100644
--- a/src/scout_apm/instruments/jinja2.py
+++ b/src/scout_apm/instruments/jinja2.py
@@ -69,12 +69,9 @@ def ensure_installed():
 @wrapt.decorator
 def wrapped_render(wrapped, instance, args, kwargs):
     tracked_request = TrackedRequest.instance()
-    span = tracked_request.start_span(operation="Template/Render")
-    span.tag("name", instance.name)
-    try:
+    with tracked_request.span(operation="Template/Render") as span:
+        span.tag("name", instance.name)
         return wrapped(*args, **kwargs)
-    finally:
-        tracked_request.stop_span()
 
 
 @wrapt.decorator
diff --git a/src/scout_apm/instruments/pymongo.py b/src/scout_apm/instruments/pymongo.py
index 991abe1..80474be 100644
--- a/src/scout_apm/instruments/pymongo.py
+++ b/src/scout_apm/instruments/pymongo.py
@@ -88,10 +88,6 @@ def wrap_collection_method(wrapped, instance, args, kwargs):
     tracked_request = TrackedRequest.instance()
     camel_name = "".join(c.title() for c in wrapped.__name__.split("_"))
     operation = "MongoDB/{}.{}".format(instance.name, camel_name)
-    span = tracked_request.start_span(operation=operation, ignore_children=True)
-    span.tag("name", instance.name)
-
-    try:
+    with tracked_request.span(operation=operation, ignore_children=True) as span:
+        span.tag("name", instance.name)
         return wrapped(*args, **kwargs)
-    finally:
-        tracked_request.stop_span()
diff --git a/src/scout_apm/instruments/redis.py b/src/scout_apm/instruments/redis.py
index e3baf4b..1262ab0 100644
--- a/src/scout_apm/instruments/redis.py
+++ b/src/scout_apm/instruments/redis.py
@@ -67,20 +67,12 @@ def wrapped_execute_command(wrapped, instance, args, kwargs):
         op = "Unknown"
 
     tracked_request = TrackedRequest.instance()
-    tracked_request.start_span(operation="Redis/{}".format(op))
-
-    try:
+    with tracked_request.span(operation="Redis/{}".format(op)):
         return wrapped(*args, **kwargs)
-    finally:
-        tracked_request.stop_span()
 
 
 @wrapt.decorator
 def wrapped_execute(wrapped, instance, args, kwargs):
     tracked_request = TrackedRequest.instance()
-    tracked_request.start_span(operation="Redis/MULTI")
-
-    try:
+    with tracked_request.span(operation="Redis/MULTI"):
         return wrapped(*args, **kwargs)
-    finally:
-        tracked_request.stop_span()
diff --git a/src/scout_apm/instruments/urllib3.py b/src/scout_apm/instruments/urllib3.py
index 1019041..67b6985 100644
--- a/src/scout_apm/instruments/urllib3.py
+++ b/src/scout_apm/instruments/urllib3.py
@@ -58,10 +58,6 @@ def wrapped_urlopen(wrapped, instance, args, kwargs):
         url = "Unknown"
 
     tracked_request = TrackedRequest.instance()
-    span = tracked_request.start_span(operation="HTTP/{}".format(method))
-    span.tag("url", text_type(url))
-
-    try:
+    with tracked_request.span(operation="HTTP/{}".format(method)) as span:
+        span.tag("url", text_type(url))
         return wrapped(*args, **kwargs)
-    finally:
-        tracked_request.stop_span()
diff --git a/src/scout_apm/pyramid.py b/src/scout_apm/pyramid.py
index 53081d3..962dfb0 100644
--- a/src/scout_apm/pyramid.py
+++ b/src/scout_apm/pyramid.py
@@ -29,11 +29,10 @@ def includeme(config):
 def instruments(handler, registry):
     def scout_tween(request):
         tracked_request = TrackedRequest.instance()
-        span = tracked_request.start_span(
-            operation="Controller/Pyramid", should_capture_backtrace=False
-        )
 
-        try:
+        with tracked_request.span(
+            operation="Controller/Pyramid", should_capture_backtrace=False
+        ) as span:
             path = request.path
             # mixed() returns values as *either* single items or lists
             url_params = [
@@ -78,8 +77,6 @@ def instruments(handler, registry):
 
             if 500 <= response.status_code <= 599:
                 tracked_request.tag("error", "true")
-        finally:
-            tracked_request.stop_span()
 
         return response
 
diff --git a/src/scout_apm/rq.py b/src/scout_apm/rq.py
index 588498b..164e5a6 100644
--- a/src/scout_apm/rq.py
+++ b/src/scout_apm/rq.py
@@ -68,11 +68,10 @@ def wrap_perform(wrapped, instance, args, kwargs):
     tracked_request.tag("queue", instance.origin)
     queue_time = (dt.datetime.utcnow() - instance.enqueued_at).total_seconds()
     tracked_request.tag("queue_time", queue_time)
-    tracked_request.start_span(operation="Job/{}".format(instance.func_name))
-    try:
-        return wrapped(*args, **kwargs)
-    except Exception:
-        tracked_request.tag("error", "true")
-        raise
-    finally:
-        tracked_request.stop_span()
+
+    with tracked_request.span(operation="Job/{}".format(instance.func_name)):
+        try:
+            return wrapped(*args, **kwargs)
+        except Exception:
+            tracked_request.tag("error", "true")
+            raise


Test patch:
diff --git a/tests/unit/core/test_tracked_request.py b/tests/unit/core/test_tracked_request.py
index 2db95e7..fcb9062 100644
--- a/tests/unit/core/test_tracked_request.py
+++ b/tests/unit/core/test_tracked_request.py
@@ -54,9 +54,8 @@ def test_span_repr(tracked_request):
 
 
 def test_tag_span(tracked_request):
-    span = tracked_request.start_span(operation="myoperation")
-    span.tag("foo", "bar")
-    tracked_request.stop_span()
+    with tracked_request.span(operation="myoperation") as span:
+        span.tag("foo", "bar")
 
     assert tracked_request.complete_spans[0].tags["foo"] == "bar"
 
@@ -94,16 +93,15 @@ def test_tags_allocations_for_spans_no_objtrace_extension(tracked_request):
 
 
 def test_start_span_does_not_ignore_children(tracked_request):
-    tracked_request.start_span(operation="parent")
-    child1 = tracked_request.start_span(operation="myoperation")
+    with tracked_request.span(operation="parent"):
+        with tracked_request.span(operation="myoperation") as child1:
+            with tracked_request.span(operation="myoperation") as child2:
+                pass
+
     assert not child1.ignore
     assert not child1.ignore_children
-    child2 = tracked_request.start_span(operation="myoperation")
     assert not child2.ignore
     assert not child2.ignore_children
-    tracked_request.stop_span()
-    tracked_request.stop_span()
-    tracked_request.stop_span()
     assert len(tracked_request.complete_spans) == 3
     assert tracked_request.complete_spans[2].operation == "parent"