Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Add new config parameter private_key #260

Merged
merged 4 commits into from
Sep 19, 2024
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
@@ -17,7 +17,8 @@ Built with the [Meltano Singer SDK](https://sdk.meltano.com).
|:---------------------------|:---------|:------------------------------|:-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| 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 file containing private key. |
| private_key | False | None | The private key contents. For KeyPair authentication either private_key or private_key_path must be provided. |
| private_key_path | False | None | Path to file containing private key. For KeyPair authentication either private_key or private_key_path must be provided. |
| 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. |
80 changes: 57 additions & 23 deletions target_snowflake/connector.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,9 @@
from __future__ import annotations

from enum import Enum
from functools import cached_property
from operator import contains, eq
from pathlib import Path
from typing import TYPE_CHECKING, Any, Iterable, Sequence, cast

import snowflake.sqlalchemy.custom_types as sct
@@ -10,6 +13,7 @@
from singer_sdk import typing as th
from singer_sdk.connectors import SQLConnector
from singer_sdk.connectors.sql import FullyQualifiedName
from singer_sdk.exceptions import ConfigValidationError
from snowflake.sqlalchemy import URL
from snowflake.sqlalchemy.base import SnowflakeIdentifierPreparer
from snowflake.sqlalchemy.snowdialect import SnowflakeDialect
@@ -62,6 +66,14 @@ def prepare_part(self, part: str) -> str:
return self.dialect.identifier_preparer.quote(part)


class SnowflakeAuthMethod(Enum):
"""Supported methods to authenticate to snowflake"""

BROWSER = 1
PASSWORD = 2
KEY_PAIR = 3


class SnowflakeConnector(SQLConnector):
"""Snowflake Target Connector.

@@ -124,6 +136,47 @@ def _convert_type(sql_type): # noqa: ANN205, ANN001

return sql_type

def get_private_key(self):
"""Get private key from the right location."""
phrase = self.config.get("private_key_passphrase")
encoded_passphrase = phrase.encode() if phrase else None
if "private_key_path" in self.config:
with Path.open(self.config["private_key_path"], "rb") as key:
key_content = key.read()
else:
key_content = self.config["private_key"].encode()

p_key = serialization.load_pem_private_key(
key_content,
password=encoded_passphrase,
backend=default_backend(),
)

return p_key.private_bytes(
encoding=serialization.Encoding.DER,
format=serialization.PrivateFormat.PKCS8,
encryption_algorithm=serialization.NoEncryption(),
)

@cached_property
def auth_method(self) -> SnowflakeAuthMethod:
"""Validate & return the authentication method based on config."""
if self.config.get("use_browser_authentication"):
return SnowflakeAuthMethod.BROWSER

valid_auth_methods = {"private_key", "private_key_path", "password"}
config_auth_methods = [x for x in self.config if x in valid_auth_methods]
if len(config_auth_methods) != 1:
msg = (
"Neither password nor private key was provided for "
"authentication. For password-less browser authentication via SSO, "
"set use_browser_authentication config option to True."
)
raise ConfigValidationError(msg)
if config_auth_methods[0] in ["private_key", "private_key_path"]:
return SnowflakeAuthMethod.KEY_PAIR
return SnowflakeAuthMethod.PASSWORD

def get_sqlalchemy_url(self, config: dict) -> str:
"""Generates a SQLAlchemy URL for Snowflake.

@@ -136,17 +189,10 @@ def get_sqlalchemy_url(self, config: dict) -> str:
"database": config["database"],
}

if config.get("use_browser_authentication"):
if self.auth_method == SnowflakeAuthMethod.BROWSER:
params["authenticator"] = "externalbrowser"
elif "password" in config:
elif self.auth_method == SnowflakeAuthMethod.PASSWORD:
params["password"] = config["password"]
elif "private_key_path" not in config:
msg = (
"Neither password nor private_key_path was provided for "
"authentication. For password-less browser authentication via SSO, "
"set use_browser_authentication config option to True."
)
raise Exception(msg) # noqa: TRY002

for option in ["warehouse", "role"]:
if config.get(option):
@@ -173,20 +219,8 @@ def create_engine(self) -> Engine:
"QUOTED_IDENTIFIERS_IGNORE_CASE": "TRUE",
},
}
if "private_key_path" in self.config:
with open(self.config["private_key_path"], "rb") as private_key_file: # noqa: PTH123
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(),
)
if self.auth_method == SnowflakeAuthMethod.KEY_PAIR:
connect_args["private_key"] = self.get_private_key()
engine = sqlalchemy.create_engine(
self.sqlalchemy_url,
connect_args=connect_args,
15 changes: 14 additions & 1 deletion target_snowflake/target.py
Original file line number Diff line number Diff line change
@@ -30,11 +30,24 @@ class TargetSnowflake(SQLTarget):
required=False,
description="The password for your Snowflake user.",
),
th.Property(
"private_key",
th.StringType,
required=False,
secret=True,
description=(
"The private key contents. For KeyPair authentication either "
"private_key or private_key_path must be provided."
),
),
th.Property(
"private_key_path",
th.StringType,
required=False,
description="Path to file containing private key.",
description=(
"Path to file containing private key. For KeyPair authentication either "
"private_key or private_key_path must be provided."
),
),
th.Property(
"private_key_passphrase",