Rd 3761 support pymongo (#152)

* support pymongo

Author: saartochner-lumigo

requirements.txt

@@ -6,3 +6,4 @@ pytest-cov==2.6.1
 capturer==2.4
 attrs==19.1.0
 requests==2.24.0
+pymongo==3.11.0

src/lumigo_tracer/spans_container.py

@@ -144,6 +144,12 @@ def get_last_span(self) -> Optional[dict]:
             return None
         return self.spans[-1]
 
+    def get_span_by_id(self, span_id: str) -> Optional[dict]:
+        for span in self.spans:
+            if span.get("id") == span_id:
+                return span
+        return None
+
     def remove_last_span(self) -> Optional[dict]:
         return self.spans.pop() if self.spans else None
 

src/lumigo_tracer/wrappers/__init__.py

@@ -1,4 +1,5 @@
 from .http.sync_http_wrappers import wrap_http_calls
+from .pymongo.pymongo_wrapper import wrap_pymongo
 
 
 already_wrapped = False
@@ -8,4 +9,5 @@ def wrap():
     global already_wrapped
     if not already_wrapped:
         wrap_http_calls()
+        wrap_pymongo()
         already_wrapped = True

src/lumigo_tracer/wrappers/pymongo/__init__.py

@@ -0,0 +1,73 @@
+from typing import Dict
+
+import time
+import uuid
+
+from lumigo_tracer.lumigo_utils import lumigo_safe_execute, get_logger, lumigo_dumps
+from lumigo_tracer.spans_container import SpansContainer
+
+try:
+    from pymongo import monitoring
+except Exception:
+    monitoring = None
+
+if not monitoring:
+    LumigoMongoMonitoring = None
+else:
+
+    class LumigoMongoMonitoring(monitoring.CommandListener):  # type: ignore
+        request_to_span_id: Dict[str, str] = {}
+        MONGO_SPAN = "mongoDb"
+
+        def started(self, event):
+            with lumigo_safe_execute("pymongo started"):
+                span_id = str(uuid.uuid4())
+                LumigoMongoMonitoring.request_to_span_id[event.request_id] = span_id
+                SpansContainer.get_span().add_span(
+                    {
+                        "id": span_id,
+                        "type": self.MONGO_SPAN,
+                        "started": int(time.time() * 1000),
+                        "databaseName": event.database_name,
+                        "commandName": event.command_name,
+                        "request": lumigo_dumps(event.command),
+                        "mongoRequestId": event.request_id,
+                        "mongoOperationId": event.operation_id,
+                        "mongoConnectionId": event.connection_id,
+                    }
+                )
+
+        def succeeded(self, event):
+            with lumigo_safe_execute("pymongo succeed"):
+                if event.request_id not in LumigoMongoMonitoring.request_to_span_id:
+                    get_logger().warning("Mongo span ended without a record on its start")
+                    return
+                span_id = LumigoMongoMonitoring.request_to_span_id.pop(event.request_id)
+                span = SpansContainer.get_span().get_span_by_id(span_id)
+                span.update(
+                    {
+                        "ended": span["started"] + (event.duration_micros / 1000),
+                        "response": lumigo_dumps(event.reply),
+                    }
+                )
+
+        def failed(self, event):
+            with lumigo_safe_execute("pymongo failed"):
+                if event.request_id not in LumigoMongoMonitoring.request_to_span_id:
+                    get_logger().warning("Mongo span ended without a record on its start")
+                    return
+                span_id = LumigoMongoMonitoring.request_to_span_id.pop(event.request_id)
+                span = SpansContainer.get_span().get_span_by_id(span_id)
+                span.update(
+                    {
+                        "ended": span["started"] + (event.duration_micros / 1000),
+                        "error": lumigo_dumps(event.failure),
+                    }
+                )
+
+
+def wrap_pymongo():
+    with lumigo_safe_execute("wrap pymogno"):
+        if monitoring:
+            get_logger().debug("wrapping pymongo")
+            monitoring.register(LumigoMongoMonitoring())

src/lumigo_tracer/wrappers/pymongo/pymongo_wrapper.py

@@ -170,3 +170,12 @@ def test_get_tags_len():
     SpansContainer.get_span().add_tag("k0", "v0")
     SpansContainer.get_span().add_tag("k1", "v1")
     assert SpansContainer.get_span().get_tags_len() == 2
+
+
+def test_get_span_by_id():
+    container = SpansContainer.get_span()
+    container.add_span({"id": 1, "extra": "a"})
+    container.add_span({"id": 2, "extra": "b"})
+    container.add_span({"id": 3, "extra": "c"})
+    assert SpansContainer.get_span().get_span_by_id(2)["extra"] == "b"
+    assert SpansContainer.get_span().get_span_by_id(5) is None

src/test/unit/test_spans_container.py

@@ -0,0 +1,89 @@
+from types import SimpleNamespace
+
+import pytest
+
+from lumigo_tracer.spans_container import SpansContainer
+from lumigo_tracer.wrappers.pymongo.pymongo_wrapper import LumigoMongoMonitoring
+
+
+@pytest.fixture
+def start_event():
+    return SimpleNamespace(
+        database_name="dname",
+        command_name="cname",
+        command="cmd",
+        request_id="rid",
+        operation_id="oid",
+        connection_id="cid",
+    )
+
+
+@pytest.fixture
+def success_event():
+    return SimpleNamespace(duration_micros=5000, reply={"code": 200}, request_id="rid")
+
+
+@pytest.fixture
+def fail_event():
+    return SimpleNamespace(duration_micros=5000, failure={"code": 500}, request_id="rid")
+
+
+def test_pymongo_happy_flow(monkeypatch, start_event, success_event):
+    monitor = LumigoMongoMonitoring()
+    monitor.started(start_event)
+    monitor.succeeded(success_event)
+
+    spans = SpansContainer.get_span().spans
+    assert len(spans) == 1
+    assert spans[0]["request"] == '"cmd"'
+    assert spans[0]["ended"] > spans[0]["started"]
+    assert spans[0]["response"] == '{"code": 200}'
+    assert "error" not in spans[0]
+
+
+def test_pymongo_only_start(monkeypatch, start_event):
+    monitor = LumigoMongoMonitoring()
+    monitor.started(start_event)
+
+    spans = SpansContainer.get_span().spans
+    assert len(spans) == 1
+    assert spans[0]["request"] == '"cmd"'
+    assert "duration" not in spans[0]
+
+
+def test_pymongo_error(monkeypatch, start_event, fail_event):
+    monitor = LumigoMongoMonitoring()
+    monitor.started(start_event)
+    monitor.failed(fail_event)
+
+    spans = SpansContainer.get_span().spans
+    assert len(spans) == 1
+    assert spans[0]["request"] == '"cmd"'
+    assert spans[0]["ended"] > spans[0]["started"]
+    assert spans[0]["error"] == '{"code": 500}'
+    assert "response" not in spans[0]
+
+
+def test_pymongo_concurrent_events(monkeypatch, start_event, success_event):
+    monitor = LumigoMongoMonitoring()
+    monitor.started(start_event)
+    monitor.started(
+        SimpleNamespace(
+            database_name="dname",
+            command_name="cname",
+            command="cmd",
+            request_id="rid2",
+            operation_id="oid",
+            connection_id="cid",
+        )
+    )
+    monitor.succeeded(success_event)
+
+    spans = SpansContainer.get_span().spans
+    assert len(spans) == 2
+    assert spans[0]["mongoRequestId"] == "rid"
+    assert spans[0]["ended"] > spans[0]["started"]
+    assert spans[0]["response"] == '{"code": 200}'
+
+    assert spans[1]["mongoRequestId"] == "rid2"
+    assert "ended" not in spans[1]

src/test/unit/wrappers/pymongo/test_pymongo_wrapper.py

@@ -0,0 +1,16 @@
+import importlib
+import sys
+
+import lumigo_tracer
+
+
+def test_wrapping_without_libraries(monkeypatch):
+    # remove library from path and reload module
+    monkeypatch.setitem(sys.modules, "pymongo", None)
+    wrapper = importlib.reload(lumigo_tracer.wrappers.pymongo.pymongo_wrapper)
+    assert wrapper.LumigoMongoMonitoring is None
+
+    lumigo_tracer.wrappers.wrap()  # should succeed
+
+    monkeypatch.undo()
+    importlib.reload(lumigo_tracer.wrappers.pymongo.pymongo_wrapper)