[ASYNC] apply #2258 to async code

Author: sfc-gh-pczajka

src/snowflake/connector/aio/_file_transfer_agent.py

@@ -195,7 +195,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"
             )
             if task.exception():
                 done_client.failed_transfers += 1

src/snowflake/connector/aio/_s3_storage_client.py

@@ -127,6 +127,9 @@ def generate_authenticated_url_and_args_v4() -> tuple[str, 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/aio/_storage_client.py

@@ -193,6 +193,7 @@ async 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)
@@ -208,10 +209,14 @@ async def _send_request_with_retry(
                     )
 
                 if await self._has_expired_presigned_url(response):
+                    logger.debug(
+                        "presigned url expired. trying to update presigned url."
+                    )
                     await self._update_presigned_url()
                 else:
                     self.last_err_is_presigned_url = False
                     if response.status in self.TRANSIENT_HTTP_ERR:
+                        logger.debug(f"transient error: {response.status}")
                         await asyncio.sleep(
                             min(
                                 # TODO should SLEEP_UNIT come from the parent
@@ -222,7 +227,9 @@ async def _send_request_with_retry(
                         )
                         self.retry_count[retry_id] += 1
                     elif await 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: