Repository: scoutapp/scout_apm_python

Issue:
Rewrite exec calls
Currently the elasticsearch and pymongo instruments rely on `exec()` to create functions. This is orders of magnitude slower than using inner functions, more fragile, and prevents inspection by our code quality tools. We should rewrite to use inner functions for their generation.


Discussion:

Resolving patch:
diff --git a/src/scout_apm/instruments/elasticsearch.py b/src/scout_apm/instruments/elasticsearch.py
--- a/src/scout_apm/instruments/elasticsearch.py
+++ b/src/scout_apm/instruments/elasticsearch.py
@@ -70,11 +70,11 @@ def uninstall(self):
             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
--- a/src/scout_apm/instruments/pymongo.py
+++ b/src/scout_apm/instruments/pymongo.py
@@ -3,10 +3,14 @@
 
 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 @@ def install(self):
 
         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 @@ def uninstall(self):
             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/instruments/test_elasticsearch.py b/tests/integration/instruments/test_elasticsearch.py
--- 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
--- a/tests/integration/instruments/test_pymongo.py
+++ b/tests/integration/instruments/test_pymongo.py
@@ -9,11 +9,9 @@
 
 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():