jprdonnelly/redash
1
0
Fork 0
mirror of https://github.com/getredash/redash.git synced 2025-12-25 10:00:45 -05:00
Code Issues Projects Releases Wiki Activity

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

Browse Source 
* 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 <yoshiokatsuneo@gmail.com>
This commit is contained in:
Adrian Oesch
2025-08-08 19:38:22 +02:00
committed by GitHub
parent 35afe880a1
commit 00a97d9266
1 changed files with 33 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
redash/query_runner/snowflake.py
View File

@@ -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 @@ class Snowflake(BaseSQLQueryRunner):
"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 @@ class Snowflake(BaseSQLQueryRunner):
"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,7 +95,7 @@ class Snowflake(BaseSQLQueryRunner):
if region == "us-west":
region = None
if self.configuration.__contains__("host"):
if self.configuration.get("host"):
host = self.configuration.get("host")
else:
if region:
@@ -96,14 +103,29 @@ class Snowflake(BaseSQLQueryRunner):
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