[pymongo] Add support for PyMongo 3.9 (#1023)

* Test most recent versions of mongoengine and pymongo

* limit the scope of what we test

* [pymongo] Add support for PyMongo 3.9

* fix trailing comma issue

* fix version spec

* last missing trailing comma

Author: brettlangdon

ddtrace/contrib/pymongo/client.py

@@ -90,7 +90,7 @@ class TracedServer(ObjectProxy):
     def __init__(self, server):
         super(TracedServer, self).__init__(server)
 
-    def send_message_with_response(self, operation, *args, **kwargs):
+    def _datadog_trace_operation(self, operation):
         cmd = None
         # Only try to parse something we think is a query.
         if self._is_query(operation):
@@ -100,40 +100,71 @@ def send_message_with_response(self, operation, *args, **kwargs):
                 log.exception('error parsing query')
 
         pin = ddtrace.Pin.get_from(self)
-
         # if we couldn't parse or shouldn't trace the message, just go.
         if not cmd or not pin or not pin.enabled():
-            return self.__wrapped__.send_message_with_response(
+            return None
+
+        span = pin.tracer.trace('pymongo.cmd', span_type=mongox.TYPE, service=pin.service)
+        span.set_tag(mongox.DB, cmd.db)
+        span.set_tag(mongox.COLLECTION, cmd.coll)
+        span.set_tags(cmd.tags)
+
+        # set `mongodb.query` tag and resource for span
+        _set_query_metadata(span, cmd)
+
+        # set analytics sample rate
+        sample_rate = config.pymongo.get_analytics_sample_rate()
+        if sample_rate is not None:
+            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)
-
-        with pin.tracer.trace(
-                'pymongo.cmd',
-                span_type=mongox.TYPE,
-                service=pin.service) as span:
+                **kwargs
+            )
 
-            span.set_tag(mongox.DB, cmd.db)
-            span.set_tag(mongox.COLLECTION, cmd.coll)
-            span.set_tags(cmd.tags)
+        try:
+            result = self.__wrapped__.run_operation_with_response(
+                sock_info,
+                operation,
+                *args,
+                **kwargs
+            )
 
-            # set `mongodb.query` tag and resource for span
-            _set_query_metadata(span, cmd)
+            if result and result.address:
+                _set_address_tags(span, result.address)
+            return result
+        finally:
+            span.finish()
 
-            # set analytics sample rate
-            span.set_tag(
-                ANALYTICS_SAMPLE_RATE_KEY,
-                config.pymongo.get_analytics_sample_rate()
+    # 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
             )
 
+        try:
             result = self.__wrapped__.send_message_with_response(
                 operation,
                 *args,
-                **kwargs)
+                **kwargs
+            )
 
             if result and result.address:
                 _set_address_tags(span, result.address)
             return result
+        finally:
+            span.finish()
 
     @contextlib.contextmanager
     def get_socket(self, *args, **kwargs):

tox.ini

@@ -81,7 +81,7 @@ envlist =
     jinja2_contrib-{py27,py34,py35,py36,py37}-jinja{27,28,29,210}
     mako_contrib-{py27,py34,py35,py36,py37}-mako{010,100}
     molten_contrib-py{36,37}-molten{070,072}
-    mongoengine_contrib-{py27,py34,py35,py36,py37}-mongoengine{015}
+    mongoengine_contrib-{py27,py34,py35,py36,py37}-mongoengine{015,016,017,018,latest}-pymongo{latest}
     mysql_contrib-{py27,py34,py35,py36,py37}-mysqlconnector
     mysqldb_contrib-{py27}-mysqldb{12}
     mysqldb_contrib-{py27,py34,py35,py36,py37}-mysqlclient{13}
@@ -90,7 +90,7 @@ envlist =
     pylibmc_contrib-{py27,py34,py35,py36,py37}-pylibmc{140,150}
     pylons_contrib-{py27}-pylons{096,097,010,10}
     pymemcache_contrib{,_autopatch}-{py27,py34,py35,py36,py37}-pymemcache{130,140}
-    pymongo_contrib-{py27,py34,py35,py36,py37}-pymongo{30,31,32,33,34,36,37,38}-mongoengine{015,016,017}
+    pymongo_contrib-{py27,py34,py35,py36,py37}-pymongo{30,31,32,33,34,35,36,37,38,39,latest}-mongoengine{latest}
     pymysql_contrib-{py27,py34,py35,py36,py37}-pymysql{07,08,09}
     pyramid_contrib{,_autopatch}-{py27,py34,py35,py36,py37}-pyramid{17,18,19}-webtest
     redis_contrib-{py27,py34,py35,py36,py37}-redis{26,27,28,29,210,300}
@@ -274,6 +274,8 @@ deps =
     mongoengine015: mongoengine>=0.15<0.16
     mongoengine016: mongoengine>=0.16<0.17
     mongoengine017: mongoengine>=0.17<0.18
+    mongoengine018: mongoengine>=0.18<0.19
+    mongoenginelatest: mongoengine>=0.18
     mysqlconnector: mysql-connector-python
     mysqldb12: mysql-python>=1.2,<1.3
     mysqlclient13: mysqlclient>=1.3,<1.4
@@ -298,6 +300,8 @@ deps =
     pymongo36: pymongo>=3.6,<3.7
     pymongo37: pymongo>=3.7,<3.8
     pymongo38: pymongo>=3.8,<3.9
+    pymongo39: pymongo>=3.9,<3.10
+    pymongolatest: pymongo>=3.9
     pymysql07: pymysql>=0.7,<0.8
     pymysql08: pymysql>=0.8,<0.9
     pymysql09: pymysql>=0.9,<0.10