Merge pull request #843 from DataDog/majorgreys/trace-search-config-complete

Second batch of trace search client configuration

Author: majorgreys

ddtrace/contrib/aiopg/connection.py

@@ -4,8 +4,10 @@
 from aiopg.utils import _ContextManager
 
 from .. import dbapi
-from ...pin import Pin
+from ...constants import ANALYTICS_SAMPLE_RATE_KEY
 from ...ext import sql, AppTypes
+from ...pin import Pin
+from ...settings import config
 
 
 class AIOTracedCursor(wrapt.ObjectProxy):
@@ -32,6 +34,12 @@ def _trace_method(self, method, resource, extra_tags, *args, **kwargs):
             s.set_tags(pin.tags)
             s.set_tags(extra_tags)
 
+            # set analytics sample rate
+            s.set_tag(
+                ANALYTICS_SAMPLE_RATE_KEY,
+                config.aiopg.get_analytics_sample_rate()
+            )
+
             try:
                 result = yield from method(*args, **kwargs)
                 return result

ddtrace/contrib/elasticsearch/patch.py

@@ -5,9 +5,11 @@
 from .quantize import quantize
 
 from ...compat import urlencode
+from ...constants import ANALYTICS_SAMPLE_RATE_KEY
 from ...ext import elasticsearch as metadata, http, AppTypes
 from ...pin import Pin
 from ...utils.wrappers import unwrap as _u
+from ...settings import config
 
 
 def _es_modules():
@@ -68,6 +70,12 @@ def _perform_request(func, instance, args, kwargs):
                 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:

ddtrace/contrib/flask_cache/tracers.py

@@ -7,6 +7,8 @@
 
 # project
 from .utils import _extract_conn_tags, _resource_from_cache_prefix
+from ...constants import ANALYTICS_SAMPLE_RATE_KEY
+from ...settings import config
 
 # 3rd party
 from flask.ext.cache import Cache
@@ -51,6 +53,11 @@ def __trace(self, cmd):
             # set span tags
             s.set_tag(CACHE_BACKEND, self.config.get("CACHE_TYPE"))
             s.set_tags(self._datadog_meta)
+            # set analytics sample rate
+            s.set_tag(
+                ANALYTICS_SAMPLE_RATE_KEY,
+                config.flask_cache.get_analytics_sample_rate()
+            )
             # add connection meta if there is one
             if getattr(self.cache, "_client", None):
                 try:

ddtrace/contrib/grpc/client_interceptor.py

@@ -2,6 +2,8 @@
 
 from ddtrace import Pin
 from .propagation import inject_span
