feat: add private key authentication (#109)

This adds two additional properties to config allowing to use key pair authentication as described in [Key pair authentication](https://docs.snowflake.com/en/user-guide/key-pair-auth) It can be used instead of password and decrypts the key with a passphrase if needed. Co-authored-by: Pat Nadolny <[email protected]>
MeltanoLabs · Sep 6, 2023 · 4eb60b8 · 4eb60b8
1 parent 2a05455
commit 4eb60b8
Show file tree

Hide file tree

Showing 4 changed files with 68 additions and 28 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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,