Pymongo instrumentation (#217)

* Add mongodb to the CircleCI build stack

* Add basic pymongo command events listener implementation

* Register MongoCommandTracer as a global event handler on startup

* Preserve MongoDB command within tracer until execution is complete

* Register mongo as an exit span

* Initiate a new span each time a command is sent to MongoDB

* Attach mongo command json to the span

* Add MongoDBData span type

* Send mongo spans and MongoDBData

* Send mongo.json and mongo.filter span tags as JSON strings

* Address mapreduce command case change in pymongo-3.9.0+

* Add myself to the contributors list and update the copyright year

Author: Andrew Slotin

.circleci/config.yml

@@ -16,6 +16,7 @@ jobs:
       - image: circleci/redis:5.0.4
       - image: rabbitmq:3.5.4
       - image: couchbase/server-sandbox:5.5.0
+      - image: circleci/mongo:4.2.3-ram
 
     working_directory: ~/repo
 
@@ -70,6 +71,7 @@ jobs:
       - image: circleci/redis:5.0.4
       - image: rabbitmq:3.5.4
       - image: couchbase/server-sandbox:5.5.0
+      - image: circleci/mongo:4.2.3-ram
 
     working_directory: ~/repo
 
@@ -122,6 +124,7 @@ jobs:
       - image: circleci/redis:5.0.4
       - image: rabbitmq:3.5.4
       - image: couchbase/server-sandbox:5.5.0
+      - image: circleci/mongo:4.2.3-ram
 
     working_directory: ~/repo
 

instana/__init__.py

@@ -25,8 +25,8 @@
 import pkg_resources
 
 __author__ = 'Instana Inc.'
-__copyright__ = 'Copyright 2019 Instana Inc.'
-__credits__ = ['Pavlo Baron', 'Peter Giacomo Lombardo']
+__copyright__ = 'Copyright 2020 Instana Inc.'
+__credits__ = ['Pavlo Baron', 'Peter Giacomo Lombardo', 'Andrey Slotin']
 __license__ = 'MIT'
 __maintainer__ = 'Peter Giacomo Lombardo'
 __email__ = '[email protected]'
@@ -82,6 +82,7 @@ def boot_agent():
         from .instrumentation import sudsjurko
         from .instrumentation import urllib3
         from .instrumentation.django import middleware
+        from .instrumentation import pymongo
 
     # Hooks
     from .hooks import hook_uwsgi

instana/instrumentation/pymongo.py

@@ -0,0 +1,95 @@
+from __future__ import absolute_import
+
+from ..log import logger
+from ..singletons import tracer
+
+try:
+    import pymongo
+    from pymongo import monitoring
+    from bson import json_util
+
+    class MongoCommandTracer(monitoring.CommandListener):
+        def __init__(self):
+            self.__active_commands = {}
+
+        def started(self, event):
+            parent_span = tracer.active_span
+
+            # return early if we're not tracing
+            if parent_span is None:
+                return
+
+            with tracer.start_active_span("mongo", child_of=parent_span) as scope:
+                self._collect_connection_tags(scope.span, event)
+                self._collect_command_tags(scope.span, event)
+
+                # include collection name into the namespace if provided
+                if event.command.has_key(event.command_name):
+                    scope.span.set_tag("collection", event.command.get(event.command_name))
+
+                self.__active_commands[event.request_id] = scope
+
+        def succeeded(self, event):
+            active_span = self.__active_commands.pop(event.request_id, None)
+
+            # return early if we're not tracing
+            if active_span is None:
+                return
+
+        def failed(self, event):
+            active_span = self.__active_commands.pop(event.request_id, None)
+
+            # return early if we're not tracing
+            if active_span is None:
+                return
+
+            active_span.log_exception(event.failure)
+
+        def _collect_connection_tags(self, span, event):
+            (host, port) = event.connection_id
+
+            span.set_tag("driver", "pymongo")
+            span.set_tag("host", host)
+            span.set_tag("port", str(port))
+            span.set_tag("db", event.database_name)
+
+        def _collect_command_tags(self, span, event):
+            """
+            Extract MongoDB command name and arguments and attach it to the span
+            """
+            cmd = event.command_name
+            span.set_tag("command", cmd)
+
+            for key in ["filter", "query"]:
+                if event.command.has_key(key):
+                    span.set_tag("filter", json_util.dumps(event.command.get(key)))
+                    break
+
+            # The location of command documents within the command object depends on the name
+            # of this command. This is the name -> command object key mapping
+            cmd_doc_locations = {
+                "insert": "documents",
+                "update": "updates",
+                "delete": "deletes",
+                "aggregate": "pipeline"
+            }
+
+            cmd_doc = None
+            if cmd in cmd_doc_locations:
+                cmd_doc = event.command.get(cmd_doc_locations[cmd])
+            elif cmd.lower() == "mapreduce": # mapreduce command was renamed to mapReduce in pymongo 3.9.0
+                # mapreduce command consists of two mandatory parts: map and reduce
+                cmd_doc = {
+                    "map": event.command.get("map"),
+                    "reduce": event.command.get("reduce")
+                }
+
+            if cmd_doc is not None:
+                span.set_tag("json", json_util.dumps(cmd_doc))
+
+    monitoring.register(MongoCommandTracer())
+
+    logger.debug("Instrumenting pymongo")
+
+except ImportError:
+    pass

instana/json_span.py

@@ -97,6 +97,15 @@ class MySQLData(BaseSpan):
     error = None
 
 
+class MongoDBData(BaseSpan):
+    service = None
+    namespace = None
+    command = None
+    filter = None
+    json = None
+    error = None
+
+
 class PostgresData(BaseSpan):
     db = None
     host = None

instana/recorder.py

@@ -10,7 +10,7 @@
 import instana.singletons
 
 from .json_span import (CassandraData, CouchbaseData, CustomData, Data, HttpData, JsonSpan, LogData,
-                        MySQLData, PostgresData, RabbitmqData, RedisData, RenderData,
+                        MongoDBData, MySQLData, PostgresData, RabbitmqData, RedisData, RenderData,
                         RPCData, SDKData, SoapData, SQLAlchemyData)
 
 from .log import logger
@@ -25,15 +25,16 @@
 class InstanaRecorder(SpanRecorder):
     THREAD_NAME = "Instana Span Reporting"
     registered_spans = ("aiohttp-client", "aiohttp-server", "cassandra", "couchbase", "django", "log",
-                        "memcache", "mysql", "postgres", "rabbitmq", "redis", "render", "rpc-client",
+                        "memcache", "mongo", "mysql", "postgres", "rabbitmq", "redis", "render", "rpc-client",
                         "rpc-server", "sqlalchemy", "soap", "tornado-client", "tornado-server",
                         "urllib3", "wsgi")
 
     http_spans = ("aiohttp-client", "aiohttp-server", "django", "http", "soap", "tornado-client",
                   "tornado-server", "urllib3", "wsgi")
 
-    exit_spans = ("aiohttp-client", "cassandra", "couchbase", "log", "memcache", "mysql", "postgres",
-                  "rabbitmq", "redis", "rpc-client", "sqlalchemy", "soap", "tornado-client", "urllib3")
+    exit_spans = ("aiohttp-client", "cassandra", "couchbase", "log", "memcache", "mongo", "mysql", "postgres",
+                  "rabbitmq", "redis", "rpc-client", "sqlalchemy", "soap", "tornado-client", "urllib3",
+                  "pymongo")
 
     entry_spans = ("aiohttp-server", "django", "wsgi", "rabbitmq", "rpc-server", "tornado-server")
 
@@ -237,6 +238,16 @@ def build_registered_span(self, span):
                 tskey = list(data.custom.logs.keys())[0]
                 data.pg.error = data.custom.logs[tskey]['message']
 
+        elif span.operation_name == "mongo":
+            service = "%s:%s" % (span.tags.pop('host', None), span.tags.pop('port', None))
+            namespace = "%s.%s" % (span.tags.pop('db', "?"), span.tags.pop('collection', "?"))
+            data.mongo = MongoDBData(service=service,
+                                     namespace=namespace,
+                                     command=span.tags.pop('command', None),
+                                     filter=span.tags.pop('filter', None),
+                                     json=span.tags.pop('json', None),
+                                     error=span.tags.pop('command', None))
+
         elif span.operation_name == "log":
             data.log = {}
             # use last special key values

setup.py

@@ -83,6 +83,7 @@ def check_setuptools():
               'pyOpenSSL>=16.1.0;python_version<="2.7"',
               'pytest>=3.0.1',
               'psycopg2>=2.7.1',
+              'pymongo>=3.7.0',
               'redis>3.0.0',
               'requests>=2.17.1',
               'sqlalchemy>=1.1.15',

tests/helpers.py

@@ -48,6 +48,14 @@
 testenv['redis_host'] = os.environ.get('REDIS_HOST', '127.0.0.1')
 
 
+"""
+MongoDB Environment
+"""
+testenv['mongodb_host'] = os.environ.get('MONGO_HOST', '127.0.0.1')
+testenv['mongodb_port'] = os.environ.get('MONGO_PORT', '27017')
+testenv['mongodb_user'] = os.environ.get('MONGO_USER', None)
+testenv['mongodb_pw'] = os.environ.get('MONGO_PW', None)
+
 def get_first_span_by_name(spans, name):
     for span in spans:
         if span.n == name: