SNOW-1955965: Fix expired S3 credentials update (#2258)

Author: sfc-gh-pbulawa

src/snowflake/connector/file_transfer_agent.py

@@ -319,6 +319,9 @@ def __init__(
     def update(self, cur_timestamp) -> None:
         with self.lock:
             if cur_timestamp < self.timestamp:
+                logger.debug(
+                    "Omitting renewal of storage token, as it already happened."
+                )
                 return
             logger.debug("Renewing expired storage token.")
             ret = self.connection.cursor()._execute_helper(self._command)
@@ -540,7 +543,7 @@ def transfer_done_cb(
         ) -> None:
             # Note: chunk_id is 0 based while num_of_chunks is count
             logger.debug(
-                f"Chunk {chunk_id}/{done_client.num_of_chunks} of file {done_client.meta.name} reached callback"
+                f"Chunk(id: {chunk_id}) {chunk_id+1}/{done_client.num_of_chunks} of file {done_client.meta.name} reached callback"
             )
             with cv_chunk_process:
                 transfer_metadata.chunks_in_queue -= 1

src/snowflake/connector/s3_storage_client.py

@@ -333,6 +333,9 @@ def generate_authenticated_url_and_args_v4() -> tuple[bytes, dict[str, bytes]]:
             amzdate = t.strftime("%Y%m%dT%H%M%SZ")
             short_amzdate = amzdate[:8]
             x_amz_headers["x-amz-date"] = amzdate
+            x_amz_headers["x-amz-security-token"] = self.credentials.creds.get(
+                "AWS_TOKEN", ""
+            )
 
             (
                 canonical_request,

src/snowflake/connector/storage_client.py

@@ -286,6 +286,7 @@ def _send_request_with_retry(
             conn = self.meta.sfagent._cursor.connection
 
         while self.retry_count[retry_id] < self.max_retry:
+            logger.debug(f"retry #{self.retry_count[retry_id]}")
             cur_timestamp = self.credentials.timestamp
             url, rest_kwargs = get_request_args()
             rest_kwargs["timeout"] = (REQUEST_CONNECTION_TIMEOUT, REQUEST_READ_TIMEOUT)
@@ -299,10 +300,14 @@ def _send_request_with_retry(
                     response = rest_call(url, **rest_kwargs)
 
                 if self._has_expired_presigned_url(response):
+                    logger.debug(
+                        "presigned url expired. trying to update presigned url."
+                    )
                     self._update_presigned_url()
                 else:
                     self.last_err_is_presigned_url = False
                     if response.status_code in self.TRANSIENT_HTTP_ERR:
+                        logger.debug(f"transient error: {response.status_code}")
                         time.sleep(
                             min(
                                 # TODO should SLEEP_UNIT come from the parent
@@ -313,7 +318,9 @@ def _send_request_with_retry(
                         )
                         self.retry_count[retry_id] += 1
                     elif self._has_expired_token(response):
+                        logger.debug("token is expired. trying to update token")
                         self.credentials.update(cur_timestamp)
+                        self.retry_count[retry_id] += 1
                     else:
                         return response
             except self.TRANSIENT_ERRORS as e: