Working protocols with debugging prints

Author: NoahStapp

pymongo/asynchronous/network.py

@@ -53,7 +53,7 @@
     from bson import CodecOptions
     from pymongo.asynchronous.client_session import AsyncClientSession
     from pymongo.asynchronous.mongo_client import AsyncMongoClient
-    from pymongo.asynchronous.pool import AsyncConnection, AsyncStreamConnection, AsyncConnectionStream
+    from pymongo.asynchronous.pool import AsyncConnection, AsyncStreamConnection, AsyncConnectionProtocol
     from pymongo.compression_support import SnappyContext, ZlibContext, ZstdContext
     from pymongo.monitoring import _EventListeners
     from pymongo.read_concern import ReadConcern
@@ -70,7 +70,7 @@
 
 
 async def command_stream(
-    conn: AsyncConnectionStream,
+    conn: AsyncConnectionProtocol,
     dbname: str,
     spec: MutableMapping[str, Any],
     is_mongos: bool,
@@ -317,7 +317,7 @@ async def command_stream(
 
 
 async def receive_message(
-    conn: AsyncConnectionStream, request_id: Optional[int], max_message_size: int = MAX_MESSAGE_SIZE
+    conn: AsyncConnectionProtocol, request_id: Optional[int], max_message_size: int = MAX_MESSAGE_SIZE
 ) -> Union[_OpReply, _OpMsg]:
     """Receive a raw BSON message or raise socket.error."""
     # if _csot.get_timeout():
@@ -330,28 +330,28 @@ async def receive_message(
     #         deadline = None
     deadline = None
     # Ignore the response's request id.
-    length, _, response_to, op_code = _UNPACK_HEADER(await async_receive_data(conn, 16, deadline))
+    data, op_code = await async_receive_data(conn, 0, deadline)
+    # length, _, response_to, op_code = _UNPACK_HEADER(await async_receive_data(conn, 16, deadline))
     # No request_id for exhaust cursor "getMore".
-    if request_id is not None:
-        if request_id != response_to:
-            raise ProtocolError(f"Got response id {response_to!r} but expected {request_id!r}")
-    if length <= 16:
-        raise ProtocolError(
-            f"Message length ({length!r}) not longer than standard message header size (16)"
-        )
-    if length > max_message_size:
-        raise ProtocolError(
-            f"Message length ({length!r}) is larger than server max "
-            f"message size ({max_message_size!r})"
-        )
-    if op_code == 2012:
-        op_code, _, compressor_id = _UNPACK_COMPRESSION_HEADER(
-            await async_receive_data(conn, 9, deadline)
-        )
-        data = decompress(await async_receive_data(conn, length - 25, deadline), compressor_id)
-    else:
-        data = await async_receive_data(conn, length - 16, deadline)
-
+    # if request_id is not None:
+    #     if request_id != response_to:
+    #         raise ProtocolError(f"Got response id {response_to!r} but expected {request_id!r}")
+    # if length <= 16:
+    #     raise ProtocolError(
+    #         f"Message length ({length!r}) not longer than standard message header size (16)"
+    #     )
+    # if length > max_message_size:
+    #     raise ProtocolError(
+    #         f"Message length ({length!r}) is larger than server max "
+    #         f"message size ({max_message_size!r})"
+    #     )
+    # if op_code == 2012:
+    #     op_code, _, compressor_id = _UNPACK_COMPRESSION_HEADER(
+    #         await async_receive_data(conn, 9, deadline)
+    #     )
+    #     data = decompress(await async_receive_data(conn, length - 25, deadline), compressor_id)
+    # else:
+    # data = await async_receive_data(conn, length - 16, deadline)
     try:
         unpack_reply = _UNPACK_REPLY[op_code]
     except KeyError:

pymongo/asynchronous/pool.py

@@ -578,6 +578,7 @@ async def send_message(self, message: bytes, max_doc_size: int) -> None:
         try:
             await async_sendall(self.conn, message)
         except BaseException as error:
+            print(error)
             self._raise_connection_failure(error)
 
     async def receive_message(self, request_id: Optional[int]) -> Union[_OpReply, _OpMsg]:
@@ -784,7 +785,7 @@ def __repr__(self) -> str:
         )
 
 
-class AsyncConnectionStream:
+class AsyncConnectionProtocol:
     """Store a connection with some metadata.
 
     :param conn: a raw connection object
@@ -1818,7 +1819,7 @@ async def remove_stale_sockets(self, reference_generation: int) -> None:
                     self.requests -= 1
                     self.size_cond.notify()
 
-    async def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> AsyncConnectionStream:
+    async def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> AsyncConnectionProtocol:
         """Connect to Mongo and return a new AsyncConnection.
 
         Can raise ConnectionFailure.
@@ -1874,7 +1875,7 @@ async def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> A
 
             raise
 
-        conn = AsyncConnectionStream(sock, self, self.address, conn_id)  # type: ignore[arg-type]
+        conn = AsyncConnectionProtocol(sock, self, self.address, conn_id)  # type: ignore[arg-type]
         async with self.lock:
             self.active_contexts.add(conn.cancel_context)
             self.active_contexts.discard(tmp_context)
@@ -1899,7 +1900,7 @@ async def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> A
     @contextlib.asynccontextmanager
     async def checkout(
         self, handler: Optional[_MongoClientErrorHandler] = None
-    ) -> AsyncGenerator[AsyncConnectionStream, None]:
+    ) -> AsyncGenerator[AsyncConnectionProtocol, None]:
         """Get a connection from the pool. Use with a "with" statement.
 
         Returns a :class:`AsyncConnection` object wrapping a connected
@@ -2002,7 +2003,7 @@ def _raise_if_not_ready(self, checkout_started_time: float, emit_event: bool) ->
 
     async def _get_conn(
         self, checkout_started_time: float, handler: Optional[_MongoClientErrorHandler] = None
-    ) -> AsyncConnectionStream:
+    ) -> AsyncConnectionProtocol:
         """Get or create a AsyncConnection. Can raise ConnectionFailure."""
         # We use the pid here to avoid issues with fork / multiprocessing.
         # See test.test_client:TestClient.test_fork for an example of

pymongo/message.py

@@ -1547,7 +1547,7 @@ def unpack(cls, msg: bytes) -> _OpMsg:
             raise ProtocolError(f"Unsupported OP_MSG payload type: 0x{first_payload_type:x}")
 
         if len(msg) != first_payload_size + 5:
-            raise ProtocolError("Unsupported OP_MSG reply: >1 section")
+            raise ProtocolError(f"Unsupported OP_MSG reply: >1 section, {len(msg)} vs {first_payload_size + 5}")
 
         payload_document = msg[5:]
         return cls(flags, payload_document)