fix(pymongo): use generic arguments for write_command (#3371) (#3372)

Pymongo 4.0.2 released which introduced a breaking change to the method signature of Socket.write_command: mongodb/mongo-python-driver@4.0.1...4.0.2diff-43b0ab0dab45c7b3e339879a4452c7dcb43a659ef6603532348a745e6405afd4R819

The patching attempted to match the signature exactly so the fix is trivial in just accepting args/kwargs and extracting what is needed from them.

(cherry picked from commit da60df9)

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

Author: mergify[bot]

ddtrace/contrib/pymongo/client.py

@@ -17,6 +17,7 @@
 from ...ext import net as netx
 from ...internal.compat import iteritems
 from ...internal.logger import get_logger
+from ...internal.utils import get_argument_value
 from .parse import parse_msg
 from .parse import parse_query
 from .parse import parse_spec
@@ -190,7 +191,8 @@ def command(self, dbname, spec, *args, **kwargs):
         with self.__trace(cmd):
             return self.__wrapped__.command(dbname, spec, *args, **kwargs)
 
-    def write_command(self, request_id, msg):
+    def write_command(self, *args, **kwargs):
+        msg = get_argument_value(args, kwargs, 1, "msg")
         cmd = None
         try:
             cmd = parse_msg(msg)
@@ -200,10 +202,10 @@ def write_command(self, request_id, msg):
         pin = ddtrace.Pin.get_from(self)
         # if we couldn't parse it, don't try to trace it.
         if not cmd or not pin or not pin.enabled():
-            return self.__wrapped__.write_command(request_id, msg)
+            return self.__wrapped__.write_command(*args, **kwargs)
 
         with self.__trace(cmd) as s:
-            result = self.__wrapped__.write_command(request_id, msg)
+            result = self.__wrapped__.write_command(*args, **kwargs)
             if result:
                 s.set_metric(mongox.ROWS, result.get("n", -1))
             return result

releasenotes/notes/pymongo-4.0.2-1f5d2b6af5c158d2.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    pymongo: fix ``write_command`` being patched with the wrong method signature.