SNOW-27443: Fixed the error handling in network connection

Author: smtakeda

connection.py

@@ -74,6 +74,7 @@
     u'ocsp_response_cache_filename': None,  # snowflake internal
     u'converter_class': SnowflakeConverter,  # snowflake internal
     u'chunk_downloader_class': SnowflakeChunkDownloader,  # snowflake internal
+    u'retry_connection_auth': True,  # snowflake internal
 }
 
 APPLICATION_RE = re.compile(r'[\w\d_]+')
@@ -451,6 +452,7 @@ def __open_connection(self, mfa_callback, password_callback):
             mfa_callback=mfa_callback,
             password_callback=password_callback,
             session_parameters=self._session_parameters,
+            retry_connection_auth=self._retry_connection_auth,
         )
         self._password = None
         self._token = self._con.token

errors.py

@@ -162,47 +162,47 @@ class NotSupportedError(DatabaseError):
 class InternalServerError(Error):
     u"""Exception for 500 HTTP code for retry"""
 
-    def __init__(self):
+    def __init__(self, **_):
         Error.__init__(self, msg=u'HTTP 500: InternalServerError')
 
 
 class ServiceUnavailableError(Error):
     u"""Exception for 503 HTTP code for retry"""
 
-    def __init__(self):
+    def __init__(self, **_):
         Error.__init__(self, msg=u'HTTP 503: ServiceUnavailable')
 
 
 class GatewayTimeoutError(Error):
     u"""Exception for 504 HTTP error for retry"""
 
-    def __init__(self):
+    def __init__(self, **_):
         Error.__init__(self, msg=u'HTTP 504: GatewayTimeout')
 
 
 class ForbiddenError(Error):
     """Exception for 403 HTTP error for retry"""
 
-    def __init__(self):
+    def __init__(self, **_):
         Error.__init__(self, msg=u'HTTP 403: Forbidden')
 
 
 class RequestTimeoutError(Error):
     u"""Exception for 408 HTTP error for retry"""
 
-    def __init__(self):
+    def __init__(self, **_):
         Error.__init__(self, msg=u'HTTP 408: RequestTimeout')
 
 
 class BadRequest(Error):
     u"""Exception for 400 HTTP error for retry"""
 
-    def __init__(self):
+    def __init__(self, **_):
         Error.__init__(self, msg=u'HTTP 400: BadRequest')
 
 
 class BadGatewayError(Error):
     u"""Exception for 502 HTTP error for retry"""
 
-    def __init__(self):
+    def __init__(self, **_):
         Error.__init__(self, msg=u'HTTP 502: BadGateway')

network.py

@@ -46,7 +46,7 @@
 from .version import VERSION
 
 """
-Moneky patch for PyOpenSSL Socket wrapper
+Monkey patch for PyOpenSSL Socket wrapper
 """
 ssl_wrap_socket.inject_into_urllib3()
 
@@ -209,7 +209,7 @@ def authenticate(self, account, user, password, master_token=None,
                      warehouse=None, role=None, passcode=None,
                      passcode_in_password=False, saml_response=None,
                      mfa_callback=None, password_callback=None,
-                     session_parameters=None):
+                     session_parameters=None, retry_connection_auth=True):
         self.logger.info(u'authenticate')
 
         if token and master_token:
@@ -296,9 +296,10 @@ def authenticate(self, account, user, password, master_token=None,
             "body['data']: %s",
             {k: v for (k, v) in body[u'data'].items() if k != u'PASSWORD'})
 
-        # retry 10 times for authentication
+        max_retry = 10 if retry_connection_auth else 1
+        # max_retry times for authentication
         ret = self._post_request(
-            url, headers, json.dumps(body), retry=10)
+            url, headers, json.dumps(body), retry=max_retry)
         # this means we are waiting for MFA authentication
         if ret[u'data'].get(u'nextAction') and ret[u'data'][
             u'nextAction'] == u'EXT_AUTHN_DUO_ALL':
