mongoengine: first pass at mongoengine

ddtrace/contrib/mongoengine/__init__.py

@@ -0,0 +1,30 @@
+"""
+To trace mongoengine queries, we patch it's connect method::
+
+    # to patch all mongoengine connections, do the following
+    # before you import mongoengine yourself.
+
+    from ddtrace import tracer
+    from ddtrace.contrib.mongoengine import trace_mongoengine
+    trace_mongoengine(tracer, service="my-mongo-db", patch=True)
+
+
+    # to patch a single mongoengine connection, do this:
+    connect = trace_mongoengine(tracer, service="my-mongo-db", patch=False()
+    connect()
+
+    # now use mongoengine ....
+    User.objects(name="Mongo")
+"""
+
+
+from ..util import require_modules
+
+
+required_modules = ['mongoengine']
+
+with require_modules(required_modules) as missing_modules:
+    if not missing_modules:
+        from .trace import trace_mongoengine
+
+        __all__ = ['trace_mongoengine']

ddtrace/contrib/mongoengine/trace.py

@@ -0,0 +1,47 @@
+
+# 3p
+import mongoengine
+import wrapt
+
+# project
+from ddtrace.ext import mongo as mongox
+from ddtrace.contrib.pymongo import trace_mongo_client
+
+
+def trace_mongoengine(tracer, service=mongox.TYPE, patch=False):
+    connect = mongoengine.connect
+    wrapped = WrappedConnect(connect, tracer, service)
+    if patch:
+        mongoengine.connect = wrapped
+    return wrapped
+
+
+class WrappedConnect(wrapt.ObjectProxy):
+    """ WrappedConnect wraps mongoengines 'connect' function to ensure
+        that all returned connections are wrapped for tracing.
+    """
+
+    _service = None
+    _tracer = None
+
+    def __init__(self, connect, tracer, service):
+        super(WrappedConnect, self).__init__(connect)
+        self._service = service
+        self._tracer = tracer
+
+    def __call__(self, *args, **kwargs):
+        client = self.__wrapped__(*args, **kwargs)
+        if _is_traced(client):
+            return client
+        # mongoengine uses pymongo internally, so we can just piggyback on the
+        # existing pymongo integration and make sure that the connections it
+        # uses internally are traced.
+        return trace_mongo_client(
+            client,
+            tracer=self._tracer,
+            service=self._service)
+
+
+def _is_traced(client):
+    return isinstance(client, wrapt.ObjectProxy)
+

setup.py

@@ -11,6 +11,7 @@
     'django',
     'elasticsearch',
     'flask',
+    'mongoengine',
     'psycopg2',
     'pymongo',
     'redis',

tests/contrib/mongoengine/__init__.py

@@ -0,0 +1,127 @@
+
+# stdib
+import time
+
+# 3p
+from nose.tools import eq_
+from mongoengine import (
+    connect,
+    Document,
+    StringField
+)
+
+
+# project
+from ddtrace import Tracer
+from ddtrace.contrib.mongoengine import trace_mongoengine
+from ...test_tracer import DummyWriter
+
+
+class Artist(Document):
+    first_name = StringField(max_length=50)
+    last_name = StringField(max_length=50)
+
+
+def test_insert_update_delete_query():
+    tracer = Tracer()
+    tracer.writer = DummyWriter()
+
+    # patch the mongo db connection
+    traced_connect = trace_mongoengine(tracer, service='my-mongo')
+    traced_connect()
+
+    start = time.time()
+    Artist.drop_collection()
+    end = time.time()
+
+    # ensure we get a drop collection span
+    spans = tracer.writer.pop()
+    eq_(len(spans), 1)
+    span = spans[0]
+    eq_(span.resource, 'drop artist')
+    eq_(span.span_type, 'mongodb')
+    eq_(span.service, 'my-mongo')
+    _assert_timing(span, start, end)
+
+    start = end
+    joni = Artist()
+    joni.first_name = 'Joni'
+    joni.last_name = 'Mitchell'
+    joni.save()
+    end = time.time()
+
+    # ensure we get an insert span
+    spans = tracer.writer.pop()
+    eq_(len(spans), 1)
+    span = spans[0]
+    eq_(span.resource, 'insert artist')
+    eq_(span.span_type, 'mongodb')
+    eq_(span.service, 'my-mongo')
+    _assert_timing(span, start, end)
+
+    # ensure full scans work
+    start = time.time()
+    artists = [a for a in Artist.objects]
+    end = time.time()
+    eq_(len(artists), 1)
+    eq_(artists[0].first_name, 'Joni')
+    eq_(artists[0].last_name, 'Mitchell')
+
+    spans = tracer.writer.pop()
+    eq_(len(spans), 1)
+    span = spans[0]
+    eq_(span.resource, 'query artist {}')
+    eq_(span.span_type, 'mongodb')
+    eq_(span.service, 'my-mongo')
+    _assert_timing(span, start, end)
+
+    # ensure filtered queries work
+    start = time.time()
+    artists = [a for a in Artist.objects(first_name="Joni")]
+    end = time.time()
+    eq_(len(artists), 1)
+    joni = artists[0]
+    eq_(artists[0].first_name, 'Joni')
+    eq_(artists[0].last_name, 'Mitchell')
+
+    spans = tracer.writer.pop()
+    eq_(len(spans), 1)
+    span = spans[0]
+    eq_(span.resource, "query artist {'first_name': '?'}")
+    eq_(span.span_type, 'mongodb')
+    eq_(span.service, 'my-mongo')
+    _assert_timing(span, start, end)
+
+    # ensure updates work
+    start = time.time()
+    joni.last_name = 'From Saskatoon'
+    joni.save()
+    end = time.time()
+
+    spans = tracer.writer.pop()
+    eq_(len(spans), 1)
+    span = spans[0]
+    eq_(span.resource, "update artist {'_id': '?'}")
+    eq_(span.span_type, 'mongodb')
+    eq_(span.service, 'my-mongo')
+    _assert_timing(span, start, end)
+
+    # ensure deletes
+    start = time.time()
+    joni.delete()
+    end = time.time()
+
+    spans = tracer.writer.pop()
+    eq_(len(spans), 1)
+    span = spans[0]
+    eq_(span.resource, "delete artist {'_id': '?'}")
+    eq_(span.span_type, 'mongodb')
+    eq_(span.service, 'my-mongo')
+    _assert_timing(span, start, end)
+
+
+
+
+def _assert_timing(span, start, end):
+    assert start < span.start < end
+    assert span.duration < end - start