pylibmc: add query metadata

Author: clutchski

ddtrace/contrib/pylibmc/client.py

@@ -8,7 +8,7 @@
 
 # project
 import ddtrace
-from ddtrace.ext import AppTypes
+from ddtrace.ext import memcached
 from ddtrace.ext import net
 from .addrs import parse_addresses
 
@@ -22,7 +22,7 @@ class TracedClient(ObjectProxy):
     _service = None
     _tracer = None
 
-    def __init__(self, client, service="memcached", tracer=None):
+    def __init__(self, client, service=memcached.SERVICE, tracer=None):
         """ Create a traced client that wraps the given memcached client. """
         super(TracedClient, self).__init__(client)
         self._service = service
@@ -38,8 +38,8 @@ def __init__(self, client, service="memcached", tracer=None):
         try:
             self._tracer.set_service_info(
                 service=service,
-                app="memcached",
-                app_type=AppTypes.cache)
+                app=memcached.SERVICE,
+                app_type=memcached.TYPE)
         except Exception:
             log.exception("error setting service info")
 
@@ -49,48 +49,65 @@ def clone(self, *args, **kwargs):
         return TracedClient(cloned, tracer=self._tracer, service=self._service)
 
     def get(self, *args, **kwargs):
-        return self._trace("get", *args, **kwargs)
-
-    def get_multi(self, *args, **kwargs):
-        return self._trace("get_multi", *args, **kwargs)
-
-    def set_multi(self, *args, **kwargs):
-        return self._trace("set_multi", *args, **kwargs)
-
-    def delete_multi(self, *args, **kwargs):
-        self._trace("delete_multi", *args, **kwargs)
+        return self._trace_cmd("get", *args, **kwargs)
 
     def set(self, *args, **kwargs):
-        return self._trace("set", *args, **kwargs)
+        return self._trace_cmd("set", *args, **kwargs)
 
     def delete(self, *args, **kwargs):
-        return self._trace("delete", *args, **kwargs)
+        return self._trace_cmd("delete", *args, **kwargs)
 
     def gets(self, *args, **kwargs):
-        return self._trace("gets", *args, **kwargs)
+        return self._trace_cmd("gets", *args, **kwargs)
 
     def touch(self, *args, **kwargs):
-        return self._trace("touch", *args, **kwargs)
+        return self._trace_cmd("touch", *args, **kwargs)
 
     def cas(self, *args, **kwargs):
-        return self._trace("cas", *args, **kwargs)
+        return self._trace_cmd("cas", *args, **kwargs)
 
     def incr(self, *args, **kwargs):
-        return self._trace("incr", *args, **kwargs)
+        return self._trace_cmd("incr", *args, **kwargs)
 
     def decr(self, *args, **kwargs):
-        return self._trace("decr", *args, **kwargs)
+        return self._trace_cmd("decr", *args, **kwargs)
 
     def append(self, *args, **kwargs):
-        return self._trace("append", *args, **kwargs)
+        return self._trace_cmd("append", *args, **kwargs)
 
     def prepend(self, *args, **kwargs):
-        return self._trace("prepend", *args, **kwargs)
+        return self._trace_cmd("prepend", *args, **kwargs)
+
+    def get_multi(self, *args, **kwargs):
+        return self._trace_multi_cmd("get_multi", *args, **kwargs)
+
+    def set_multi(self, *args, **kwargs):
+        return self._trace_multi_cmd("set_multi", *args, **kwargs)
 
-    def _trace(self, method_name, *args, **kwargs):
-        """ trace the execution of the method with the given name. """
+    def delete_multi(self, *args, **kwargs):
+        return self._trace_multi_cmd("delete_multi", *args, **kwargs)
+
+    def _trace_cmd(self, method_name, *args, **kwargs):
+        """ trace the execution of the method with the given name and will
+            patch the first arg.
+        """
+        method = getattr(self.__wrapped__, method_name)
+        with self._span(method_name) as span:
+
+            if args:
+                span.set_tag(memcached.QUERY, "%s %s" % (method_name, args[0]))
+
+            return method(*args, **kwargs)
+
+    def _trace_multi_cmd(self, method_name, *args, **kwargs):
+        """ trace the execution of the multi command with the given name. """
         method = getattr(self.__wrapped__, method_name)
-        with self._span(method_name):
+        with self._span(method_name) as span:
+
+            pre = kwargs.get('key_prefix')
+            if pre:
+                span.set_tag(memcached.QUERY, "%s %s" % (method_name, pre))
+
             return method(*args, **kwargs)
 
     def _span(self, cmd_name):
@@ -105,8 +122,8 @@ def _span(self, cmd_name):
             self._tag_span(span)
         except Exception:
             log.exception("error tagging span")
-        finally:
-            return span
+
+        return span
 
     def _tag_span(self, span):
         # FIXME[matt] the host selection is buried in c code. we can't tell what it's actually

ddtrace/ext/memcached.py

@@ -0,0 +1,7 @@
+
+from ddtrace.ext import AppTypes
+
+SERVICE = "memcached"
+TYPE = AppTypes.cache
+
+QUERY = "memcached.query"

tests/contrib/pylibmc/test.py

@@ -100,10 +100,29 @@ def test_get_set_multi(self):
         resources = sorted(s.resource for s in spans)
         eq_(expected_resources, resources)
 
+    def test_get_set_multi_prefix(self):
+        client, tracer = _setup()
+        # test
+        start = time.time()
+        client.set_multi({"a":1, "b":2}, key_prefix='foo')
+        out = client.get_multi(["a", "c"], key_prefix='foo')
+        eq_(out, {"a":1})
+        client.delete_multi(["a", "c"], key_prefix='foo')
+        end = time.time()
+        # verify
+        spans = tracer.writer.pop()
+        for s in spans:
+            _verify_cache_span(s, start, end)
+            eq_(s.get_tag("memcached.query"), "%s foo" % s.resource,)
+        expected_resources = sorted(["get_multi", "set_multi", "delete_multi"])
+        resources = sorted(s.resource for s in spans)
+        eq_(expected_resources, resources)
+
+
     def test_get_set_delete(self):
         client, tracer = _setup()
         # test
-        k = "key-foo"
+        k = u'cafe'
         v = "val-foo"
         start = time.time()
         client.delete(k) # just in case
@@ -117,10 +136,12 @@ def test_get_set_delete(self):
         spans = tracer.writer.pop()
         for s in spans:
             _verify_cache_span(s, start, end)
+            eq_(s.get_tag("memcached.query"), "%s %s" % (s.resource, k))
         expected_resources = sorted(["get", "get", "delete", "set"])
         resources = sorted(s.resource for s in spans)
         eq_(expected_resources, resources)
 
+
 def _verify_cache_span(s, start, end):
     assert s.start > start
     assert s.start + s.duration < end