fix small bugs in get private key

Signed-off-by: HH <hhcs9527@gmail.com>
flyteorg · Sep 27, 2023 · 62f9170 · 62f9170
1 parent 8346fb9
commit 62f9170
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/flytekit/types/structured/snowflake.py b/flytekit/types/structured/snowflake.py
@@ -24,9 +24,16 @@ def get_private_key():
 
     import flytekit
 
-    pk_string = flytekit.current_context().secrets.get(SNOWFLAKE, "private_key")
+    pk_string = flytekit.current_context().secrets.get(TASK_TYPE, "private_key", encode_mode="rb")
+    p_key = serialization.load_pem_private_key(pk_string, password=None, backend=default_backend())
 
-    return serialization.load_pem_private_key(pk_string, password=None, backend=default_backend())
+    pkb = p_key.private_bytes(
+        encoding=serialization.Encoding.DER,
+        format=serialization.PrivateFormat.PKCS8,
+        encryption_algorithm=serialization.NoEncryption(),
+    )
+
+    return pkb
 
 
 def _write_to_sf(structured_dataset: StructuredDataset):