Added support for encrypted key-pair authentication on snowflake (#474)

* Added support for encrypted key-pair authentication on snowflake
Current repo should support non encrypted p8 file too, but pywht needs a fix in sending that

* PR comments addressed
Added cryptography in snowflake import packages too

src/predictions/profiles_mlcorelib/connectors/SnowflakeConnector.py

@@ -8,7 +8,8 @@
 import pandas as pd
 from datetime import datetime
 
-from typing import Any, Iterable, List, Tuple, Union, Optional, Sequence, Dict
+
+from typing import Any, Iterable, List, Union, Sequence, Dict
 
 import snowflake.snowpark
 import snowflake.snowpark.types as T
@@ -22,6 +23,8 @@
 from ..utils.logger import logger
 from ..connectors.Connector import Connector
 from ..wht.rudderPB import MATERIAL_PREFIX
+from cryptography.hazmat.primitives.serialization import load_pem_private_key
+from cryptography.hazmat.backends import default_backend
 
 local_folder = constants.SF_LOCAL_STORAGE_DIR
 
@@ -74,7 +77,20 @@ def __init__(self, creds: dict) -> None:
     def build_session(self, credentials: dict) -> snowflake.snowpark.Session:
         self.schema = credentials.get("schema", None)
         self.connection_parameters = self.remap_credentials(credentials)
+        if "privateKey" in credentials:
+            private_key = load_pem_private_key(
+                credentials["privateKey"].encode(),
+                password=(
+                    credentials["privateKeyPassphrase"].encode()
+                    if credentials.get("privateKeyPassphrase")
+                    else None
+                ),
+                backend=default_backend(),
+            )
+            self.connection_parameters["private_key"] = private_key
         session = Session.builder.configs(self.connection_parameters).create()
+        # Removing the private key to prevent serialisation error in the snowflake stored procedure
+        _ = self.connection_parameters.pop("private_key", None)
         return session
 
     def join_file_path(self, file_name: str) -> str:

src/predictions/profiles_mlcorelib/utils/constants.py

@@ -83,6 +83,7 @@
     "seaborn==0.12.0",
     "scikit-plot==0.3.7",
     "pycaret<=3.3.0",
+    "cryptography==42.0.2",
 ]
 
 

src/predictions/setup.py

@@ -39,6 +39,7 @@
         "pycaret==3.3.1",
         "boto3>=1.34.153",
         "google-auth-oauthlib>=1.0.0",
+        "cryptography>=42.0.2",
         "plotly>=5.24.1",
         "networkx>=3.3",
     ],

tests/unit/pythonWHT.py

@@ -77,7 +77,6 @@ def test_all_data_present_and_valid(self, mock_validate_historical_materials_has
         self.assertEqual(len(materials), 1)
 
     def test_missing_sequence_number(self):
-
         self.pythonWHT.connector.check_table_entry_in_material_registry = Mock(
             return_value=True
         )