SNOW-294266 make connection object exit autocommit aware (#641)

Author: sfc-gh-mkeller

src/snowflake/connector/auth.py

@@ -15,6 +15,7 @@
 from os import getenv, makedirs, mkdir, path, remove, removedirs, rmdir
 from os.path import expanduser
 from threading import Lock, Thread
+from typing import Dict, Union
 
 from .auth_keypair import AuthByKeyPair
 from .auth_usrpwdmfa import AuthByUsrPwdMfa
@@ -124,7 +125,7 @@ def authenticate(self, auth_instance, account, user,
                      warehouse=None, role=None, passcode=None,
                      passcode_in_password=False,
                      mfa_callback=None, password_callback=None,
-                     session_parameters=None, timeout=120):
+                     session_parameters=None, timeout=120) -> Dict[str, Union[str, int, bool]]:
         logger.debug('authenticate')
 
         if session_parameters is None:
@@ -356,9 +357,10 @@ def post_request_wrapper(self, url, headers, body):
                 self._rest._connection._schema = session_info.get('schemaName')
                 self._rest._connection._warehouse = session_info.get('warehouseName')
                 self._rest._connection._role = session_info.get('roleName')
-            self._rest._connection._set_parameters(ret, session_parameters)
-
-        return session_parameters
+            if 'parameters' in ret['data']:
+                session_parameters.update({p['name']: p['value'] for p in ret['data']['parameters']})
+            self._rest._connection._update_parameters(session_parameters)
+            return session_parameters
 
     def _read_temporary_credential(self, host, user, cred_type):
         cred = None

src/snowflake/connector/connection.py

@@ -214,6 +214,7 @@ def __init__(self, **kwargs):
         self._async_sfqids = set()
         self._done_async_sfqids = set()
         self.telemetry_enabled = False
+        self._session_parameters: Dict[str, Union[str, int, bool]] = {}
         logger.info(
             "Snowflake Connector for Python Version: %s, "
             "Python Version: %s, Platform: %s",
@@ -889,7 +890,7 @@ def __authenticate(self, auth_instance):
             auth_instance, 'consent_cache_id_token', True)
 
         auth = Auth(self.rest)
-        self._session_parameters = auth.authenticate(
+        auth.authenticate(
             auth_instance=auth_instance,
             account=self.account,
             user=self.user,
@@ -1114,17 +1115,15 @@ def _validate_client_prefetch_threads(self):
             self.client_prefetch_threads)
         return self.client_prefetch_threads
 
-    def _set_parameters(self, ret, session_parameters):
-        """Sets session parameters."""
-        if 'parameters' not in ret['data']:
-            return
-        parameters = ret['data']['parameters']
+    def _update_parameters(
+            self,
+            parameters: Dict[str, Union[str, int, bool]],
+    ) -> None:
+        """Update session parameters."""
         with self._lock_converter:
             self.converter.set_parameters(parameters)
-        for kv in parameters:
-            name = kv['name']
-            value = kv['value']
-            session_parameters[name] = value
+        for name, value in parameters.items():
+            self._session_parameters[name] = value
             if PARAMETER_CLIENT_TELEMETRY_ENABLED == name:
                 self.telemetry_enabled = value
             elif PARAMETER_CLIENT_TELEMETRY_OOB_ENABLED == name:
@@ -1153,10 +1152,12 @@ def __enter__(self):
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         """Context manager with commit or rollback teardown."""
-        if exc_tb is None:
-            self.commit()
-        else:
-            self.rollback()
+        if not self._session_parameters.get('AUTOCOMMIT', False):
+            # Either AUTOCOMMIT is turned off, or is not set so we default to old behavior
+            if exc_tb is None:
+                self.commit()
+            else:
+                self.rollback()
         self.close()
 
     def _get_query_status(self, sf_qid: str) -> Tuple[QueryStatus, Dict[str, Any]]:

src/snowflake/connector/converter.py

@@ -11,7 +11,7 @@
 from datetime import time as dt_t
 from datetime import timedelta
 from logging import getLogger
-from typing import Any, Tuple, Union
+from typing import Any, Dict, Optional, Tuple, Union
 
 import pytz
 
@@ -129,24 +129,22 @@ def _generate_tzinfo_from_tzoffset(tzoffset_minutes: int) -> pytz._FixedOffset:
 
 class SnowflakeConverter(object):
     def __init__(self, **kwargs):
-        self._parameters = {}
+        self._parameters: Dict[str, Union[str, int, bool]] = {}
         self._use_numpy = kwargs.get('use_numpy', False) and numpy is not None
 
         logger.debug('use_numpy: %s', self._use_numpy)
 
-    def set_parameters(self, parameters):
-        self._parameters = {}
-        for kv in parameters:
-            self._parameters[kv['name']] = kv['value']
+    def set_parameters(self, new_parameters: Dict) -> None:
+        self._parameters = new_parameters
 
-    def set_parameter(self, param, value):
+    def set_parameter(self, param: Any, value: Any) -> None:
         self._parameters[param] = value
 
-    def get_parameters(self):
+    def get_parameters(self) -> Dict[str, Union[str, int, bool]]:
         return self._parameters
 
-    def get_parameter(self, param):
-        return self._parameters[param] if param in self._parameters else None
+    def get_parameter(self, param: str) -> Optional[Union[str, int, bool]]:
+        return self._parameters.get(param)
 
     #
     # FROM Snowflake to Python Objects

src/snowflake/connector/cursor.py

@@ -415,7 +415,9 @@ def interrupt_handler(*_):  # pragma: no cover
                 logger.debug('cancelled timebomb in finally')
 
         if 'data' in ret and 'parameters' in ret['data']:
-            for kv in ret['data']['parameters']:
+            parameters = ret['data']['parameters']
+            # Set session parameters for cursor object
+            for kv in parameters:
                 if 'TIMESTAMP_OUTPUT_FORMAT' in kv['name']:
                     self._timestamp_output_format = kv['value']
                 if 'TIMESTAMP_NTZ_OUTPUT_FORMAT' in kv['name']:
@@ -432,8 +434,8 @@ def interrupt_handler(*_):  # pragma: no cover
                     self._timezone = kv['value']
                 if 'BINARY_OUTPUT_FORMAT' in kv['name']:
                     self._binary_output_format = kv['value']
-            self._connection._set_parameters(
-                ret, self._connection._session_parameters)
+            # Set session parameters for connection object
+            self._connection._update_parameters({p['name']: p['value'] for p in parameters})
 
         self._sequence_counter = -1
         return ret

test/integ/test_connection.py

@@ -926,3 +926,16 @@ def test_process_param_error(conn_cnx):
                             side_effect=Exception('test')):
                 conn._process_params(mock.Mock())
             assert pe.errno == ER_FAILED_PROCESSING_PYFORMAT
+
+
+@pytest.mark.parametrize('auto_commit', [pytest.param(True, marks=pytest.mark.skipolddriver), False])
+def test_autocommit(conn_cnx, db_parameters, auto_commit):
+    conn = snowflake.connector.connect(**db_parameters)
+    with mock.patch.object(conn, 'commit') as mocked_commit:
+        with conn:
+            with conn.cursor() as cur:
+                cur.execute(f"alter session set autocommit = {auto_commit}")
+    if auto_commit:
+        assert not mocked_commit.called
+    else:
+        assert mocked_commit.called