chore(django): migrate template render wrappers to WrappingContext

ddtrace/_trace/trace_handlers.py

@@ -1,10 +1,12 @@
 import functools
 import sys
+from types import TracebackType
 from typing import Any
 from typing import Callable
 from typing import Dict
 from typing import List
 from typing import Optional
+from typing import Tuple
 from urllib import parse
 
 import wrapt
@@ -144,6 +146,24 @@ def _start_span(ctx: core.ExecutionContext, call_trace: bool = True, **kwargs) -
     return span
 
 
+def _finish_span(
+    ctx: core.ExecutionContext,
+    exc_info: Tuple[Optional[type], Optional[BaseException], Optional[TracebackType]],
+):
+    """
+    Finish the span in the context.
+    If no span is present, do nothing.
+    """
+    span = ctx.span
+    if not span:
+        return
+
+    exc_type, exc_value, exc_traceback = exc_info
+    if exc_type and exc_value and exc_traceback:
+        span.set_exc_info(exc_type, exc_value, exc_traceback)
+    span.finish()
+
+
 def _set_web_frameworks_tags(ctx, span, int_config):
     span.set_tag_str(COMPONENT, int_config.integration_name)
     span.set_tag_str(SPAN_KIND, SpanKind.SERVER)
@@ -963,5 +983,8 @@ def listen():
     ):
         core.on(f"context.started.{context_name}", _start_span)
 
+    for name in ("django.template.render",):
+        core.on(f"context.ended.{name}", _finish_span)
+
 
 listen()

ddtrace/contrib/internal/django/patch.py

@@ -35,7 +35,6 @@
 from ddtrace.ext import sql as sqlx
 from ddtrace.internal import core
 from ddtrace.internal._exceptions import BlockingException
-from ddtrace.internal.compat import maybe_stringify
 from ddtrace.internal.constants import COMPONENT
 from ddtrace.internal.core.event_hub import ResultType
 from ddtrace.internal.logger import get_logger
@@ -573,36 +572,6 @@ def blocked_response():
                     return response  # noqa: B012
 
 
