ext/pymongo: Small code rework

Pop span instead of getting and then removing it.
open-telemetry · Apr 23, 2020 · 8f7b1ca · 8f7b1ca
1 parent 8eaae6d
commit 8f7b1ca
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 24 deletions.
diff --git a/ext/opentelemetry-ext-pymongo/src/opentelemetry/ext/pymongo/__init__.py b/ext/opentelemetry-ext-pymongo/src/opentelemetry/ext/pymongo/__init__.py
@@ -129,39 +129,32 @@ def started(self, event: monitoring.CommandStartedEvent):
             if span is not None:
                 span.set_status(Status(StatusCanonicalCode.INTERNAL, str(ex)))
                 span.end()
-                self._remove_span(event)
+                self._pop_span(event)
 
     def succeeded(self, event: monitoring.CommandSucceededEvent):
         """ Method to handle a pymongo CommandSucceededEvent """
         if not self.enable:
             return
-        span = self._get_span(event)
-        if span is not None:
-            span.set_attribute(
-                "db.mongo.duration_micros", event.duration_micros
-            )
-            span.set_status(Status(StatusCanonicalCode.OK, event.reply))
-            span.end()
-            self._remove_span(event)
+        span = self._pop_span(event)
+        if span is None:
+            return
+        span.set_attribute("db.mongo.duration_micros", event.duration_micros)
+        span.set_status(Status(StatusCanonicalCode.OK, event.reply))
+        span.end()
 
     def failed(self, event: monitoring.CommandFailedEvent):
         """ Method to handle a pymongo CommandFailedEvent """
         if not self.enable:
             return
-        span = self._get_span(event)
-        if span is not None:
-            span.set_attribute(
-                "db.mongo.duration_micros", event.duration_micros
-            )
-            span.set_status(Status(StatusCanonicalCode.UNKNOWN, event.failure))
-            span.end()
-            self._remove_span(event)
-
-    def _get_span(self, event):
-        return self._span_dict.get(_get_span_dict_key(event))
-
-    def _remove_span(self, event):
-        self._span_dict.pop(_get_span_dict_key(event))
+        span = self._pop_span(event)
+        if span is None:
+            return
+        span.set_attribute("db.mongo.duration_micros", event.duration_micros)
+        span.set_status(Status(StatusCanonicalCode.UNKNOWN, event.failure))
+        span.end()
+
+    def _pop_span(self, event):
+        return self._span_dict.pop(_get_span_dict_key(event), None)
 
 
 def _get_span_dict_key(event):

diff --git a/ext/opentelemetry-ext-pymongo/tests/test_pymongo.py b/ext/opentelemetry-ext-pymongo/tests/test_pymongo.py
@@ -50,7 +50,7 @@ def test_started(self):
         # the memory exporter can't be used here because the span isn't ended
         # yet
         # pylint: disable=protected-access
-        span = command_tracer._get_span(mock_event)
+        span = command_tracer._pop_span(mock_event)
         self.assertIs(span.kind, trace_api.SpanKind.CLIENT)
         self.assertEqual(span.name, "mongodb.command_name.find")
         self.assertEqual(span.attributes["component"], "mongodb")