Apply #2184+#2413 to async code

Author: sfc-gh-pczajka

src/snowflake/connector/aio/_azure_storage_client.py

@@ -50,13 +50,15 @@ def __init__(
         chunk_size: int,
         stage_info: dict[str, Any],
         use_s3_regional_url: bool = False,
+        unsafe_file_write: bool = False,
     ) -> None:
         SnowflakeAzureRestClientSync.__init__(
             self,
             meta=meta,
             stage_info=stage_info,
             chunk_size=chunk_size,
             credentials=credentials,
+            unsafe_file_write=unsafe_file_write,
         )
 
     async def _has_expired_token(self, response: aiohttp.ClientResponse) -> bool:

src/snowflake/connector/aio/_cursor.py

@@ -662,6 +662,7 @@ async def execute(
                     source_from_stream=file_stream,
                     multipart_threshold=data.get("threshold"),
                     use_s3_regional_url=self._connection.enable_stage_s3_privatelink_for_us_east_1,
+                    unsafe_file_write=self._connection.unsafe_file_write,
                 )
                 await sf_file_transfer_agent.execute()
                 data = sf_file_transfer_agent.result()

src/snowflake/connector/aio/_file_transfer_agent.py

@@ -62,24 +62,26 @@ def __init__(
         multipart_threshold: int | None = None,
         source_from_stream: IO[bytes] | None = None,
         use_s3_regional_url: bool = False,
+        unsafe_file_write: bool = False,
     ) -> None:
         super().__init__(
-            cursor,
-            command,
-            ret,
-            put_callback,
-            put_azure_callback,
-            put_callback_output_stream,
-            get_callback,
-            get_azure_callback,
-            get_callback_output_stream,
-            show_progress_bar,
-            raise_put_get_error,
-            force_put_overwrite,
-            skip_upload_on_content_match,
-            multipart_threshold,
-            source_from_stream,
-            use_s3_regional_url,
+            cursor=cursor,
+            command=command,
+            ret=ret,
+            put_callback=put_callback,
+            put_azure_callback=put_azure_callback,
+            put_callback_output_stream=put_callback_output_stream,
+            get_callback=get_callback,
+            get_azure_callback=get_azure_callback,
+            get_callback_output_stream=get_callback_output_stream,
+            show_progress_bar=show_progress_bar,
+            raise_put_get_error=raise_put_get_error,
+            force_put_overwrite=force_put_overwrite,
+            skip_upload_on_content_match=skip_upload_on_content_match,
+            multipart_threshold=multipart_threshold,
+            source_from_stream=source_from_stream,
+            use_s3_regional_url=use_s3_regional_url,
+            unsafe_file_write=unsafe_file_write,
         )
 
     async def execute(self) -> None:
@@ -271,6 +273,7 @@ async def _create_file_transfer_client(
                 meta,
                 self._stage_info,
                 4 * megabyte,
+                unsafe_file_write=self._unsafe_file_write,
             )
         elif self._stage_location_type == AZURE_FS:
             return SnowflakeAzureRestClient(
@@ -279,6 +282,7 @@ async def _create_file_transfer_client(
                 AZURE_CHUNK_SIZE,
                 self._stage_info,
                 use_s3_regional_url=self._use_s3_regional_url,
+                unsafe_file_write=self._unsafe_file_write,
             )
         elif self._stage_location_type == S3_FS:
             client = SnowflakeS3RestClient(
@@ -288,6 +292,7 @@ async def _create_file_transfer_client(
                 chunk_size=_chunk_size_calculator(meta.src_file_size),
                 use_accelerate_endpoint=self._use_accelerate_endpoint,
                 use_s3_regional_url=self._use_s3_regional_url,
+                unsafe_file_write=self._unsafe_file_write,
             )
             await client.transfer_accelerate_config(self._use_accelerate_endpoint)
             return client
@@ -299,6 +304,7 @@ async def _create_file_transfer_client(
                 self._cursor._connection,
                 self._command,
                 use_s3_regional_url=self._use_s3_regional_url,
+                unsafe_file_write=self._unsafe_file_write,
             )
             if client.security_token:
                 logger.debug(f"len(GCS_ACCESS_TOKEN): {len(client.security_token)}")

src/snowflake/connector/aio/_gcs_storage_client.py

@@ -39,6 +39,7 @@ def __init__(
         cnx: SnowflakeConnection,
         command: str,
         use_s3_regional_url: bool = False,
+        unsafe_file_write: bool = False,
     ) -> None:
         """Creates a client object with given stage credentials.
 
@@ -55,6 +56,7 @@ def __init__(
             chunk_size=-1,
             credentials=credentials,
             chunked_transfer=False,
+            unsafe_file_write=unsafe_file_write,
         )
         self.stage_info = stage_info
         self._command = command

src/snowflake/connector/aio/_s3_storage_client.py

@@ -48,6 +48,7 @@ def __init__(
         chunk_size: int,
         use_accelerate_endpoint: bool | None = None,
         use_s3_regional_url: bool = False,
+        unsafe_file_write: bool = False,
     ) -> None:
         """Rest client for S3 storage.
 
@@ -60,6 +61,7 @@ def __init__(
             stage_info=stage_info,
             chunk_size=chunk_size,
             credentials=credentials,
+            unsafe_file_write=unsafe_file_write,
         )
         # Signature version V4
         # Addressing style Virtual Host

src/snowflake/connector/aio/_storage_client.py

@@ -37,6 +37,7 @@ def __init__(
         chunked_transfer: bool | None = True,
         credentials: StorageCredential | None = None,
         max_retry: int = 5,
+        unsafe_file_write: bool = False,
     ) -> None:
         SnowflakeStorageClientSync.__init__(
             self,
@@ -46,6 +47,7 @@ def __init__(
             chunked_transfer=chunked_transfer,
             credentials=credentials,
             max_retry=max_retry,
+            unsafe_file_write=unsafe_file_write,
         )
 
     @abstractmethod
@@ -162,6 +164,7 @@ async def finish_download(self) -> None:
                     meta.encryption_material,
                     str(self.intermediate_dst_path),
                     tmp_dir=self.tmp_dir,
+                    unsafe_file_write=self.unsafe_file_write,
                 )
                 shutil.move(tmp_dst_file_name, self.full_dst_file_name)
                 self.intermediate_dst_path.unlink()