chore(pymongo): update pymongo integration for compatibility with 4.5…

….x (#6736)

Pymongo 4.5.0 [changed the
name](mongodb/mongo-python-driver#1329) of the
`get_socket` function on which our integration is based. This pull
request adapts the integration to work with version 4.5.0 by changing
the expected function name based on the version.

Fixes #6723

## Checklist

- [x] Change(s) are motivated and described in the PR description.
- [x] Testing strategy is described if automated tests are not included
in the PR.
- [x] Risk is outlined (performance impact, potential for breakage,
maintainability, etc).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] [Library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
are followed. If no release note is required, add label
`changelog/no-changelog`.
- [x] Documentation is included (in-code, generated user docs, [public
corp docs](https://github.com/DataDog/documentation/)).
- [x] Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist

- [x] Title is accurate.
- [x] No unnecessary changes are introduced.
- [x] Description motivates each change.
- [x] Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes unless absolutely necessary.
- [x] Testing strategy adequately addresses listed risk(s).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] Release note makes sense to a user of the library.
- [x] Reviewer has explicitly acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment.
- [x] Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)

---------

Co-authored-by: Teague Bick <[email protected]>

.riot/requirements/1854ba5.txt

@@ -2,20 +2,20 @@
 # This file is autogenerated by pip-compile with Python 3.11
 # by the following command:
 #
-#    pip-compile --no-annotate --resolver=backtracking .riot/requirements/1854ba5.in
+#    pip-compile --no-annotate .riot/requirements/1854ba5.in
 #
-attrs==22.2.0
-coverage[toml]==7.2.2
-dnspython==2.3.0
+attrs==23.1.0
+coverage[toml]==7.3.0
+dnspython==2.4.2
 hypothesis==6.45.0
 iniconfig==2.0.0
-mock==5.0.1
+mock==5.1.0
 mongoengine==0.27.0
 opentracing==2.4.0
-packaging==23.0
-pluggy==1.0.0
-pymongo==4.3.3
-pytest==7.2.2
-pytest-cov==4.0.0
-pytest-mock==3.10.0
+packaging==23.1
+pluggy==1.2.0
+pymongo==4.5.0
+pytest==7.4.0
+pytest-cov==4.1.0
+pytest-mock==3.11.1
 sortedcontainers==2.4.0

.riot/requirements/1fd1158.txt

@@ -2,22 +2,22 @@
 # This file is autogenerated by pip-compile with Python 3.10
 # by the following command:
 #
-#    pip-compile --no-annotate --resolver=backtracking .riot/requirements/1fd1158.in
+#    pip-compile --no-annotate .riot/requirements/1fd1158.in
 #
-attrs==22.2.0
-coverage[toml]==7.2.2
-dnspython==2.3.0
-exceptiongroup==1.1.1
+attrs==23.1.0
+coverage[toml]==7.3.0
+dnspython==2.4.2
+exceptiongroup==1.1.3
 hypothesis==6.45.0
 iniconfig==2.0.0
-mock==5.0.1
+mock==5.1.0
 mongoengine==0.27.0
 opentracing==2.4.0
-packaging==23.0
-pluggy==1.0.0
-pymongo==4.3.3
-pytest==7.2.2
-pytest-cov==4.0.0
-pytest-mock==3.10.0
+packaging==23.1
+pluggy==1.2.0
+pymongo==4.5.0
+pytest==7.4.0
+pytest-cov==4.1.0
+pytest-mock==3.11.1
 sortedcontainers==2.4.0
 tomli==2.0.1

.riot/requirements/41bf6ef.txt

@@ -0,0 +1,23 @@
+#
+# This file is autogenerated by pip-compile with Python 3.9
+# by the following command:
+#
+#    pip-compile --no-annotate .riot/requirements/41bf6ef.in
+#
+attrs==23.1.0
+coverage[toml]==7.3.0
+dnspython==2.4.2
+exceptiongroup==1.1.3
+hypothesis==6.45.0
+iniconfig==2.0.0
+mock==5.1.0
+mongoengine==0.27.0
+opentracing==2.4.0
+packaging==23.1
+pluggy==1.2.0
+pymongo==4.5.0
+pytest==7.4.0
+pytest-cov==4.1.0
+pytest-mock==3.11.1
+sortedcontainers==2.4.0
+tomli==2.0.1

.riot/requirements/4bf8418.txt

@@ -0,0 +1,26 @@
+#
+# This file is autogenerated by pip-compile with Python 3.7
+# by the following command:
+#
+#    pip-compile --config=pyproject.toml --no-annotate --resolver=backtracking .riot/requirements/4bf8418.in
+#
+attrs==23.1.0
+coverage[toml]==7.2.7
+dnspython==2.3.0
+exceptiongroup==1.1.3
+hypothesis==6.45.0
+importlib-metadata==6.7.0
+iniconfig==2.0.0
+mock==5.1.0
+mongoengine==0.27.0
+opentracing==2.4.0
+packaging==23.1
+pluggy==1.2.0
+pymongo==4.5.0
+pytest==7.4.0
+pytest-cov==4.1.0
+pytest-mock==3.11.1
+sortedcontainers==2.4.0
+tomli==2.0.1
+typing-extensions==4.7.1
+zipp==3.15.0

.riot/requirements/c01e3b4.txt

@@ -2,7 +2,7 @@
 # This file is autogenerated by pip-compile with Python 3.10
 # by the following command:
 #
-#    pip-compile --no-annotate --resolver=backtracking .riot/requirements/c01e3b4.in
+#    pip-compile --no-annotate .riot/requirements/c01e3b4.in
 #
 attrs==22.2.0
 coverage[toml]==7.2.2

.riot/requirements/f92d9dc.txt

@@ -0,0 +1,23 @@
+#
+# This file is autogenerated by pip-compile with Python 3.8
+# by the following command:
+#
+#    pip-compile --no-annotate .riot/requirements/f92d9dc.in
+#
+attrs==23.1.0
+coverage[toml]==7.3.0
+dnspython==2.4.2
+exceptiongroup==1.1.3
+hypothesis==6.45.0
+iniconfig==2.0.0
+mock==5.1.0
+mongoengine==0.27.0
+opentracing==2.4.0
+packaging==23.1
+pluggy==1.2.0
+pymongo==4.5.0
+pytest==7.4.0
+pytest-cov==4.1.0
+pytest-mock==3.11.1
+sortedcontainers==2.4.0
+tomli==2.0.1

ddtrace/contrib/pymongo/client.py

@@ -143,7 +143,26 @@ def _datadog_trace_operation(self, operation):
             span.set_tag(ANALYTICS_SAMPLE_RATE_KEY, sample_rate)
         return span
 
-    # Pymongo >= 3.12
+    if VERSION >= (4, 5, 0):
+
+        @contextlib.contextmanager
+        def checkout(self, *args, **kwargs):
+            with self.__wrapped__.checkout(*args, **kwargs) as s:
+                if not isinstance(s, TracedSocket):
+                    s = TracedSocket(s)
+                ddtrace.Pin.get_from(self).onto(s)
+                yield s
+
+    else:
+
+        @contextlib.contextmanager
+        def get_socket(self, *args, **kwargs):
+            with self.__wrapped__.get_socket(*args, **kwargs) as s:
+                if not isinstance(s, TracedSocket):
+                    s = TracedSocket(s)
+                ddtrace.Pin.get_from(self).onto(s)
+                yield s
+
     if VERSION >= (3, 12, 0):
 
         def run_operation(self, sock_info, operation, *args, **kwargs):
@@ -159,7 +178,6 @@ def run_operation(self, sock_info, operation, *args, **kwargs):
                         set_query_rowcount(docs=result.docs, span=span)
                 return result
 
-    # Pymongo >= 3.9, <3.12
     elif (3, 9, 0) <= VERSION < (3, 12, 0):
 
         def run_operation_with_response(self, sock_info, operation, *args, **kwargs):
@@ -175,7 +193,6 @@ def run_operation_with_response(self, sock_info, operation, *args, **kwargs):
                         set_query_rowcount(docs=result.docs, span=span)
                 return result
 
-    # Pymongo < 3.9
     else:
 
         def send_message_with_response(self, operation, *args, **kwargs):
@@ -200,14 +217,6 @@ def send_message_with_response(self, operation, *args, **kwargs):
                                     set_query_rowcount(docs=docs, span=span)
                 return result
 
-    @contextlib.contextmanager
-    def get_socket(self, *args, **kwargs):
-        with self.__wrapped__.get_socket(*args, **kwargs) as s:
-            if not isinstance(s, TracedSocket):
-                s = TracedSocket(s)
-            ddtrace.Pin.get_from(self).onto(s)
-            yield s
-
     @staticmethod
     def _is_query(op):
         # NOTE: _Query should always have a spec field

ddtrace/contrib/pymongo/patch.py

@@ -28,6 +28,9 @@
 # Original Client class
 _MongoClient = pymongo.MongoClient
 
+_VERSION = pymongo.version_tuple
+_CHECKOUT_FN_NAME = "get_socket" if _VERSION < (4, 5) else "checkout"
+
 
 def patch():
     patch_pymongo_module()
@@ -50,15 +53,15 @@ def patch_pymongo_module():
     # Whenever a pymongo command is invoked, the lib either:
     # - Creates a new socket & performs a TCP handshake
     # - Grabs a socket already initialized before
-    _w("pymongo.server", "Server.get_socket", traced_get_socket)
+    _w("pymongo.server", "Server.%s" % _CHECKOUT_FN_NAME, traced_get_socket)
 
 
 def unpatch_pymongo_module():
     if not getattr(pymongo, "_datadog_patch", False):
         return
     pymongo._datadog_patch = False
 
-    _u(pymongo.server.Server, "get_socket")
+    _u(pymongo.server.Server, _CHECKOUT_FN_NAME)
 
 
 @contextlib.contextmanager
@@ -70,7 +73,9 @@ def traced_get_socket(wrapped, instance, args, kwargs):
             return
 
     with pin.tracer.trace(
-        "pymongo.get_socket", service=trace_utils.int_service(pin, config.pymongo), span_type=SpanTypes.MONGODB
+        "pymongo.%s" % _CHECKOUT_FN_NAME,
+        service=trace_utils.int_service(pin, config.pymongo),
+        span_type=SpanTypes.MONGODB,
     ) as span:
         span.set_tag_str(COMPONENT, config.pymongo.integration_name)
         span.set_tag_str(db.SYSTEM, mongo.SERVICE)

releasenotes/notes/pymongo-get-socket-808421288343ab81.yaml

@@ -0,0 +1,5 @@
+---
+fixes:
+  - |
+    pymongo: This upgrades the PyMongo integration to work with PyMongo versions 4.5.0 and above by choosing the root function
+    of the integration on the basis of the PyMongo version.

riotfile.py

@@ -670,7 +670,7 @@ def select_pys(min_version=MIN_PYTHON_VERSION, max_version=MAX_PYTHON_VERSION):
                     pkgs={"pymongo": ["~=3.4", "~=3.11", "~=3.13"]},
                 ),
                 Venv(
-                    pys=select_pys(min_version="3.7", max_version="3.9"), pkgs={"pymongo": ["~=3.11", "~=4.0", latest]}
+                    pys=select_pys(min_version="3.7", max_version="3.9"), pkgs={"pymongo": ["~=3.11", "~=4.5", latest]}
                 ),
                 Venv(
                     # pymongo added support for Python 3.10 in 3.12.1

tests/contrib/pymongo/test.py

@@ -8,6 +8,7 @@
 from ddtrace import Pin
 from ddtrace.constants import ANALYTICS_SAMPLE_RATE_KEY
 from ddtrace.contrib.pymongo.client import normalize_filter
+from ddtrace.contrib.pymongo.patch import _CHECKOUT_FN_NAME
 from ddtrace.contrib.pymongo.patch import patch
 from ddtrace.contrib.pymongo.patch import unpatch
 from ddtrace.ext import SpanTypes
@@ -752,7 +753,7 @@ def tearDown(self):
 
     @staticmethod
     def check_socket_metadata(span):
-        assert span.name == "pymongo.get_socket"
+        assert span.name == "pymongo.%s" % _CHECKOUT_FN_NAME
         assert span.service == "pymongo"
         assert span.span_type == SpanTypes.MONGODB
         assert span.get_tag("out.host") == "localhost"
@@ -820,8 +821,8 @@ def test_multiple_ops(self):
             # run_operation_with_response() which takes as an argument a sock_info. The
             # lib now first calls get_socket() and then run_operation_with_response(sock_info),
             # which makes more sense and also allows us to trace the function correctly.
-            assert {first_span.name, second_span.name} == {"pymongo.cmd", "pymongo.get_socket"}
-            if first_span.name == "pymongo.get_socket":
+            assert {first_span.name, second_span.name} == {"pymongo.cmd", "pymongo.%s" % _CHECKOUT_FN_NAME}
+            if first_span.name == "pymongo.%s" % _CHECKOUT_FN_NAME:
                 self.check_socket_metadata(first_span)
             else:
                 self.check_socket_metadata(second_span)