feat: implement private key loading functionality in AsyncConnection and Connection classes

Author: ya7010

turu-snowflake/src/turu/snowflake/_key_pair.py

@@ -0,0 +1,56 @@
+import base64
+from typing import TYPE_CHECKING, Optional
+from cryptography.hazmat.primitives.serialization import (
+    load_der_private_key,
+    load_pem_private_key,
+    Encoding,
+    PrivateFormat,
+    NoEncryption,
+)
+
+if TYPE_CHECKING:
+    from asn1crypto.keys import RSAPrivateKey
+
+
+def load_private_key(
+    data: bytes, private_key_passphrase: Optional[bytes] = None
+) -> "RSAPrivateKey":
+    if data.startswith(b"-----BEGIN "):
+        return _load_private_key_from_pem(data, private_key_passphrase)
+    else:
+        return _load_private_key_from_der(data, private_key_passphrase)
+
+
+def _load_private_key_from_pem(
+    data: bytes, private_key_passphrase: Optional[bytes] = None
+) -> "RSAPrivateKey":
+    """
+    Load a private key from a PEM encoded byte string.
+    """
+    p_key = load_pem_private_key(data, password=private_key_passphrase)
+    private_key = p_key.private_bytes(
+        encoding=Encoding.DER,
+        format=PrivateFormat.PKCS8,
+        encryption_algorithm=NoEncryption(),
+    )
+    return load_der_private_key(
+        data=private_key,
+        password=None,
+    )  # type: ignore[assignment]
+
+
+def _load_private_key_from_der(
+    data: bytes, private_key_passphrase: Optional[bytes] = None
+) -> "RSAPrivateKey":
+    """
+    Load a private key from a DER encoded byte string.
+    """
+    try:
+        data = base64.b64decode(data)
+    except Exception:
+        pass
+
+    return load_der_private_key(
+        data=data,
+        password=private_key_passphrase,
+    )  # type: ignore[assignment]

turu-snowflake/src/turu/snowflake/async_connection.py

@@ -55,40 +55,19 @@ async def connect(  # type: ignore[override]
         private_key_passphrase: Union[str, bytes, None] = None,
         **kwargs,
     ) -> Self:
-        if isinstance(private_key_passphrase, str):
-            private_key_passphrase = private_key_passphrase.encode("utf-8")
-
         if isinstance(private_key_file, (str, Path)):
-            from cryptography.hazmat.primitives.serialization import (
-                load_der_private_key, load_pem_private_key, Encoding, PrivateFormat, NoEncryption
-            )
             with open(private_key_file, "rb") as key:
-                p_key = load_pem_private_key(
-                    key.read(), password=private_key_passphrase
-                )
-            private_key = p_key.private_bytes(
-                encoding=Encoding.DER,
-                format=PrivateFormat.PKCS8,
-                encryption_algorithm=NoEncryption(),
-            )
-            private_key = load_der_private_key(
-                data=private_key,
-                password=None,
-            ) # type: ignore[assignment]
-        else:
-            if isinstance(private_key, str):
-                private_key = private_key.encode("utf-8")
-
-            if isinstance(private_key, bytes):
-                import base64
-                from cryptography.hazmat.primitives.serialization import (
-                    load_der_private_key,
-                )
-
-                private_key = load_der_private_key(
-                    data=base64.b64decode(private_key),
-                    password=private_key_passphrase,
-                )  # type: ignore[assignment]
+                private_key = key.read()
+        elif isinstance(private_key, str):
+            private_key = private_key.encode("utf-8")
+
+        if isinstance(private_key, bytes):
+            from turu.snowflake._key_pair import load_private_key
+
+            if isinstance(private_key_passphrase, str):
+                private_key_passphrase = private_key_passphrase.encode("utf-8")
+
+            private_key = load_private_key(private_key, private_key_passphrase)
 
         return cls(
             snowflake.connector.SnowflakeConnection(

turu-snowflake/src/turu/snowflake/connection.py

@@ -67,40 +67,19 @@ def connect(  # type: ignore[override]
         private_key_passphrase: Union[str, bytes, None] = None,
         **kwargs,
     ) -> Self:
-        if isinstance(private_key_passphrase, str):
-            private_key_passphrase = private_key_passphrase.encode("utf-8")
-
         if isinstance(private_key_file, (str, Path)):
-            from cryptography.hazmat.primitives.serialization import (
-                load_der_private_key, load_pem_private_key, Encoding, PrivateFormat, NoEncryption
-            )
             with open(private_key_file, "rb") as key:
-                p_key = load_pem_private_key(
-                    key.read(), password=private_key_passphrase
-                )
-            private_key = p_key.private_bytes(
-                encoding=Encoding.DER,
-                format=PrivateFormat.PKCS8,
-                encryption_algorithm=NoEncryption(),
-            )
-            private_key = load_der_private_key(
-                data=private_key,
-                password=None,
-            ) # type: ignore[assignment]
-        else:
-            if isinstance(private_key, str):
-                private_key = private_key.encode("utf-8")
-
-            if isinstance(private_key, bytes):
-                import base64
-                from cryptography.hazmat.primitives.serialization import (
-                    load_der_private_key,
-                )
-
-                private_key = load_der_private_key(
-                    data=base64.b64decode(private_key),
-                    password=private_key_passphrase,
-                )  # type: ignore[assignment]
+                private_key = key.read()
+        elif isinstance(private_key, str):
+            private_key = private_key.encode("utf-8")
+
+        if isinstance(private_key, bytes):
+            from turu.snowflake._key_pair import load_private_key
+
+            if isinstance(private_key_passphrase, str):
+                private_key_passphrase = private_key_passphrase.encode("utf-8")
+
+            private_key = load_private_key(private_key, private_key_passphrase)
 
         return cls(
             snowflake.connector.SnowflakeConnection(