Pymongo instrumentation: Capturing statements for additional command types

CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+### Added
+- `opentelemetry-instrumentation-pymongo` Add full command statement capturing for additional command types
+  ([#1955](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1955))
+
 ### Fixed
 
 - Fix version of Flask dependency `werkzeug`

instrumentation/opentelemetry-instrumentation-pymongo/src/opentelemetry/instrumentation/pymongo/__init__.py

@@ -83,7 +83,7 @@ def failed_hook(span, event):
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from opentelemetry.instrumentation.pymongo.package import _instruments
 from opentelemetry.instrumentation.pymongo.utils import (
-    COMMAND_TO_ATTRIBUTE_MAPPING,
+    COMMAND_TO_ATTRIBUTES_MAPPING,
 )
 from opentelemetry.instrumentation.pymongo.version import __version__
 from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
@@ -207,10 +207,11 @@ def _pop_span(self, event):
 
     def _get_statement_by_command_name(self, command_name, event):
         statement = command_name
-        command_attribute = COMMAND_TO_ATTRIBUTE_MAPPING.get(command_name)
-        command = event.command.get(command_attribute)
-        if command and self.capture_statement:
-            statement += " " + str(command)
+        command_attributes = COMMAND_TO_ATTRIBUTES_MAPPING.get(command_name)
+        if self.capture_statement:
+            for attribute in command_attributes:
+                if command := event.command.get(attribute):
+                    statement += " " + str(command)
         return statement
 
 

instrumentation/opentelemetry-instrumentation-pymongo/src/opentelemetry/instrumentation/pymongo/utils.py

@@ -12,9 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-COMMAND_TO_ATTRIBUTE_MAPPING = {
-    "insert": "documents",
-    "delete": "deletes",
-    "update": "updates",
-    "find": "filter",
+COMMAND_TO_ATTRIBUTES_MAPPING = {
+    "insert": ["documents"],
+    "delete": ["deletes"],
+    "update": ["updates"],
+    "find": ["filter"],
+    "findAndModify": ["query", "update"],
+    "aggregate": ["pipeline"],
+    "createIndexes": ["indexes"],
 }

tests/opentelemetry-docker-tests/tests/pymongo/test_pymongo_functional.py

@@ -14,7 +14,7 @@
 
 import os
 
-from pymongo import MongoClient
+from pymongo import MongoClient, IndexModel, ASCENDING
 
 from opentelemetry import trace as trace_api
 from opentelemetry.instrumentation.pymongo import PymongoInstrumentor
@@ -113,6 +113,30 @@ def test_find(self):
         expected_db_statement = "find {'name': 'testName'}"
         self.validate_spans(expected_db_statement)
 
+    def test_find_one_and_update(self):
+        """Should create a child span for find_one_and_update"""
+        with self._tracer.start_as_current_span("rootSpan"):
+            self._collection.find_one_and_update(
+                {"name": "testName"}, {"$set": {"value": "someOtherValue"}}
+            )
+
+        expected_db_statement = "findAndModify {'name': 'testName'} {'$set': {'value': 'someOtherValue'}}"
+        self.validate_spans(expected_db_statement)
+
+    def test_aggregate(self):
+        """Should create a child span for aggregate"""
+        with self._tracer.start_as_current_span("rootSpan"):
+            self._collection.aggregate(
+                [
+                    {"$match": {"field": "value"}},
+                    {"$group": {"_id": "$anotherField", "count": {"$sum": 1}}},
+                ]
+            )
+
+
+        expected_db_statement = "aggregate [{'$match': {'field': 'value'}}, {'$group': {'_id': '$anotherField', 'count': {'$sum': 1}}}]"
+        self.validate_spans(expected_db_statement)
+
     def test_delete(self):
         """Should create a child span for delete"""
         with self._tracer.start_as_current_span("rootSpan"):
@@ -123,6 +147,15 @@ def test_delete(self):
         )
         self.validate_spans(expected_db_statement)
 
+    def test_create_indexes(self):
+        """Should create a child span for create_indexes"""
+        some_index = IndexModel([("anotherField", ASCENDING)], name="some_index")
+        with self._tracer.start_as_current_span("rootSpan"):
+            self._collection.create_indexes([some_index])
+
+        expected_db_statement = "createIndexes [{'name': 'some_index', 'key': SON([('anotherField', 1)])}]"
+        self.validate_spans(expected_db_statement)
+
     def test_find_without_capture_statement(self):
         """Should create a child span for find"""
         self.instrumentor._commandtracer_instance.capture_statement = False