From a62ade864ddc07f1c0ee2782ef07ecfbf07fefd7 Mon Sep 17 00:00:00 2001 From: Noah Stapp Date: Thu, 17 Oct 2024 11:32:39 -0400 Subject: [PATCH] PYTHON-4874 - Add KMS support for async Windows (#1939) --- 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 d14a21f41..7a325853c 100644 --- a/pymongo/network_layer.py +++ b/pymongo/network_layer.py @@ -205,7 +205,7 @@ else: 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 @@ else: 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: