feat(pymongo): add PyMongo integration

Adds breadcrumbs and performance traces for PyMongo queries using an
official monitoring API. Integration is similar to the one available in
OpenTelemetry, tags set to values recommended for attributes by OT as
specified in `Span Operations` guidelines.

PyMongo version selection explanation:
* 3.1 - introduction of monitoring API. Only Python 2.7 and 3.6
supported.
* 3.12 - latest 3.x release, support for 2.7, 3.6-3.9 (3.7-3.9 added in
various minor releases between 3.1 and 3.12).
* 4.0 - no support for 2.7, added support for 3.10.
* 4.1 - no support for 3.6.0-3.6.1.
* 4.2 - no support for any 3.6.

Author: Agalin

sentry_sdk/integrations/pymongo.py

@@ -0,0 +1,120 @@
+from __future__ import absolute_import
+
+from sentry_sdk import Hub
+from sentry_sdk.integrations import DidNotEnable, Integration
+from sentry_sdk.tracing import Span
+from sentry_sdk.utils import capture_internal_exceptions
+
+from sentry_sdk._types import MYPY
+
+try:
+    from pymongo import monitoring
+except ImportError:
+    raise DidNotEnable("Pymongo not installed")
+
+if MYPY:
+    from typing import Dict
+
+    from pymongo.monitoring import (
+        CommandFailedEvent,
+        CommandStartedEvent,
+        CommandSucceededEvent,
+    )
+
+
+class CommandTracer(monitoring.CommandListener):
+    def __init__(self):
+        self._ongoing_operations = {}  # type: Dict[int, Span]
+
+    def _operation_key(self, event):
+        # type: (CommandFailedEvent | CommandStartedEvent | CommandSucceededEvent) -> int
+        return event.request_id
+
+    def started(self, event):
+        # type: (CommandStartedEvent) -> None
+        hub = Hub.current
+        if hub.get_integration(PyMongoIntegration) is None:
+            return
+        with capture_internal_exceptions():
+            command = dict(event.command)
+
+            command.pop("$db", None)
+            command.pop("$clusterTime", None)
+            command.pop("$signature", None)
+
+            if event.command_name:
+                op = "db." + event.command_name
+            else:
+                op = "db"
+
+            tags = {
+                "db.name": event.database_name,
+                "db.system": "mongodb",
+                "db.operation": event.command_name,
+            }
+
+            try:
+                tags["net.peer.name"] = event.connection_id[0]
+                tags["net.peer.port"] = str(event.connection_id[1])
+            except TypeError:
+                pass
+
+            data = {"operation_ids": {}}
+
+            data["operation_ids"]["operation"] = event.operation_id
+            data["operation_ids"]["request"] = event.request_id
+
+            try:
+                lsid = command.pop("lsid")["id"]
+                data["operation_ids"]["session"] = str(lsid)
+            except KeyError:
+                pass
+
+            query = "{} {}".format(event.command_name, command)
+            span = hub.start_span(op=op, description=query)
+
+            for tag, value in tags.items():
+                span.set_tag(tag, value)
+
+            for key, value in data.items():
+                span.set_data(key, value)
+
+            with capture_internal_exceptions():
+                hub.add_breadcrumb(message=query, category="query", type=op, data=tags)
+
+            self._ongoing_operations[self._operation_key(event)] = span.__enter__()
+
+    def failed(self, event):
+        # type: (CommandFailedEvent) -> None
+        hub = Hub.current
+        if hub.get_integration(PyMongoIntegration) is None:
+            return
+
+        try:
+            span = self._ongoing_operations.pop(self._operation_key(event))
+            span.set_status("internal_error")
+            span.__exit__(None, None, None)
+        except KeyError:
+            return
+
+    def succeeded(self, event):
+        # type: (CommandSucceededEvent) -> None
+        hub = Hub.current
+        if hub.get_integration(PyMongoIntegration) is None:
+            return
+
+        try:
+            span = self._ongoing_operations.pop(self._operation_key(event))
+            span.set_status("ok")
+            span.__exit__(None, None, None)
+        except KeyError:
+            pass
+
+
+class PyMongoIntegration(Integration):
+    identifier = "pymongo"
+
+    @staticmethod
+    def setup_once():
+        # type: () -> None
+        monitoring.register(CommandTracer())

setup.py

@@ -62,6 +62,7 @@ def get_file_text(file_name):
         "httpx": ["httpx>=0.16.0"],
         "starlette": ["starlette>=0.19.1"],
         "fastapi": ["fastapi>=0.79.0"],
