diff --git a/pymongo/srv_resolver.py b/pymongo/srv_resolver.py index cf9ccffc43..45a4f66611 100644 --- a/pymongo/srv_resolver.py +++ b/pymongo/srv_resolver.py @@ -16,6 +16,13 @@ try: from dns import resolver + + try: + # dnspython >= 2 + from dns.resolver import resolve as _resolve + except ImportError: + # dnspython 1.X + from dns.resolver import query as _resolve _HAVE_DNSPYTHON = True except ImportError: _HAVE_DNSPYTHON = False @@ -48,8 +55,8 @@ def __init__(self, fqdn, connect_timeout=None): def get_options(self): try: - results = resolver.query(self.__fqdn, 'TXT', - lifetime=self.__connect_timeout) + results = _resolve(self.__fqdn, 'TXT', + lifetime=self.__connect_timeout) except (resolver.NoAnswer, resolver.NXDOMAIN): # No TXT records return None @@ -63,8 +70,8 @@ def get_options(self): def _resolve_uri(self, encapsulate_errors): try: - results = resolver.query('_mongodb._tcp.' + self.__fqdn, 'SRV', - lifetime=self.__connect_timeout) + results = _resolve('_mongodb._tcp.' + self.__fqdn, 'SRV', + lifetime=self.__connect_timeout) except Exception as exc: if not encapsulate_errors: # Raise the original error. diff --git a/test/test_client.py b/test/test_client.py index 2544bcef4c..e62516bd58 100644 --- a/test/test_client.py +++ b/test/test_client.py @@ -392,11 +392,11 @@ def test_uri_option_precedence(self): _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 + from pymongo.srv_resolver import _resolve + patched_resolver = FunctionCallRecorder(_resolve) + pymongo.srv_resolver._resolve = patched_resolver def reset_resolver(): - pymongo.srv_resolver.resolver.query = resolver.query + pymongo.srv_resolver._resolve = _resolve self.addCleanup(reset_resolver) # Setup.