Add private_key auth method to snowflake query runner (#7371)

* add private_key auth method

* fix casing

* fix private_key parsing

* use params and add optional pwd

* use private_key_b64

* add file option

* remove __contains__

* fix pem pwd

* fix lint issues

* fix black

---------

Co-authored-by: Tsuneo Yoshioka <[email protected]>

Author: adrianoesch

redash/query_runner/snowflake.py

@@ -1,11 +1,14 @@
 try:
     import snowflake.connector
+    from cryptography.hazmat.primitives.serialization import load_pem_private_key
 
     enabled = True
 except ImportError:
     enabled = False
 
 
+from base64 import b64decode
+
 from redash import __version__
 from redash.query_runner import (
     TYPE_BOOLEAN,
@@ -43,6 +46,8 @@ def configuration_schema(cls):
                 "account": {"type": "string"},
                 "user": {"type": "string"},
                 "password": {"type": "string"},
+                "private_key_File": {"type": "string"},
+                "private_key_pwd": {"type": "string"},
                 "warehouse": {"type": "string"},
                 "database": {"type": "string"},
                 "region": {"type": "string", "default": "us-west"},
@@ -57,13 +62,15 @@ def configuration_schema(cls):
                 "account",
                 "user",
                 "password",
+                "private_key_File",
+                "private_key_pwd",
                 "warehouse",
                 "database",
                 "region",
                 "host",
             ],
-            "required": ["user", "password", "account", "database", "warehouse"],
-            "secret": ["password"],
+            "required": ["user", "account", "database", "warehouse"],
+            "secret": ["password", "private_key_File", "private_key_pwd"],
             "extra_options": [
                 "host",
             ],
@@ -88,22 +95,37 @@ def _get_connection(self):
         if region == "us-west":
             region = None
 
-        if self.configuration.__contains__("host"):
+        if self.configuration.get("host"):
             host = self.configuration.get("host")
         else:
             if region:
                 host = "{}.{}.snowflakecomputing.com".format(account, region)
             else:
                 host = "{}.snowflakecomputing.com".format(account)
 
-        connection = snowflake.connector.connect(
-            user=self.configuration["user"],
-            password=self.configuration["password"],
-            account=account,
-            region=region,
-            host=host,
-            application="Redash/{} (Snowflake)".format(__version__.split("-")[0]),
-        )
+        params = {
+            "user": self.configuration["user"],
+            "account": account,
+            "region": region,
+            "host": host,
+            "application": "Redash/{} (Snowflake)".format(__version__.split("-")[0]),
+        }
+
+        if self.configuration.get("password"):
+            params["password"] = self.configuration["password"]
+        elif self.configuration.get("private_key_File"):
+            private_key_b64 = self.configuration.get("private_key_File")
+            private_key_bytes = b64decode(private_key_b64)
+            if self.configuration.get("private_key_pwd"):
+                private_key_pwd = self.configuration.get("private_key_pwd").encode()
+            else:
+                private_key_pwd = None
+            private_key_pem = load_pem_private_key(private_key_bytes, private_key_pwd)
+            params["private_key"] = private_key_pem
+        else:
+            raise Exception("Neither password nor private_key_b64 is set.")
+
+        connection = snowflake.connector.connect(**params)
 
         return connection