Merge pull request #1224 from convoyinc/adriank/add_snowflake_sso_sup…

…port

Add support for Snowflake SSO authentication round 2

plugins/snowflake/dbt/adapters/snowflake/connections.py

@@ -25,6 +25,10 @@
         'password': {
             'type': 'string',
         },
+        'authenticator': {
+            'type': 'string',
+            'description': "Either 'externalbrowser', or a valid Okta url"
+        },
         'database': {
             'type': 'string',
         },
@@ -41,7 +45,7 @@
             'type': 'boolean',
         }
     },
-    'required': ['account', 'user', 'password', 'database', 'schema'],
+    'required': ['account', 'user', 'database', 'schema'],
 }
 
 
@@ -95,17 +99,21 @@ def open(cls, connection):
 
         try:
             credentials = connection.credentials
+            # Pull all of the optional authentication args for the connector,
+            # let connector handle the actual arg validation
+            auth_args = {auth_key: credentials[auth_key]
+                         for auth_key in ['user', 'password', 'authenticator']
+                         if auth_key in credentials}
             handle = snowflake.connector.connect(
                 account=credentials.account,
-                user=credentials.user,
-                password=credentials.password,
                 database=credentials.database,
                 schema=credentials.schema,
                 warehouse=credentials.warehouse,
                 role=credentials.get('role', None),
                 autocommit=False,
                 client_session_keep_alive=credentials.get(
-                    'client_session_keep_alive', False)
+                    'client_session_keep_alive', False),
+                **auth_args
             )
 
             connection.handle = handle

test/unit/test_snowflake_adapter.py

@@ -11,6 +11,7 @@
 
 from .utils import config_from_parts_or_dicts
 
+
 class TestSnowflakeAdapter(unittest.TestCase):
     def setUp(self):
         flags.STRICT_MODE = False
@@ -21,7 +22,6 @@ def setUp(self):
                     'type': 'snowflake',
                     'account': 'test_account',
                     'user': 'test_user',
-                    'password': 'test_password',
                     'database': 'test_databse',
                     'warehouse': 'test_warehouse',
                     'schema': 'public',
@@ -42,10 +42,12 @@ def setUp(self):
         }
         self.config = config_from_parts_or_dicts(project_cfg, profile_cfg)
 
-        self.handle = mock.MagicMock(spec=snowflake_connector.SnowflakeConnection)
+        self.handle = mock.MagicMock(
+            spec=snowflake_connector.SnowflakeConnection)
         self.cursor = self.handle.cursor.return_value
         self.mock_execute = self.cursor.execute
-        self.patcher = mock.patch('dbt.adapters.snowflake.connections.snowflake.connector.connect')
+        self.patcher = mock.patch(
+            'dbt.adapters.snowflake.connections.snowflake.connector.connect')
         self.snowflake = self.patcher.start()
 
         self.snowflake.return_value = self.handle
@@ -74,7 +76,9 @@ def test_quoting_on_drop(self):
         )
         self.adapter.drop_relation(relation)
         self.mock_execute.assert_has_calls([
-            mock.call('drop table if exists "test_schema".test_table cascade', None)
+            mock.call(
+                'drop table if exists "test_schema".test_table cascade',
+                None)
         ])
 
     def test_quoting_on_truncate(self):
@@ -108,7 +112,9 @@ def test_quoting_on_rename(self):
             to_relation=to_relation
         )
         self.mock_execute.assert_has_calls([
-            mock.call('alter table "test_schema".table_a rename to table_b', None)
+            mock.call(
+                'alter table "test_schema".table_a rename to table_b',
+                None)
         ])
 
     def test_cancel_open_connections_empty(self):
@@ -131,18 +137,20 @@ def test_cancel_open_connections_single(self):
             query_result = mock.MagicMock()
             add_query.return_value = (None, query_result)
 
-            self.assertEqual(len(list(self.adapter.cancel_open_connections())), 1)
+            self.assertEqual(
+                len(list(self.adapter.cancel_open_connections())), 1)
 
-            add_query.assert_called_once_with('select system$abort_session(42)', 'master')
+            add_query.assert_called_once_with(
+                'select system$abort_session(42)', 'master')
 
     def test_client_session_keep_alive_false_by_default(self):
         self.adapter.connections.get(name='new_connection_with_new_config')
         self.snowflake.assert_has_calls([
             mock.call(
                 account='test_account', autocommit=False,
                 client_session_keep_alive=False, database='test_databse',
-                password='test_password', role=None, schema='public',
-                user='test_user', warehouse='test_warehouse')
+                role=None, schema='public', user='test_user',
+                warehouse='test_warehouse')
         ])
 
     def test_client_session_keep_alive_true(self):
@@ -155,6 +163,49 @@ def test_client_session_keep_alive_true(self):
             mock.call(
                 account='test_account', autocommit=False,
                 client_session_keep_alive=True, database='test_databse',
+                role=None, schema='public', user='test_user',
+                warehouse='test_warehouse')
+        ])
+
+    def test_user_pass_authentication(self):
+        self.config.credentials = self.config.credentials.incorporate(
+            password='test_password')
+        self.adapter = SnowflakeAdapter(self.config)
+        self.adapter.connections.get(name='new_connection_with_new_config')
+
+        self.snowflake.assert_has_calls([
+            mock.call(
+                account='test_account', autocommit=False,
+                client_session_keep_alive=False, database='test_databse',
                 password='test_password', role=None, schema='public',
                 user='test_user', warehouse='test_warehouse')
         ])
+
+    def test_authenticator_user_pass_authentication(self):
+        self.config.credentials = self.config.credentials.incorporate(
+            password='test_password', authenticator='test_sso_url')
+        self.adapter = SnowflakeAdapter(self.config)
+        self.adapter.connections.get(name='new_connection_with_new_config')
+
+        self.snowflake.assert_has_calls([
+            mock.call(
+                account='test_account', autocommit=False,
+                client_session_keep_alive=False, database='test_databse',
+                password='test_password', role=None, schema='public',
+                user='test_user', warehouse='test_warehouse',
+                authenticator='test_sso_url')
+        ])
+
+    def test_authenticator_externalbrowser_authentication(self):
+        self.config.credentials = self.config.credentials.incorporate(
+            authenticator='externalbrowser')
+        self.adapter = SnowflakeAdapter(self.config)
+        self.adapter.connections.get(name='new_connection_with_new_config')
+
+        self.snowflake.assert_has_calls([
+            mock.call(
+                account='test_account', autocommit=False,
+                client_session_keep_alive=False, database='test_databse',
+                role=None, schema='public', user='test_user',
+                warehouse='test_warehouse', authenticator='externalbrowser')
+        ])