SNOW-27430: added retry for BAD_GATEWAY and increased the retry counter

Author: smtakeda

compat.py

@@ -42,6 +42,7 @@
     PKCS5_OFFSET = lambda v: ord(v[-1])
     IS_BINARY = lambda v: isinstance(v, bytearray)
 
+    BAD_GATEWAY = httplib.BAD_GATEWAY
     BAD_REQUEST = httplib.BAD_REQUEST
     SERVICE_UNAVAILABLE = httplib.SERVICE_UNAVAILABLE
     GATEWAY_TIMEOUT = httplib.GATEWAY_TIMEOUT
@@ -74,6 +75,7 @@
     PKCS5_OFFSET = lambda v: v[-1]
     IS_BINARY = lambda v: isinstance(v, (bytes, bytearray))
 
+    BAD_GATEWAY = http.client.BAD_GATEWAY
     BAD_REQUEST = http.client.BAD_REQUEST
     SERVICE_UNAVAILABLE = http.client.SERVICE_UNAVAILABLE
     GATEWAY_TIMEOUT = http.client.GATEWAY_TIMEOUT

connection.py

@@ -559,11 +559,12 @@ def _cmd_query(self, sql, sequence_counter, request_id,
 
         url_parameters = {u'requestId': request_id}
 
+        # retry 1000 times/4.5 hours for general queries
         ret = self._con.request(
             u'/queries/v1/query-request?' + urlencode(url_parameters),
             data,
             client=client,
-            _no_results=_no_results)
+            _no_results=_no_results, retry=1000)
 
         if ret is not None and u'data' in ret and ret[u'data'] is None:
             ret[u'data'] = {}

cursor.py

@@ -578,7 +578,9 @@ def chunk_info(self, data, use_ijson=False):
 
     def query_result(self, qid, _use_ijson=False):
         url = ('/queries/{qid}/result').format(qid=qid)
-        ret = self._connection._con.request(url=url, method='get')
+        # retry 1000 times for getting the results
+        ret = self._connection._con.request(
+            url=url, method='get', retry=1000)
         if ret.get(u'success'):
             data = ret.get(u'data')
             self.chunk_info(data, use_ijson=_use_ijson)
@@ -602,7 +604,8 @@ def query_result(self, qid, _use_ijson=False):
 
     def abort_query(self, qid):
         url = ('/queries/{qid}/abort-request').format(qid=qid)
-        ret = self._connection._con.request(url=url, method='post')
+        ret = self._connection._con.request(
+            url=url, method='post', retry=100)
         return ret.get(u'success')
 
     def executemany(self, command, seqparams):

errors.py

@@ -199,3 +199,10 @@ class BadRequest(Error):
 
     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):
+        Error.__init__(self, msg=u'HTTP 502: BadGateway')

network.py

@@ -26,7 +26,7 @@
 from . import ssl_wrap_socket
 from .compat import (
     BAD_REQUEST, SERVICE_UNAVAILABLE, GATEWAY_TIMEOUT,
-    FORBIDDEN,
+    FORBIDDEN, BAD_GATEWAY,
     UNAUTHORIZED, INTERNAL_SERVER_ERROR, OK, BadStatusLine)
 from .compat import (Queue, EmptyQueue)
 from .compat import (TO_UNICODE, urlencode)
@@ -37,7 +37,7 @@
 from .errors import (Error, OperationalError, DatabaseError, ProgrammingError,
                      GatewayTimeoutError, ServiceUnavailableError,
                      InterfaceError, InternalServerError, ForbiddenError,
-                     BadRequest)
+                     BadGatewayError, BadRequest)
 from .gzip_decoder import (decompress_raw_data)
 from .sqlstate import (SQLSTATE_CONNECTION_NOT_EXISTS,
                        SQLSTATE_CONNECTION_WAS_NOT_ESTABLISHED,
@@ -95,6 +95,7 @@
     SERVICE_UNAVAILABLE: ServiceUnavailableError,
     GATEWAY_TIMEOUT: GatewayTimeoutError,
     BAD_REQUEST: BadRequest,
+    BAD_GATEWAY: BadGatewayError,
 }
 
 
@@ -295,7 +296,9 @@ 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'})
 
-        ret = self._post_request(url, headers, json.dumps(body))
+        # retry 10 times for authentication
+        ret = self._post_request(
+            url, headers, json.dumps(body), retry=10)
         # 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':
@@ -304,7 +307,8 @@ def authenticate(self, account, user, password, master_token=None,
             self.ret = None
 
             def post_request_wrapper(self, url, headers, body):
-                self.ret = self._post_request(url, headers, body)
+                # retry 10 times for MFA approval
+                self.ret = self._post_request(url, headers, body, retry=10)
 
             # send new request to wait until MFA is approved
             t = Thread(target=post_request_wrapper,
@@ -316,14 +320,16 @@ def post_request_wrapper(self, url, headers, body):
                 while not self.ret:
                     next(c)
             else:
-                t.join()
+                t.join(timeout=120)
             ret = self.ret
             if ret[u'data'].get(u'nextAction') and ret[u'data'][
                 u'nextAction'] == u'EXT_AUTHN_SUCCESS':
                 body = copy.deepcopy(body_template)
                 body[u'inFlightCtx'] = ret[u'data'][u'inFlightCtx']
                 # final request to get tokens
-                ret = self._post_request(url, headers, json.dumps(body))
+                # retry 10 times
+                ret = self._post_request(
+                    url, headers, json.dumps(body), retry=10)
 
         elif ret[u'data'].get(u'nextAction') and ret[u'data'][
             u'nextAction'] == u'PWD_CHANGE':
@@ -333,7 +339,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()
-                ret = self._post_request(url, headers, json.dumps(body))
+                # retry 10 times for New Password input
+                ret = self._post_request(
+                    url, headers, json.dumps(body), retry=10)
 
         self.logger.debug(u'completed authentication')
         if not ret[u'success']:
@@ -373,7 +381,7 @@ def post_request_wrapper(self, url, headers, body):
                     self._connection._warehouse = session_info[u'warehouseName']
 
     def request(self, url, body=None, method=u'post', client=u'sfsql',
-                _no_results=False):
+                _no_results=False, retry=10):
         if body is None:
             body = {}
         if not hasattr(self, u'_master_token'):
@@ -396,11 +404,12 @@ def request(self, url, body=None, method=u'post', client=u'sfsql',
             u"User-Agent": PYTHON_CONNECTOR_USER_AGENT,
         }
         if method == u'post':
-            return self._post_request(url, headers, json.dumps(body),
-                                      token=self._token,
-                                      _no_results=_no_results)
+            return self._post_request(
+                url, headers, json.dumps(body),
+                token=self._token, _no_results=_no_results, retry=retry)
         else:
-            return self._get_request(url, headers, token=self._token)
+            return self._get_request(
+                url, headers, token=self._token, retry=retry)
 
     def _renew_session(self):
         if not hasattr(self, u'_master_token'):