fix(pymongo): support 3.12.0 (#2648) (#2649)

* fix(pymongo): support 3.12.0

With pymongo 3.12 an API that was used for hooking span creation [was
renamed](mongodb/mongo-python-driver@3.11.4...3.12.0#diff-d5e3bba0b330986914b87fdbdbb9fe506eaf7f0d0cdd8b2f92081ff23620deaaL70).

The code was also cleaned up a bit by using the span context manager
rather than a try..finally which also means that exception information
will be included on spans.

* Only install methods in versions necessary

(cherry picked from commit 97f81ca)

Co-authored-by: Kyle Verhoog <[email protected]>

Author: mergify[bot]

ddtrace/contrib/pymongo/client.py

@@ -28,6 +28,9 @@
 log = get_logger(__name__)
 
 
+VERSION = pymongo.version_tuple
+
+
 class TracedMongoClient(ObjectProxy):
     def __init__(self, client=None, *args, **kwargs):
         # To support the former trace_mongo_client interface, we have to keep this old interface
@@ -114,35 +117,35 @@ def _datadog_trace_operation(self, operation):
             span.set_tag(ANALYTICS_SAMPLE_RATE_KEY, sample_rate)
         return span
 
-    # Pymongo >= 3.9
-    def run_operation_with_response(self, sock_info, operation, *args, **kwargs):
-        span = self._datadog_trace_operation(operation)
-        if not span:
-            return self.__wrapped__.run_operation_with_response(sock_info, operation, *args, **kwargs)
+    # Pymongo >= 3.12
+    if VERSION >= (3, 12, 0):
 
-        try:
-            result = self.__wrapped__.run_operation_with_response(sock_info, operation, *args, **kwargs)
+        def run_operation(self, sock_info, operation, *args, **kwargs):
+            with self._datadog_trace_operation(operation) as span:
+                result = self.__wrapped__.run_operation(sock_info, operation, *args, **kwargs)
+                if result and result.address:
+                    set_address_tags(span, result.address)
+                return result
 
-            if result and result.address:
-                set_address_tags(span, result.address)
-            return result
-        finally:
-            span.finish()
+    # Pymongo >= 3.9, <3.12
+    elif (3, 9, 0) <= VERSION < (3, 12, 0):
 
-    # Pymongo < 3.9
-    def send_message_with_response(self, operation, *args, **kwargs):
-        span = self._datadog_trace_operation(operation)
-        if not span:
-            return self.__wrapped__.send_message_with_response(operation, *args, **kwargs)
+        def run_operation_with_response(self, sock_info, operation, *args, **kwargs):
+            with self._datadog_trace_operation(operation) as span:
+                result = self.__wrapped__.run_operation_with_response(sock_info, operation, *args, **kwargs)
+                if result and result.address:
+                    set_address_tags(span, result.address)
+                return result
 
-        try:
-            result = self.__wrapped__.send_message_with_response(operation, *args, **kwargs)
+    # Pymongo < 3.9
+    else:
 
-            if result and result.address:
-                set_address_tags(span, result.address)
-            return result
-        finally:
-            span.finish()
+        def send_message_with_response(self, operation, *args, **kwargs):
+            with self._datadog_trace_operation(operation) as span:
+                result = self.__wrapped__.send_message_with_response(operation, *args, **kwargs)
+                if result and result.address:
+                    set_address_tags(span, result.address)
+                return result
 
     @contextlib.contextmanager
     def get_socket(self, *args, **kwargs):

releasenotes/notes/mongo-2f086c64dbbc8be1.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    Fix pymongo 3.12.0+ spans not being generated.

riotfile.py

@@ -414,6 +414,7 @@ def select_pys(min_version=MIN_PYTHON_VERSION, max_version=MAX_PYTHON_VERSION):
                             ">=3.8,<3.9",
                             ">=3.9,<3.10",
                             ">=3.10,<3.11",
+                            ">=3.12,<3.13",
                             latest,
                         ],
                         "mongoengine": latest,
@@ -433,6 +434,7 @@ def select_pys(min_version=MIN_PYTHON_VERSION, max_version=MAX_PYTHON_VERSION):
                             ">=3.8,<3.9",
                             ">=3.9,<3.10",
                             ">=3.10,<3.11",
+                            ">=3.12,<3.13",
                             latest,
                         ],
                         "mongoengine": latest,