getaddrinfo helper method

Author: NoahStapp

pymongo/asynchronous/auth.py

@@ -15,7 +15,6 @@
 """Authentication helpers."""
 from __future__ import annotations
 
-import asyncio
 import functools
 import hashlib
 import hmac
@@ -39,6 +38,7 @@
     _authenticate_oidc,
     _get_authenticator,
 )
+from pymongo.asynchronous.helpers import getaddrinfo
 from pymongo.auth_shared import (
     MongoCredential,
     _authenticate_scram_start,
@@ -184,22 +184,16 @@ async def _canonicalize_hostname(hostname: str, option: str | bool) -> str:
     if option in [False, "none"]:
         return hostname
 
-    if not _IS_SYNC:
-        loop = asyncio.get_running_loop()
-        af, socktype, proto, canonname, sockaddr = (
-            await loop.getaddrinfo(
-                hostname,
-                None,
-                family=0,
-                type=0,
-                proto=socket.IPPROTO_TCP,
-                flags=socket.AI_CANONNAME,
-            )
-        )[0]  # type: ignore[index]
-    else:
-        af, socktype, proto, canonname, sockaddr = socket.getaddrinfo(
-            hostname, None, 0, 0, socket.IPPROTO_TCP, socket.AI_CANONNAME
-        )[0]
+    af, socktype, proto, canonname, sockaddr = (
+        await getaddrinfo(
+            hostname,
+            None,
+            family=0,
+            type=0,
+            proto=socket.IPPROTO_TCP,
+            flags=socket.AI_CANONNAME,
+        )
+    )[0]  # type: ignore[index]
 
     # For forward just to resolve the cname as dns.lookup() will not return it.
     if option == "forward":

pymongo/asynchronous/helpers.py

@@ -15,7 +15,9 @@
 """Miscellaneous pieces that need to be synchronized."""
 from __future__ import annotations
 
+import asyncio
 import builtins
+import socket
 import sys
 from typing import (
     Any,
@@ -68,6 +70,16 @@ async def inner(*args: Any, **kwargs: Any) -> Any:
     return cast(F, inner)
 
 
+async def getaddrinfo(host, port, **kwargs):
+    if not _IS_SYNC:
+        loop = asyncio.get_running_loop()
+        return await loop.getaddrinfo(  # type: ignore[assignment]
+            host, port, **kwargs
+        )
+    else:
+        return socket.getaddrinfo(host, port, **kwargs)  # type: ignore[assignment]
+
+
 if sys.version_info >= (3, 10):
     anext = builtins.anext
     aiter = builtins.aiter

pymongo/asynchronous/pool.py

@@ -41,7 +41,7 @@
 from pymongo import _csot, helpers_shared
 from pymongo._asyncio_executor import _PYMONGO_EXECUTOR
 from pymongo.asynchronous.client_session import _validate_session_write_concern
-from pymongo.asynchronous.helpers import _handle_reauth
+from pymongo.asynchronous.helpers import _handle_reauth, getaddrinfo
 from pymongo.asynchronous.network import command, receive_message
 from pymongo.common import (
     MAX_BSON_SIZE,
@@ -815,14 +815,7 @@ async def _create_connection(address: _Address, options: PoolOptions) -> socket.
         family = socket.AF_UNSPEC
 
     err = None
-    if not _IS_SYNC:
-        loop = asyncio.get_running_loop()
-        results = await loop.getaddrinfo(  # type: ignore[assignment]
-            host, port, family=family, type=socket.SOCK_STREAM
-        )
-    else:
-        results = socket.getaddrinfo(host, port, family, socket.SOCK_STREAM)  # type: ignore[assignment]
-    for res in results:  # type: ignore[attr-defined]
+    for res in await getaddrinfo(host, port, family=family, type=socket.SOCK_STREAM):  # type: ignore[attr-defined]
         af, socktype, proto, dummy, sa = res
         # SOCK_CLOEXEC was new in CPython 3.2, and only available on a limited
         # number of platforms (newer Linux and *BSD). Starting with CPython 3.4

pymongo/synchronous/auth.py

@@ -15,7 +15,6 @@
 """Authentication helpers."""
 from __future__ import annotations
 
-import asyncio
 import functools
 import hashlib
 import hmac
@@ -46,6 +45,7 @@
     _authenticate_oidc,
     _get_authenticator,
 )
+from pymongo.synchronous.helpers import getaddrinfo
 
 if TYPE_CHECKING:
     from pymongo.hello import Hello
@@ -181,22 +181,16 @@ def _canonicalize_hostname(hostname: str, option: str | bool) -> str:
     if option in [False, "none"]:
         return hostname
 
-    if not _IS_SYNC:
-        loop = asyncio.get_running_loop()
-        af, socktype, proto, canonname, sockaddr = (
-            loop.getaddrinfo(
-                hostname,
-                None,
-                family=0,
-                type=0,
-                proto=socket.IPPROTO_TCP,
-                flags=socket.AI_CANONNAME,
-            )
-        )[0]  # type: ignore[index]
-    else:
-        af, socktype, proto, canonname, sockaddr = socket.getaddrinfo(
-            hostname, None, 0, 0, socket.IPPROTO_TCP, socket.AI_CANONNAME
-        )[0]
+    af, socktype, proto, canonname, sockaddr = (
+        getaddrinfo(
+            hostname,
+            None,
+            family=0,
+            type=0,
+            proto=socket.IPPROTO_TCP,
+            flags=socket.AI_CANONNAME,
+        )
+    )[0]  # type: ignore[index]
 
     # For forward just to resolve the cname as dns.lookup() will not return it.
     if option == "forward":

pymongo/synchronous/helpers.py

@@ -15,7 +15,9 @@
 """Miscellaneous pieces that need to be synchronized."""
 from __future__ import annotations
 
+import asyncio
 import builtins
+import socket
 import sys
 from typing import (
     Any,
@@ -68,6 +70,16 @@ def inner(*args: Any, **kwargs: Any) -> Any:
     return cast(F, inner)
 
 
+def getaddrinfo(host, port, **kwargs):
+    if not _IS_SYNC:
+        loop = asyncio.get_running_loop()
+        return loop.getaddrinfo(  # type: ignore[assignment]
+            host, port, **kwargs
+        )
+    else:
+        return socket.getaddrinfo(host, port, **kwargs)  # type: ignore[assignment]
+
+
 if sys.version_info >= (3, 10):
     next = builtins.next
     iter = builtins.iter

pymongo/synchronous/pool.py

@@ -85,7 +85,7 @@
 from pymongo.socket_checker import SocketChecker
 from pymongo.ssl_support import HAS_SNI, SSLError
 from pymongo.synchronous.client_session import _validate_session_write_concern
-from pymongo.synchronous.helpers import _handle_reauth
+from pymongo.synchronous.helpers import _handle_reauth, getaddrinfo
 from pymongo.synchronous.network import command, receive_message
 
 if TYPE_CHECKING:
@@ -813,14 +813,7 @@ def _create_connection(address: _Address, options: PoolOptions) -> socket.socket
         family = socket.AF_UNSPEC
 
     err = None
-    if not _IS_SYNC:
-        loop = asyncio.get_running_loop()
-        results = loop.getaddrinfo(  # type: ignore[assignment]
-            host, port, family=family, type=socket.SOCK_STREAM
-        )
-    else:
-        results = socket.getaddrinfo(host, port, family, socket.SOCK_STREAM)  # type: ignore[assignment]
-    for res in results:  # type: ignore[attr-defined]
+    for res in getaddrinfo(host, port, family=family, type=socket.SOCK_STREAM):  # type: ignore[attr-defined]
         af, socktype, proto, dummy, sa = res
         # SOCK_CLOEXEC was new in CPython 3.2, and only available on a limited
         # number of platforms (newer Linux and *BSD). Starting with CPython 3.4