[misc] Adds optional port envar for local tests (#378)

snowflakedb · Jan 28, 2020 · 70c78bf · 70c78bf
1 parent 395c1fe
commit 70c78bf
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 4 deletions.
diff --git a/src/runners/helpers/db.py b/src/runners/helpers/db.py
@@ -24,6 +24,8 @@
     DATABASE,
     USER,
     WAREHOUSE,
+    PORT,
+    PROTOCOL,
     PRIVATE_KEY,
     PRIVATE_KEY_PASSWORD,
     TIMEOUT,
@@ -97,6 +99,8 @@ def connect():
             # 2) If using OAuth and OAUTH_CONNECTION_* env vars have been set, use these
             # 3) Else, use OAuth'd user's default namespace (see ALTER USER docs)
             account=oauth_account or ACCOUNT,
+            port=PORT,
+            protocol=PROTOCOL,
             database=environ.get('OAUTH_CONNECTION_DATABASE', None)
             if oauth_account
             else DATABASE,

diff --git a/src/runners/helpers/dbconfig.py b/src/runners/helpers/dbconfig.py
@@ -31,6 +31,8 @@ def format_p8_from_key(key, encrypted=False):
 REGION = environ.get('SA_REGION', environ.get('REGION', "us-west-2"))
 REGION_SUBDOMAIN_POSTFIX = '' if REGION == 'us-west-2' else f'.{REGION}'
 ACCOUNT = environ.get('SNOWFLAKE_ACCOUNT', '') + REGION_SUBDOMAIN_POSTFIX
+PORT = environ.get('SNOWFLAKE_PORT', '443')
+PROTOCOL = environ.get('SNOWFLAKE_PROTOCOL', 'https')
 
 USER = environ.get('SA_USER', "snowalert") + tail
 PRIVATE_KEY_PASSWORD = environ.get('PRIVATE_KEY_PASSWORD', '').encode('utf-8')

diff --git a/src/scripts/install.py b/src/scripts/install.py
@@ -32,7 +32,7 @@
 from runners.config import DATABASE, DATA_SCHEMA, RULES_SCHEMA, RESULTS_SCHEMA
 
 from runners.helpers import log
-from runners.helpers.dbconfig import USER, ROLE, WAREHOUSE
+from runners.helpers.dbconfig import USER, ROLE, WAREHOUSE, PORT, PROTOCOL
 from runners.helpers.dbconnect import snowflake_connect
 
 
@@ -218,7 +218,13 @@ def login(configuration=None):
     else:
         print(f"Loaded password: {'*' * len(password)}")
 
-    connect_kwargs = {'user': username, 'account': account}
+    connect_kwargs = {
+        'user': username,
+        'account': account,
+        'port': configuration.get('port', PORT),
+        'protocol': configuration.get('protocol', PROTOCOL),
+    }
+
     if password == '':
         connect_kwargs['authenticator'] = 'externalbrowser'
     else:
@@ -493,6 +499,8 @@ def main(
     jira=None,
     config_account=None,
     config_region=None,
+    config_port=None,
+    config_protocol=None,
     config_username=None,
     config_password=None,
     connection_name=None,
@@ -510,7 +518,11 @@ def main(
         connection_name
         if connection_name is not None
         else {
-            'region': config_region or environ.get('REGION'),
+            'region': (
+                config_region or environ.get('SNOWFLAKE_REGION', environ.get('REGION'))
+            ),
+            'port': config_port or environ.get('SNOWFLAKE_PORT'),
+            'protocol': config_protocol or environ.get('SNOWFLAKE_PROTOCOL'),
             'accountname': config_account or environ.get('SNOWFLAKE_ACCOUNT'),
             'username': config_username or environ.get('SA_ADMIN_USER'),
             'password': config_password or environ.get('SA_ADMIN_PASSWORD'),
@@ -571,7 +583,10 @@ def main(
             f"ALTER USER {USER} SET rsa_public_key='{rsa_public_key}'",
         )
     else:
-        print("Please ask account admin to run:\n\n", f"  ALTER USER {USER} SET rsa_public_key='{rsa_public_key}'")
+        print(
+            "Please ask account admin to run:\n\n",
+            f"  ALTER USER {USER} SET rsa_public_key='{rsa_public_key}'",
+        )
 
     aws_key, aws_secret = load_aws_config()