From 5965d3c930b2d569a1b8c849921d8d5b3c61bce8 Mon Sep 17 00:00:00 2001 From: Noah Stapp Date: Thu, 17 Oct 2024 10:18:37 -0400 Subject: [PATCH] PYTHON-4874 - Add KMS support for async Windows --- pymongo/network_layer.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/pymongo/network_layer.py b/pymongo/network_layer.py index d14a21f41d..7a325853c8 100644 --- a/pymongo/network_layer.py +++ b/pymongo/network_layer.py @@ -205,7 +205,7 @@ async def _async_sendall_ssl( total_sent += sent async def _async_receive_ssl( - conn: _sslConn, length: int, dummy: AbstractEventLoop + conn: _sslConn, length: int, dummy: AbstractEventLoop, once: Optional[bool] = False ) -> memoryview: mv = memoryview(bytearray(length)) total_read = 0 @@ -215,6 +215,9 @@ async def _async_receive_ssl( while total_read < length: try: read = conn.recv_into(mv[total_read:]) + # KMS responses update their expected size after the first batch, stop reading after one loop + if once: + return mv[:read] if read == 0: raise OSError("connection closed") except BLOCKING_IO_ERRORS: