Repository: scoutapp/scout_apm_python
Pull request #360 (2019-11-05 15:37:59)

Convert Bottle decorator to wrapt
Should have been done as part of #311. Defining the decorator as an outer function also provides a small performance improvement.


Patch:
diff --git a/src/scout_apm/bottle.py b/src/scout_apm/bottle.py
index 746f7bf..ca16689 100644
--- a/src/scout_apm/bottle.py
+++ b/src/scout_apm/bottle.py
@@ -1,6 +1,7 @@
 # coding=utf-8
 from __future__ import absolute_import, division, print_function, unicode_literals
 
+import wrapt
 from bottle import request
 
 import scout_apm.core
@@ -34,57 +35,55 @@ class ScoutPlugin(object):
     def apply(self, callback, context):
         if not scout_config.value("monitor"):
             return callback
+        return wrap_callback(callback)
 
-        def wrapper(*args, **kwargs):
-            tracked_request = TrackedRequest.instance()
-            tracked_request.is_real_request = True
 
-            path = request.path
-            # allitems() is an undocumented bottle internal
-            tracked_request.tag(
-                "path", create_filtered_path(path, request.query.allitems())
-            )
-            if ignore_path(path):
-                tracked_request.tag("ignore_transaction", True)
+@wrapt.decorator
+def wrap_callback(wrapped, instance, args, kwargs):
+    tracked_request = TrackedRequest.instance()
+    tracked_request.is_real_request = True
 
-            if request.route.name is not None:
-                controller_name = request.route.name
-            else:
-                controller_name = request.route.rule
-            if controller_name == "/":
-                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,
-            )
+    path = request.path
+    # allitems() is an undocumented bottle internal
+    tracked_request.tag("path", create_filtered_path(path, request.query.allitems()))
+    if ignore_path(path):
+        tracked_request.tag("ignore_transaction", True)
 
-            # Determine a remote IP to associate with the request. The
-            # value is spoofable by the requester so this is not suitable
-            # to use in any security sensitive context.
-            user_ip = (
-                request.headers.get("x-forwarded-for", "").split(",")[0]
-                or request.headers.get("client-ip", "").split(",")[0]
-                or request.environ.get("REMOTE_ADDR")
-            )
-            tracked_request.tag("user_ip", user_ip)
+    if request.route.name is not None:
+        controller_name = request.route.name
+    else:
+        controller_name = request.route.rule
+    if controller_name == "/":
+        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
+    )
 
-            tracked_queue_time = False
-            queue_time = request.headers.get(
-                "x-queue-start", ""
-            ) or request.headers.get("x-request-start", "")
-            tracked_queue_time = track_request_queue_time(queue_time, tracked_request)
-            if not tracked_queue_time:
-                amazon_queue_time = request.headers.get("x-amzn-trace-id", "")
-                track_amazon_request_queue_time(amazon_queue_time, tracked_request)
+    # Determine a remote IP to associate with the request. The
+    # value is spoofable by the requester so this is not suitable
+    # to use in any security sensitive context.
+    user_ip = (
+        request.headers.get("x-forwarded-for", "").split(",")[0]
+        or request.headers.get("client-ip", "").split(",")[0]
+        or request.environ.get("REMOTE_ADDR")
+    )
+    tracked_request.tag("user_ip", user_ip)
 
-            try:
-                return callback(*args, **kwargs)
-            except Exception:
-                tracked_request.tag("error", "true")
-                raise
-            finally:
-                tracked_request.stop_span()
+    tracked_queue_time = False
+    queue_time = request.headers.get("x-queue-start", "") or request.headers.get(
+        "x-request-start", ""
+    )
+    tracked_queue_time = track_request_queue_time(queue_time, tracked_request)
+    if not tracked_queue_time:
+        amazon_queue_time = request.headers.get("x-amzn-trace-id", "")
+        track_amazon_request_queue_time(amazon_queue_time, tracked_request)
 
