From be3ecb2e324fe7122c75d5eb5ffb40e5bd91b878 Mon Sep 17 00:00:00 2001 From: Artur Date: Mon, 5 Aug 2024 18:01:27 -0400 Subject: [PATCH 1/8] allow feast snowflake to read in byte string for private-key authentication Signed-off-by: Artur --- .../infra/materialization/snowflake_engine.py | 4 ++-- .../feast/infra/offline_stores/snowflake.py | 4 ++-- .../feast/infra/online_stores/snowflake.py | 4 ++-- sdk/python/feast/infra/registry/snowflake.py | 4 ++-- .../infra/utils/snowflake/snowflake_utils.py | 20 +++++++++++++++---- 5 files changed, 24 insertions(+), 12 deletions(-) diff --git a/sdk/python/feast/infra/materialization/snowflake_engine.py b/sdk/python/feast/infra/materialization/snowflake_engine.py index f77239398e6..ce3d555f97a 100644 --- a/sdk/python/feast/infra/materialization/snowflake_engine.py +++ b/sdk/python/feast/infra/materialization/snowflake_engine.py @@ -67,8 +67,8 @@ class SnowflakeMaterializationEngineConfig(FeastConfigBaseModel): authenticator: Optional[str] = None """ Snowflake authenticator name """ - private_key: Optional[str] = None - """ Snowflake private key file path""" + private_key: Optional[str | bytes] = None + """ Snowflake private key stored as bytes or file path""" private_key_passphrase: Optional[str] = None """ Snowflake private key file passphrase""" diff --git a/sdk/python/feast/infra/offline_stores/snowflake.py b/sdk/python/feast/infra/offline_stores/snowflake.py index 96552ff87ec..c214c961730 100644 --- a/sdk/python/feast/infra/offline_stores/snowflake.py +++ b/sdk/python/feast/infra/offline_stores/snowflake.py @@ -104,8 +104,8 @@ class SnowflakeOfflineStoreConfig(FeastConfigBaseModel): authenticator: Optional[str] = None """ Snowflake authenticator name """ - private_key: Optional[str] = None - """ Snowflake private key file path""" + private_key: Optional[str | bytes] = None + """ Snowflake private key stored as bytes or file path""" private_key_passphrase: Optional[str] = None """ Snowflake private key file passphrase""" diff --git a/sdk/python/feast/infra/online_stores/snowflake.py b/sdk/python/feast/infra/online_stores/snowflake.py index fef804a3773..b015283929f 100644 --- a/sdk/python/feast/infra/online_stores/snowflake.py +++ b/sdk/python/feast/infra/online_stores/snowflake.py @@ -50,8 +50,8 @@ class SnowflakeOnlineStoreConfig(FeastConfigBaseModel): authenticator: Optional[str] = None """ Snowflake authenticator name """ - private_key: Optional[str] = None - """ Snowflake private key file path""" + private_key: Optional[str | bytes] = None + """ Snowflake private key stored as bytes or file path""" private_key_passphrase: Optional[str] = None """ Snowflake private key file passphrase""" diff --git a/sdk/python/feast/infra/registry/snowflake.py b/sdk/python/feast/infra/registry/snowflake.py index f2bc09e7e42..183c247759b 100644 --- a/sdk/python/feast/infra/registry/snowflake.py +++ b/sdk/python/feast/infra/registry/snowflake.py @@ -93,8 +93,8 @@ class SnowflakeRegistryConfig(RegistryConfig): authenticator: Optional[str] = None """ Snowflake authenticator name """ - private_key: Optional[str] = None - """ Snowflake private key file path""" + private_key: Optional[str | bytes] = None + """ Snowflake private key stored as bytes or file path""" private_key_passphrase: Optional[str] = None """ Snowflake private key file passphrase""" diff --git a/sdk/python/feast/infra/utils/snowflake/snowflake_utils.py b/sdk/python/feast/infra/utils/snowflake/snowflake_utils.py index dd965c4bed1..1baabb2ba43 100644 --- a/sdk/python/feast/infra/utils/snowflake/snowflake_utils.py +++ b/sdk/python/feast/infra/utils/snowflake/snowflake_utils.py @@ -6,7 +6,7 @@ from logging import getLogger from pathlib import Path from tempfile import TemporaryDirectory -from typing import Any, Dict, Iterator, List, Optional, Tuple, cast +from typing import Any, Dict, Iterator, List, Optional, Tuple, Union, cast import pandas as pd import pyarrow @@ -510,13 +510,25 @@ def chunk_helper(lst: pd.DataFrame, n: int) -> Iterator[Tuple[int, pd.DataFrame] 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: +def parse_private_key_path( + private_key: Union[bytes | str], private_key_passphrase: str +) -> bytes: + """Returns snowflake pkb by parsing and reading either from private key path or as byte string.""" + if isinstance(private_key, str): + with open(private_key, "rb") as key: + p_key = serialization.load_pem_private_key( + key.read(), + password=private_key_passphrase.encode(), + backend=default_backend(), + ) + elif isinstance(private_key, bytes): p_key = serialization.load_pem_private_key( - key.read(), + private_key, password=private_key_passphrase.encode(), backend=default_backend(), ) + else: + raise ValueError("private_key must be either type of str or bytes.") pkb = p_key.private_bytes( encoding=serialization.Encoding.DER, From a9eac3ec01697830afca1c2ccd3673a0862315a1 Mon Sep 17 00:00:00 2001 From: Artur Date: Thu, 8 Aug 2024 08:31:25 -0400 Subject: [PATCH 2/8] Update type hint for to use Union instead of | syntax Signed-off-by: Artur --- sdk/python/feast/infra/materialization/snowflake_engine.py | 2 +- sdk/python/feast/infra/offline_stores/snowflake.py | 2 +- sdk/python/feast/infra/online_stores/snowflake.py | 4 ++-- sdk/python/feast/infra/registry/snowflake.py | 2 +- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/sdk/python/feast/infra/materialization/snowflake_engine.py b/sdk/python/feast/infra/materialization/snowflake_engine.py index ce3d555f97a..0636dcbb51d 100644 --- a/sdk/python/feast/infra/materialization/snowflake_engine.py +++ b/sdk/python/feast/infra/materialization/snowflake_engine.py @@ -67,7 +67,7 @@ class SnowflakeMaterializationEngineConfig(FeastConfigBaseModel): authenticator: Optional[str] = None """ Snowflake authenticator name """ - private_key: Optional[str | bytes] = None + private_key: Optional[Union[str,bytes]] = None """ Snowflake private key stored as bytes or file path""" private_key_passphrase: Optional[str] = None diff --git a/sdk/python/feast/infra/offline_stores/snowflake.py b/sdk/python/feast/infra/offline_stores/snowflake.py index c214c961730..bb550cb1fad 100644 --- a/sdk/python/feast/infra/offline_stores/snowflake.py +++ b/sdk/python/feast/infra/offline_stores/snowflake.py @@ -104,7 +104,7 @@ class SnowflakeOfflineStoreConfig(FeastConfigBaseModel): authenticator: Optional[str] = None """ Snowflake authenticator name """ - private_key: Optional[str | bytes] = None + private_key: Optional[Union[str,bytes]] = None """ Snowflake private key stored as bytes or file path""" private_key_passphrase: Optional[str] = None diff --git a/sdk/python/feast/infra/online_stores/snowflake.py b/sdk/python/feast/infra/online_stores/snowflake.py index b015283929f..4bb3f88d4fc 100644 --- a/sdk/python/feast/infra/online_stores/snowflake.py +++ b/sdk/python/feast/infra/online_stores/snowflake.py @@ -2,7 +2,7 @@ import os from binascii import hexlify from datetime import datetime -from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Tuple +from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Tuple, Union import pandas as pd from pydantic import ConfigDict, Field, StrictStr @@ -50,7 +50,7 @@ class SnowflakeOnlineStoreConfig(FeastConfigBaseModel): authenticator: Optional[str] = None """ Snowflake authenticator name """ - private_key: Optional[str | bytes] = None + private_key: Optional[Union[str,bytes]] = None """ Snowflake private key stored as bytes or file path""" private_key_passphrase: Optional[str] = None diff --git a/sdk/python/feast/infra/registry/snowflake.py b/sdk/python/feast/infra/registry/snowflake.py index 183c247759b..ad5dfa72ae6 100644 --- a/sdk/python/feast/infra/registry/snowflake.py +++ b/sdk/python/feast/infra/registry/snowflake.py @@ -93,7 +93,7 @@ class SnowflakeRegistryConfig(RegistryConfig): authenticator: Optional[str] = None """ Snowflake authenticator name """ - private_key: Optional[str | bytes] = None + private_key: Optional[Union[str,bytes]] = None """ Snowflake private key stored as bytes or file path""" private_key_passphrase: Optional[str] = None From 625d2a6467ec0afcb55aed1d466dc2131258018d Mon Sep 17 00:00:00 2001 From: Artur Date: Thu, 8 Aug 2024 08:33:12 -0400 Subject: [PATCH 3/8] Update type hint for private_key to use Union instead of | syntax Signed-off-by: Artur --- sdk/python/feast/infra/materialization/snowflake_engine.py | 2 +- sdk/python/feast/infra/offline_stores/snowflake.py | 2 +- sdk/python/feast/infra/online_stores/snowflake.py | 2 +- sdk/python/feast/infra/registry/snowflake.py | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/sdk/python/feast/infra/materialization/snowflake_engine.py b/sdk/python/feast/infra/materialization/snowflake_engine.py index 0636dcbb51d..dab2e9640f2 100644 --- a/sdk/python/feast/infra/materialization/snowflake_engine.py +++ b/sdk/python/feast/infra/materialization/snowflake_engine.py @@ -67,7 +67,7 @@ class SnowflakeMaterializationEngineConfig(FeastConfigBaseModel): authenticator: Optional[str] = None """ Snowflake authenticator name """ - private_key: Optional[Union[str,bytes]] = None + private_key: Optional[Union[str, bytes]] = None """ Snowflake private key stored as bytes or file path""" private_key_passphrase: Optional[str] = None diff --git a/sdk/python/feast/infra/offline_stores/snowflake.py b/sdk/python/feast/infra/offline_stores/snowflake.py index bb550cb1fad..081c05658f4 100644 --- a/sdk/python/feast/infra/offline_stores/snowflake.py +++ b/sdk/python/feast/infra/offline_stores/snowflake.py @@ -104,7 +104,7 @@ class SnowflakeOfflineStoreConfig(FeastConfigBaseModel): authenticator: Optional[str] = None """ Snowflake authenticator name """ - private_key: Optional[Union[str,bytes]] = None + private_key: Optional[Union[str, bytes]] = None """ Snowflake private key stored as bytes or file path""" private_key_passphrase: Optional[str] = None diff --git a/sdk/python/feast/infra/online_stores/snowflake.py b/sdk/python/feast/infra/online_stores/snowflake.py index 4bb3f88d4fc..560b24b8fac 100644 --- a/sdk/python/feast/infra/online_stores/snowflake.py +++ b/sdk/python/feast/infra/online_stores/snowflake.py @@ -50,7 +50,7 @@ class SnowflakeOnlineStoreConfig(FeastConfigBaseModel): authenticator: Optional[str] = None """ Snowflake authenticator name """ - private_key: Optional[Union[str,bytes]] = None + private_key: Optional[Union[str, bytes]] = None """ Snowflake private key stored as bytes or file path""" private_key_passphrase: Optional[str] = None diff --git a/sdk/python/feast/infra/registry/snowflake.py b/sdk/python/feast/infra/registry/snowflake.py index ad5dfa72ae6..3e58a859a9a 100644 --- a/sdk/python/feast/infra/registry/snowflake.py +++ b/sdk/python/feast/infra/registry/snowflake.py @@ -93,7 +93,7 @@ class SnowflakeRegistryConfig(RegistryConfig): authenticator: Optional[str] = None """ Snowflake authenticator name """ - private_key: Optional[Union[str,bytes]] = None + private_key: Optional[Union[str, bytes]] = None """ Snowflake private key stored as bytes or file path""" private_key_passphrase: Optional[str] = None From f48638422e386f5d73b8bb24b8bdef49a451ecab Mon Sep 17 00:00:00 2001 From: Artur Date: Thu, 8 Aug 2024 08:47:25 -0400 Subject: [PATCH 4/8] Update type hint in parse_private_key_path Signed-off-by: Artur --- sdk/python/feast/infra/utils/snowflake/snowflake_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/python/feast/infra/utils/snowflake/snowflake_utils.py b/sdk/python/feast/infra/utils/snowflake/snowflake_utils.py index 1baabb2ba43..e917c797cd5 100644 --- a/sdk/python/feast/infra/utils/snowflake/snowflake_utils.py +++ b/sdk/python/feast/infra/utils/snowflake/snowflake_utils.py @@ -511,7 +511,7 @@ def chunk_helper(lst: pd.DataFrame, n: int) -> Iterator[Tuple[int, pd.DataFrame] def parse_private_key_path( - private_key: Union[bytes | str], private_key_passphrase: str + private_key: Union[str, bytes], private_key_passphrase: str ) -> bytes: """Returns snowflake pkb by parsing and reading either from private key path or as byte string.""" if isinstance(private_key, str): From 636089fe720b0ef4b15e59c9ab59f36d432672dc Mon Sep 17 00:00:00 2001 From: Artur Date: Mon, 12 Aug 2024 15:58:18 -0400 Subject: [PATCH 5/8] added private_key_content in Snowflake configs to support key-pair auth by reading in byte string Signed-off-by: Artur --- .../infra/materialization/snowflake_engine.py | 7 +++-- .../feast/infra/offline_stores/snowflake.py | 7 +++-- .../feast/infra/online_stores/snowflake.py | 7 +++-- sdk/python/feast/infra/registry/snowflake.py | 7 +++-- .../infra/utils/snowflake/snowflake_utils.py | 26 +++++++++---------- 5 files changed, 33 insertions(+), 21 deletions(-) diff --git a/sdk/python/feast/infra/materialization/snowflake_engine.py b/sdk/python/feast/infra/materialization/snowflake_engine.py index dab2e9640f2..5d0f08c2f59 100644 --- a/sdk/python/feast/infra/materialization/snowflake_engine.py +++ b/sdk/python/feast/infra/materialization/snowflake_engine.py @@ -67,8 +67,11 @@ class SnowflakeMaterializationEngineConfig(FeastConfigBaseModel): authenticator: Optional[str] = None """ Snowflake authenticator name """ - private_key: Optional[Union[str, bytes]] = None - """ Snowflake private key stored as bytes or file path""" + private_key: Optional[str] = None + """ Snowflake private key file path""" + + private_key_content: Optional[bytes] = None + """ Snowflake private key stored as bytes""" private_key_passphrase: Optional[str] = None """ Snowflake private key file passphrase""" diff --git a/sdk/python/feast/infra/offline_stores/snowflake.py b/sdk/python/feast/infra/offline_stores/snowflake.py index 081c05658f4..ada6c99c985 100644 --- a/sdk/python/feast/infra/offline_stores/snowflake.py +++ b/sdk/python/feast/infra/offline_stores/snowflake.py @@ -104,8 +104,11 @@ class SnowflakeOfflineStoreConfig(FeastConfigBaseModel): authenticator: Optional[str] = None """ Snowflake authenticator name """ - private_key: Optional[Union[str, bytes]] = None - """ Snowflake private key stored as bytes or file path""" + private_key: Optional[str] = None + """ Snowflake private key file path""" + + private_key_content: Optional[bytes] = None + """ Snowflake private key stored as bytes""" private_key_passphrase: Optional[str] = None """ Snowflake private key file passphrase""" diff --git a/sdk/python/feast/infra/online_stores/snowflake.py b/sdk/python/feast/infra/online_stores/snowflake.py index 560b24b8fac..926563d0835 100644 --- a/sdk/python/feast/infra/online_stores/snowflake.py +++ b/sdk/python/feast/infra/online_stores/snowflake.py @@ -50,8 +50,11 @@ class SnowflakeOnlineStoreConfig(FeastConfigBaseModel): authenticator: Optional[str] = None """ Snowflake authenticator name """ - private_key: Optional[Union[str, bytes]] = None - """ Snowflake private key stored as bytes or file path""" + private_key: Optional[str] = None + """ Snowflake private key file path""" + + private_key_content: Optional[bytes] = None + """ Snowflake private key stored as bytes""" private_key_passphrase: Optional[str] = None """ Snowflake private key file passphrase""" diff --git a/sdk/python/feast/infra/registry/snowflake.py b/sdk/python/feast/infra/registry/snowflake.py index 3e58a859a9a..ac4f52dc06e 100644 --- a/sdk/python/feast/infra/registry/snowflake.py +++ b/sdk/python/feast/infra/registry/snowflake.py @@ -93,8 +93,11 @@ class SnowflakeRegistryConfig(RegistryConfig): authenticator: Optional[str] = None """ Snowflake authenticator name """ - private_key: Optional[Union[str, bytes]] = None - """ Snowflake private key stored as bytes or file path""" + private_key: Optional[str] = None + """ Snowflake private key file path""" + + private_key_content: Optional[bytes] = None + """ Snowflake private key stored as bytes""" private_key_passphrase: Optional[str] = None """ Snowflake private key file passphrase""" diff --git a/sdk/python/feast/infra/utils/snowflake/snowflake_utils.py b/sdk/python/feast/infra/utils/snowflake/snowflake_utils.py index e917c797cd5..06db4fa7326 100644 --- a/sdk/python/feast/infra/utils/snowflake/snowflake_utils.py +++ b/sdk/python/feast/infra/utils/snowflake/snowflake_utils.py @@ -84,9 +84,9 @@ def __enter__(self): # 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: + if "private_key" in kwargs or "private_key_content" in kwargs: kwargs["private_key"] = parse_private_key_path( - kwargs["private_key"], kwargs["private_key_passphrase"] + kwargs.get("private_key_passphrase"), kwargs.get("private_key"), kwargs.get("private_key_content") ) try: @@ -511,24 +511,24 @@ def chunk_helper(lst: pd.DataFrame, n: int) -> Iterator[Tuple[int, pd.DataFrame] def parse_private_key_path( - private_key: Union[str, bytes], private_key_passphrase: str + private_key_passphrase: str, key_path: str = None, private_key_content: bytes = None ) -> bytes: - """Returns snowflake pkb by parsing and reading either from private key path or as byte string.""" - if isinstance(private_key, str): - with open(private_key, "rb") as key: + """Returns snowflake pkb by parsing and reading either from key path or private_key_content as byte string.""" + if private_key_content: + p_key = serialization.load_pem_private_key( + private_key_content, + password=private_key_passphrase.encode(), + backend=default_backend(), + ) + elif key_path: + with open(key_path, "rb") as key: p_key = serialization.load_pem_private_key( key.read(), password=private_key_passphrase.encode(), backend=default_backend(), ) - elif isinstance(private_key, bytes): - p_key = serialization.load_pem_private_key( - private_key, - password=private_key_passphrase.encode(), - backend=default_backend(), - ) else: - raise ValueError("private_key must be either type of str or bytes.") + raise ValueError("Please provide key_path or private_key_content.") pkb = p_key.private_bytes( encoding=serialization.Encoding.DER, From fb175dfdc57c63dec9b8885a61125bd54aafa969 Mon Sep 17 00:00:00 2001 From: Artur Date: Mon, 12 Aug 2024 16:18:28 -0400 Subject: [PATCH 6/8] fix incompatible linting types Signed-off-by: Artur --- sdk/python/feast/infra/utils/snowflake/snowflake_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/python/feast/infra/utils/snowflake/snowflake_utils.py b/sdk/python/feast/infra/utils/snowflake/snowflake_utils.py index 06db4fa7326..72cba383036 100644 --- a/sdk/python/feast/infra/utils/snowflake/snowflake_utils.py +++ b/sdk/python/feast/infra/utils/snowflake/snowflake_utils.py @@ -511,7 +511,7 @@ def chunk_helper(lst: pd.DataFrame, n: int) -> Iterator[Tuple[int, pd.DataFrame] def parse_private_key_path( - private_key_passphrase: str, key_path: str = None, private_key_content: bytes = None + private_key_passphrase: str, key_path: Optional[str] = None, private_key_content: Optional[bytes] = None ) -> bytes: """Returns snowflake pkb by parsing and reading either from key path or private_key_content as byte string.""" if private_key_content: From 855347251bfa7a6c73a8285119b5482cd831b744 Mon Sep 17 00:00:00 2001 From: Artur Date: Mon, 12 Aug 2024 17:12:23 -0400 Subject: [PATCH 7/8] remove unused Union import Signed-off-by: Artur --- sdk/python/feast/infra/online_stores/snowflake.py | 2 +- sdk/python/feast/infra/utils/snowflake/snowflake_utils.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/sdk/python/feast/infra/online_stores/snowflake.py b/sdk/python/feast/infra/online_stores/snowflake.py index 926563d0835..6f39bdd0f65 100644 --- a/sdk/python/feast/infra/online_stores/snowflake.py +++ b/sdk/python/feast/infra/online_stores/snowflake.py @@ -2,7 +2,7 @@ import os from binascii import hexlify from datetime import datetime -from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Tuple, Union +from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Tuple import pandas as pd from pydantic import ConfigDict, Field, StrictStr diff --git a/sdk/python/feast/infra/utils/snowflake/snowflake_utils.py b/sdk/python/feast/infra/utils/snowflake/snowflake_utils.py index 72cba383036..1c305acc30f 100644 --- a/sdk/python/feast/infra/utils/snowflake/snowflake_utils.py +++ b/sdk/python/feast/infra/utils/snowflake/snowflake_utils.py @@ -6,7 +6,7 @@ from logging import getLogger from pathlib import Path from tempfile import TemporaryDirectory -from typing import Any, Dict, Iterator, List, Optional, Tuple, Union, cast +from typing import Any, Dict, Iterator, List, Optional, Tuple, cast import pandas as pd import pyarrow From af277a45ea41aeba91b4fdeafb2d2bf37bc59948 Mon Sep 17 00:00:00 2001 From: Artur Date: Mon, 12 Aug 2024 21:18:48 -0400 Subject: [PATCH 8/8] fix formating Signed-off-by: Artur --- sdk/python/feast/infra/utils/snowflake/snowflake_utils.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/sdk/python/feast/infra/utils/snowflake/snowflake_utils.py b/sdk/python/feast/infra/utils/snowflake/snowflake_utils.py index 1c305acc30f..b9035b40dbf 100644 --- a/sdk/python/feast/infra/utils/snowflake/snowflake_utils.py +++ b/sdk/python/feast/infra/utils/snowflake/snowflake_utils.py @@ -86,7 +86,9 @@ def __enter__(self): # https://docs.snowflake.com/en/user-guide/key-pair-auth.html#configuring-key-pair-authentication if "private_key" in kwargs or "private_key_content" in kwargs: kwargs["private_key"] = parse_private_key_path( - kwargs.get("private_key_passphrase"), kwargs.get("private_key"), kwargs.get("private_key_content") + kwargs.get("private_key_passphrase"), + kwargs.get("private_key"), + kwargs.get("private_key_content"), ) try: @@ -511,7 +513,9 @@ def chunk_helper(lst: pd.DataFrame, n: int) -> Iterator[Tuple[int, pd.DataFrame] def parse_private_key_path( - private_key_passphrase: str, key_path: Optional[str] = None, private_key_content: Optional[bytes] = None + private_key_passphrase: str, + key_path: Optional[str] = None, + private_key_content: Optional[bytes] = None, ) -> bytes: """Returns snowflake pkb by parsing and reading either from key path or private_key_content as byte string.""" if private_key_content: