SNOW-1944208 add unsafe write flag (#2184)

Author: sfc-gh-mkubik

src/snowflake/connector/azure_storage_client.py

@@ -63,8 +63,15 @@ def __init__(
         chunk_size: int,
         stage_info: dict[str, Any],
         use_s3_regional_url: bool = False,
+        unsafe_file_write: bool = False,
     ) -> None:
-        super().__init__(meta, stage_info, chunk_size, credentials=credentials)
+        super().__init__(
+            meta,
+            stage_info,
+            chunk_size,
+            credentials=credentials,
+            unsafe_file_write=unsafe_file_write,
+        )
         end_point: str = stage_info["endPoint"]
         if end_point.startswith("blob."):
             end_point = end_point[len("blob.") :]

src/snowflake/connector/connection.py

@@ -381,6 +381,7 @@ class SnowflakeConnection:
         server_session_keep_alive: When true, the connector does not destroy the session on the Snowflake server side
           before the connector shuts down. Default value is false.
         token_file_path: The file path of the token file. If both token and token_file_path are provided, the token in token_file_path will be used.
+        unsafe_file_write: When true, files downloaded by GET will be saved with 644 permissions. Otherwise, files will be saved with safe - owner-only permissions: 600.
     """
 
     OCSP_ENV_LOCK = Lock()
@@ -761,6 +762,14 @@ def is_query_context_cache_disabled(self) -> bool:
     def iobound_tpe_limit(self) -> int | None:
         return self._iobound_tpe_limit
 
+    @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
+
     def connect(self, **kwargs) -> None:
         """Establishes connection to Snowflake."""
         logger.debug("connect")
@@ -1232,6 +1241,11 @@ def __config(self, **kwargs):
             if "protocol" not in kwargs:
                 self._protocol = "https"
 
+        if "unsafe_file_write" in kwargs:
+            self._unsafe_file_write = kwargs["unsafe_file_write"]
+        else:
+            self._unsafe_file_write = False
+
         logger.info(
             f"Connecting to {_DOMAIN_NAME_MAP.get(extract_top_level_domain_from_hostname(self._host), 'GLOBAL')} Snowflake domain"
         )

src/snowflake/connector/cursor.py

@@ -1058,6 +1058,7 @@ def execute(
                     multipart_threshold=data.get("threshold"),
                     use_s3_regional_url=self._connection.enable_stage_s3_privatelink_for_us_east_1,
                     iobound_tpe_limit=self._connection.iobound_tpe_limit,
+                    unsafe_file_write=self._connection.unsafe_file_write,
                 )
                 sf_file_transfer_agent.execute()
                 data = sf_file_transfer_agent.result()

src/snowflake/connector/encryption_util.py

@@ -194,6 +194,7 @@ def decrypt_file(
         in_filename: str,
         chunk_size: int = 64 * kilobyte,
         tmp_dir: str | None = None,
+        unsafe_file_write: bool = False,
     ) -> str:
         """Decrypts a file and stores the output in the temporary directory.
 
@@ -212,8 +213,10 @@ def decrypt_file(
             temp_output_file = os.path.join(tmp_dir, temp_output_file)
 
         logger.debug("encrypted file: %s, tmp file: %s", in_filename, temp_output_file)
+
+        file_opener = None if unsafe_file_write else owner_rw_opener
         with open(in_filename, "rb") as infile:
-            with open(temp_output_file, "wb") as outfile:
+            with open(temp_output_file, "wb", opener=file_opener) as outfile:
                 SnowflakeEncryptionUtil.decrypt_stream(
                     metadata, encryption_material, infile, outfile, chunk_size
                 )

src/snowflake/connector/file_transfer_agent.py

@@ -355,6 +355,7 @@ def __init__(
         source_from_stream: IO[bytes] | None = None,
         use_s3_regional_url: bool = False,
         iobound_tpe_limit: int | None = None,
+        unsafe_file_write: bool = False,
     ) -> None:
         self._cursor = cursor
         self._command = command
@@ -386,6 +387,7 @@ def __init__(
         self._use_s3_regional_url = use_s3_regional_url
         self._credentials: StorageCredential | None = None
         self._iobound_tpe_limit = iobound_tpe_limit
+        self._unsafe_file_write = unsafe_file_write
 
     def execute(self) -> None:
         self._parse_command()
@@ -673,6 +675,7 @@ 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(
@@ -681,6 +684,7 @@ 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:
             return SnowflakeS3RestClient(
@@ -690,6 +694,7 @@ def _create_file_transfer_client(
                 _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,
             )
         elif self._stage_location_type == GCS_FS:
             return SnowflakeGCSRestClient(
@@ -699,6 +704,7 @@ 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,
             )
         raise Exception(f"{self._stage_location_type} is an unknown stage type")
 

src/snowflake/connector/gcs_storage_client.py

@@ -54,6 +54,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.
 
@@ -64,7 +65,12 @@ def __init__(
             The client to communicate with GCS.
         """
         super().__init__(
-            meta, stage_info, -1, credentials=credentials, chunked_transfer=False
+            meta,
+            stage_info,
+            -1,
+            credentials=credentials,
+            chunked_transfer=False,
+            unsafe_file_write=unsafe_file_write,
         )
         self.stage_info = stage_info
         self._command = command

src/snowflake/connector/local_storage_client.py

@@ -26,8 +26,11 @@ def __init__(
         meta: SnowflakeFileMeta,
         stage_info: dict[str, Any],
         chunk_size: int,
+        unsafe_file_write: bool = False,
     ) -> None:
-        super().__init__(meta, stage_info, chunk_size)
+        super().__init__(
+            meta, stage_info, chunk_size, unsafe_file_write=unsafe_file_write
+        )
         self.data_file = meta.src_file_name
         self.full_dst_file_name: str = os.path.join(
             stage_info["location"], os.path.basename(meta.dst_file_name)

src/snowflake/connector/s3_storage_client.py

@@ -61,13 +61,20 @@ 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.
 
         Args:
             stage_info:
         """
-        super().__init__(meta, stage_info, chunk_size, credentials=credentials)
+        super().__init__(
+            meta,
+            stage_info,
+            chunk_size,
+            credentials=credentials,
+            unsafe_file_write=unsafe_file_write,
+        )
         # Signature version V4
         # Addressing style Virtual Host
         self.region_name: str = stage_info["region"]

src/snowflake/connector/storage_client.py

@@ -77,6 +77,7 @@ def __init__(
         chunked_transfer: bool | None = True,
         credentials: StorageCredential | None = None,
         max_retry: int = 5,
+        unsafe_file_write: bool = False,
     ) -> None:
         self.meta = meta
         self.stage_info = stage_info
@@ -115,6 +116,7 @@ def __init__(
         self.failed_transfers: int = 0
         # only used when PRESIGNED_URL expires
         self.last_err_is_presigned_url = False
+        self.unsafe_file_write = unsafe_file_write
 
     def compress(self) -> None:
         if self.meta.require_compress:
@@ -371,7 +373,7 @@ def finish_download(self) -> None:
                 # For storage utils that do not have the privilege of
                 # getting the metadata early, both object and metadata
                 # are downloaded at once. In which case, the file meta will
-                # be updated with all the metadata that we need and
+                # be updated with all the metadata that we need, and
                 # then we can call get_file_header to get just that and also
                 # preserve the idea of getting metadata in the first place.
                 # One example of this is the utils that use presigned url
@@ -385,6 +387,7 @@ 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()

test/integ/test_put_get.py

@@ -20,6 +20,13 @@
 
 from snowflake.connector import OperationalError
 
+try:
+    from src.snowflake.connector.compat import IS_WINDOWS
+except ImportError:
+    import platform
+
+    IS_WINDOWS = platform.system() == "Windows"
+
 try:
     from snowflake.connector.util_text import random_string
 except ImportError:
@@ -740,16 +747,44 @@ def test_get_empty_file(tmp_path, conn_cnx):
 
 
 @pytest.mark.skipolddriver
+@pytest.mark.skipif(IS_WINDOWS, reason="not supported on Windows")
 def test_get_file_permission(tmp_path, conn_cnx, 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_")
     with conn_cnx() as cnx:
         with cnx.cursor() as cur:
             cur.execute(f"create temporary stage {stage_name}")
             filename_in_put = str(test_file).replace("\\", "/")
             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):
+                cur.execute(f"GET @{stage_name}/data.csv file://{tmp_path}")
+            assert "FileNotFoundError" not in caplog.text
+
+            default_mask = os.umask(0)
+            os.umask(default_mask)
+
+            assert (
+                oct(os.stat(test_file).st_mode)[-3:] == oct(0o600 & ~default_mask)[-3:]
+            )
+
+
+@pytest.mark.skipolddriver
+@pytest.mark.skipif(IS_WINDOWS, reason="not supported on Windows")
+def test_get_unsafe_file_permission_when_flag_set(tmp_path, conn_cnx, caplog):
+    test_file = tmp_path / "data.csv"
+    test_file.write_text("1,2,3\n")
+    stage_name = random_string(5, "test_get_file_permission_")
+    with conn_cnx() as cnx:
+        cnx.unsafe_file_write = True
+        with cnx.cursor() as cur:
+            cur.execute(f"create temporary stage {stage_name}")
+            filename_in_put = str(test_file).replace("\\", "/")
+            cur.execute(
+                f"PUT 'file://{filename_in_put}' @{stage_name} AUTO_COMPRESS=FALSE",
             )
 
             with caplog.at_level(logging.ERROR):
@@ -764,6 +799,7 @@ def test_get_file_permission(tmp_path, conn_cnx, caplog):
             assert (
                 oct(os.stat(test_file).st_mode)[-3:] == oct(0o666 & ~default_mask)[-3:]
             )
+        cnx.unsafe_file_write = False
 
 
 @pytest.mark.skipolddriver