-        return wrapper
+    try:
+        return wrapped(*args, **kwargs)
+    except Exception:
+        tracked_request.tag("error", "true")
+        raise
+    finally:
+        tracked_request.stop_span()
diff --git a/src/scout_apm/instruments/elasticsearch.py b/src/scout_apm/instruments/elasticsearch.py
index b255134..4450e04 100644
--- a/src/scout_apm/instruments/elasticsearch.py
+++ b/src/scout_apm/instruments/elasticsearch.py
@@ -70,11 +70,11 @@ class Instrument(object):
             logger.info("Elasticsearch instruments are not installed. Skipping.")
             return False
 
-        self.__class__.installed = False
-
         self.uninstrument_client()
         self.uninstrument_transport()
 
+        self.__class__.installed = False
+
     def instrument_client(self):
         for name in self.__class__.CLIENT_METHODS:
             method = getattr(Elasticsearch, name, None)
diff --git a/src/scout_apm/instruments/pymongo.py b/src/scout_apm/instruments/pymongo.py
index 0119895..2e53c5d 100644
--- a/src/scout_apm/instruments/pymongo.py
+++ b/src/scout_apm/instruments/pymongo.py
@@ -3,10 +3,14 @@ from __future__ import absolute_import, division, print_function, unicode_litera
 
 import logging
 
-# Used in the exec() call below.
-import wrapt  # noqa: F401
+import wrapt
 
-from scout_apm.core.tracked_request import TrackedRequest  # noqa: F401
+from scout_apm.core.tracked_request import TrackedRequest
+
+try:
+    from pymongo.collection import Collection
+except ImportError:
+    Collection = None
 
 logger = logging.getLogger(__name__)
 
@@ -49,9 +53,7 @@ class Instrument(object):
     ]
 
     def installable(self):
-        try:
-            from pymongo.collection import Collection  # noqa: F401
-        except ImportError:
+        if Collection is None:
             logger.info("Unable to import for PyMongo instruments")
             return False
         if self.installed:
@@ -66,44 +68,10 @@ class Instrument(object):
 
         self.__class__.installed = True
 
-        try:
-            from pymongo.collection import Collection  # noqa: F401
-        # There is no way the import can fail if self.installable() succeeded.
-        except ImportError:  # pragma: no cover
-            logger.info(
-                "Unable to import for PyMongo instruments. Instrument install failed."
-            )
-            return False
-
-        for method_str in self.__class__.PYMONGO_METHODS:
-            try:
-                code_str = """
-@wrapt.decorator
-def wrapped_{method_str}(wrapped, instance, args, kwargs):
-    tracked_request = TrackedRequest.instance()
-    name = '/'.join(['MongoDB', instance.name, '{camel_name}'])
-    span = tracked_request.start_span(operation=name, ignore_children=True)
-    span.tag('name', instance.name)
-
-    try:
-        return wrapped(*args, **kwargs)
-    finally:
-        tracked_request.stop_span()
-
-Collection.{method_str} = wrapped_{method_str}(Collection.{method_str})
-""".format(
-                    method_str=method_str,
-                    camel_name="".join(c.title() for c in method_str.split("_")),
-                )
-
-                exec(code_str)
-                logger.info("Instrumented PyMongo Collection.%s", method_str)
-
-            except Exception as e:
-                logger.warning(
-                    "Unable to instrument for PyMongo Collection.%s: %r", method_str, e
-                )
-                return False
+        for name in self.__class__.PYMONGO_METHODS:
+            method = getattr(Collection, name, None)
+            if method is not None:
+                setattr(Collection, name, wrap_collection_method(method))
         return True
 
     def uninstall(self):
@@ -111,9 +79,23 @@ Collection.{method_str} = wrapped_{method_str}(Collection.{method_str})
             logger.info("PyMongo instruments are not installed. Skipping.")
             return False
 
+        for name in self.__class__.PYMONGO_METHODS:
+            method = getattr(Collection, name, None)
+            if method is not None:
+                setattr(Collection, name, method.__wrapped__)
+
         self.__class__.installed = False
 
