Merge pull request #970 from borisuvarov/client_session_keep_alive_sn…

…owflake Add client_session_keep_alive parameter for Snowflake connection
dbt-labs · Sep 28, 2018 · 95f3064 · 95f3064
2 parents 6454a81 + 4da156f
commit 95f3064
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 1 deletion.
diff --git a/dbt/adapters/snowflake/impl.py b/dbt/adapters/snowflake/impl.py
@@ -84,7 +84,9 @@ def open_connection(cls, connection):
                 schema=credentials.schema,
                 warehouse=credentials.warehouse,
                 role=credentials.get('role', None),
-                autocommit=False
+                autocommit=False,
+                client_session_keep_alive=credentials.get(
+                    'client_session_keep_alive', False)
             )
 
             connection.handle = handle

diff --git a/dbt/contracts/connection.py b/dbt/contracts/connection.py
@@ -110,6 +110,9 @@
         'role': {
             'type': 'string',
         },
+        'client_session_keep_alive': {
+            'type': 'boolean',
+        }
     },
     'required': ['account', 'user', 'password', 'database', 'schema'],
 }

diff --git a/test/unit/test_snowflake_adapter.py b/test/unit/test_snowflake_adapter.py
@@ -94,3 +94,26 @@ def test_quoting_on_rename(self):
         self.mock_execute.assert_has_calls([
             mock.call('alter table "test_schema".table_a rename to table_b', None)
         ])
+
+    def test_client_session_keep_alive_false_by_default(self):
+        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_client_session_keep_alive_true(self):
+        self.config.credentials = self.config.credentials.incorporate(
+            client_session_keep_alive=True)
+        self.adapter = SnowflakeAdapter(self.config)
+        self.adapter.get_connection(name='new_connection_with_new_config')
+
+        self.snowflake.assert_has_calls([
+            mock.call(
+                account='test_account', autocommit=False,
+                client_session_keep_alive=True, database='test_databse',
+                password='test_password', role=None, schema='public',
+                user='test_user', warehouse='test_warehouse')
+        ])