feat: support private_key access.

Author: ya7010

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

@@ -1,7 +1,10 @@
 import os
 from pathlib import Path
-from typing import Any, Optional, Sequence, Tuple, Type, Union, cast, overload
+from typing import TYPE_CHECKING, Any, Optional, Sequence, Tuple, Type, Union, cast, overload
 
+from typing_extensions import Never, Unpack, override
+
+import snowflake.connector
 import turu.core.connection
 import turu.core.cursor
 import turu.core.mock
@@ -13,9 +16,6 @@
     PanderaDataFrame,
     PyArrowTable,
 )
-from typing_extensions import Never, Unpack, override
-
-import snowflake.connector
 
 from .cursor import (
     Cursor,
@@ -25,6 +25,9 @@
     Self,
 )
 
+if TYPE_CHECKING:
+    from asn1crypto.keys import RSAPrivateKey
+
 
 class Connection(turu.core.connection.Connection):
     """
@@ -49,8 +52,25 @@ def connect(  # type: ignore[override]
         schema: Optional[str] = None,
         warehouse: Optional[str] = None,
         role: Optional[str] = None,
+        private_key: "Union[str ,bytes ,RSAPrivateKey, None]" = None,
+        private_key_passphrase: Union[str, bytes, None] = None,
         **kwargs,
     ) -> Self:
+        if isinstance(private_key, str):
+            private_key = private_key.encode("utf-8")
+
+        if isinstance(private_key_passphrase, str):
+            private_key_passphrase = private_key_passphrase.encode("utf-8")
+
+        if isinstance(private_key, bytes) and  private_key_passphrase is not None:
+            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]
+
         return cls(
             snowflake.connector.SnowflakeConnection(
                 connection_name,
@@ -62,6 +82,7 @@ def connect(  # type: ignore[override]
                 schema=schema,
                 warehouse=warehouse,
                 role=role,
+                private_key=private_key,
                 **kwargs,
             )
         )
@@ -80,6 +101,8 @@ def connect_from_env(  # type: ignore[override]
         warehouse_envname: str = "SNOWFLAKE_WAREHOUSE",
         role_envname: str = "SNOWFLAKE_ROLE",
         authenticator_envname: str = "SNOWFLAKE_AUTHENTICATOR",
+        private_key_envname: str = "SNOWFLAKE_PRIVATE_KEY",
+        private_key_passphrase_envname: str = "SNOWFLAKE_PRIVATE_KEY_PASSPHRASE",
         **kwargs: Any,
     ) -> Self:
         if (
@@ -100,6 +123,10 @@ def connect_from_env(  # type: ignore[override]
             schema=kwargs.get("schema", os.environ.get(schema_envname)),
             warehouse=kwargs.get("warehouse", os.environ.get(warehouse_envname)),
             role=kwargs.get("role", os.environ.get(role_envname)),
+            private_key=kwargs.get("private_key", os.environ.get(private_key_envname)),
+            private_key_passphrase=kwargs.get(
+                "private_key_passphrase", os.environ.get(private_key_passphrase_envname)
+            ),
             **kwargs,
         )