address review

mongodb · Jan 10, 2025 · 321b3e3 · 321b3e3
1 parent 8232f26
commit 321b3e3
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 8 deletions.
diff --git a/pymongo/asynchronous/encryption.py b/pymongo/asynchronous/encryption.py
@@ -213,16 +213,20 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
                     if not data:
                         raise OSError("KMS connection closed")
                     kms_context.feed(data)
-            # Async raises an OSError instead of returning empty bytes
-            except OSError as err:
-                raise OSError("KMS connection closed") from err
             except MongoCryptError:
                 raise  # Propagate MongoCryptError errors directly.
             except Exception as exc:
                 # Wrap I/O errors in PyMongo exceptions.
                 if isinstance(exc, BLOCKING_IO_ERRORS):
                     exc = socket.timeout("timed out")
-                _raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts))
+                # Async raises an OSError instead of returning empty bytes.
+                if isinstance(exc, OSError):
+                    msg_prefix = "KMS connection closed"
+                else:
+                    msg_prefix = None
+                _raise_connection_failure(
+                    address, exc, msg_prefix=msg_prefix, timeout_details=_get_timeout_details(opts)
+                )
             finally:
                 conn.close()
         except MongoCryptError:

diff --git a/pymongo/synchronous/encryption.py b/pymongo/synchronous/encryption.py
@@ -213,16 +213,20 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
                     if not data:
                         raise OSError("KMS connection closed")
                     kms_context.feed(data)
-            # Async raises an OSError instead of returning empty bytes
-            except OSError as err:
-                raise OSError("KMS connection closed") from err
             except MongoCryptError:
                 raise  # Propagate MongoCryptError errors directly.
             except Exception as exc:
                 # Wrap I/O errors in PyMongo exceptions.
                 if isinstance(exc, BLOCKING_IO_ERRORS):
                     exc = socket.timeout("timed out")
-                _raise_connection_failure(address, exc, timeout_details=_get_timeout_details(opts))
+                # Async raises an OSError instead of returning empty bytes.
+                if isinstance(exc, OSError):
+                    msg_prefix = "KMS connection closed"
+                else:
+                    msg_prefix = None
+                _raise_connection_failure(
+                    address, exc, msg_prefix=msg_prefix, timeout_details=_get_timeout_details(opts)
+                )
             finally:
                 conn.close()
         except MongoCryptError: