PYTHON-5544 [v4.15] Revert changes to base protocol layer (#2538)

Author: blink1073

doc/changelog.rst

@@ -1,7 +1,7 @@
 Changelog
 =========
 
-Changes in Version 4.15.1 (2025/09/11)
+Changes in Version 4.15.1 (2025/09/16)
 --------------------------------------
 
 Version 4.15.1 is a bug fix release.

pymongo/asynchronous/encryption.py

@@ -64,7 +64,6 @@
 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, TextOpts
@@ -77,11 +76,11 @@
     ServerSelectionTimeoutError,
 )
 from pymongo.helpers_shared import _get_timeout_details
-from pymongo.network_layer import PyMongoKMSProtocol, async_receive_kms, async_sendall
+from pymongo.network_layer import async_socket_sendall
 from pymongo.operations import UpdateOne
 from pymongo.pool_options import PoolOptions
 from pymongo.pool_shared import (
-    _configured_protocol_interface,
+    _async_configured_socket,
     _raise_connection_failure,
 )
 from pymongo.read_concern import ReadConcern
@@ -94,8 +93,10 @@
 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
@@ -110,10 +111,9 @@
 _KEY_VAULT_OPTS = CodecOptions(document_class=RawBSONDocument)
 
 
-async def _connect_kms(address: _Address, opts: PoolOptions) -> AsyncBaseConnection:
+async def _connect_kms(address: _Address, opts: PoolOptions) -> Union[socket.socket, _sslConn]:
     try:
-        interface = await _configured_protocol_interface(address, opts, PyMongoKMSProtocol)
-        return AsyncBaseConnection(interface, opts)
+        return await _async_configured_socket(address, opts)
     except Exception as exc:
         _raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts))
 
@@ -198,11 +198,19 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
         try:
             conn = await _connect_kms(address, opts)
             try:
-                await async_sendall(conn.conn.get_conn, message)
+                await async_socket_sendall(conn, message)
                 while kms_context.bytes_needed > 0:
                     # CSOT: update timeout.
-                    conn.set_conn_timeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0))
-                    data = await async_receive_kms(conn, kms_context.bytes_needed)
+                    conn.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0))
+                    data: memoryview | bytes
+                    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)
                     if not data:
                         raise OSError("KMS connection closed")
                     kms_context.feed(data)
@@ -221,7 +229,7 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
                     address, exc, msg_prefix=msg_prefix, timeout_details=_get_timeout_details(opts)
                 )
             finally:
-                await conn.close_conn(None)
+                conn.close()
         except MongoCryptError:
             raise  # Propagate MongoCryptError errors directly.
         except Exception as exc:

pymongo/asynchronous/pool.py

@@ -123,89 +123,7 @@ def _set_non_inheritable_non_atomic(fd: int) -> None:  # noqa: ARG001
 _IS_SYNC = False
 
 
-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):
+class AsyncConnection:
     """Store a connection with some metadata.
 
     :param conn: a raw connection object
@@ -223,27 +141,29 @@ def __init__(
         id: int,
         is_sdam: bool,
     ):
-        super().__init__(conn, pool.opts)
         self.pool_ref = weakref.ref(pool)
-        self.address: tuple[str, int] = address
-        self.id: int = id
+        self.conn = conn
+        self.address = address
+        self.id = 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: int = MAX_BSON_SIZE
-        self.max_message_size: int = MAX_MESSAGE_SIZE
-        self.max_write_batch_size: int = MAX_WRITE_BATCH_SIZE
+        self.max_bson_size = MAX_BSON_SIZE
+        self.max_message_size = MAX_MESSAGE_SIZE
+        self.max_write_batch_size = MAX_WRITE_BATCH_SIZE
         self.supports_sessions = False
         self.hello_ok: bool = False
-        self.is_mongos: bool = False
+        self.is_mongos = 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
@@ -254,6 +174,9 @@ 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
@@ -269,6 +192,44 @@ 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
@@ -612,6 +573,26 @@ 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],