respond to comments and move srv_resolver to async

sleepyStick · sleepyStick · commit 8c2b36852ff3 · 2025-03-13T11:20:30.000-07:00
diff --git a/pymongo/asynchronous/mongo_client.py b/pymongo/asynchronous/mongo_client.py
@@ -1216,10 +1216,12 @@ def options(self) -> ClientOptions:
 
     def __eq__(self, other: Any) -> bool:
         if isinstance(other, self.__class__):
-            if hasattr(self, "_topology"):
+            if hasattr(self, "_topology") and hasattr(other, "_topology"):
                 return self._topology == other._topology
             else:
-                raise InvalidOperation("Cannot compare client equality until both clients are connected")
+                raise InvalidOperation(
+                    "Cannot compare client equality until both clients are connected"
+                )
         return NotImplemented
 
     def __ne__(self, other: Any) -> bool:
@@ -1245,13 +1247,16 @@ def option_repr(option: str, value: Any) -> str:
             return f"{option}={value!r}"
 
         # Host first...
-        options = [
-            "host=%r"
-            % [
-                "%s:%d" % (host, port) if port is not None else host
-                for host, port in self._topology_settings.seeds
+        if hasattr(self, "_topology"):
+            options = [
+                "host=%r"
+                % [
+                    "%s:%d" % (host, port) if port is not None else host
+                    for host, port in self._topology_settings.seeds
+                ]
             ]
-        ]
+        else:
+            options = []
         # ... then everything in self._constructor_args...
         options.extend(
             option_repr(key, self._options._options[key]) for key in self._constructor_args
@@ -1265,9 +1270,7 @@ def option_repr(option: str, value: Any) -> str:
         return ", ".join(options)
 
     def __repr__(self) -> str:
-        if hasattr(self, "_topology"):
-            return f"{type(self).__name__}({self._repr_helper()})"
-        raise InvalidOperation("Cannot perform operation until client is connected")
+        return f"{type(self).__name__}({self._repr_helper()})"
 
     def __getattr__(self, name: str) -> database.AsyncDatabase[_DocumentType]:
         """Get a database by name.
diff --git a/pymongo/asynchronous/monitor.py b/pymongo/asynchronous/monitor.py
@@ -25,6 +25,7 @@
 
 from pymongo import common, periodic_executor
 from pymongo._csot import MovingMinimum
+from pymongo.asynchronous.srv_resolver import _SrvResolver
 from pymongo.errors import NetworkTimeout, _OperationCancelled
 from pymongo.hello import Hello
 from pymongo.lock import _async_create_lock
@@ -33,7 +34,6 @@
 from pymongo.pool_options import _is_faas
 from pymongo.read_preferences import MovingAverage
 from pymongo.server_description import ServerDescription
-from pymongo.srv_resolver import _SrvResolver
 
 if TYPE_CHECKING:
     from pymongo.asynchronous.pool import AsyncConnection, Pool, _CancellationContext
diff --git a/pymongo/asynchronous/srv_resolver.py b/pymongo/asynchronous/srv_resolver.py
@@ -0,0 +1,155 @@
+# Copyright 2019-present MongoDB, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); you
+# may not use this file except in compliance with the License.  You
+# may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+# implied.  See the License for the specific language governing
+# permissions and limitations under the License.
+
+"""Support for resolving hosts and options from mongodb+srv:// URIs."""
+from __future__ import annotations
+
+import ipaddress
+import random
+from typing import TYPE_CHECKING, Any, Optional, Union
+
+from pymongo.common import CONNECT_TIMEOUT
+from pymongo.errors import ConfigurationError
+
+if TYPE_CHECKING:
+    from dns import resolver
+
+_IS_SYNC = False
+
+
+def _have_dnspython() -> bool:
+    try:
+        import dns  # noqa: F401
+
+        return True
+    except ImportError:
+        return False
+
+
+# dnspython can return bytes or str from various parts
+# of its API depending on version. We always want str.
+def maybe_decode(text: Union[str, bytes]) -> str:
+    if isinstance(text, bytes):
+        return text.decode()
+    return text
+
+
+# PYTHON-2667 Lazily call dns.resolver methods for compatibility with eventlet.
+async def _resolve(*args: Any, **kwargs: Any) -> resolver.Answer:
+    if _IS_SYNC:
+        from dns import resolver
+
+        if hasattr(resolver, "resolve"):
+            # dnspython >= 2
+            return resolver.resolve(*args, **kwargs)
+        # dnspython 1.X
+        return resolver.query(*args, **kwargs)
+    else:
+        from dns.asyncresolver import Resolver
+
+        return await Resolver.resolve(*args, **kwargs)
+
+
+_INVALID_HOST_MSG = (
+    "Invalid URI host: %s is not a valid hostname for 'mongodb+srv://'. "
+    "Did you mean to use 'mongodb://'?"
+)
+
+
+class _SrvResolver:
+    def __init__(
+        self,
+        fqdn: str,
+        connect_timeout: Optional[float],
+        srv_service_name: str,
+        srv_max_hosts: int = 0,
+    ):
+        self.__fqdn = fqdn
+        self.__srv = srv_service_name
+        self.__connect_timeout = connect_timeout or CONNECT_TIMEOUT
+        self.__srv_max_hosts = srv_max_hosts or 0
+        # Validate the fully qualified domain name.
+        try:
+            ipaddress.ip_address(fqdn)
+            raise ConfigurationError(_INVALID_HOST_MSG % ("an IP address",))
+        except ValueError:
+            pass
+
+        try:
+            self.__plist = self.__fqdn.split(".")[1:]
+        except Exception:
+            raise ConfigurationError(_INVALID_HOST_MSG % (fqdn,)) from None
+        self.__slen = len(self.__plist)
+        if self.__slen < 2:
+            raise ConfigurationError(_INVALID_HOST_MSG % (fqdn,))
+
+    async def get_options(self) -> Optional[str]:
+        from dns import resolver
+
+        try:
+            results = await _resolve(self.__fqdn, "TXT", lifetime=self.__connect_timeout)
+        except (resolver.NoAnswer, resolver.NXDOMAIN):
+            # No TXT records
+            return None
+        except Exception as exc:
+            raise ConfigurationError(str(exc)) from None
+        if len(results) > 1:
+            raise ConfigurationError("Only one TXT record is supported")
+        return (b"&".join([b"".join(res.strings) for res in results])).decode("utf-8")  # type: ignore[attr-defined]
+
+    async def _resolve_uri(self, encapsulate_errors: bool) -> resolver.Answer:
+        try:
+            results = await _resolve(
+                "_" + self.__srv + "._tcp." + self.__fqdn, "SRV", lifetime=self.__connect_timeout
+            )
+        except Exception as exc:
+            if not encapsulate_errors:
+                # Raise the original error.
+                raise
+            # Else, raise all errors as ConfigurationError.
+            raise ConfigurationError(str(exc)) from None
+        return results
+
+    async def _get_srv_response_and_hosts(
+        self, encapsulate_errors: bool
+    ) -> tuple[resolver.Answer, list[tuple[str, Any]]]:
+        results = await self._resolve_uri(encapsulate_errors)
+
+        # Construct address tuples
+        nodes = [
+            (maybe_decode(res.target.to_text(omit_final_dot=True)), res.port)  # type: ignore[attr-defined]
+            for res in results
+        ]
+
+        # Validate hosts
+        for node in nodes:
+            try:
+                nlist = node[0].lower().split(".")[1:][-self.__slen :]
+            except Exception:
+                raise ConfigurationError(f"Invalid SRV host: {node[0]}") from None
+            if self.__plist != nlist:
+                raise ConfigurationError(f"Invalid SRV host: {node[0]}")
+        if self.__srv_max_hosts:
+            nodes = random.sample(nodes, min(self.__srv_max_hosts, len(nodes)))
+        return results, nodes
+
+    async def get_hosts(self) -> list[tuple[str, Any]]:
+        _, nodes = await self._get_srv_response_and_hosts(True)
+        return nodes
+
+    async def get_hosts_and_min_ttl(self) -> tuple[list[tuple[str, Any]], int]:
+        results, nodes = await self._get_srv_response_and_hosts(False)
+        rrset = results.rrset
+        ttl = rrset.ttl if rrset else 0
+        return nodes, ttl
diff --git a/pymongo/synchronous/mongo_client.py b/pymongo/synchronous/mongo_client.py
@@ -1214,10 +1214,12 @@ def options(self) -> ClientOptions:
 
     def __eq__(self, other: Any) -> bool:
         if isinstance(other, self.__class__):
-            if hasattr(self, "_topology"):
+            if hasattr(self, "_topology") and hasattr(other, "_topology"):
                 return self._topology == other._topology
             else:
-                raise InvalidOperation("Cannot perform operation until client is connected")
+                raise InvalidOperation(
+                    "Cannot compare client equality until both clients are connected"
+                )
         return NotImplemented
 
     def __ne__(self, other: Any) -> bool:
@@ -1227,7 +1229,7 @@ def __hash__(self) -> int:
         if hasattr(self, "_topology"):
             return hash(self._topology)
         else:
-            raise InvalidOperation("Cannot perform operation until client is connected")
+            raise InvalidOperation("Cannot hash client until it is connected")
 
     def _repr_helper(self) -> str:
         def option_repr(option: str, value: Any) -> str:
@@ -1243,13 +1245,16 @@ def option_repr(option: str, value: Any) -> str:
             return f"{option}={value!r}"
 
         # Host first...
-        options = [
-            "host=%r"
-            % [
-                "%s:%d" % (host, port) if port is not None else host
-                for host, port in self._topology_settings.seeds
+        if hasattr(self, "_topology"):
+            options = [
+                "host=%r"
+                % [
+                    "%s:%d" % (host, port) if port is not None else host
+                    for host, port in self._topology_settings.seeds
+                ]
             ]
-        ]
+        else:
+            options = []
         # ... then everything in self._constructor_args...
         options.extend(
             option_repr(key, self._options._options[key]) for key in self._constructor_args
@@ -1263,9 +1268,7 @@ def option_repr(option: str, value: Any) -> str:
         return ", ".join(options)
 
     def __repr__(self) -> str:
-        if hasattr(self, "_topology"):
-            return f"{type(self).__name__}({self._repr_helper()})"
-        raise InvalidOperation("Cannot perform operation until client is connected")
+        return f"{type(self).__name__}({self._repr_helper()})"
 
     def __getattr__(self, name: str) -> database.Database[_DocumentType]:
         """Get a database by name.
diff --git a/pymongo/synchronous/monitor.py b/pymongo/synchronous/monitor.py
@@ -33,7 +33,7 @@
 from pymongo.pool_options import _is_faas
 from pymongo.read_preferences import MovingAverage
 from pymongo.server_description import ServerDescription
-from pymongo.srv_resolver import _SrvResolver
+from pymongo.synchronous.srv_resolver import _SrvResolver
 
 if TYPE_CHECKING:
     from pymongo.synchronous.pool import Connection, Pool, _CancellationContext
diff --git a/pymongo/synchronous/srv_resolver.py b/pymongo/synchronous/srv_resolver.py
@@ -25,6 +25,8 @@
 if TYPE_CHECKING:
     from dns import resolver
 
+_IS_SYNC = True
+
 
 def _have_dnspython() -> bool:
     try:
@@ -45,13 +47,18 @@ def maybe_decode(text: Union[str, bytes]) -> str:
 
 # PYTHON-2667 Lazily call dns.resolver methods for compatibility with eventlet.
 def _resolve(*args: Any, **kwargs: Any) -> resolver.Answer:
-    from dns import resolver
+    if _IS_SYNC:
+        from dns import resolver
+
+        if hasattr(resolver, "resolve"):
+            # dnspython >= 2
+            return resolver.resolve(*args, **kwargs)
+        # dnspython 1.X
+        return resolver.query(*args, **kwargs)
+    else:
+        from dns.asyncresolver import Resolver
 
-    if hasattr(resolver, "resolve"):
-        # dnspython >= 2
-        return resolver.resolve(*args, **kwargs)
-    # dnspython 1.X
-    return resolver.query(*args, **kwargs)
+        return Resolver.resolve(*args, **kwargs)
 
 
 _INVALID_HOST_MSG = (
diff --git a/pymongo/uri_parser.py b/pymongo/uri_parser.py
@@ -34,6 +34,7 @@
 )
 from urllib.parse import unquote_plus
 
+from pymongo.asynchronous.srv_resolver import _have_dnspython, _SrvResolver
 from pymongo.client_options import _parse_ssl_options
 from pymongo.common import (
     INTERNAL_URI_OPTION_NAME_MAP,
@@ -43,7 +44,6 @@
     get_validated_options,
 )
 from pymongo.errors import ConfigurationError, InvalidURI
-from pymongo.srv_resolver import _have_dnspython, _SrvResolver
 from pymongo.typings import _Address
 
 if TYPE_CHECKING:
diff --git a/test/asynchronous/test_client.py b/test/asynchronous/test_client.py
@@ -512,13 +512,13 @@ async def test_uri_option_precedence(self):
 
     async def test_connection_timeout_ms_propagates_to_DNS_resolver(self):
         # Patch the resolver.
-        from pymongo.srv_resolver import _resolve
+        from pymongo.asynchronous.srv_resolver import _resolve
 
         patched_resolver = FunctionCallRecorder(_resolve)
-        pymongo.srv_resolver._resolve = patched_resolver
+        pymongo.asynchronous.srv_resolver._resolve = patched_resolver
 
         def reset_resolver():
-            pymongo.srv_resolver._resolve = _resolve
+            pymongo.asynchronous.srv_resolver._resolve = _resolve
 
         self.addCleanup(reset_resolver)
 
diff --git a/test/asynchronous/test_srv_polling.py b/test/asynchronous/test_srv_polling.py
diff --git a/test/test_client.py b/test/test_client.py
diff --git a/test/test_srv_polling.py b/test/test_srv_polling.py

Original file line number	Diff line number	Diff line change
`@@ -34,6 +34,7 @@`
`34`	`34`	`)`
`35`	`35`	`from urllib.parse import unquote_plus`
`36`	`36`
	`37`	`+from pymongo.asynchronous.srv_resolver import _have_dnspython, _SrvResolver`
`37`	`38`	`from pymongo.client_options import _parse_ssl_options`
`38`	`39`	`from pymongo.common import (`
`39`	`40`	`INTERNAL_URI_OPTION_NAME_MAP,`
`@@ -43,7 +44,6 @@`
`43`	`44`	`get_validated_options,`
`44`	`45`	`)`
`45`	`46`	`from pymongo.errors import ConfigurationError, InvalidURI`
`46`		`-from pymongo.srv_resolver import _have_dnspython, _SrvResolver`
`47`	`47`	`from pymongo.typings import _Address`
`48`	`48`
`49`	`49`	`if TYPE_CHECKING:`