Fix for an issue where specifying a private_key_file_pwd in a connections.toml file causes an error. (#1878)

* Fix for issue SNOW-1045815 where specifying a private_key_file_pwd in a connections.toml file causes an error.

* Fix lint

---------

Co-authored-by: Sophie Tan <[email protected]>

Author: rosspb3

src/snowflake/connector/connection.py

@@ -124,8 +124,10 @@ def DefaultConverterClass() -> type:
 
 def _get_private_bytes_from_file(
     private_key_file: str | bytes | os.PathLike[str] | os.PathLike[bytes],
-    private_key_file_pwd: bytes | None = None,
+    private_key_file_pwd: bytes | str | None = None,
 ) -> bytes:
+    if private_key_file_pwd is not None and isinstance(private_key_file_pwd, str):
+        private_key_file_pwd = private_key_file_pwd.encode("utf-8")
     with open(private_key_file, "rb") as key:
         private_key = serialization.load_pem_private_key(
             key.read(),
@@ -178,7 +180,7 @@ def _get_private_bytes_from_file(
     "passcode": (None, (type(None), str)),  # Snowflake MFA
     "private_key": (None, (type(None), str, RSAPrivateKey)),
     "private_key_file": (None, (type(None), str)),
-    "private_key_file_pwd": (None, (type(None), str)),
+    "private_key_file_pwd": (None, (type(None), str, bytes)),
     "token": (None, (type(None), str)),  # OAuth or JWT Token
     "authenticator": (DEFAULT_AUTHENTICATOR, (type(None), str)),
     "mfa_callback": (None, (type(None), Callable)),

test/unit/test_connection.py

@@ -9,6 +9,7 @@
 import os
 import sys
 from pathlib import Path
+from secrets import token_urlsafe
 from textwrap import dedent
 from unittest import mock
 from unittest.mock import MagicMock, patch
@@ -428,6 +429,49 @@ def test_private_key_file_reading(tmp_path: Path):
     assert m.call_args_list[0].kwargs["private_key"] == pkb
 
 
+def test_encrypted_private_key_file_reading(tmp_path: Path):
+    key_file = tmp_path / "key.pem"
+    private_key_password = token_urlsafe(25)
+    private_key = rsa.generate_private_key(
+        backend=default_backend(), public_exponent=65537, key_size=2048
+    )
+
+    private_key_pem = private_key.private_bytes(
+        encoding=serialization.Encoding.PEM,
+        format=serialization.PrivateFormat.PKCS8,
+        encryption_algorithm=serialization.BestAvailableEncryption(
+            private_key_password.encode("utf-8")
+        ),
+    )
+
+    key_file.write_bytes(private_key_pem)
+
+    pkb = private_key.private_bytes(
+        encoding=serialization.Encoding.DER,
+        format=serialization.PrivateFormat.PKCS8,
+        encryption_algorithm=serialization.NoEncryption(),
+    )
+
+    exc_msg = "stop execution"
+
+    with mock.patch(
+        "snowflake.connector.auth.keypair.AuthByKeyPair.__init__",
+        side_effect=Exception(exc_msg),
+    ) as m:
+        with pytest.raises(
+            Exception,
+            match=exc_msg,
+        ):
+            snowflake.connector.connect(
+                account="test_account",
+                user="test_user",
+                private_key_file=str(key_file),
+                private_key_file_pwd=private_key_password,
+            )
+    assert m.call_count == 1
+    assert m.call_args_list[0].kwargs["private_key"] == pkb
+
+
 def test_expired_detection():
     with mock.patch(
         "snowflake.connector.network.SnowflakeRestful._post_request",