From ad041368476e50ed6816fc3565b152f57bdf19f9 Mon Sep 17 00:00:00 2001 From: Onkar Ravgan Date: Tue, 15 Mar 2022 18:19:02 +0530 Subject: [PATCH] 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 --- ingestion/setup.py | 2 +- .../src/metadata/ingestion/source/snowflake.py | 17 +++++++++++++++++ .../src/metadata/ingestion/source/sql_source.py | 15 ++++++++++++++- 3 files changed, 32 insertions(+), 2 deletions(-) diff --git a/ingestion/setup.py b/ingestion/setup.py index 7e0bb20f29f..df21f7278d9 100644 --- a/ingestion/setup.py +++ b/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": {}, diff --git a/ingestion/src/metadata/ingestion/source/snowflake.py b/ingestion/src/metadata/ingestion/source/snowflake.py index e56a138faa9..5e8fcd0ac69 100644 --- a/ingestion/src/metadata/ingestion/source/snowflake.py +++ b/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]: diff --git a/ingestion/src/metadata/ingestion/source/sql_source.py b/ingestion/src/metadata/ingestion/source/sql_source.py index bace8616338..5b8983474f4 100644 --- a/ingestion/src/metadata/ingestion/source/sql_source.py +++ b/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