-@trace_utils.with_traced_module
-def traced_template_render(django, pin, wrapped, instance, args, kwargs):
-    # DEV: Check here in case this setting is configured after a template has been instrumented
-    if not config_django.instrument_templates:
-        return wrapped(*args, **kwargs)
-
-    template_name = maybe_stringify(getattr(instance, "name", None))
-    if template_name:
-        resource = template_name
-    else:
-        resource = "{0}.{1}".format(func_name(instance), wrapped.__name__)
-
-    tags = {COMPONENT: config_django.integration_name}
-    if template_name:
-        tags["django.template.name"] = template_name
-    engine = getattr(instance, "engine", None)
-    if engine:
-        tags["django.template.engine.class"] = func_name(engine)
-
-    with core.context_with_data(
-        "django.template.render",
-        span_name="django.template.render",
-        resource=resource,
-        span_type=http.TEMPLATE,
-        tags=tags,
-        pin=pin,
-    ) as ctx, ctx.span:
-        return wrapped(*args, **kwargs)
-
-
 def instrument_view(django, view, path=None):
     """
     Helper to wrap Django views.
@@ -985,9 +954,9 @@ def _(m):
         )
 
     if config_django.instrument_templates:
-        when_imported("django.template.base")(
-            lambda m: trace_utils.wrap(m, "Template.render", traced_template_render(django))
-        )
+        from .templates import DjangoTemplateWrappingContext
+
+        when_imported("django.template.base")(DjangoTemplateWrappingContext.instrument_module)
 
     if django.VERSION < (4, 0, 0):
         when_imported("django.conf.urls")(lambda m: trace_utils.wrap(m, "url", traced_urls_path(django)))
@@ -1060,6 +1029,11 @@ def _unpatch(django):
         trace_utils.unwrap(conn, "cursor")
     trace_utils.unwrap(django.db.utils.ConnectionHandler, "__getitem__")
 
+    if config.django.instrument_templates:
+        from .templates import DjangoTemplateWrappingContext
+
+        DjangoTemplateWrappingContext.uninstrument_module(django.template.base)
+
 
 def unpatch():
     import django

ddtrace/contrib/internal/django/templates.py

@@ -0,0 +1,134 @@
+from types import FunctionType
+from types import ModuleType
+from types import TracebackType
+import typing
+from typing import Optional
+from typing import Type
+from typing import TypeVar
+
+import ddtrace
+from ddtrace import config
+from ddtrace._trace.pin import Pin
+from ddtrace.ext import http
+from ddtrace.internal import core
+from ddtrace.internal.compat import maybe_stringify
+from ddtrace.internal.constants import COMPONENT
+from ddtrace.internal.logger import get_logger
+from ddtrace.internal.utils.importlib import func_name
+from ddtrace.internal.wrapping.context import WrappingContext
+from ddtrace.settings.integration import IntegrationConfig
+
+
+T = TypeVar("T")
+
+log = get_logger(__name__)
+
+
+# PERF: cache the getattr lookup for the Django config
+config_django: IntegrationConfig = typing.cast(IntegrationConfig, config.django)
+
+
+class DjangoTemplateWrappingContext(WrappingContext):
+    """
+    A context for wrapping django.template.base:Template.render method.
+    """
+
+    def __init__(self, f: FunctionType, django: ModuleType) -> None:
+        super().__init__(f)
+        self._django = django
+
+    @classmethod
+    def instrument_module(cls, django_template_base: ModuleType) -> None:
+        """
+        Instrument the django template base module to wrap the render method.
+        """
+        if not config_django.instrument_templates:
+            return
+
+        # Wrap the render method of the Template class
+        Template = getattr(django_template_base, "Template", None)
+        if not Template or not hasattr(Template, "render"):
+            return
+
+        import django
+
+        cls(typing.cast(FunctionType, Template.render), django).wrap()
+
+    @classmethod
+    def uninstrument_module(cls, django_template_base: ModuleType) -> None:
+        # Unwrap the render method of the Template class
+        Template = getattr(django_template_base, "Template", None)
+        if not Template or not hasattr(Template, "render"):
+            return
+
+        if cls.is_wrapped(Template.render):
+            ctx = cls.extract(Template.render)
+            ctx.unwrap()
+
+    def __enter__(self) -> "DjangoTemplateWrappingContext":
+        super().__enter__()
+
+        if not config_django.instrument_templates:
+            return self
+
+        # Get the template instance (self parameter of the render method)
+        # Note: instance is a django.template.base.Template
+        instance = self.get_local("self")
+
+        # Extract template name
+        template_name = maybe_stringify(getattr(instance, "name", None))
+        if template_name:
+            resource = template_name
+        else:
+            resource = "{0}.{1}".format(func_name(instance), self.__wrapped__.__name__)
+
+        # Build tags
+        tags = {COMPONENT: config_django.integration_name}
+        if template_name:
+            tags["django.template.name"] = template_name
+
+        engine = getattr(instance, "engine", None)
+        if engine:
+            tags["django.template.engine.class"] = func_name(engine)
+
+        # Create the span context
+        pin = Pin.get_from(self._django)
+        tracer = ddtrace.tracer
+        if pin:
+            tracer = pin.tracer or tracer
+        ctx = core.context_with_data(
+            "django.template.render",
+            span_name="django.template.render",
+            resource=resource,
+            span_type=http.TEMPLATE,
+            tags=tags,
+            tracer=tracer,
+        )
+
+        # Enter the context and store it
+        ctx.__enter__()
+        self.set("ctx", ctx)
+
+        return self
+
+    def _close_ctx(
+        self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
+    ) -> None:
+        try:
+            # Close the context and any open span
+            ctx = self.get("ctx")
+            ctx.__exit__(exc_type, exc_val, exc_tb)
+        except Exception:
+            log.exception("Failed to close Django template render wrapping context")
+
+    def __return__(self, value: T) -> T:
+        if config_django.instrument_templates:
+            self._close_ctx(None, None, None)
+        return super().__return__(value)
+
+    def __exit__(
+        self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
+    ) -> None:
+        if config_django.instrument_templates:
+            self._close_ctx(exc_type, exc_val, exc_tb)
+        return super().__exit__(exc_type, exc_val, exc_tb)

tests/contrib/django/test_django_patch.py

@@ -1,5 +1,6 @@
 from ddtrace.contrib.internal.django.patch import get_version
 from ddtrace.contrib.internal.django.patch import patch
+from ddtrace.contrib.internal.django.templates import DjangoTemplateWrappingContext
 from tests.contrib.patch import PatchTestCase
 
 
@@ -22,7 +23,7 @@ def assert_module_patched(self, django):
 
         import django.template.base
 
-        self.assert_wrapped(django.template.base.Template.render)
+        assert DjangoTemplateWrappingContext.is_wrapped(django.template.base.Template.render)
         if django.VERSION >= (2, 0, 0):
             self.assert_wrapped(django.urls.path)
             self.assert_wrapped(django.urls.re_path)
@@ -34,7 +35,7 @@ def assert_not_module_patched(self, django):
 
         self.assert_not_wrapped(django.core.handlers.base.BaseHandler.load_middleware)
         self.assert_not_wrapped(django.core.handlers.base.BaseHandler.get_response)
-        self.assert_not_wrapped(django.template.base.Template.render)
+        assert not DjangoTemplateWrappingContext.is_wrapped(django.template.base.Template.render)
         if django.VERSION >= (2, 0, 0):
             self.assert_not_wrapped(django.urls.path)
             self.assert_not_wrapped(django.urls.re_path)
@@ -46,7 +47,12 @@ def assert_not_module_double_patched(self, django):
         self.assert_not_double_wrapped(django.apps.registry.Apps.populate)
         self.assert_not_double_wrapped(django.core.handlers.base.BaseHandler.load_middleware)
         self.assert_not_double_wrapped(django.core.handlers.base.BaseHandler.get_response)
-        self.assert_not_double_wrapped(django.template.base.Template.render)
+
+        assert DjangoTemplateWrappingContext.is_wrapped(django.template.base.Template.render)
+        ctx = DjangoTemplateWrappingContext.extract(django.template.base.Template.render)
+        ctx.unwrap()
+        assert not DjangoTemplateWrappingContext.is_wrapped(django.template.base.Template.render)
+
         if django.VERSION >= (2, 0, 0):
             self.assert_not_double_wrapped(django.urls.path)
             self.assert_not_double_wrapped(django.urls.re_path)