PYTHON-1902 DNS resolution should honor connectTimeoutMS

Author: prashantmital

doc/changelog.rst

@@ -78,6 +78,11 @@ Unavoidable breaking changes:
   to avoid running into :class:`~pymongo.errors.OperationFailure` exceptions
   during write operations. The MMAPv1 storage engine is deprecated and does
   not support retryable writes which are now turned on by default.
+- In order to ensure that the ``connectTimeoutMS`` URI option is honored when
+  connecting to clusters with a ``mongodb+srv://`` connection string, the
+  minimum required version of the optional ``dnspython`` dependency has been
+  bumped to 1.16.0. This is a breaking change for applications that use
+  PyMongo's SRV support with a version of ``dnspython`` older than 1.16.0.
 
 .. _URI options specification: https://github.com/mongodb/specifications/blob/master/source/uri-options/uri-options.rst
 

pymongo/mongo_client.py

@@ -592,6 +592,16 @@ def __init__(
         if not isinstance(port, int):
             raise TypeError("port must be an instance of int")
 
+        # _pool_class, _monitor_class, and _condition_class are for deep
+        # customization of PyMongo, e.g. Motor.
+        pool_class = kwargs.pop('_pool_class', None)
+        monitor_class = kwargs.pop('_monitor_class', None)
+        condition_class = kwargs.pop('_condition_class', None)
+
+        # Parse options passed as kwargs.
+        keyword_opts = common._CaseInsensitiveDictionary(kwargs)
+        keyword_opts['document_class'] = document_class
+
         seeds = set()
         username = None
         password = None
@@ -600,8 +610,14 @@ def __init__(
         fqdn = None
         for entity in host:
             if "://" in entity:
+                # Determine connection timeout from kwargs.
+                timeout = keyword_opts.get("connecttimeoutms")
+                if timeout is not None:
+                    timeout = common.validate_timeout_or_none(
+                        keyword_opts.cased_key("connecttimeoutms"), timeout)
                 res = uri_parser.parse_uri(
-                    entity, port, validate=True, warn=True, normalize=False)
+                    entity, port, validate=True, warn=True, normalize=False,
+                    connect_timeout=timeout)
                 seeds.update(res["nodelist"])
                 username = res["username"] or username
                 password = res["password"] or password
@@ -613,14 +629,7 @@ def __init__(
         if not seeds:
             raise ConfigurationError("need to specify at least one host")
 
-        # _pool_class, _monitor_class, and _condition_class are for deep
-        # customization of PyMongo, e.g. Motor.
-        pool_class = kwargs.pop('_pool_class', None)
-        monitor_class = kwargs.pop('_monitor_class', None)
-        condition_class = kwargs.pop('_condition_class', None)
-
-        keyword_opts = common._CaseInsensitiveDictionary(kwargs)
-        keyword_opts['document_class'] = document_class
+        # Add options with named keyword arguments to the parsed kwarg options.
         if type_registry is not None:
             keyword_opts['type_registry'] = type_registry
         if tz_aware is None:

pymongo/srv_resolver.py

@@ -22,6 +22,7 @@
 
 from bson.py3compat import PY3
 
+from pymongo.common import CONNECT_TIMEOUT
 from pymongo.errors import ConfigurationError
 
 
@@ -38,8 +39,9 @@ def maybe_decode(text):
 
 
 class _SrvResolver(object):
-    def __init__(self, fqdn):
+    def __init__(self, fqdn, connect_timeout=None):
         self.__fqdn = fqdn
+        self.__connect_timeout = connect_timeout or CONNECT_TIMEOUT
 
         # Validate the fully qualified domain name.
         try:
@@ -52,7 +54,8 @@ def __init__(self, fqdn):
 
     def get_options(self):
         try:
-            results = resolver.query(self.__fqdn, 'TXT')
+            results = resolver.query(self.__fqdn, 'TXT',
+                                     lifetime=self.__connect_timeout)
         except (resolver.NoAnswer, resolver.NXDOMAIN):
             # No TXT records
             return None
@@ -66,7 +69,8 @@ def get_options(self):
 
     def _resolve_uri(self, encapsulate_errors):
         try:
-            results = resolver.query('_mongodb._tcp.' + self.__fqdn, 'SRV')
+            results = resolver.query('_mongodb._tcp.' + self.__fqdn, 'SRV',
+                                     lifetime=self.__connect_timeout)
         except Exception as exc:
             if not encapsulate_errors:
                 # Raise the original error.

pymongo/uri_parser.py

@@ -325,7 +325,7 @@ def split_hosts(hosts, default_port=DEFAULT_PORT):
 
 
 def parse_uri(uri, default_port=DEFAULT_PORT, validate=True, warn=False,
-              normalize=True):
+              normalize=True, connect_timeout=None):
     """Parse and validate a MongoDB URI.
 
     Returns a dict of the form::
@@ -355,6 +355,8 @@ def parse_uri(uri, default_port=DEFAULT_PORT, validate=True, warn=False,
           invalid. Default: ``False``.
         - `normalize` (optional): If ``True``, convert names of URI options
           to their internally-used names. Default: ``True``.
+        - `connect_timeout` (optional): The maximum time in milliseconds to
+          wait for a response from the DNS server.
 
     .. versionchanged:: 3.9
         Added the ``normalize`` parameter.
@@ -400,6 +402,25 @@ def parse_uri(uri, default_port=DEFAULT_PORT, validate=True, warn=False,
         raise InvalidURI("A '/' is required between "
                          "the host list and any options.")
 
+    if path_part:
+        if path_part[0] == '?':
+            opts = unquote_plus(path_part[1:])
+        else:
+            dbase, _, opts = map(unquote_plus, path_part.partition('?'))
+            if '.' in dbase:
+                dbase, collection = dbase.split('.', 1)
+
+            if _BAD_DB_CHARS.search(dbase):
+                raise InvalidURI('Bad database name "%s"' % dbase)
+
+        if opts:
+            options.update(split_options(opts, validate, warn, normalize))
+
+    if dbase is not None:
+        dbase = unquote_plus(dbase)
+    if collection is not None:
+        collection = unquote_plus(collection)
+
     if '@' in host_part:
         userinfo, _, hosts = host_part.rpartition('@')
         user, passwd = parse_userinfo(userinfo)
@@ -424,37 +445,26 @@ def parse_uri(uri, default_port=DEFAULT_PORT, validate=True, warn=False,
             raise InvalidURI(
                 "%s URIs must not include a port number" % (SRV_SCHEME,))
 
-        dns_resolver = _SrvResolver(fqdn)
+        # Use the connection timeout. connectTimeoutMS passed as a keyword
+        # argument overrides the same option passed in the connection string.
+        connect_timeout = connect_timeout or options.get("connectTimeoutMS")
+        dns_resolver = _SrvResolver(fqdn, connect_timeout=connect_timeout)
         nodes = dns_resolver.get_hosts()
         dns_options = dns_resolver.get_options()
         if dns_options:
-            options = split_options(dns_options, validate, warn, normalize)
-            if set(options) - _ALLOWED_TXT_OPTS:
+            parsed_dns_options = split_options(
+                dns_options, validate, warn, normalize)
+            if set(parsed_dns_options) - _ALLOWED_TXT_OPTS:
                 raise ConfigurationError(
                     "Only authSource and replicaSet are supported from DNS")
-        options["ssl"] = True if validate else 'true'
+            for opt, val in parsed_dns_options.items():
+                if opt not in options:
+                    options[opt] = val
+        if "ssl" not in options:
+            options["ssl"] = True if validate else 'true'
     else:
         nodes = split_hosts(hosts, default_port=default_port)
 
-    if path_part:
-        if path_part[0] == '?':
-            opts = unquote_plus(path_part[1:])
-        else:
-            dbase, _, opts = map(unquote_plus, path_part.partition('?'))
-            if '.' in dbase:
-                dbase, collection = dbase.split('.', 1)
-
-            if _BAD_DB_CHARS.search(dbase):
-                raise InvalidURI('Bad database name "%s"' % dbase)
-
-        if opts:
-            options.update(split_options(opts, validate, warn, normalize))
-
-    if dbase is not None:
-        dbase = unquote_plus(dbase)
-    if collection is not None:
-        collection = unquote_plus(collection)
-
     return {
         'nodelist': nodes,
         'username': user,

setup.py

@@ -318,17 +318,16 @@ def build_extension(self, ext):
                                   'bson/buffer.c'])]
 
 extras_require = {
+    'encryption': ['pymongocrypt'],  # For client side field level encryption.
     'snappy': ['python-snappy'],
+    'srv': ["dnspython>=1.16.0,<2.0.0"],
     'zstd': ['zstandard'],
-    'encryption': ['pymongocrypt'],  # For client side field level encryption.
 }
 vi = sys.version_info
 if vi[0] == 2:
-    extras_require.update(
-        {'tls': ["ipaddress"], 'srv': ["dnspython>=1.8.0,<2.0.0"]})
+    extras_require.update({'tls': ["ipaddress"]})
 else:
-    extras_require.update(
-        {'tls': [], 'srv': ["dnspython>=1.13.0,<2.0.0"]})
+    extras_require.update({'tls': []})
 if sys.platform == 'win32':
     extras_require['gssapi'] = ["winkerberos>=0.5.0"]
     if vi < (2, 7, 9):

test/test_client.py

@@ -33,8 +33,9 @@
 from bson.py3compat import thread
 from bson.son import SON
 from bson.tz_util import utc
+import pymongo
 from pymongo import auth, message
-from pymongo.common import _UUID_REPRESENTATIONS
+from pymongo.common import CONNECT_TIMEOUT, _UUID_REPRESENTATIONS
 from pymongo.command_cursor import CommandCursor
 from pymongo.compression_support import _HAVE_SNAPPY, _HAVE_ZSTD
 from pymongo.cursor import Cursor, CursorType
@@ -56,6 +57,7 @@
 from pymongo.server_selectors import (any_server_selector,
                                       writable_server_selector)
 from pymongo.server_type import SERVER_TYPE
+from pymongo.srv_resolver import _HAVE_DNSPYTHON
 from pymongo.write_concern import WriteConcern
 from test import (client_context,
                   client_knobs,
@@ -70,6 +72,7 @@
 from test.utils import (assertRaisesExactly,
                         connected,
                         delay,
+                        FunctionCallRecorder,
                         get_pool,
                         gevent_monkey_patched,
                         ignore_deprecations,
@@ -358,6 +361,43 @@ def test_uri_option_precedence(self):
         self.assertEqual(
             clopts.read_preference, ReadPreference.SECONDARY_PREFERRED)
 
+    @unittest.skipUnless(
+        _HAVE_DNSPYTHON, "DNS-related tests need dnspython to be installed")
+    def test_connection_timeout_ms_propagates_to_DNS_resolver(self):
+        # Patch the resolver.
+        from pymongo.srv_resolver import resolver
+        patched_resolver = FunctionCallRecorder(resolver.query)
+        pymongo.srv_resolver.resolver.query = patched_resolver
+        def reset_resolver():
+            pymongo.srv_resolver.resolver.query = resolver.query
+        self.addCleanup(reset_resolver)
+
+        # Setup.
+        base_uri = "mongodb+srv://test5.test.build.10gen.cc"
+        connectTimeoutMS = 5000
+        expected_kw_value = 5.0
+        uri_with_timeout = base_uri + "/?connectTimeoutMS=6000"
+        expected_uri_value = 6.0
+
+        def test_scenario(args, kwargs, expected_value):
+            patched_resolver.reset()
+            MongoClient(*args, **kwargs)
+            for _, kw in patched_resolver.call_list():
+                self.assertAlmostEqual(kw['lifetime'], expected_value)
+
+        # No timeout specified.
+        test_scenario((base_uri,), {}, CONNECT_TIMEOUT)
+
+        # Timeout only specified in connection string.
+        test_scenario((uri_with_timeout,), {}, expected_uri_value)
+
+        # Timeout only specified in keyword arguments.
+        kwarg = {'connectTimeoutMS': connectTimeoutMS}
+        test_scenario((base_uri,), kwarg, expected_kw_value)
+
+        # Timeout specified in both kwargs and connection string.
+        test_scenario((uri_with_timeout,), kwarg, expected_kw_value)
+
     def test_uri_security_options(self):
         # Ensure that we don't silently override security-related options.
         with self.assertRaises(InvalidURI):

test/test_server_selection.py

@@ -28,7 +28,7 @@
 
 from test import client_context, unittest, IntegrationTest
 from test.utils import (rs_or_single_client, wait_until, EventListener,
-                        FunctionCallCounter)
+                        FunctionCallRecorder)
 from test.utils_selection_tests import (
     create_selection_tests, get_addresses, get_topology_settings_dict,
     make_server_description)
@@ -105,7 +105,7 @@ def test_invalid_server_selector(self):
 
     @client_context.require_replica_set
     def test_selector_called(self):
-        selector = FunctionCallCounter(lambda x: x)
+        selector = FunctionCallRecorder(lambda x: x)
 
         # Client setup.
         mongo_client = rs_or_single_client(server_selector=selector)
@@ -169,7 +169,7 @@ def test_latency_threshold_application(self):
 
     @client_context.require_replica_set
     def test_server_selector_bypassed(self):
-        selector = FunctionCallCounter(lambda x: x)
+        selector = FunctionCallRecorder(lambda x: x)
 
         scenario_def = {
             'topology_description': {

test/test_srv_polling.py

@@ -27,7 +27,7 @@
 from pymongo.srv_resolver import _HAVE_DNSPYTHON
 from pymongo.mongo_client import MongoClient
 from test import client_knobs, unittest
-from test.utils import wait_until, FunctionCallCounter
+from test.utils import wait_until, FunctionCallRecorder
 
 
 WAIT_TIME = 0.1
@@ -62,7 +62,7 @@ def mock_get_hosts_and_min_ttl(resolver, *args):
             return nodes, ttl
 
         if self.count_resolver_calls:
-            patch_func = FunctionCallCounter(mock_get_hosts_and_min_ttl)
+            patch_func = FunctionCallRecorder(mock_get_hosts_and_min_ttl)
         else:
             patch_func = mock_get_hosts_and_min_ttl
 

test/utils.py

@@ -216,19 +216,28 @@ def __ne__(self, other):
         return not self.__eq__(other)
 
 
-class FunctionCallCounter(object):
-    """Class that wraps a function and keeps count of invocations."""
+class FunctionCallRecorder(object):
+    """Utility class to wrap a callable and record its invocations."""
     def __init__(self, function):
         self._function = function
-        self._call_count = 0
+        self._call_list = []
 
     def __call__(self, *args, **kwargs):
-        self._call_count += 1
+        self._call_list.append((args, kwargs))
         return self._function(*args, **kwargs)
 
+    def reset(self):
+        """Wipes the call list."""
+        self._call_list = []
+
+    def call_list(self):
+        """Returns a copy of the call list."""
+        return self._call_list[:]
+
     @property
     def call_count(self):
-        return self._call_count
+        """Returns the number of times the function has been called."""
+        return len(self._call_list)
 
 
 class TestCreator(object):