yujunjun/datahub
2
0
Fork 0
mirror of https://github.com/datahub-project/datahub.git synced 2025-09-26 17:45:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(ingest): snowflake - config variable for specifying a direct private key (#6609)

Browse Source
This commit is contained in:
Mayuri Nehate 2022-12-05 19:09:08 +05:30 committed by GitHub
parent ed9110ea44
commit fdcb731e29
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 30 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37
metadata-ingestion/src/datahub/ingestion/source_config/sql/snowflake.py
View File

@ -105,15 +105,21 @@ class BaseSnowflakeConfig(BaseTimeWindowConfig):
password: Optional[pydantic.SecretStr] = pydantic.Field( password: Optional[pydantic.SecretStr] = pydantic.Field(
default=None, exclude=True, description="Snowflake password." default=None, exclude=True, description="Snowflake password."
) )
private_key: Optional[str] = pydantic.Field(
default=None,
description="Private key in a form of '-----BEGIN PRIVATE KEY-----\\nprivate-key\\n-----END PRIVATE KEY-----\\n' if using key pair authentication. Encrypted version of private key will be in a form of '-----BEGIN ENCRYPTED PRIVATE KEY-----\\nencrypted-private-key\\n-----END ECNCRYPTED PRIVATE KEY-----\\n' See: https://docs.snowflake.com/en/user-guide/key-pair-auth.html",
)
private_key_path: Optional[str] = pydantic.Field( private_key_path: Optional[str] = pydantic.Field(
default=None, default=None,
description="The path to the private key if using key pair authentication. See: https://docs.snowflake.com/en/user-guide/key-pair-auth.html", description="The path to the private key if using key pair authentication. Ignored if `private_key` is set. See: https://docs.snowflake.com/en/user-guide/key-pair-auth.html",
) )
private_key_password: Optional[pydantic.SecretStr] = pydantic.Field( private_key_password: Optional[pydantic.SecretStr] = pydantic.Field(
default=None, default=None,
exclude=True, exclude=True,
description="Password for your private key if using key pair authentication.", description="Password for your private key. Required if using key pair authentication with encrypted private key.",
) )
oauth_config: Optional[OauthConfiguration] = pydantic.Field( oauth_config: Optional[OauthConfiguration] = pydantic.Field(
default=None, default=None,
description="oauth configuration - https://docs.snowflake.com/en/user-guide/python-connector-example.html#connecting-with-oauth", description="oauth configuration - https://docs.snowflake.com/en/user-guide/python-connector-example.html#connecting-with-oauth",
@ -182,10 +188,13 @@ class BaseSnowflakeConfig(BaseTimeWindowConfig):
) )
if v == "KEY_PAIR_AUTHENTICATOR": if v == "KEY_PAIR_AUTHENTICATOR":
# If we are using key pair auth, we need the private key path and password to be set # If we are using key pair auth, we need the private key path and password to be set
if values.get("private_key_path") is None: if (
values.get("private_key") is None
and values.get("private_key_path") is None
):
raise ValueError( raise ValueError(
f"'private_key_path' was none " f"Both `private_key` and `private_key_path` are none. "
f"but should be set when using {v} authentication" f"At least one should be set when using {v} authentication"
) )
elif v == "OAUTH_AUTHENTICATOR": elif v == "OAUTH_AUTHENTICATOR":
if values.get("oauth_config") is None: if values.get("oauth_config") is None:
@ -275,14 +284,20 @@ class BaseSnowflakeConfig(BaseTimeWindowConfig):
if self.authentication_type != "KEY_PAIR_AUTHENTICATOR": if self.authentication_type != "KEY_PAIR_AUTHENTICATOR":
return {} return {}
if self.connect_args is None: if self.connect_args is None:
if self.private_key_path is None: if self.private_key is not None:
raise ValueError("missing required private key path to read key from") pkey_bytes = self.private_key.replace("\\n", "\n").encode()
if self.private_key_password is None: else:
raise ValueError("missing required private key password") assert (
self.private_key_path
), "missing required private key path to read key from"
with open(self.private_key_path, "rb") as key: with open(self.private_key_path, "rb") as key:
pkey_bytes = key.read()
p_key = serialization.load_pem_private_key( p_key = serialization.load_pem_private_key(
key.read(), pkey_bytes,
password=self.private_key_password.get_secret_value().encode(), password=self.private_key_password.get_secret_value().encode()
if self.private_key_password is not None
else None,
backend=default_backend(), backend=default_backend(),
) )