Skip to content

Commit 5215a21

Browse files
arturkolakowskiarturkolakowski
authored
feat: Allow feast snowflake to read in byte string for private-key authentication (feast-dev#4384)
* allow feast snowflake to read in byte string for private-key authentication Signed-off-by: Artur <artur.kolakowski@medely.com> * Update type hint for to use Union instead of | syntax Signed-off-by: Artur <artur.kolakowski@medely.com> * Update type hint for private_key to use Union instead of | syntax Signed-off-by: Artur <artur.kolakowski@medely.com> * Update type hint in parse_private_key_path Signed-off-by: Artur <artur.kolakowski@medely.com> * added private_key_content in Snowflake configs to support key-pair auth by reading in byte string Signed-off-by: Artur <artur.kolakowski@medely.com> * fix incompatible linting types Signed-off-by: Artur <artur.kolakowski@medely.com> * remove unused Union import Signed-off-by: Artur <artur.kolakowski@medely.com> * fix formating Signed-off-by: Artur <artur.kolakowski@medely.com> --------- Signed-off-by: Artur <artur.kolakowski@medely.com> Co-authored-by: Artur <artur.kolakowski@medely.com>
1 parent 419ca5e commit 5215a21

5 files changed

Lines changed: 33 additions & 5 deletions

File tree

‎sdk/python/feast/infra/materialization/snowflake_engine.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,9 @@ class SnowflakeMaterializationEngineConfig(FeastConfigBaseModel):
7070
private_key: Optional[str] = None
7171
""" Snowflake private key file path"""
7272

73+
private_key_content: Optional[bytes] = None
74+
""" Snowflake private key stored as bytes"""
75+
7376
private_key_passphrase: Optional[str] = None
7477
""" Snowflake private key file passphrase"""
7578

‎sdk/python/feast/infra/offline_stores/snowflake.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,9 @@ class SnowflakeOfflineStoreConfig(FeastConfigBaseModel):
107107
private_key: Optional[str] = None
108108
""" Snowflake private key file path"""
109109

110+
private_key_content: Optional[bytes] = None
111+
""" Snowflake private key stored as bytes"""
112+
110113
private_key_passphrase: Optional[str] = None
111114
""" Snowflake private key file passphrase"""
112115

‎sdk/python/feast/infra/online_stores/snowflake.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,9 @@ class SnowflakeOnlineStoreConfig(FeastConfigBaseModel):
5353
private_key: Optional[str] = None
5454
""" Snowflake private key file path"""
5555

56+
private_key_content: Optional[bytes] = None
57+
""" Snowflake private key stored as bytes"""
58+
5659
private_key_passphrase: Optional[str] = None
5760
""" Snowflake private key file passphrase"""
5861

‎sdk/python/feast/infra/registry/snowflake.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,9 @@ class SnowflakeRegistryConfig(RegistryConfig):
9696
private_key: Optional[str] = None
9797
""" Snowflake private key file path"""
9898

99+
private_key_content: Optional[bytes] = None
100+
""" Snowflake private key stored as bytes"""
101+
99102
private_key_passphrase: Optional[str] = None
100103
""" Snowflake private key file passphrase"""
101104

‎sdk/python/feast/infra/utils/snowflake/snowflake_utils.py‎

Lines changed: 21 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -84,9 +84,11 @@ def __enter__(self):
8484

8585
# https://docs.snowflake.com/en/user-guide/python-connector-example.html#using-key-pair-authentication-key-pair-rotation
8686
# https://docs.snowflake.com/en/user-guide/key-pair-auth.html#configuring-key-pair-authentication
87-
if "private_key" in kwargs:
87+
if "private_key" in kwargs or "private_key_content" in kwargs:
8888
kwargs["private_key"] = parse_private_key_path(
89-
kwargs["private_key"], kwargs["private_key_passphrase"]
89+
kwargs.get("private_key_passphrase"),
90+
kwargs.get("private_key"),
91+
kwargs.get("private_key_content"),
9092
)
9193

9294
try:
@@ -510,13 +512,27 @@ def chunk_helper(lst: pd.DataFrame, n: int) -> Iterator[Tuple[int, pd.DataFrame]
510512
yield int(i / n), lst[i : i + n]
511513

512514

513-
def parse_private_key_path(key_path: str, private_key_passphrase: str) -> bytes:
514-
with open(key_path, "rb") as key:
515+
def parse_private_key_path(
516+
private_key_passphrase: str,
517+
key_path: Optional[str] = None,
518+
private_key_content: Optional[bytes] = None,
519+
) -> bytes:
520+
"""Returns snowflake pkb by parsing and reading either from key path or private_key_content as byte string."""
521+
if private_key_content:
515522
p_key = serialization.load_pem_private_key(
516-
key.read(),
523+
private_key_content,
517524
password=private_key_passphrase.encode(),
518525
backend=default_backend(),
519526
)
527+
elif key_path:
528+
with open(key_path, "rb") as key:
529+
p_key = serialization.load_pem_private_key(
530+
key.read(),
531+
password=private_key_passphrase.encode(),
532+
backend=default_backend(),
533+
)
534+
else:
535+
raise ValueError("Please provide key_path or private_key_content.")
520536

521537
pkb = p_key.private_bytes(
522538
encoding=serialization.Encoding.DER,

0 commit comments

Comments
 (0)