Files
2020-04-08 10:39:44 -07:00

67 lines
2.7 KiB
Python

# DEV: This will import the first available module from:
# `elasticsearch`, `elasticsearch1`, `elasticsearch2`, `elasticsearch5`, 'elasticsearch6'
from .elasticsearch import elasticsearch
from .quantize import quantize
from ...utils.deprecation import deprecated
from ...compat import urlencode
from ...ext import SpanTypes, http, elasticsearch as metadata
from ...settings import config
DEFAULT_SERVICE = 'elasticsearch'
@deprecated(message='Use patching instead (see the docs).', version='1.0.0')
def get_traced_transport(datadog_tracer, datadog_service=DEFAULT_SERVICE):
class TracedTransport(elasticsearch.Transport):
""" Extend elasticseach transport layer to allow Datadog
tracer to catch any performed request.
"""
_datadog_tracer = datadog_tracer
_datadog_service = datadog_service
def perform_request(self, method, url, params=None, body=None):
with self._datadog_tracer.trace('elasticsearch.query', span_type=SpanTypes.ELASTICSEARCH) as s:
# Don't instrument if the trace is not sampled
if not s.sampled:
return super(TracedTransport, self).perform_request(
method, url, params=params, body=body)
s.service = self._datadog_service
s.set_tag(metadata.METHOD, method)
s.set_tag(metadata.URL, url)
s.set_tag(metadata.PARAMS, urlencode(params))
if config.elasticsearch.trace_query_string:
s.set_tag(http.QUERY_STRING, urlencode(params))
if method == 'GET':
s.set_tag(metadata.BODY, self.serializer.dumps(body))
s = quantize(s)
try:
result = super(TracedTransport, self).perform_request(method, url, params=params, body=body)
except elasticsearch.exceptions.TransportError as e:
s.set_tag(http.STATUS_CODE, e.status_code)
raise
status = None
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
if status:
s.set_tag(http.STATUS_CODE, status)
took = data.get('took')
if took:
s.set_metric(metadata.TOOK, int(took))
return result
return TracedTransport