change key path to SNOWFLAKE_PRIVATE_KEY (#1879)

Signed-off-by: HH <[email protected]> Signed-off-by: Kevin Su <[email protected]> Co-authored-by: Kevin Su <[email protected]>
flyteorg · Oct 9, 2023 · b35af10 · b35af10
1 parent 1f0d120
commit b35af10
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/plugins/flytekit-snowflake/flytekitplugins/snowflake/agent.py b/plugins/flytekit-snowflake/flytekitplugins/snowflake/agent.py
@@ -23,6 +23,7 @@
 from flytekit.models.types import LiteralType, StructuredDatasetType
 
 TASK_TYPE = "snowflake"
+SNOWFLAKE_PRIVATE_KEY = "snowflake_private_key"
 
 
 @dataclass
@@ -46,7 +47,7 @@ def get_private_key(self):
 
         import flytekit
 
-        pk_string = flytekit.current_context().secrets.get(TASK_TYPE, "private_key", encode_mode="rb")
+        pk_string = flytekit.current_context().secrets.get(SNOWFLAKE_PRIVATE_KEY, encode_mode="rb")
         p_key = serialization.load_pem_private_key(pk_string, password=None, backend=default_backend())
 
         pkb = p_key.private_bytes(