+        "pymongo": ["pymongo>=3.1"],
     },
     classifiers=[
         "Development Status :: 5 - Production/Stable",

tests/integrations/pymongo/__init__.py

@@ -0,0 +1,3 @@
+import pytest
+
+pytest.importorskip("pymongo")

tests/integrations/pymongo/test_pymongo.py

@@ -0,0 +1,96 @@
+from sentry_sdk import capture_message, start_transaction
+from sentry_sdk.integrations.pymongo import PyMongoIntegration
+
+from mockupdb import MockupDB, OpQuery
+from pymongo import MongoClient
+import pytest
+
+
+@pytest.fixture(scope="session")
+def mongo_server():
+    server = MockupDB(verbose=True)
+    server.autoresponds("ismaster", maxWireVersion=6)
+    server.run()
+    server.autoresponds(
+        {"find": "test_collection"}, cursor={"id": 123, "firstBatch": []}
+    )
+    # Find query changed somewhere between PyMongo 3.1 and 3.12.
+    # This line is to respond to "find" queries sent by old PyMongo the same way it's done above.
+    server.autoresponds(OpQuery({"foobar": 1}), cursor={"id": 123, "firstBatch": []})
+    server.autoresponds({"insert": "test_collection"}, ok=1)
+    server.autoresponds({"insert": "erroneous"}, ok=0, errmsg="test error")
+    yield server
+    server.stop()
+
+
+def test_transactions(sentry_init, capture_events, mongo_server):
+    sentry_init(integrations=[PyMongoIntegration()], traces_sample_rate=1.0)
+    events = capture_events()
+
+    connection = MongoClient(mongo_server.uri)
+
+    with start_transaction():
+        list(
+            connection["test_db"]["test_collection"].find({"foobar": 1})
+        )  # force query execution
+        connection["test_db"]["test_collection"].insert_one({"foo": 2})
+        try:
+            connection["test_db"]["erroneous"].insert_many([{"bar": 3}, {"baz": 3}])
+            pytest.fail("Request should raise")
+        except Exception:
+            pass
+
+    (event,) = events
+    (find, insert_success, insert_fail) = event["spans"]
+
+    common_tags = {
+        "db.name": "test_db",
+        "db.system": "mongodb",
+        "net.peer.name": mongo_server.host,
+        "net.peer.port": str(mongo_server.port),
+    }
+    for span in find, insert_success, insert_fail:
+        for field, value in common_tags.items():
+            assert span["tags"][field] == value
+
+    assert find["op"] == "db.find"
+    assert insert_success["op"] == "db.insert"
+    assert insert_fail["op"] == "db.insert"
+
+    assert find["tags"]["db.operation"] == "find"
+    assert insert_success["tags"]["db.operation"] == "insert"
+    assert insert_fail["tags"]["db.operation"] == "insert"
+
+    assert find["description"].startswith("find {")
+    assert insert_success["description"].startswith("insert {")
+    assert insert_fail["description"].startswith("insert {")
+
+    assert find["tags"]["status"] == "ok"
+    assert insert_success["tags"]["status"] == "ok"
+    assert insert_fail["tags"]["status"] == "internal_error"
+
+
+def test_breadcrumbs(sentry_init, capture_events, mongo_server):
+    sentry_init(integrations=[PyMongoIntegration()], traces_sample_rate=1.0)
+    events = capture_events()
+
+    connection = MongoClient(mongo_server.uri)
+
+    list(
+        connection["test_db"]["test_collection"].find({"foobar": 1})
+    )  # force query execution
+    capture_message("hi")
+
+    (event,) = events
+    (crumb,) = event["breadcrumbs"]["values"]
+
+    assert crumb["category"] == "query"
+    assert crumb["message"].startswith("find {")
+    assert crumb["type"] == "db.find"
+    assert crumb["data"] == {
+        "db.name": "test_db",
+        "db.system": "mongodb",
+        "db.operation": "find",
+        "net.peer.name": mongo_server.host,
+        "net.peer.port": str(mongo_server.port),
+    }

tox.ini

@@ -89,6 +89,11 @@ envlist =
 
     {py3.6,py3.7,py3.8,py3.9,py3.10}-httpx-{0.16,0.17}
 
+    {py2.7,py3.6}-pymongo-{3.1}
+    {py2.7,py3.6,py3.7,py3.8,py3.9}-pymongo-{3.12}
+    {py3.6,py3.7,py3.8,py3.9,py3.10}-pymongo-{4.0}
+    {py3.7,py3.8,py3.9,py3.10}-pymongo-{4.1,4.2}
+
 [testenv]
 deps =
     # if you change test-requirements.txt and your change is not being reflected
@@ -257,6 +262,13 @@ deps =
     httpx-0.16: httpx>=0.16,<0.17
     httpx-0.17: httpx>=0.17,<0.18
 
+    pymongo: mockupdb
+    pymongo-3.1: pymongo>=3.1,<3.2
+    pymongo-3.12: pymongo>=3.12,<4.0
+    pymongo-4.0: pymongo>=4.0,<4.1
+    pymongo-4.1: pymongo>=4.1,<4.2
+    pymongo-4.2: pymongo>=4.2,<4.3
+
 setenv =
     PYTHONDONTWRITEBYTECODE=1
     TESTPATH=tests
@@ -286,6 +298,7 @@ setenv =
     chalice: TESTPATH=tests/integrations/chalice
     boto3: TESTPATH=tests/integrations/boto3
     httpx: TESTPATH=tests/integrations/httpx
+    pymongo: TESTPATH=tests/integrations/pymongo
 
     COVERAGE_FILE=.coverage-{envname}
 passenv =
@@ -301,6 +314,7 @@ extras =
     bottle: bottle
     falcon: falcon
     quart: quart
+    pymongo: pymongo
 
 basepython =
     py2.7: python2.7