diff --git a/ddtrace/contrib/mongoengine/__init__.py b/ddtrace/contrib/mongoengine/__init__.py new file mode 100644 index 00000000000..9c15a12e606 --- /dev/null +++ b/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'] diff --git a/ddtrace/contrib/mongoengine/trace.py b/ddtrace/contrib/mongoengine/trace.py new file mode 100644 index 00000000000..fa1febbbc4c --- /dev/null +++ b/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) + diff --git a/setup.py b/setup.py index 96dc987d63a..b9b4ccc0086 100644 --- a/setup.py +++ b/setup.py @@ -11,6 +11,7 @@ 'django', 'elasticsearch', 'flask', + 'mongoengine', 'psycopg2', 'pymongo', 'redis', diff --git a/tests/contrib/mongoengine/__init__.py b/tests/contrib/mongoengine/__init__.py new file mode 100644 index 00000000000..8bbebc53218 --- /dev/null +++ b/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