Merge pull request #51 from palazzem/flask-cache-integration

flask-cache integration

Author: Emanuele Palazzetti

ddtrace/contrib/flask_cache/__init__.py

@@ -0,0 +1,33 @@
+"""
+The flask cache tracer will track any access to a cache backend.
+You can this tracer together with the Flask tracer middleware.
+
+To install the tracer, do the following::
+
+    from flask import Flask
+
+    from ddtrace import tracer
+    from ddtrace.contrib.flask_cache import get_traced_cache
+
+    app = Flask(__name__)
+
+    # get the traced Cache class
+    Cache = get_traced_cache(tracer, service='flask-cache-experiments')
+
+    # use the Cache as usual
+    cache = Cache(app, config={'CACHE_TYPE': 'simple'})
+
+    def counter():
+        # this access is traced
+        conn_counter = cache.get("conn_counter")
+"""
+
+from ..util import require_modules
+
+required_modules = ['flask_cache']
+
+with require_modules(required_modules) as missing_modules:
+    if not missing_modules:
+        from .tracers import get_traced_cache
+
+        __all__ = ['get_traced_cache']

ddtrace/contrib/flask_cache/tracers.py

@@ -0,0 +1,147 @@
+"""
+Datadog trace code for flask_cache
+"""
+
+# stdlib
+import logging
+
+# project
+from .utils import _extract_conn_tags, _resource_from_cache_prefix
+from ...ext import AppTypes
+
+# 3rd party
+from flask.ext.cache import Cache
+
+
+log = logging.Logger(__name__)
+
+TYPE = "cache"
+DEFAULT_SERVICE = "flask-cache"
+
+# standard tags
+COMMAND_KEY = "flask_cache.key"
+CACHE_BACKEND = "flask_cache.backend"
+CONTACT_POINTS = "flask_cache.contact_points"
+
+
+def get_traced_cache(ddtracer, service=DEFAULT_SERVICE, meta=None):
+    """
+    Return a traced Cache object that behaves exactly as the ``flask.ext.cache.Cache class``
+    """
+
+    # set the Tracer info
+    ddtracer.set_service_info(
+        app="flask",
+        app_type=AppTypes.cache,
+        service=service,
+    )
+
+    class TracedCache(Cache):
+        """
+        Traced cache backend that monitors any operations done by flask_cache. Observed actions are:
+            * get, set, add, delete, clear
+            * all many_ operations
+        """
+        _datadog_tracer = ddtracer
+        _datadog_service = service
+        _datadog_meta = meta
+
+        def __trace(self, cmd):
+            """
+            Start a tracing with default attributes and tags
+            """
+            # create a new span
+            s = self._datadog_tracer.trace(
+                cmd,
+                span_type=TYPE,
+                service=self._datadog_service
+            )
+            # set span tags
+            s.set_tag(CACHE_BACKEND, self.config.get("CACHE_TYPE"))
+            s.set_tags(self._datadog_meta)
+            # add connection meta if there is one
+            if getattr(self.cache, "_client", None):
+                try:
+                    s.set_tags(_extract_conn_tags(self.cache._client))
+                except Exception:
+                    log.exception("error parsing connection tags")
+
+            return s
+
+        def get(self, *args, **kwargs):
+            """
+            Track ``get`` operation
+            """
+            with self.__trace("flask_cache.cmd") as span:
+                span.resource = _resource_from_cache_prefix("GET", self.config)
+                if len(args) > 0:
+                    span.set_tag(COMMAND_KEY, args[0])
+                return super(TracedCache, self).get(*args, **kwargs)
+
+        def set(self, *args, **kwargs):
+            """
+            Track ``set`` operation
+            """
+            with self.__trace("flask_cache.cmd") as span:
+                span.resource = _resource_from_cache_prefix("SET", self.config)
+                if len(args) > 0:
+                    span.set_tag(COMMAND_KEY, args[0])
+                return super(TracedCache, self).set(*args, **kwargs)
+
+        def add(self, *args, **kwargs):
+            """
+            Track ``add`` operation
+            """
+            with self.__trace("flask_cache.cmd") as span:
+                span.resource = _resource_from_cache_prefix("ADD", self.config)
+                if len(args) > 0:
+                    span.set_tag(COMMAND_KEY, args[0])
+                return super(TracedCache, self).add(*args, **kwargs)
+
+        def delete(self, *args, **kwargs):
+            """
+            Track ``delete`` operation
+            """
+            with self.__trace("flask_cache.cmd") as span:
+                span.resource = _resource_from_cache_prefix("DELETE", self.config)
+                if len(args) > 0:
+                    span.set_tag(COMMAND_KEY, args[0])
+                return super(TracedCache, self).delete(*args, **kwargs)
+
+        def delete_many(self, *args, **kwargs):
+            """
+            Track ``delete_many`` operation
+            """
+            with self.__trace("flask_cache.cmd") as span:
+                span.resource = _resource_from_cache_prefix("DELETE_MANY", self.config)
+                span.set_tag(COMMAND_KEY, list(args))
+                return super(TracedCache, self).delete_many(*args, **kwargs)
+
+        def clear(self, *args, **kwargs):
+            """
+            Track ``clear`` operation
+            """
+            with self.__trace("flask_cache.cmd") as span:
+                span.resource = _resource_from_cache_prefix("CLEAR", self.config)
+                return super(TracedCache, self).clear(*args, **kwargs)
+
+        def get_many(self, *args, **kwargs):
+            """
+            Track ``get_many`` operation
+            """
+            with self.__trace("flask_cache.cmd") as span:
+                span.resource = _resource_from_cache_prefix("GET_MANY", self.config)
+                span.set_tag(COMMAND_KEY, list(args))
+                return super(TracedCache, self).get_many(*args, **kwargs)
+
+        def set_many(self, *args, **kwargs):
+            """
+            Track ``set_many`` operation
+            """
+            with self.__trace("flask_cache.cmd") as span:
+                span.resource = _resource_from_cache_prefix("SET_MANY", self.config)
+                if len(args) > 0:
+                    span.set_tag(COMMAND_KEY, list(args[0].keys()))
+                return super(TracedCache, self).set_many(*args, **kwargs)
+
+    return TracedCache

ddtrace/contrib/flask_cache/utils.py

@@ -0,0 +1,39 @@
+# project
+from ...ext import net
+from ..redis.util import _extract_conn_tags as extract_redis_tags
+
+
+def _resource_from_cache_prefix(resource, cache):
+    """
+    Combine the resource name with the cache prefix (if any)
+    """
+    if getattr(cache, "key_prefix", None):
+        name = "{} {}".format(resource, cache.key_prefix)
+    else:
+        name = resource
+
+    # enforce lowercase to make the output nicer to read
+    return name.lower()
+
+
+def _extract_conn_tags(client):
+    """
+    For the given client extracts connection tags
+    """
+    tags = {}
+
+    if getattr(client, "servers", None):
+        # Memcached backend supports an address pool
+        if isinstance(client.servers, list) and len(client.servers) > 0:
+            # use the first address of the pool as a host because
+            # the code doesn't expose more information
+            contact_point = client.servers[0].address
+            tags[net.TARGET_HOST] = contact_point[0]
+            tags[net.TARGET_PORT] = contact_point[1]
+
+    if getattr(client, "connection_pool", None):
+        # Redis main connection
+        redis_tags = extract_redis_tags(client.connection_pool.connection_kwargs)
+        tags.update(**redis_tags)
+
+    return tags