opentelemetry-python-contrib/reference/ddtrace/contrib/jinja2/patch.py

94 lines
2.9 KiB
Python

import jinja2
from ddtrace.vendor.wrapt import wrap_function_wrapper as _w
from ddtrace import config
from ...ext import SpanTypes
from ...utils.formats import get_env
from ...pin import Pin
from ...utils.wrappers import unwrap as _u
from .constants import DEFAULT_TEMPLATE_NAME
# default settings
config._add('jinja2', {
'service_name': get_env('jinja2', 'service_name', None),
})
def patch():
if getattr(jinja2, '__datadog_patch', False):
# already patched
return
setattr(jinja2, '__datadog_patch', True)
Pin(
service=config.jinja2['service_name'],
_config=config.jinja2,
).onto(jinja2.environment.Environment)
_w(jinja2, 'environment.Template.render', _wrap_render)
_w(jinja2, 'environment.Template.generate', _wrap_render)
_w(jinja2, 'environment.Environment.compile', _wrap_compile)
_w(jinja2, 'environment.Environment._load_template', _wrap_load_template)
def unpatch():
if not getattr(jinja2, '__datadog_patch', False):
return
setattr(jinja2, '__datadog_patch', False)
_u(jinja2.Template, 'render')
_u(jinja2.Template, 'generate')
_u(jinja2.Environment, 'compile')
_u(jinja2.Environment, '_load_template')
def _wrap_render(wrapped, instance, args, kwargs):
"""Wrap `Template.render()` or `Template.generate()`
"""
pin = Pin.get_from(instance.environment)
if not pin or not pin.enabled():
return wrapped(*args, **kwargs)
template_name = instance.name or DEFAULT_TEMPLATE_NAME
with pin.tracer.trace('jinja2.render', pin.service, span_type=SpanTypes.TEMPLATE) as span:
try:
return wrapped(*args, **kwargs)
finally:
span.resource = template_name
span.set_tag('jinja2.template_name', template_name)
def _wrap_compile(wrapped, instance, args, kwargs):
pin = Pin.get_from(instance)
if not pin or not pin.enabled():
return wrapped(*args, **kwargs)
if len(args) > 1:
template_name = args[1]
else:
template_name = kwargs.get('name', DEFAULT_TEMPLATE_NAME)
with pin.tracer.trace('jinja2.compile', pin.service, span_type=SpanTypes.TEMPLATE) as span:
try:
return wrapped(*args, **kwargs)
finally:
span.resource = template_name
span.set_tag('jinja2.template_name', template_name)
def _wrap_load_template(wrapped, instance, args, kwargs):
pin = Pin.get_from(instance)
if not pin or not pin.enabled():
return wrapped(*args, **kwargs)
template_name = kwargs.get('name', args[0])
with pin.tracer.trace('jinja2.load', pin.service, span_type=SpanTypes.TEMPLATE) as span:
template = None
try:
template = wrapped(*args, **kwargs)
return template
finally:
span.resource = template_name
span.set_tag('jinja2.template_name', template_name)
if template:
span.set_tag('jinja2.template_path', template.filename)