From b24dc8e157de70e5a15a6eda80952ccb3262a84d Mon Sep 17 00:00:00 2001
From: Ofek Weiss <ofek1weiss@gmail.com>
Date: Sun, 4 Feb 2024 18:03:16 +0200
Subject: [PATCH] added snowflake keyfile support

---
 README.md     |  6 ++++++
 action.yml    |  3 +++
 entrypoint.py | 18 ++++++++++++++----
 3 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index dd5a204..7ae0e45 100644
--- a/README.md
+++ b/README.md
@@ -149,6 +149,12 @@ If you're using BigQuery with a key file,
 supply the `bigquery-keyfile` argument to the action and make sure your `keyfile` in the `profiles-yml`
 is `/tmp/bigquery_keyfile.json`.
 
+### Snowflake Keyfile Authentication
+
+If you're using Snowflake with a key file,
+supply the `snowflake-keyfile` argument to the action and make sure your `private_key_path` in the `profiles-yml`
+is `tmp/snowflake_keyfile.key`.
+
 ### Google Cloud Storage Keyfile
 
 If you want to upload your report to a Google Cloud Storage bucket using `send-report`,
diff --git a/action.yml b/action.yml
index 90d43d8..feb4280 100644
--- a/action.yml
+++ b/action.yml
@@ -40,6 +40,9 @@ inputs:
   bigquery-keyfile:
     description: The content of your BigQuery keyfile.
     required: false
+  snowflake-keyfile:
+    description: The content of your Snowflake keyfile.
+    required: false
   gcs-keyfile:
     description: The content of your Google service account keyfile, for usage with Google Cloud Storage.
     required: false
diff --git a/entrypoint.py b/entrypoint.py
index 6371343..38d4552 100644
--- a/entrypoint.py
+++ b/entrypoint.py
@@ -27,8 +27,9 @@ def setup_env(
     profiles_yml: Optional[str],
     bigquery_keyfile: Optional[str],
     gcs_keyfile: Optional[str],
+    snowflake_keyfile: Optional[str],
 ):
-    logging.info(f"Setting up the environment.")
+    logging.info("Setting up the environment.")
     dbt_dir = Path.home() / ".dbt"
     dbt_dir.mkdir(parents=True, exist_ok=True)
     if profiles_yml:
@@ -37,6 +38,8 @@ def setup_env(
         Path("/tmp/bigquery_keyfile.json").write_text(bigquery_keyfile)
     if gcs_keyfile:
         Path("/tmp/gcs_keyfile.json").write_text(gcs_keyfile)
+    if snowflake_keyfile:
+        Path("tmp/snowflake_keyfile.key").write_text(snowflake_keyfile)
 
 
 def install_edr(
@@ -108,7 +111,7 @@ def install_edr(
 
 
 def run_edr(edr_command: str, project_dir: Optional[str]):
-    logging.info(f"Running the edr command.")
+    logging.info("Running the edr command.")
     subprocess.run(edr_command, shell=True, check=True, cwd=project_dir)
 
 
@@ -119,6 +122,7 @@ class Args(BaseModel):
     profiles_yml: Optional[str]
     edr_command: str
     bigquery_keyfile: Optional[str]
+    snowflake_keyfile: Optional[str]
     gcs_keyfile: Optional[str]
     profile_target: Optional[str]
 
@@ -131,16 +135,22 @@ def main():
         profile_target=os.getenv("INPUT_PROFILE-TARGET") or None,
         project_dir=os.getenv("INPUT_PROJECT-DIR") or None,
         bigquery_keyfile=os.getenv("INPUT_BIGQUERY-KEYFILE") or None,
+        snowflake_keyfile=os.getenv("INPUT_SNOWFLAKE-KEYFILE") or None,
         gcs_keyfile=os.getenv("INPUT_GCS-KEYFILE") or None,
         adapter_version=os.getenv("INPUT_ADAPTER-VERSION") or None,
     )
     install_dbt(args.adapter, args.adapter_version)
-    setup_env(args.profiles_yml, args.bigquery_keyfile, args.gcs_keyfile)
+    setup_env(
+        args.profiles_yml,
+        args.bigquery_keyfile,
+        args.gcs_keyfile,
+        args.snowflake_keyfile,
+    )
     install_edr(args.adapter, args.project_dir, args.profile_target)
     try:
         run_edr(args.edr_command, args.project_dir)
     except subprocess.CalledProcessError:
-        logging.exception(f"Failed to run the edr command.")
+        logging.exception("Failed to run the edr command.")
         raise