@@ -307,8 +308,9 @@ def authenticate(self, account, user, password, master_token=None,
             self.ret = None
 
             def post_request_wrapper(self, url, headers, body):
-                # retry 10 times for MFA approval
-                self.ret = self._post_request(url, headers, body, retry=10)
+                # max_retry times for MFA approval
+                self.ret = self._post_request(
+                    url, headers, body, retry=max_retry)
 
             # send new request to wait until MFA is approved
             t = Thread(target=post_request_wrapper,
@@ -327,9 +329,9 @@ def post_request_wrapper(self, url, headers, body):
                 body = copy.deepcopy(body_template)
                 body[u'inFlightCtx'] = ret[u'data'][u'inFlightCtx']
                 # final request to get tokens
-                # retry 10 times
+                # max_retry times
                 ret = self._post_request(
-                    url, headers, json.dumps(body), retry=10)
+                    url, headers, json.dumps(body), retry=max_retry)
 
         elif ret[u'data'].get(u'nextAction') and ret[u'data'][
             u'nextAction'] == u'PWD_CHANGE':
@@ -339,9 +341,9 @@ def post_request_wrapper(self, url, headers, body):
                 body[u'data'][u"LOGIN_NAME"] = user
                 body[u'data'][u"PASSWORD"] = password
                 body[u'data'][u'CHOSEN_NEW_PASSWORD'] = password_callback()
-                # retry 10 times for New Password input
+                # max_retry times for New Password input
                 ret = self._post_request(
-                    url, headers, json.dumps(body), retry=10)
+                    url, headers, json.dumps(body), retry=max_retry)
 
         self.logger.debug(u'completed authentication')
         if not ret[u'success']:
@@ -560,7 +562,7 @@ def _post_request(self, url, headers, body, token=None, retry=10,
             max_connection_pool=self._max_connection_pool)
         self.logger.debug(
             u'ret[code] = {code}, after post request'.format(
-                code=(ret[u'code'] if u'code' in ret else u'N/A')))
+                code=(ret.get(u'code', u'N/A'))))
 
         if u'code' in ret and ret[u'code'] == SESSION_EXPIRED_GS_CODE:
             ret = self._renew_session()
@@ -685,7 +687,7 @@ def request_thread(result_queue):
                 elif raw_ret.status_code in STATUS_TO_EXCEPTION:
                     # retryable exceptions
                     result_queue.put(
-                        (STATUS_TO_EXCEPTION[raw_ret.status_code], True))
+                        (STATUS_TO_EXCEPTION[raw_ret.status_code](), True))
                 elif raw_ret.status_code == UNAUTHORIZED and \
                         catch_okta_unauthorized_error:
                     # OKTA Unauthorized errors
@@ -757,6 +759,7 @@ def request_thread(result_queue):
         sleeping_time = 1
         return_object = None
         for retry_cnt in range(retry):
+            return_object = None
             request_result_queue = Queue()
             th = Thread(name='request_thread', target=request_thread,
                         args=(request_result_queue,))
@@ -801,8 +804,17 @@ def request_thread(result_queue):
                     retry,
                     sleeping_time)
             time.sleep(sleeping_time)
-            return_object = None
 
+        if isinstance(return_object, Error):
+            Error.errorhandler_wrapper(conn, None, return_object)
+        elif isinstance(return_object, Exception):
+            Error.errorhandler_wrapper(
+                conn, None, OperationalError,
+                {
+                    u'msg': u'Failed to execute request: {0}'.format(
+                        return_object),
+                    u'errno': ER_FAILED_TO_REQUEST,
+                })
         return return_object
 
     def authenticate_by_saml(self, authenticator, account, user, password):

test/test_connection.py

@@ -7,8 +7,10 @@
 import pytest
 
 import snowflake.connector
-from snowflake.connector import (DatabaseError,
-                                 ProgrammingError)
+from snowflake.connector import (
+    DatabaseError,
+    ProgrammingError)
+from snowflake.connector.errors import ForbiddenError
 
 
 def test_basic(conn_testaccount):
@@ -120,7 +122,9 @@ def test_bogus(db_parameters):
             password='bogus',
             host=db_parameters['host'],
             port=db_parameters['port'],
-            account=db_parameters['account'])
+            account=db_parameters['account'],
+            retry_connection_auth=False
+        )
 
     with pytest.raises(DatabaseError):
         snowflake.connector.connect(
@@ -130,7 +134,8 @@ def test_bogus(db_parameters):
             account='testaccount123',
             host=db_parameters['host'],
             port=db_parameters['port'],
-            insecure_mode=True)
+            insecure_mode=True,
+            retry_connection_auth=False)
 
     with pytest.raises(DatabaseError):
         snowflake.connector.connect(
@@ -140,6 +145,7 @@ def test_bogus(db_parameters):
             account='testaccount123',
             host=db_parameters['host'],
             port=db_parameters['port'],
+            retry_connection_auth=False
         )
 
     with pytest.raises(ProgrammingError):
@@ -150,6 +156,7 @@ def test_bogus(db_parameters):
             account='testaccount123',
             host=db_parameters['host'],
             port=db_parameters['port'],
+            retry_connection_auth=False
         )
 
 
@@ -224,3 +231,13 @@ def exe(sql):
                 db_parameters['database']))
         exe('drop role snowdog_role')
         exe('drop user if exists snowdog')
+
+
+def test_invalid_account():
+    with pytest.raises(ForbiddenError):
+        snowflake.connector.connect(
+            account='bogus',
+            user='test',
+            password='test',
+            retry_connection_auth=False
+        )