Skip to content

Commit c42c9b0

Browse files
fix: Added private_key auth for Snowflake (#2508)
* added private_key auth Signed-off-by: Miles Adkins <miles.adkins@snowflake.com> * added private_key auth comment to docs Signed-off-by: Miles Adkins <miles.adkins@snowflake.com> * added private_key auth comment to docs Signed-off-by: Miles Adkins <miles.adkins@snowflake.com> * type kwargs to be str -> any Signed-off-by: Achal Shah <achals@gmail.com> * type kwargs to be str -> any and some other fixes Signed-off-by: Achal Shah <achals@gmail.com> Co-authored-by: Achal Shah <achals@gmail.com>
1 parent 83f3e0d commit c42c9b0

1 file changed

Lines changed: 31 additions & 5 deletions

File tree

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

Lines changed: 31 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,11 @@
44
import string
55
from logging import getLogger
66
from tempfile import TemporaryDirectory
7-
from typing import Dict, Iterator, List, Optional, Tuple, cast
7+
from typing import Any, Dict, Iterator, List, Optional, Tuple, cast
88

99
import pandas as pd
10+
from cryptography.hazmat.backends import default_backend
11+
from cryptography.hazmat.primitives import serialization
1012
from tenacity import (
1113
retry,
1214
retry_if_exception_type,
@@ -40,18 +42,17 @@ def execute_snowflake_statement(conn: SnowflakeConnection, query) -> SnowflakeCu
4042

4143

4244
def get_snowflake_conn(config, autocommit=True) -> SnowflakeConnection:
43-
if config.type == "snowflake.offline":
44-
config_header = "connections.feast_offline_store"
45+
assert config.type == "snowflake.offline"
46+
config_header = "connections.feast_offline_store"
4547

4648
config_dict = dict(config)
4749

4850
# read config file
4951
config_reader = configparser.ConfigParser()
5052
config_reader.read([config_dict["config_path"]])
53+
kwargs: Dict[str, Any] = {}
5154
if config_reader.has_section(config_header):
5255
kwargs = dict(config_reader[config_header])
53-
else:
54-
kwargs = {}
5556

5657
if "schema" in kwargs:
5758
kwargs["schema_"] = kwargs.pop("schema")
@@ -67,6 +68,13 @@ def get_snowflake_conn(config, autocommit=True) -> SnowflakeConnection:
6768
else:
6869
kwargs["schema"] = '"PUBLIC"'
6970

71+
# https://docs.snowflake.com/en/user-guide/python-connector-example.html#using-key-pair-authentication-key-pair-rotation
72+
# https://docs.snowflake.com/en/user-guide/key-pair-auth.html#configuring-key-pair-authentication
73+
if "private_key" in kwargs:
74+
kwargs["private_key"] = parse_private_key_path(
75+
kwargs["private_key"], kwargs["private_key_passphrase"]
76+
)
77+
7078
try:
7179
conn = snowflake.connector.connect(
7280
application="feast", autocommit=autocommit, **kwargs
@@ -288,3 +296,21 @@ def chunk_helper(lst: pd.DataFrame, n: int) -> Iterator[Tuple[int, pd.DataFrame]
288296
"""Helper generator to chunk a sequence efficiently with current index like if enumerate was called on sequence."""
289297
for i in range(0, len(lst), n):
290298
yield int(i / n), lst[i : i + n]
299+
300+
301+
def parse_private_key_path(key_path: str, private_key_passphrase: str) -> bytes:
302+
303+
with open(key_path, "rb") as key:
304+
p_key = serialization.load_pem_private_key(
305+
key.read(),
306+
password=private_key_passphrase.encode(),
307+
backend=default_backend(),
308+
)
309+
310+
pkb = p_key.private_bytes(
311+
encoding=serialization.Encoding.DER,
312+
format=serialization.PrivateFormat.PKCS8,
313+
encryption_algorithm=serialization.NoEncryption(),
314+
)
315+
316+
return pkb

0 commit comments

Comments
 (0)