-        from pymongo.collection import Collection
 
-        for method_str in self.__class__.PYMONGO_METHODS:
-            setattr(Collection, method_str, getattr(Collection, method_str).__wrapped__)
+@wrapt.decorator
+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:
+        return wrapped(*args, **kwargs)
+    finally:
+        tracked_request.stop_span()


Test patch:
diff --git a/tests/integration/django_app.py b/tests/integration/django_app.py
index bdc8e81..728fbf3 100644
--- a/tests/integration/django_app.py
+++ b/tests/integration/django_app.py
@@ -1,9 +1,8 @@
 # coding: utf-8
 from __future__ import absolute_import, division, print_function, unicode_literals
 
-from functools import wraps
-
 import django
+import wrapt
 from django.conf import settings
 from django.template.response import TemplateResponse
 
@@ -135,14 +134,11 @@ def template(request):
     return HttpResponse(template.render(context))
 
 
-def exclaimify_template_response_name(func):
-    @wraps(func)
-    def wrapper(*args, **kwargs):
-        response = func(*args, **kwargs)
-        response.context_data["name"] = response.context_data["name"] + "!!"
-        return response
-
-    return wrapper
+@wrapt.decorator
+def exclaimify_template_response_name(wrapped, instance, args, kwargs):
+    response = wrapped(*args, **kwargs)
+    response.context_data["name"] = response.context_data["name"] + "!!"
+    return response
 
 
 @exclaimify_template_response_name
diff --git a/tests/integration/instruments/test_elasticsearch.py b/tests/integration/instruments/test_elasticsearch.py
index b22bfac..c2b9616 100644
--- a/tests/integration/instruments/test_elasticsearch.py
+++ b/tests/integration/instruments/test_elasticsearch.py
@@ -170,7 +170,10 @@ def test_instrument_client_install_missing_attribute():
         "scout_apm.instruments.elasticsearch.Elasticsearch"
     ) as mock_elasticsearch:
         del mock_elasticsearch.bulk
-    instrument.instrument_client()  # no crash
+    try:
+        instrument.instrument_client()  # no crash
+    finally:
+        instrument.uninstrument_client()
 
 
 @mock.patch(
diff --git a/tests/integration/instruments/test_pymongo.py b/tests/integration/instruments/test_pymongo.py
index 8204b8c..207dc90 100644
--- a/tests/integration/instruments/test_pymongo.py
+++ b/tests/integration/instruments/test_pymongo.py
@@ -9,11 +9,9 @@ import pytest
 
 from scout_apm.instruments.pymongo import Instrument
 from tests.compat import mock
-from tests.tools import pretend_package_unavailable
 
 # e.g. export MONGODB_URL="mongodb://localhost:27017/"
 MONGODB_URL = os.environ.get("MONGODB_URL")
-
 skip_if_mongodb_not_running = pytest.mark.skipif(
     MONGODB_URL is None, reason="MongoDB isn't available"
 )
@@ -57,22 +55,22 @@ def test_installable():
 
 
 def test_installable_no_pymongo_module():
-    with pretend_package_unavailable("pymongo"):
+    with mock.patch("scout_apm.instruments.pymongo.Collection", new=None):
         assert not instrument.installable()
 
 
 def test_install_no_pymongo_module():
-    with pretend_package_unavailable("pymongo"):
+    with mock.patch("scout_apm.instruments.pymongo.Collection", new=None):
         assert not instrument.install()
 
 
-@mock.patch("scout_apm.instruments.pymongo.wrapt.decorator", side_effect=RuntimeError)
-def test_install_failure(mock_decorator):
+def test_install_missing_attribute():
+    with mock.patch("scout_apm.instruments.pymongo.Collection") as mock_collection:
+        del mock_collection.aggregate
     try:
-        assert not instrument.install()  # doesn't crash
+        instrument.install()  # no crash
     finally:
-        # Currently installed = True even if installing failed.
-        Instrument.installed = False
+        instrument.uninstall()
 
 
 def test_install_is_idempotent():