Test 1731 (#1808)

Co-authored-by: Benny Lu <[email protected]>
Co-authored-by: Benny Lu <[email protected]>

Author: sfc-gh-mkeller

DESCRIPTION.md

@@ -13,6 +13,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
 
   - Added support for Vector types
   - Changed urllib3 version pin to only affect Python versions < 3.10.
+  - Support for `private_key_file` and `private_key_file_pwd` connection parameters
 
 - v3.5.0(November 13,2023)
 

src/snowflake/connector/connection.py

@@ -26,6 +26,8 @@
 from typing import Any, Callable, Generator, Iterable, Iterator, NamedTuple, Sequence
 from uuid import UUID
 
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey
 
 from . import errors, proxy
@@ -120,6 +122,26 @@ def DefaultConverterClass() -> type:
         return SnowflakeConverter
 
 
+def _get_private_bytes_from_file(
+    private_key_file: str | bytes | os.PathLike[str] | os.PathLike[bytes],
+    private_key_file_pwd: bytes | None = None,
+) -> bytes:
+    with open(private_key_file, "rb") as key:
+        private_key = serialization.load_pem_private_key(
+            key.read(),
+            password=private_key_file_pwd,
+            backend=default_backend(),
+        )
+
+    pkb = private_key.private_bytes(
+        encoding=serialization.Encoding.DER,
+        format=serialization.PrivateFormat.PKCS8,
+        encryption_algorithm=serialization.NoEncryption(),
+    )
+
+    return pkb
+
+
 SUPPORTED_PARAMSTYLES = {
     "qmark",
     "numeric",
@@ -155,6 +177,8 @@ def DefaultConverterClass() -> type:
     "passcode_in_password": (False, bool),  # Snowflake MFA
     "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)),
     "token": (None, (type(None), str)),  # OAuth or JWT Token
     "authenticator": (DEFAULT_AUTHENTICATOR, (type(None), str)),
     "mfa_callback": (None, (type(None), Callable)),
@@ -935,8 +959,16 @@ def __open_connection(self):
                 )
 
         elif self._authenticator == KEY_PAIR_AUTHENTICATOR:
+            private_key = self._private_key
+
+            if self._private_key_file:
+                private_key = _get_private_bytes_from_file(
+                    self._private_key_file,
+                    self._private_key_file_pwd,
+                )
+
             self.auth_class = AuthByKeyPair(
-                private_key=self._private_key,
+                private_key=private_key,
                 timeout=self._login_timeout,
                 backoff_generator=self._backoff_generator,
             )
@@ -1091,7 +1123,7 @@ def __config(self, **kwargs):
                 {"msg": "User is empty", "errno": ER_NO_USER},
             )
 
-        if self._private_key:
+        if self._private_key or self._private_key_file:
             self._authenticator = KEY_PAIR_AUTHENTICATOR
 
         if (

test/unit/test_connection.py

@@ -8,10 +8,15 @@
 import json
 import os
 import sys
+from pathlib import Path
 from textwrap import dedent
+from unittest import mock
 from unittest.mock import MagicMock, patch
 
 import pytest
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
 
 import snowflake.connector
 from snowflake.connector.errors import (
@@ -353,3 +358,71 @@ def test_handle_timeout(mockSessionRequest, next_action):
     # 9 seconds should be enough for authenticator to attempt twice
     # however, loosen restrictions to avoid thread scheduling causing failure
     assert 1 < mockSessionRequest.call_count < 4
+
+
+def test__get_private_bytes_from_file(tmp_path: Path):
+    private_key_file = tmp_path / "key.pem"
+
+    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.NoEncryption(),
+    )
+
+    pkb = private_key.private_bytes(
+        encoding=serialization.Encoding.DER,
+        format=serialization.PrivateFormat.PKCS8,
+        encryption_algorithm=serialization.NoEncryption(),
+    )
+
+    private_key_file.write_bytes(private_key_pem)
+
+    private_key = snowflake.connector.connection._get_private_bytes_from_file(
+        private_key_file=str(private_key_file)
+    )
+
+    assert pkb == private_key
+
+
+def test_private_key_file_reading(tmp_path: Path):
+    key_file = tmp_path / "key.pem"
+
+    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.NoEncryption(),
+    )
+
+    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),
+            )
+    assert m.call_count == 1
+    assert m.call_args_list[0].kwargs["private_key"] == pkb