Fix #3095: Added support for private key (#3408)

* ISSUE-3095: Add support for private key in snowflake

rebase main branch

* Fix #3095: Added support for private key

* Fix #3095: Added support for private key and review changes

* Fix #3095: Added support for private key and review changes

* changed direct assignment of dictionary

Co-authored-by: Onkar Ravgan <onkarravgan@Onkars-MacBook-Pro.local>

ingestion/setup.py

@@ -109,7 +109,7 @@ plugins: Dict[str, Set[str]] = {
         "psycopg2-binary",
         "GeoAlchemy2",
     },
-    "snowflake": {"snowflake-sqlalchemy<=1.3.2"},
+    "snowflake": {"snowflake-sqlalchemy<=1.3.2", "cryptography"},
     "snowflake-usage": {"snowflake-sqlalchemy<=1.3.2"},
     "sample-entity": {"faker~=8.1.1"},
     "superset": {},

ingestion/src/metadata/ingestion/source/snowflake.py

@@ -9,8 +9,12 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 import logging
+import os
 from typing import Optional
 
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric import dsa, rsa
 from snowflake.sqlalchemy.custom_types import VARIANT
 from snowflake.sqlalchemy.snowdialect import SnowflakeDialect, ischema_names
 from sqlalchemy.engine import reflection
@@ -57,6 +61,19 @@ class SnowflakeConfig(SQLConnectionConfig):
 
 class SnowflakeSource(SQLSource):
     def __init__(self, config, metadata_config, ctx):
+        if config.connect_args.get("private_key"):
+            private_key = config.connect_args["private_key"]
+            p_key = serialization.load_pem_private_key(
+                bytes(private_key, "utf-8"),
+                password=os.environ["SNOWFLAKE_PRIVATE_KEY_PASSPHRASE"].encode(),
+                backend=default_backend(),
+            )
+            pkb = p_key.private_bytes(
+                encoding=serialization.Encoding.DER,
+                format=serialization.PrivateFormat.PKCS8,
+                encryption_algorithm=serialization.NoEncryption(),
+            )
+            config.connect_args["private_key"] = pkb
         super().__init__(config, metadata_config, ctx)
 
     def fetch_sample_data(self, schema: str, table: str) -> Optional[TableData]:

ingestion/src/metadata/ingestion/source/sql_source.py

@@ -11,6 +11,7 @@
 """
 Generic source to build SQL connectors.
 """
+import copy
 import json
 import logging
 import re
@@ -71,6 +72,16 @@ def _get_table_description(schema: str, table: str, inspector: Inspector) -> str
     return description
 
 
+def _get_private_key_config(config: SQLConnectionConfig) -> SQLConnectionConfig:
+    new_config = copy.deepcopy(config)
+    if new_config.connect_args.get("private_key"):
+        new_config.connect_args["private_key"] = str(
+            new_config.connect_args["private_key"]
+        )
+        return new_config
+    return config
+
+
 class SQLSource(Source[OMetaDatabaseAndTable]):
     """
     Source Connector implementation to extract
@@ -87,7 +98,9 @@ class SQLSource(Source[OMetaDatabaseAndTable]):
         super().__init__(ctx)
         self.config = config
         self.metadata_config = metadata_config
-        self.service = get_database_service_or_create(config, metadata_config)
+        self.service = get_database_service_or_create(
+            _get_private_key_config(config), metadata_config
+        )
         self.metadata = OpenMetadata(metadata_config)
         self.status = SQLSourceStatus()
         self.sql_config = self.config