From 378b192e93929bdea6a8b8d6050642d03598fe9d Mon Sep 17 00:00:00 2001
From: Vitaly Markov <wild.desu@gmail.com>
Date: Wed, 17 Jul 2024 10:20:45 +0100
Subject: [PATCH] Allow passing raw private key with SNOWFLAKE_PRIVATE_KEY
 environment variable

---
 CHANGELOG.md        |  6 +++++-
 snowddl/app/base.py | 19 ++++++++++++++-----
 snowddl/version.py  |  2 +-
 3 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2a1f60c..48c3967 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,8 +1,12 @@
 # Changelog
 
+## [0.30.2] - 2024-07-17
+
+- Allowed passing raw private key with `SNOWFLAKE_PRIVATE_KEY` environment variable for convenience of GitHub actions. This is an addition to original `SNOWFLAKE_PRIVATE_KEY_PATH`, but does not require creation of file.
+
 ## [0.30.1] - 2024-07-15
 
-- Add missing `__init__.py` to `fernet` package. Make sure this package is included by `find:` during build process.
+- Added missing `__init__.py` to `fernet` package. Make sure this package is included by `find:` during build process.
 
 ## [0.30.0] - 2024-07-14
 
diff --git a/snowddl/app/base.py b/snowddl/app/base.py
index 8162c34..68472d9 100644
--- a/snowddl/app/base.py
+++ b/snowddl/app/base.py
@@ -1,4 +1,5 @@
 from argparse import ArgumentParser, HelpFormatter
+from cryptography.hazmat.primitives import serialization
 from importlib.util import module_from_spec, spec_from_file_location
 from json import loads as json_loads
 from json.decoder import JSONDecodeError
@@ -255,7 +256,7 @@ def init_arguments(self):
 
     def validate_auth_args(self, args):
         if args["authenticator"] == "snowflake":
-            if not args["a"] or not args["u"] or (not args["p"] and not args["k"]):
+            if not args["a"] or not args["u"] or (not args["p"] and not args["k"] and "SNOWFLAKE_PRIVATE_KEY" not in environ):
                 return False
         elif args["authenticator"] == "externalbrowser":
             if not args["a"] or not args["u"]:
@@ -445,13 +446,21 @@ def get_connection(self):
         }
 
         if self.args.get("authenticator") == "snowflake":
-            if self.args.get("k"):
-                from cryptography.hazmat.primitives import serialization
+            key_bytes = None
 
+            if self.args.get("k"):
                 key_path = Path(self.args.get("k"))
-                key_password = str(self.args.get("passphrase")).encode("utf-8") if self.args.get("passphrase") else None
 
-                pk = serialization.load_pem_private_key(data=key_path.read_bytes(), password=key_password)
+                if not key_path.is_file():
+                    raise ValueError(f"Private key file [{key_path}] does not exist or not a file")
+
+                key_bytes = key_path.read_bytes()
+            elif "SNOWFLAKE_PRIVATE_KEY" in environ:
+                key_bytes = str(environ["SNOWFLAKE_PRIVATE_KEY"]).encode("utf-8")
+
+            if key_bytes:
+                key_password = str(self.args.get("passphrase")).encode("utf-8") if self.args.get("passphrase") else None
+                pk = serialization.load_pem_private_key(data=key_bytes, password=key_password)
 
                 options["private_key"] = pk.private_bytes(
                     encoding=serialization.Encoding.DER,
diff --git a/snowddl/version.py b/snowddl/version.py
index 8925285..e89b4cf 100644
--- a/snowddl/version.py
+++ b/snowddl/version.py
@@ -1 +1 @@
-__version__ = "0.30.1"
+__version__ = "0.30.2"