Merge branch 'master' of github.com:mongodb/mongo-python-driver into PYTHON-5143

Author: blink1073

pymongo/asynchronous/encryption.py

@@ -64,6 +64,7 @@
 from pymongo.asynchronous.cursor import AsyncCursor
 from pymongo.asynchronous.database import AsyncDatabase
 from pymongo.asynchronous.mongo_client import AsyncMongoClient
+from pymongo.asynchronous.pool import AsyncBaseConnection
 from pymongo.common import CONNECT_TIMEOUT
 from pymongo.daemon import _spawn_daemon
 from pymongo.encryption_options import AutoEncryptionOpts, RangeOpts
@@ -75,12 +76,12 @@
     NetworkTimeout,
     ServerSelectionTimeoutError,
 )
-from pymongo.network_layer import async_socket_sendall
+from pymongo.helpers_shared import _get_timeout_details
+from pymongo.network_layer import PyMongoKMSProtocol, async_receive_kms, async_sendall
 from pymongo.operations import UpdateOne
 from pymongo.pool_options import PoolOptions
 from pymongo.pool_shared import (
-    _async_configured_socket,
-    _get_timeout_details,
+    _configured_protocol_interface,
     _raise_connection_failure,
 )
 from pymongo.read_concern import ReadConcern
@@ -93,10 +94,8 @@
 if TYPE_CHECKING:
     from pymongocrypt.mongocrypt import MongoCryptKmsContext
 
-    from pymongo.pyopenssl_context import _sslConn
     from pymongo.typings import _Address
 
-
 _IS_SYNC = False
 
 _HTTPS_PORT = 443
@@ -111,9 +110,10 @@
 _KEY_VAULT_OPTS = CodecOptions(document_class=RawBSONDocument)
 
 
-async def _connect_kms(address: _Address, opts: PoolOptions) -> Union[socket.socket, _sslConn]:
+async def _connect_kms(address: _Address, opts: PoolOptions) -> AsyncBaseConnection:
     try:
-        return await _async_configured_socket(address, opts)
+        interface = await _configured_protocol_interface(address, opts, PyMongoKMSProtocol)
+        return AsyncBaseConnection(interface, opts)
     except Exception as exc:
         _raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts))
 
@@ -198,18 +198,11 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
         try:
             conn = await _connect_kms(address, opts)
             try:
-                await async_socket_sendall(conn, message)
+                await async_sendall(conn.conn.get_conn, message)
                 while kms_context.bytes_needed > 0:
                     # CSOT: update timeout.
-                    conn.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0))
-                    if _IS_SYNC:
-                        data = conn.recv(kms_context.bytes_needed)
-                    else:
-                        from pymongo.network_layer import (  # type: ignore[attr-defined]
-                            async_receive_data_socket,
-                        )
-
-                        data = await async_receive_data_socket(conn, kms_context.bytes_needed)
+                    conn.set_conn_timeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0))
+                    data = await async_receive_kms(conn, kms_context.bytes_needed)
                     if not data:
                         raise OSError("KMS connection closed")
                     kms_context.feed(data)
@@ -228,7 +221,7 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
                     address, exc, msg_prefix=msg_prefix, timeout_details=_get_timeout_details(opts)
                 )
             finally:
-                conn.close()
+                await conn.close_conn(None)
         except MongoCryptError:
             raise  # Propagate MongoCryptError errors directly.
         except Exception as exc:

pymongo/asynchronous/pool.py

@@ -58,6 +58,7 @@
     WaitQueueTimeoutError,
 )
 from pymongo.hello import Hello, HelloCompat
+from pymongo.helpers_shared import _get_timeout_details, format_timeout_details
 from pymongo.lock import (
     _async_cond_wait,
     _async_create_condition,
@@ -79,9 +80,7 @@
     SSLErrors,
     _CancellationContext,
     _configured_protocol_interface,
-    _get_timeout_details,
     _raise_connection_failure,
-    format_timeout_details,
 )
 from pymongo.read_preferences import ReadPreference
 from pymongo.server_api import _add_to_command
@@ -124,7 +123,89 @@ def _set_non_inheritable_non_atomic(fd: int) -> None:  # noqa: ARG001
 _IS_SYNC = False
 
 
-class AsyncConnection:
+class AsyncBaseConnection:
+    """A base connection object for server and kms connections."""
+
+    def __init__(self, conn: AsyncNetworkingInterface, opts: PoolOptions):
+        self.conn = conn
+        self.socket_checker: SocketChecker = SocketChecker()
+        self.cancel_context: _CancellationContext = _CancellationContext()
+        self.is_sdam = False
+        self.closed = False
+        self.last_timeout: float | None = None
+        self.more_to_come = False
+        self.opts = opts
+        self.max_wire_version = -1
+
+    def set_conn_timeout(self, timeout: Optional[float]) -> None:
+        """Cache last timeout to avoid duplicate calls to conn.settimeout."""
+        if timeout == self.last_timeout:
+            return
+        self.last_timeout = timeout
+        self.conn.get_conn.settimeout(timeout)
+
+    def apply_timeout(
+        self, client: AsyncMongoClient[Any], cmd: Optional[MutableMapping[str, Any]]
+    ) -> Optional[float]:
+        # CSOT: use remaining timeout when set.
+        timeout = _csot.remaining()
+        if timeout is None:
+            # Reset the socket timeout unless we're performing a streaming monitor check.
+            if not self.more_to_come:
+                self.set_conn_timeout(self.opts.socket_timeout)
+            return None
+        # RTT validation.
+        rtt = _csot.get_rtt()
+        if rtt is None:
+            rtt = self.connect_rtt
+        max_time_ms = timeout - rtt
+        if max_time_ms < 0:
+            timeout_details = _get_timeout_details(self.opts)
+            formatted = format_timeout_details(timeout_details)
+            # CSOT: raise an error without running the command since we know it will time out.
+            errmsg = f"operation would exceed time limit, remaining timeout:{timeout:.5f} <= network round trip time:{rtt:.5f} {formatted}"
+            if self.max_wire_version != -1:
+                raise ExecutionTimeout(
+                    errmsg,
+                    50,
+                    {"ok": 0, "errmsg": errmsg, "code": 50},
+                    self.max_wire_version,
+                )
+            else:
+                raise TimeoutError(errmsg)
+        if cmd is not None:
+            cmd["maxTimeMS"] = int(max_time_ms * 1000)
+        self.set_conn_timeout(timeout)
+        return timeout
+
+    async def close_conn(self, reason: Optional[str]) -> None:
+        """Close this connection with a reason."""
+        if self.closed:
+            return
+        await self._close_conn()
+
+    async def _close_conn(self) -> None:
+        """Close this connection."""
+        if self.closed:
+            return
+        self.closed = True
+        self.cancel_context.cancel()
+        # Note: We catch exceptions to avoid spurious errors on interpreter
+        # shutdown.
+        try:
+            await self.conn.close()
+        except Exception:  # noqa: S110
+            pass
+
+    def conn_closed(self) -> bool:
+        """Return True if we know socket has been closed, False otherwise."""
+        if _IS_SYNC:
+            return self.socket_checker.socket_closed(self.conn.get_conn)
+        else:
+            return self.conn.is_closing()
+
+
+class AsyncConnection(AsyncBaseConnection):
     """Store a connection with some metadata.
 
     :param conn: a raw connection object
@@ -142,29 +223,27 @@ def __init__(
         id: int,
         is_sdam: bool,
     ):
+        super().__init__(conn, pool.opts)
         self.pool_ref = weakref.ref(pool)
-        self.conn = conn
-        self.address = address
-        self.id = id
+        self.address: tuple[str, int] = address
+        self.id: int = id
         self.is_sdam = is_sdam
-        self.closed = False
         self.last_checkin_time = time.monotonic()
         self.performed_handshake = False
         self.is_writable: bool = False
         self.max_wire_version = MAX_WIRE_VERSION
-        self.max_bson_size = MAX_BSON_SIZE
-        self.max_message_size = MAX_MESSAGE_SIZE
-        self.max_write_batch_size = MAX_WRITE_BATCH_SIZE
+        self.max_bson_size: int = MAX_BSON_SIZE
+        self.max_message_size: int = MAX_MESSAGE_SIZE
+        self.max_write_batch_size: int = MAX_WRITE_BATCH_SIZE
         self.supports_sessions = False
         self.hello_ok: bool = False
-        self.is_mongos = False
+        self.is_mongos: bool = False
         self.op_msg_enabled = False
         self.listeners = pool.opts._event_listeners
         self.enabled_for_cmap = pool.enabled_for_cmap
         self.enabled_for_logging = pool.enabled_for_logging
         self.compression_settings = pool.opts._compression_settings
         self.compression_context: Union[SnappyContext, ZlibContext, ZstdContext, None] = None
-        self.socket_checker: SocketChecker = SocketChecker()
         self.oidc_token_gen_id: Optional[int] = None
         # Support for mechanism negotiation on the initial handshake.
         self.negotiated_mechs: Optional[list[str]] = None
@@ -175,9 +254,6 @@ def __init__(
         self.pool_gen = pool.gen
         self.generation = self.pool_gen.get_overall()
         self.ready = False
-        self.cancel_context: _CancellationContext = _CancellationContext()
-        self.opts = pool.opts
-        self.more_to_come: bool = False
         # For load balancer support.
         self.service_id: Optional[ObjectId] = None
         self.server_connection_id: Optional[int] = None
@@ -193,44 +269,6 @@ def __init__(
         # For gossiping $clusterTime from the connection handshake to the client.
         self._cluster_time = None
 
-    def set_conn_timeout(self, timeout: Optional[float]) -> None:
-        """Cache last timeout to avoid duplicate calls to conn.settimeout."""
-        if timeout == self.last_timeout:
-            return
-        self.last_timeout = timeout
-        self.conn.get_conn.settimeout(timeout)
-
-    def apply_timeout(
-        self, client: AsyncMongoClient[Any], cmd: Optional[MutableMapping[str, Any]]
-    ) -> Optional[float]:
-        # CSOT: use remaining timeout when set.
-        timeout = _csot.remaining()
-        if timeout is None:
-            # Reset the socket timeout unless we're performing a streaming monitor check.
-            if not self.more_to_come:
-                self.set_conn_timeout(self.opts.socket_timeout)
-            return None
-        # RTT validation.
-        rtt = _csot.get_rtt()
-        if rtt is None:
-            rtt = self.connect_rtt
-        max_time_ms = timeout - rtt
-        if max_time_ms < 0:
-            timeout_details = _get_timeout_details(self.opts)
-            formatted = format_timeout_details(timeout_details)
-            # CSOT: raise an error without running the command since we know it will time out.
-            errmsg = f"operation would exceed time limit, remaining timeout:{timeout:.5f} <= network round trip time:{rtt:.5f} {formatted}"
-            raise ExecutionTimeout(
-                errmsg,
-                50,
-                {"ok": 0, "errmsg": errmsg, "code": 50},
-                self.max_wire_version,
-            )
-        if cmd is not None:
-            cmd["maxTimeMS"] = int(max_time_ms * 1000)
-        self.set_conn_timeout(timeout)
-        return timeout
-
     def pin_txn(self) -> None:
         self.pinned_txn = True
         assert not self.pinned_cursor
@@ -574,26 +612,6 @@ async def close_conn(self, reason: Optional[str]) -> None:
                     error=reason,
                 )
 
-    async def _close_conn(self) -> None:
-        """Close this connection."""
-        if self.closed:
-            return
-        self.closed = True
-        self.cancel_context.cancel()
-        # Note: We catch exceptions to avoid spurious errors on interpreter
-        # shutdown.
-        try:
-            await self.conn.close()
-        except Exception:  # noqa: S110
-            pass
-
-    def conn_closed(self) -> bool:
-        """Return True if we know socket has been closed, False otherwise."""
-        if _IS_SYNC:
-            return self.socket_checker.socket_closed(self.conn.get_conn)
-        else:
-            return self.conn.is_closing()
-
     def send_cluster_time(
         self,
         command: MutableMapping[str, Any],

pymongo/asynchronous/server.py

@@ -38,7 +38,6 @@
     _SDAMStatusMessage,
 )
 from pymongo.message import _convert_exception, _GetMore, _OpMsg, _Query
-from pymongo.pool_shared import _get_timeout_details, format_timeout_details
 from pymongo.response import PinnedResponse, Response
 
 if TYPE_CHECKING:
@@ -225,11 +224,7 @@ async def run_operation(
             if use_cmd:
                 first = docs[0]
                 await operation.client._process_response(first, operation.session)  # type: ignore[misc, arg-type]
-                # Append timeout details to MaxTimeMSExpired responses.
-                if first.get("code") == 50:
-                    timeout_details = _get_timeout_details(conn.opts)  # type:ignore[has-type]
-                    first["errmsg"] += format_timeout_details(timeout_details)  # type:ignore[index]
-                _check_command_response(first, conn.max_wire_version)
+                _check_command_response(first, conn.max_wire_version, pool_opts=conn.opts)  # type:ignore[has-type]
         except Exception as exc:
             duration = datetime.now() - start
             if isinstance(exc, (NotPrimaryError, OperationFailure)):

pymongo/helpers_shared.py

@@ -47,6 +47,7 @@
 if TYPE_CHECKING:
     from pymongo.cursor_shared import _Hint
     from pymongo.operations import _IndexList
+    from pymongo.pool_options import PoolOptions
     from pymongo.typings import _DocumentOut
 
 
@@ -108,6 +109,34 @@
 }
 
 
+def _get_timeout_details(options: PoolOptions) -> dict[str, float]:
+    from pymongo import _csot
+
+    details = {}
+    timeout = _csot.get_timeout()
+    socket_timeout = options.socket_timeout
+    connect_timeout = options.connect_timeout
+    if timeout:
+        details["timeoutMS"] = timeout * 1000
+    if socket_timeout and not timeout:
+        details["socketTimeoutMS"] = socket_timeout * 1000
+    if connect_timeout:
+        details["connectTimeoutMS"] = connect_timeout * 1000
+    return details
+
+
+def format_timeout_details(details: Optional[dict[str, float]]) -> str:
+    result = ""
+    if details:
+        result += " (configured timeouts:"
+        for timeout in ["socketTimeoutMS", "timeoutMS", "connectTimeoutMS"]:
+            if timeout in details:
+                result += f" {timeout}: {details[timeout]}ms,"
+        result = result[:-1]
+        result += ")"
+    return result
+
+
 def _gen_index_name(keys: _IndexList) -> str:
     """Generate an index name from the set of fields it is over."""
     return "_".join(["{}_{}".format(*item) for item in keys])
@@ -188,6 +217,7 @@ def _check_command_response(
     max_wire_version: Optional[int],
     allowable_errors: Optional[Container[Union[int, str]]] = None,
     parse_write_concern_error: bool = False,
+    pool_opts: Optional[PoolOptions] = None,
 ) -> None:
     """Check the response to a command for errors."""
     if "ok" not in response:
@@ -243,6 +273,10 @@ def _check_command_response(
     if code in (11000, 11001, 12582):
         raise DuplicateKeyError(errmsg, code, response, max_wire_version)
     elif code == 50:
+        # Append timeout details to MaxTimeMSExpired responses.
+        if pool_opts:
+            timeout_details = _get_timeout_details(pool_opts)
+            errmsg += format_timeout_details(timeout_details)
         raise ExecutionTimeout(errmsg, code, response, max_wire_version)
     elif code == 43:
         raise CursorNotFound(errmsg, code, response, max_wire_version)