Repository: scoutapp/scout_apm_python

Issue:
Fix Jinja2 async integration
Since https://github.com/pallets/jinja/issues/765 , Jinja2 only lazily imports and adds the `render_async` function that we instrument (since #398).

We need to update to only instrument at the point the `render_async` method becomes available.


Discussion:

Resolving patch:
diff --git a/src/scout_apm/instruments/jinja2.py b/src/scout_apm/instruments/jinja2.py
--- a/src/scout_apm/instruments/jinja2.py
+++ b/src/scout_apm/instruments/jinja2.py
@@ -2,11 +2,17 @@
 from __future__ import absolute_import, division, print_function, unicode_literals
 
 import logging
+import sys
 
 import wrapt
 
 from scout_apm.core.tracked_request import TrackedRequest
 
+try:
+    from jinja2 import Environment
+except ImportError:  # pragma: no cover
+    Environment = None
+
 try:
     from jinja2 import Template
 except ImportError:  # pragma: no cover
@@ -22,13 +28,14 @@
 logger = logging.getLogger(__name__)
 
 
+have_patched_environment_init = False
 have_patched_template_render = False
 have_patched_template_render_async = False
 
 
 def ensure_installed():
+    global have_patched_environment_init
     global have_patched_template_render
-    global have_patched_template_render_async
 
     logger.info("Ensuring Jinja2 instrumentation is installed.")
 
@@ -36,27 +43,27 @@ def ensure_installed():
         logger.info("Unable to import jinja2.Template")
         return
 
-    if not have_patched_template_render:
+    if not have_patched_environment_init:
         try:
-            Template.render = wrapped_render(Template.render)
+            Environment.__init__ = wrapped_environment_init(Environment.__init__)
         except Exception as exc:
             logger.warning(
-                "Unable to instrument jinja2.Template.render: %r", exc, exc_info=exc
+                "Unable to instrument jinja2.Environment.__init__: %r",
+                exc,
+                exc_info=exc,
             )
         else:
-            have_patched_template_render = True
+            have_patched_environment_init = True
 
-    if not have_patched_template_render_async and wrapped_render_async is not None:
+    if not have_patched_template_render:
         try:
-            Template.render_async = wrapped_render_async(Template.render_async)
+            Template.render = wrapped_render(Template.render)
         except Exception as exc:
             logger.warning(
-                "Unable to instrument jinja2.Template.render_async: %r",
-                exc,
-                exc_info=exc,
+                "Unable to instrument jinja2.Template.render: %r", exc, exc_info=exc
             )
         else:
-            have_patched_template_render_async = True
+            have_patched_template_render = True
 
 
 @wrapt.decorator
@@ -68,3 +75,33 @@ def wrapped_render(wrapped, instance, args, kwargs):
         return wrapped(*args, **kwargs)
     finally:
         tracked_request.stop_span()
+
+
+@wrapt.decorator
+def wrapped_environment_init(wrapped, instance, args, kwargs):
+    """
+    Delayed wrapping of render_async(), since Template won't have this method
+    until after jinja2.asyncsupport is imported, which since Jinja2 2.11.0 is
+    done conditionally in Environment.__init__:
+    https://github.com/pallets/jinja/issues/765
+    """
+    global have_patched_template_render_async
+    result = wrapped(*args, **kwargs)
+
+    if (
+        wrapped_render_async is not None
+        and not have_patched_template_render_async
+        and "jinja2.asyncsupport" in sys.modules
+    ):
+        try:
+            Template.render_async = wrapped_render_async(Template.render_async)
+        except Exception as exc:
+            logger.warning(
+                "Unable to instrument jinja2.Template.render_async: %r",
+                exc,
+                exc_info=exc,
+            )
+        else:
+            have_patched_template_render_async = True
+
+    return result


Test patch:
diff --git a/tests/integration/instruments/test_jinja2.py b/tests/integration/instruments/test_jinja2.py
--- a/tests/integration/instruments/test_jinja2.py
+++ b/tests/integration/instruments/test_jinja2.py
@@ -12,13 +12,16 @@
 
 @contextmanager
 def mock_not_attempted():
+    not_environment = mock.patch(
+        "scout_apm.instruments.jinja2.have_patched_environment_init", new=False
+    )
     not_render = mock.patch(
         "scout_apm.instruments.jinja2.have_patched_template_render", new=False
     )
     not_render_async = mock.patch(
         "scout_apm.instruments.jinja2.have_patched_template_render_async", new=False
     )
-    with not_render, not_render_async:
+    with not_environment, not_render, not_render_async:
         yield
 
 
diff --git a/tests/integration/instruments/test_jinja2_py36plus.py b/tests/integration/instruments/test_jinja2_py36plus.py
--- a/tests/integration/instruments/test_jinja2_py36plus.py
+++ b/tests/integration/instruments/test_jinja2_py36plus.py
@@ -18,18 +18,15 @@ def test_ensure_installed_fail_no_render_async_attribute(caplog):
 
         ensure_installed()
 
-    assert len(caplog.record_tuples) == 2
-    assert caplog.record_tuples[0] == (
-        "scout_apm.instruments.jinja2",
-        logging.INFO,
-        "Ensuring Jinja2 instrumentation is installed.",
-    )
-    logger, level, message = caplog.record_tuples[1]
-    assert logger == "scout_apm.instruments.jinja2"
-    assert level == logging.WARNING
-    assert message.startswith(
-        "Unable to instrument jinja2.Template.render_async: AttributeError"
-    )
+        jinja2.Environment()
+
+    assert caplog.record_tuples == [
+        (
+            "scout_apm.instruments.jinja2",
+            logging.INFO,
+            "Ensuring Jinja2 instrumentation is installed.",
+        )
+    ]
 
 
 @async_test