From e5ee1eaca5dff8d766c538ce26412199d00d95d7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mi=C5=82osz=20Szymczak?= Date: Thu, 17 Aug 2023 10:38:16 +0200 Subject: [PATCH] feat: add private key authentication --- README.md | 34 +++++++++++++++++---------------- meltano.yml | 8 ++++++++ target_snowflake/connector.py | 36 +++++++++++++++++++++++++++-------- target_snowflake/target.py | 18 ++++++++++++++---- 4 files changed, 68 insertions(+), 28 deletions(-) diff --git a/README.md b/README.md index 2924168..22f4ad4 100644 --- a/README.md +++ b/README.md @@ -13,22 +13,24 @@ Built with the [Meltano Singer SDK](https://sdk.meltano.com). ## Settings -| Setting | Required | Default | Description | -|:---------------------|:--------:|:-------:|:------------| -| user | True | None | The login name for your Snowflake user. | -| password | True | None | The password for your Snowflake user. | -| account | True | None | Your account identifier. See [Account Identifiers](https://docs.snowflake.com/en/user-guide/admin-account-identifier.html). | -| database | True | None | The initial database for the Snowflake session. | -| schema | False | None | The initial schema for the Snowflake session. | -| warehouse | False | None | The initial warehouse for the session. | -| role | False | None | The initial role for the session. | -| add_record_metadata | False | 1 | Whether to add metadata columns. | -| clean_up_batch_files | False | 1 | Whether to remove batch files after processing. | -| default_target_schema| False | None | The default target database schema name to use for all streams. | -| stream_maps | False | None | Config object for stream maps capability. For more information check out [Stream Maps](https://sdk.meltano.com/en/latest/stream_maps.html). | -| stream_map_config | False | None | User-defined config values to be used within map expressions. | -| flattening_enabled | False | None | 'True' to enable schema flattening and automatically expand nested properties. | -| flattening_max_depth | False | None | The max depth to flatten schemas. | +| Setting | Required | Default | Description | +|:-----------------------|:--------:|:-------:|:--------------------------------------------------------------------------------------------------------------------------------------------| +| user | True | None | The login name for your Snowflake user. | +| password | False | None | The password for your Snowflake user. | +| private_key_path | False | None | Path to private key file. See [Key pair auth](https://docs.snowflake.com/en/user-guide/key-pair-auth) | +| private_key_passphrase | False | None | Passphrase to decrypt private key if encrypted. | +| account | True | None | Your account identifier. See [Account Identifiers](https://docs.snowflake.com/en/user-guide/admin-account-identifier.html). | +| database | True | None | The initial database for the Snowflake session. | +| schema | False | None | The initial schema for the Snowflake session. | +| warehouse | False | None | The initial warehouse for the session. | +| role | False | None | The initial role for the session. | +| add_record_metadata | False | 1 | Whether to add metadata columns. | +| clean_up_batch_files | False | 1 | Whether to remove batch files after processing. | +| default_target_schema | False | None | The default target database schema name to use for all streams. | +| stream_maps | False | None | Config object for stream maps capability. For more information check out [Stream Maps](https://sdk.meltano.com/en/latest/stream_maps.html). | +| stream_map_config | False | None | User-defined config values to be used within map expressions. | +| flattening_enabled | False | None | 'True' to enable schema flattening and automatically expand nested properties. | +| flattening_max_depth | False | None | The max depth to flatten schemas. | A full list of supported settings and capabilities is available by running: `target-snowflake --about` diff --git a/meltano.yml b/meltano.yml index 6317132..b341ea5 100644 --- a/meltano.yml +++ b/meltano.yml @@ -49,6 +49,14 @@ plugins: kind: password label: Password name: password + - description: Path to file containing private key + kind: string + label: Private Key Path + name: private_key_path + - description: Passphrase to decrypt private key if encrypted + label: Private Key Passphrase + kind: password + name: private_key_passphrase - description: The initial role for the session. kind: string label: Role diff --git a/target_snowflake/connector.py b/target_snowflake/connector.py index 469dace..8efdc6e 100644 --- a/target_snowflake/connector.py +++ b/target_snowflake/connector.py @@ -3,6 +3,8 @@ import snowflake.sqlalchemy.custom_types as sct import sqlalchemy +from cryptography.hazmat.backends import default_backend +from cryptography.hazmat.primitives import serialization from singer_sdk import typing as th from singer_sdk.connectors import SQLConnector from snowflake.sqlalchemy import URL @@ -15,6 +17,7 @@ SNOWFLAKE_MAX_STRING_LENGTH = 16777216 + class TypeMap: def __init__(self, operator, map_value, match_value=None): self.operator = operator @@ -89,7 +92,8 @@ def get_table_columns( self.table_cache[full_table_name] = parsed_columns return parsed_columns - def _convert_type(self, sql_type): + @staticmethod + def _convert_type(sql_type): if isinstance(sql_type, sct.TIMESTAMP_NTZ): return TIMESTAMP_NTZ elif isinstance(sql_type, sct.NUMBER): @@ -108,10 +112,14 @@ def get_sqlalchemy_url(self, config: dict) -> str: params = { "account": config["account"], "user": config["user"], - "password": config["password"], "database": config["database"], } + if "password" in config: + params["password"] = config["password"] + elif "private_key_path" not in config: + raise Exception("Neither password nor private_key_path was provided for authentication.") + for option in ["warehouse", "role"]: if config.get(option): params[option] = config.get(option) @@ -132,13 +140,26 @@ def create_engine(self) -> Engine: Returns: A new SQLAlchemy Engine. """ + connect_args = { + "session_parameters": { + "QUOTED_IDENTIFIERS_IGNORE_CASE": "TRUE", + } + } + if "private_key_path" in self.config: + with open(self.config["private_key_path"], "rb") as private_key_file: + private_key = serialization.load_pem_private_key( + private_key_file.read(), + password=self.config["private_key_passphrase"].encode() if "private_key_passphrase" in self.config else None, + backend=default_backend(), + ) + connect_args["private_key"] = private_key.private_bytes( + encoding=serialization.Encoding.DER, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption(), + ) engine = sqlalchemy.create_engine( self.sqlalchemy_url, - connect_args={ - "session_parameters": { - "QUOTED_IDENTIFIERS_IGNORE_CASE": "TRUE", - } - }, + connect_args=connect_args, echo=False, ) connection = engine.connect() @@ -147,7 +168,6 @@ def create_engine(self) -> Engine: raise Exception(f"Database '{self.config['database']}' does not exist or the user/role doesn't have access to it.") return engine - def prepare_column( self, full_table_name: str, diff --git a/target_snowflake/target.py b/target_snowflake/target.py index e54de1e..acb5a9a 100644 --- a/target_snowflake/target.py +++ b/target_snowflake/target.py @@ -2,11 +2,9 @@ from __future__ import annotations -import sys - import click from singer_sdk import typing as th -from singer_sdk.target_base import SQLTarget, Target +from singer_sdk.target_base import SQLTarget from target_snowflake.initializer import initializer from target_snowflake.sinks import SnowflakeSink @@ -27,9 +25,21 @@ class TargetSnowflake(SQLTarget): th.Property( "password", th.StringType, - required=True, + required=False, description="The password for your Snowflake user.", ), + th.Property( + "private_key_path", + th.StringType, + required=False, + description="Path to file containing private key.", + ), + th.Property( + "private_key_passphrase", + th.StringType, + required=False, + description="Passphrase to decrypt private key if encrypted.", + ), th.Property( "account", th.StringType,