mirror of
https://github.com/open-telemetry/opentelemetry-python-contrib.git
synced 2025-08-01 09:13:23 +08:00
119 lines
3.9 KiB
Python
119 lines
3.9 KiB
Python
from importlib import import_module
|
|
|
|
from ddtrace.vendor.wrapt import wrap_function_wrapper as _w
|
|
|
|
from .quantize import quantize
|
|
|
|
from ...compat import urlencode
|
|
from ...constants import ANALYTICS_SAMPLE_RATE_KEY
|
|
from ...ext import SpanTypes, elasticsearch as metadata, http
|
|
from ...pin import Pin
|
|
from ...utils.wrappers import unwrap as _u
|
|
from ...settings import config
|
|
|
|
|
|
def _es_modules():
|
|
module_names = ('elasticsearch', 'elasticsearch1', 'elasticsearch2', 'elasticsearch5', 'elasticsearch6')
|
|
for module_name in module_names:
|
|
try:
|
|
yield import_module(module_name)
|
|
except ImportError:
|
|
pass
|
|
|
|
|
|
# NB: We are patching the default elasticsearch.transport module
|
|
def patch():
|
|
for elasticsearch in _es_modules():
|
|
_patch(elasticsearch)
|
|
|
|
|
|
def _patch(elasticsearch):
|
|
if getattr(elasticsearch, '_datadog_patch', False):
|
|
return
|
|
setattr(elasticsearch, '_datadog_patch', True)
|
|
_w(elasticsearch.transport, 'Transport.perform_request', _get_perform_request(elasticsearch))
|
|
Pin(service=metadata.SERVICE, app=metadata.APP).onto(elasticsearch.transport.Transport)
|
|
|
|
|
|
def unpatch():
|
|
for elasticsearch in _es_modules():
|
|
_unpatch(elasticsearch)
|
|
|
|
|
|
def _unpatch(elasticsearch):
|
|
if getattr(elasticsearch, '_datadog_patch', False):
|
|
setattr(elasticsearch, '_datadog_patch', False)
|
|
_u(elasticsearch.transport.Transport, 'perform_request')
|
|
|
|
|
|
def _get_perform_request(elasticsearch):
|
|
def _perform_request(func, instance, args, kwargs):
|
|
pin = Pin.get_from(instance)
|
|
if not pin or not pin.enabled():
|
|
return func(*args, **kwargs)
|
|
|
|
with pin.tracer.trace('elasticsearch.query', span_type=SpanTypes.ELASTICSEARCH) as span:
|
|
# Don't instrument if the trace is not sampled
|
|
if not span.sampled:
|
|
return func(*args, **kwargs)
|
|
|
|
method, url = args
|
|
params = kwargs.get('params')
|
|
body = kwargs.get('body')
|
|
|
|
span.service = pin.service
|
|
span.set_tag(metadata.METHOD, method)
|
|
span.set_tag(metadata.URL, url)
|
|
span.set_tag(metadata.PARAMS, urlencode(params))
|
|
if config.elasticsearch.trace_query_string:
|
|
span.set_tag(http.QUERY_STRING, urlencode(params))
|
|
if method == 'GET':
|
|
span.set_tag(metadata.BODY, instance.serializer.dumps(body))
|
|
status = None
|
|
|
|
# set analytics sample rate
|
|
span.set_tag(
|
|
ANALYTICS_SAMPLE_RATE_KEY,
|
|
config.elasticsearch.get_analytics_sample_rate()
|
|
)
|
|
|
|
span = quantize(span)
|
|
|
|
try:
|
|
result = func(*args, **kwargs)
|
|
except elasticsearch.exceptions.TransportError as e:
|
|
span.set_tag(http.STATUS_CODE, getattr(e, 'status_code', 500))
|
|
raise
|
|
|
|
try:
|
|
# Optional metadata extraction with soft fail.
|
|
if isinstance(result, tuple) and len(result) == 2:
|
|
# elasticsearch<2.4; it returns both the status and the body
|
|
status, data = result
|
|
else:
|
|
# elasticsearch>=2.4; internal change for ``Transport.perform_request``
|
|
# that just returns the body
|
|
data = result
|
|
|
|
took = data.get('took')
|
|
if took:
|
|
span.set_metric(metadata.TOOK, int(took))
|
|
except Exception:
|
|
pass
|
|
|
|
if status:
|
|
span.set_tag(http.STATUS_CODE, status)
|
|
|
|
return result
|
|
return _perform_request
|
|
|
|
|
|
# Backwards compatibility for anyone who decided to import `ddtrace.contrib.elasticsearch.patch._perform_request`
|
|
# DEV: `_perform_request` is a `wrapt.FunctionWrapper`
|
|
try:
|
|
# DEV: Import as `es` to not shadow loop variables above
|
|
import elasticsearch as es
|
|
_perform_request = _get_perform_request(es)
|
|
except ImportError:
|
|
pass
|