Merge pull request #286 from yassun7010/support_snowflake_private_key_async

feat: support private_key to turu.snowflake.AsyncConnection.

Author: ya7010

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

@@ -1,7 +1,20 @@
 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, Self, Unpack, override
 
+import snowflake.connector
 import turu.core.async_connection
 import turu.core.cursor
 import turu.core.mock
@@ -13,9 +26,9 @@
     PanderaDataFrame,
     PyArrowTable,
 )
-from typing_extensions import Never, Self, Unpack, override
 
-import snowflake.connector
+if TYPE_CHECKING:
+    from asn1crypto.keys import RSAPrivateKey
 
 from .async_cursor import AsyncCursor, ExecuteOptions
 
@@ -37,8 +50,27 @@ async 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,
@@ -50,6 +82,7 @@ async def connect(  # type: ignore[override]
                 schema=schema,
                 warehouse=warehouse,
                 role=role,
+                private_key=private_key,
                 **kwargs,
             )
         )
@@ -68,6 +101,8 @@ async 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,
     ) -> Self:
         if (
@@ -88,6 +123,10 @@ async 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.pop("private_key", os.environ.get(private_key_envname)),
+            private_key_passphrase=kwargs.pop(
+                "private_key_passphrase", os.environ.get(private_key_passphrase_envname)
+            ),
             **kwargs,
         )