Apply #2184 to async code

Author: sfc-gh-pczajka

src/snowflake/connector/aio/_azure_storage_client.py

@@ -47,13 +47,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/_connection.py

@@ -492,6 +492,10 @@ def _init_connection_parameters(
         for name, (value, _) in DEFAULT_CONFIGURATION.items():
             setattr(self, f"_{name}", value)
 
+        # Initialize unsafe_file_write explicitly since it's not in DEFAULT_CONFIGURATION
+        # TODO SNOW-2207863
+        self._unsafe_file_write = False
+
         self._heartbeat_task = None
         is_kwargs_empty = not connection_init_kwargs
 
@@ -704,6 +708,14 @@ def errorhandler(self, value) -> None:
     def rest(self) -> SnowflakeRestful | None:
         return self._rest
 
+    @property
+    def unsafe_file_write(self) -> bool:
+        return self._unsafe_file_write
+
+    @unsafe_file_write.setter
+    def unsafe_file_write(self, value: bool) -> None:
+        self._unsafe_file_write = value
+
     async def authenticate_with_retry(self, auth_instance) -> None:
         # make some changes if needed before real __authenticate
         try:

src/snowflake/connector/aio/_cursor.py

@@ -661,6 +661,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,6 +62,7 @@ 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,
@@ -80,6 +81,7 @@ def __init__(
             multipart_threshold,
             source_from_stream,
             use_s3_regional_url,
+            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

test/integ/aio/test_put_get_async.py

@@ -157,13 +157,13 @@ async def test_get_empty_file(tmp_path, aio_connection):
 async def test_get_file_permission(tmp_path, aio_connection, caplog):
     test_file = tmp_path / "data.csv"
     test_file.write_text("1,2,3\n")
-    stage_name = random_string(5, "test_get_empty_file_")
+    stage_name = random_string(5, "test_get_file_permission_")
     await aio_connection.connect()
     cur = aio_connection.cursor()
     await cur.execute(f"create temporary stage {stage_name}")
     filename_in_put = str(test_file).replace("\\", "/")
     await cur.execute(
-        f"PUT 'file://{filename_in_put}' @{stage_name}",
+        f"PUT 'file://{filename_in_put}' @{stage_name} AUTO_COMPRESS=FALSE",
     )
 
     with caplog.at_level(logging.ERROR):
@@ -173,9 +173,8 @@ async def test_get_file_permission(tmp_path, aio_connection, caplog):
     # get the default mask, usually it is 0o022
     default_mask = os.umask(0)
     os.umask(default_mask)
-    # files by default are given the permission 644 (Octal)
-    # umask is for denial, we need to negate
-    assert oct(os.stat(test_file).st_mode)[-3:] == oct(0o666 & ~default_mask)[-3:]
+    # files should be created with 0o600 permissions (owner read/write only)
+    assert oct(os.stat(test_file).st_mode)[-3:] == oct(0o600 & ~default_mask)[-3:]
 
 
 async def test_get_multiple_files_with_same_name(tmp_path, aio_connection, caplog):