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

Author: sfc-gh-pbulawa

DESCRIPTION.md

@@ -20,6 +20,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
   - Added `check_arrow_conversion_error_on_every_column` connection property that can be set to `False` to restore previous behaviour in which driver will ignore errors until it occurs in the last column. This flag's purpose is to unblock workflows that may be impacted by the bugfix and will be removed in later releases.
   - Lower log levels from info to debug for some of the messages to make the output easier to follow.
   - Allow the connector to inherit a UUID4 generated upstream, provided in statement parameters (field: `requestId`), rather than automatically generate a UUID4 to use for the HTTP Request ID.
+  - Fix expired S3 credentials update and increment retry when expired credentials are found.
 
 - v3.14.0(March 03, 2025)
   - Bumped pyOpenSSL dependency upper boundary from <25.0.0 to <26.0.0.

src/snowflake/connector/file_transfer_agent.py

@@ -315,6 +315,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)
@@ -536,7 +539,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

@@ -329,6 +329,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

@@ -282,6 +282,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)
@@ -295,10 +296,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
@@ -309,7 +314,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: