From 22adeabcf8beb6c3e6322b89c44550047453a76b Mon Sep 17 00:00:00 2001 From: Miles Adkins Date: Thu, 7 Apr 2022 14:30:16 -0500 Subject: [PATCH 1/5] added private_key auth Signed-off-by: Miles Adkins --- .../feast/infra/utils/snowflake_utils.py | 25 +++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/sdk/python/feast/infra/utils/snowflake_utils.py b/sdk/python/feast/infra/utils/snowflake_utils.py index 3513daa878c..e87437db3d7 100644 --- a/sdk/python/feast/infra/utils/snowflake_utils.py +++ b/sdk/python/feast/infra/utils/snowflake_utils.py @@ -7,6 +7,8 @@ from typing import Dict, Iterator, List, Optional, Tuple, cast import pandas as pd +from cryptography.hazmat.backends import default_backend +from cryptography.hazmat.primitives import serialization from tenacity import ( retry, retry_if_exception_type, @@ -67,6 +69,11 @@ def get_snowflake_conn(config, autocommit=True) -> SnowflakeConnection: else: kwargs["schema"] = '"PUBLIC"' + if "private_key" in kwargs: + kwargs["private_key"] = parse_private_key_path( + kwargs["private_key"], kwargs["private_key_passphrase"] + ) + try: conn = snowflake.connector.connect( application="feast", autocommit=autocommit, **kwargs @@ -288,3 +295,21 @@ def chunk_helper(lst: pd.DataFrame, n: int) -> Iterator[Tuple[int, pd.DataFrame] """Helper generator to chunk a sequence efficiently with current index like if enumerate was called on sequence.""" for i in range(0, len(lst), n): yield int(i / n), lst[i : i + n] + + +def parse_private_key_path(key_path: str, private_key_passphrase: str) -> bytes: + + with open(key_path, "rb") as key: + p_key = serialization.load_pem_private_key( + key.read(), + password=private_key_passphrase.encode(), + backend=default_backend(), + ) + + pkb = p_key.private_bytes( + encoding=serialization.Encoding.DER, + format=serialization.PrivateFormat.PKCS8, + encryption_algorithm=serialization.NoEncryption(), + ) + + return pkb From a8910e7778b2245bef1f8035d26ffc91290ed6f0 Mon Sep 17 00:00:00 2001 From: Miles Adkins Date: Thu, 7 Apr 2022 16:39:27 -0500 Subject: [PATCH 2/5] added private_key auth comment to docs Signed-off-by: Miles Adkins --- sdk/python/feast/infra/utils/snowflake_utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/sdk/python/feast/infra/utils/snowflake_utils.py b/sdk/python/feast/infra/utils/snowflake_utils.py index e87437db3d7..89789592aed 100644 --- a/sdk/python/feast/infra/utils/snowflake_utils.py +++ b/sdk/python/feast/infra/utils/snowflake_utils.py @@ -69,6 +69,7 @@ def get_snowflake_conn(config, autocommit=True) -> SnowflakeConnection: else: kwargs["schema"] = '"PUBLIC"' + #https://docs.snowflake.com/en/user-guide/python-connector-example.html#using-key-pair-authentication-key-pair-rotation if "private_key" in kwargs: kwargs["private_key"] = parse_private_key_path( kwargs["private_key"], kwargs["private_key_passphrase"] From b4e5e34bd03acd1dbd3240f70eb3567e15d8eff8 Mon Sep 17 00:00:00 2001 From: Miles Adkins Date: Thu, 7 Apr 2022 16:44:55 -0500 Subject: [PATCH 3/5] added private_key auth comment to docs Signed-off-by: Miles Adkins --- sdk/python/feast/infra/utils/snowflake_utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/sdk/python/feast/infra/utils/snowflake_utils.py b/sdk/python/feast/infra/utils/snowflake_utils.py index 89789592aed..3c842561971 100644 --- a/sdk/python/feast/infra/utils/snowflake_utils.py +++ b/sdk/python/feast/infra/utils/snowflake_utils.py @@ -70,6 +70,7 @@ def get_snowflake_conn(config, autocommit=True) -> SnowflakeConnection: kwargs["schema"] = '"PUBLIC"' #https://docs.snowflake.com/en/user-guide/python-connector-example.html#using-key-pair-authentication-key-pair-rotation + #https://docs.snowflake.com/en/user-guide/key-pair-auth.html#configuring-key-pair-authentication if "private_key" in kwargs: kwargs["private_key"] = parse_private_key_path( kwargs["private_key"], kwargs["private_key_passphrase"] From abdc98decfb4d752f118804816ef0be3209ea418 Mon Sep 17 00:00:00 2001 From: Achal Shah Date: Fri, 8 Apr 2022 17:44:02 -0700 Subject: [PATCH 4/5] type kwargs to be str -> any Signed-off-by: Achal Shah --- sdk/python/feast/infra/utils/snowflake_utils.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/sdk/python/feast/infra/utils/snowflake_utils.py b/sdk/python/feast/infra/utils/snowflake_utils.py index 3c842561971..20663787e3b 100644 --- a/sdk/python/feast/infra/utils/snowflake_utils.py +++ b/sdk/python/feast/infra/utils/snowflake_utils.py @@ -4,7 +4,7 @@ import string from logging import getLogger from tempfile import TemporaryDirectory -from typing import Dict, Iterator, List, Optional, Tuple, cast +from typing import Dict, Iterator, List, Optional, Tuple, cast, Any import pandas as pd from cryptography.hazmat.backends import default_backend @@ -50,10 +50,9 @@ def get_snowflake_conn(config, autocommit=True) -> SnowflakeConnection: # read config file config_reader = configparser.ConfigParser() config_reader.read([config_dict["config_path"]]) + kwargs: Dict[str, Any] = {} if config_reader.has_section(config_header): kwargs = dict(config_reader[config_header]) - else: - kwargs = {} if "schema" in kwargs: kwargs["schema_"] = kwargs.pop("schema") From 19fac9ebd61b40812a51c142ac922aa4ff6f05e4 Mon Sep 17 00:00:00 2001 From: Achal Shah Date: Fri, 8 Apr 2022 17:45:43 -0700 Subject: [PATCH 5/5] type kwargs to be str -> any and some other fixes Signed-off-by: Achal Shah --- sdk/python/feast/infra/utils/snowflake_utils.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/sdk/python/feast/infra/utils/snowflake_utils.py b/sdk/python/feast/infra/utils/snowflake_utils.py index 20663787e3b..a467a9de421 100644 --- a/sdk/python/feast/infra/utils/snowflake_utils.py +++ b/sdk/python/feast/infra/utils/snowflake_utils.py @@ -4,7 +4,7 @@ import string from logging import getLogger from tempfile import TemporaryDirectory -from typing import Dict, Iterator, List, Optional, Tuple, cast, Any +from typing import Any, Dict, Iterator, List, Optional, Tuple, cast import pandas as pd from cryptography.hazmat.backends import default_backend @@ -42,8 +42,8 @@ def execute_snowflake_statement(conn: SnowflakeConnection, query) -> SnowflakeCu def get_snowflake_conn(config, autocommit=True) -> SnowflakeConnection: - if config.type == "snowflake.offline": - config_header = "connections.feast_offline_store" + assert config.type == "snowflake.offline" + config_header = "connections.feast_offline_store" config_dict = dict(config) @@ -68,8 +68,8 @@ def get_snowflake_conn(config, autocommit=True) -> SnowflakeConnection: else: kwargs["schema"] = '"PUBLIC"' - #https://docs.snowflake.com/en/user-guide/python-connector-example.html#using-key-pair-authentication-key-pair-rotation - #https://docs.snowflake.com/en/user-guide/key-pair-auth.html#configuring-key-pair-authentication + # https://docs.snowflake.com/en/user-guide/python-connector-example.html#using-key-pair-authentication-key-pair-rotation + # https://docs.snowflake.com/en/user-guide/key-pair-auth.html#configuring-key-pair-authentication if "private_key" in kwargs: kwargs["private_key"] = parse_private_key_path( kwargs["private_key"], kwargs["private_key_passphrase"]