+from ...constants import ANALYTICS_SAMPLE_RATE_KEY
+from ...settings import config
 
 
 class GrpcClientInterceptor(
@@ -23,6 +25,11 @@ def _start_span(self, method):
             span.set_tag('grpc.port', self._port)
         if self._pin.tags:
             span.set_tags(self._pin.tags)
+        # set analytics sample rate
+        span.set_tag(
+            ANALYTICS_SAMPLE_RATE_KEY,
+            config.grpc.get_analytics_sample_rate()
+        )
         return span
 
     def intercept_unary_unary(self, continuation, client_call_details, request):

ddtrace/contrib/httplib/patch.py

@@ -3,11 +3,12 @@
 
 # Project
 from ...compat import PY2, httplib, parse
-from ddtrace import config
+from ...constants import ANALYTICS_SAMPLE_RATE_KEY
 from ...ext import http as ext_http
 from ...http import store_request_headers, store_response_headers
 from ...internal.logger import get_logger
 from ...pin import Pin
+from ...settings import config
 from ...utils.wrappers import unwrap as _u
 
 span_name = 'httplib.request' if PY2 else 'http.client.request'
@@ -78,6 +79,12 @@ def _wrap_putrequest(func, instance, args, kwargs):
 
         span.set_tag(ext_http.URL, sanitized_url)
         span.set_tag(ext_http.METHOD, method)
+
+        # set analytics sample rate
+        span.set_tag(
+            ANALYTICS_SAMPLE_RATE_KEY,
+            config.httplib.get_analytics_sample_rate()
+        )
     except Exception:
         log.debug('error applying request tags', exc_info=True)
     return func(*args, **kwargs)

ddtrace/contrib/kombu/patch.py

@@ -3,15 +3,16 @@
 from ddtrace.vendor import wrapt
 
 # project
-from ddtrace import config
-
-from ...pin import Pin
-from ...utils.formats import get_env
-from .constants import DEFAULT_SERVICE
+from ...constants import ANALYTICS_SAMPLE_RATE_KEY
 from ...ext import kombu as kombux
 from ...ext import AppTypes
-from ...utils.wrappers import unwrap
+from ...pin import Pin
 from ...propagation.http import HTTPPropagator
+from ...settings import config
+from ...utils.formats import get_env
+from ...utils.wrappers import unwrap
+
+from .constants import DEFAULT_SERVICE
 from .utils import (
     get_exchange_from_args,
     get_body_length_from_args,
@@ -88,6 +89,11 @@ def traced_receive(func, instance, args, kwargs):
 
         s.set_tags(extract_conn_tags(message.channel.connection))
         s.set_tag(kombux.ROUTING_KEY, message.delivery_info['routing_key'])
+        # set analytics sample rate
+        s.set_tag(
+            ANALYTICS_SAMPLE_RATE_KEY,
+            config.kombu.get_analytics_sample_rate()
+        )
         return func(*args, **kwargs)
 
 
@@ -105,6 +111,11 @@ def traced_publish(func, instance, args, kwargs):
         s.set_tag(kombux.ROUTING_KEY, get_routing_key_from_args(args))
         s.set_tags(extract_conn_tags(instance.channel.connection))
         s.set_metric(kombux.BODY_LEN, get_body_length_from_args(args))
+        # set analytics sample rate
+        s.set_tag(
+            ANALYTICS_SAMPLE_RATE_KEY,
+            config.kombu.get_analytics_sample_rate()
+        )
         # run the command
         propagator.inject(s.context, args[HEADER_POS])
         return func(*args, **kwargs)

ddtrace/contrib/pylibmc/client.py

@@ -6,9 +6,11 @@
 
 # project
 import ddtrace
+from ...constants import ANALYTICS_SAMPLE_RATE_KEY
 from ...ext import memcached
 from ...ext import net
 from ...internal.logger import get_logger
+from ...settings import config
 from .addrs import parse_addresses
 
 
@@ -143,3 +145,9 @@ def _tag_span(self, span):
             _, host, port, _ = random.choice(self._addresses)
             span.set_meta(net.TARGET_HOST, host)
             span.set_meta(net.TARGET_PORT, port)
+
+        # set analytics sample rate
+        span.set_tag(
+            ANALYTICS_SAMPLE_RATE_KEY,
+            config.pylibmc.get_analytics_sample_rate()
+        )

ddtrace/contrib/pymemcache/client.py

@@ -13,10 +13,12 @@
 )
 
 # project
+from ...constants import ANALYTICS_SAMPLE_RATE_KEY
 from ...compat import reraise
 from ...ext import net, memcached as memcachedx
 from ...internal.logger import get_logger
 from ...pin import Pin
+from ...settings import config
 
 log = get_logger(__name__)
 
@@ -141,6 +143,12 @@ def _traced_cmd(self, method_name, *args, **kwargs):
             resource=method_name,
             span_type=memcachedx.TYPE,
         ) as span:
+            # set analytics sample rate
+            span.set_tag(
+                ANALYTICS_SAMPLE_RATE_KEY,
+                config.pymemcache.get_analytics_sample_rate()
+            )
+
             # try to set relevant tags, catch any exceptions so we don't mess
             # with the application
             try:

ddtrace/contrib/pymongo/client.py

@@ -9,10 +9,12 @@
 # project
 import ddtrace
 from ...compat import iteritems
+from ...constants import ANALYTICS_SAMPLE_RATE_KEY
 from ...ext import AppTypes
 from ...ext import mongo as mongox
 from ...ext import net as netx
 from ...internal.logger import get_logger
+from ...settings import config
 from ...utils.deprecation import deprecated
 from .parse import parse_spec, parse_query, parse_msg
 
@@ -118,6 +120,12 @@ def send_message_with_response(self, operation, *args, **kwargs):
             # set `mongodb.query` tag and resource for span
             _set_query_metadata(span, cmd)
 
+            # set analytics sample rate
+            span.set_tag(
+                ANALYTICS_SAMPLE_RATE_KEY,
+                config.pymongo.get_analytics_sample_rate()
+            )
+
             result = self.__wrapped__.send_message_with_response(
                 operation,
                 *args,
@@ -197,6 +205,12 @@ def __trace(self, cmd):
         # set `mongodb.query` tag and resource for span
         _set_query_metadata(s, cmd)
 
+        # set analytics sample rate
+        s.set_tag(
+            ANALYTICS_SAMPLE_RATE_KEY,
+            config.pymongo.get_analytics_sample_rate()
+        )
+
         if self.address:
             _set_address_tags(s, self.address)
         return s

ddtrace/contrib/requests/connection.py

@@ -73,7 +73,7 @@ def _wrap_send(func, instance, args, kwargs):
         span.service = _extract_service_name(instance, span, hostname=hostname)
 
         # Configure trace search sample rate
-        # DEV: Not enabled by default when global analytics config is enabled
+        # DEV: analytics enabled on per-session basis
         cfg = config.get_from(instance)
         analytics_enabled = cfg.get('analytics_enabled')
         if analytics_enabled: