diff --git a/ddtrace/_trace/trace_handlers.py b/ddtrace/_trace/trace_handlers.py index d73642ddc91..6a65e0f69f8 100644 --- a/ddtrace/_trace/trace_handlers.py +++ b/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() diff --git a/ddtrace/contrib/internal/django/patch.py b/ddtrace/contrib/internal/django/patch.py index df2f0c389d2..dd9deae0057 100644 --- a/ddtrace/contrib/internal/django/patch.py +++ b/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 diff --git a/ddtrace/contrib/internal/django/templates.py b/ddtrace/contrib/internal/django/templates.py new file mode 100644 index 00000000000..692e8fcbd01 --- /dev/null +++ b/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) diff --git a/tests/contrib/django/test_django_patch.py b/tests/contrib/django/test_django_patch.py index 41e1de6b3e7..5b4578fcc56 100644 --- a/tests/contrib/django/test_django_patch.py +++ b/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)