SNOW-269926-fix-timeout-issue (#612)

Author: sfc-gh-stan

src/snowflake/connector/connection.py

@@ -829,10 +829,12 @@ def cmd_query(self,
                 is_file_transfer
             )
 
+        url_parameters = {REQUEST_ID: request_id}
+
         ret = self.rest.request(
-            '/queries/v1/query-request',
+            '/queries/v1/query-request?' + urlencode(url_parameters),
             data, client=client, _no_results=_no_results,
-            _include_retry_params=True, request_id=request_id)
+            _include_retry_params=True)
 
         if ret is None:
             ret = {'data': {}}

src/snowflake/connector/network.py

@@ -310,7 +310,7 @@ def close(self):
                 logger.info("Session cleanup failed: %s", e)
 
     def request(self, url, body=None, method='post', client='sfsql',
-                _no_results=False, timeout=None, _include_retry_params=False, request_id=None):
+                _no_results=False, timeout=None, _include_retry_params=False):
         if body is None:
             body = {}
         if self.master_token is None and self.token is None:
@@ -341,11 +341,11 @@ def request(self, url, body=None, method='post', client='sfsql',
             return self._post_request(
                 url, headers, json.dumps(body),
                 token=self.token, _no_results=_no_results,
-                timeout=timeout, _include_retry_params=_include_retry_params, request_id=request_id)
+                timeout=timeout, _include_retry_params=_include_retry_params)
         else:
             return self._get_request(
                 url, headers, token=self.token,
-                timeout=timeout, request_id=request_id)
+                timeout=timeout)
 
     def update_tokens(self, session_token, master_token,
                       master_validity_in_seconds=None,
@@ -493,7 +493,7 @@ def delete_session(self, retry=False):
 
     def _get_request(self, url, headers, token=None,
                      timeout=None,
-                     socket_timeout=DEFAULT_SOCKET_CONNECT_TIMEOUT, request_id=None):
+                     socket_timeout=DEFAULT_SOCKET_CONNECT_TIMEOUT):
         if 'Content-Encoding' in headers:
             del headers['Content-Encoding']
         if 'Content-Length' in headers:
@@ -506,7 +506,7 @@ def _get_request(self, url, headers, token=None,
             url=url,
         )
         ret = self.fetch('get', full_url, headers, timeout=timeout,
-                         token=token, socket_timeout=socket_timeout, request_id=request_id)
+                         token=token, socket_timeout=socket_timeout)
         if ret.get('code') == SESSION_EXPIRED_GS_CODE:
             try:
                 ret = self._renew_session()
@@ -526,7 +526,7 @@ def _get_request(self, url, headers, token=None,
     def _post_request(self, url, headers, body, token=None,
                       timeout=None, _no_results=False, no_retry=False,
                       socket_timeout=DEFAULT_SOCKET_CONNECT_TIMEOUT,
-                      _include_retry_params=False, request_id=None):
+                      _include_retry_params=False):
         full_url = '{protocol}://{host}:{port}{url}'.format(
             protocol=self._protocol,
             host=self._host,
@@ -541,7 +541,7 @@ def _post_request(self, url, headers, body, token=None,
         ret = self.fetch('post', full_url, headers, data=body,
                          timeout=timeout, token=token,
                          no_retry=no_retry, socket_timeout=socket_timeout,
-                         _include_retry_params=_include_retry_params, request_id=request_id)
+                         _include_retry_params=_include_retry_params)
         logger.debug(
             'ret[code] = {code}, after post request'.format(
                 code=(ret.get('code', 'N/A'))))
@@ -561,9 +561,6 @@ def _post_request(self, url, headers, body, token=None,
                 return self._post_request(
                     url, headers, body, token=self.token, timeout=timeout)
 
-        if ret.get('data') and ret['data'].get('queryId'):
-            logger.debug('Query id: {}'.format(ret['data']['queryId']))
-
         if ret.get('code') == QUERY_IN_PROGRESS_ASYNC_CODE and _no_results:
             return ret
 
@@ -583,7 +580,7 @@ def _post_request(self, url, headers, body, token=None,
 
         return ret
 
-    def fetch(self, method, full_url, headers, data=None, timeout=None, request_id=None,
+    def fetch(self, method, full_url, headers, data=None, timeout=None,
               **kwargs):
         """Carry out API request with session management."""
 
@@ -620,32 +617,28 @@ def add_retry_params(self, full_url):
             retry_ctx = RetryCtx(timeout, include_retry_params)
             while True:
                 ret = self._request_exec_wrapper(
-                    session, method, full_url, headers, data, retry_ctx, request_id=request_id,
+                    session, method, full_url, headers, data, retry_ctx,
                     **kwargs)
                 if ret is not None:
                     return ret
 
     @staticmethod
-    def add_request_id_and_guid(full_url, request_id):
+    def add_request_guid(full_url):
         """Adds request_guid parameter for HTTP request tracing."""
         parsed_url = urlparse(full_url)
         if not parsed_url.hostname.endswith(SNOWFLAKE_HOST_SUFFIX):
             return full_url
-        request_guid = uuid.uuid4()
         suffix = urlencode({
-            REQUEST_ID: request_id,
-            REQUEST_GUID: request_guid
+            REQUEST_GUID: str(uuid.uuid4())
         })
-        if request_id:
-            logger.debug(f'request_id: {str(request_id)}, request_guid: {str(request_guid)}')
         sep = '&' if parsed_url.query else '?'
         # url has query string already, just add fields
         return full_url + sep + suffix
 
     def _request_exec_wrapper(
             self,
             session, method, full_url, headers, data, retry_ctx,
-            no_retry=False, token=NO_TOKEN, request_id=None,
+            no_retry=False, token=NO_TOKEN,
             **kwargs):
 
         conn = self._connection
@@ -654,7 +647,7 @@ def _request_exec_wrapper(
 
         start_request_thread = time.time()
         full_url = retry_ctx.add_retry_params(full_url)
-        full_url = SnowflakeRestful.add_request_id_and_guid(full_url, request_id)
+        full_url = SnowflakeRestful.add_request_guid(full_url)
         try:
             return_object = self._request_exec(
                 session=session,

test/integ/lambda/test_basic_query.py

@@ -4,6 +4,7 @@
 # Copyright (c) 2012-2021 Snowflake Computing Inc. All right reserved.
 #
 
+
 def test_connection(conn_cnx):
     """Test basic connection."""
     with conn_cnx() as cnx:

test/integ/test_reuse_cursor.py

@@ -4,6 +4,7 @@
 # Copyright (c) 2012-2021 Snowflake Computing Inc. All right reserved.
 #
 
+
 def test_reuse_cursor(conn_cnx, db_parameters):
     """Ensures only the last executed command/query's result sets are returned."""
     with conn_cnx() as cnx: