SNOW-749141: Fix uploading a file > 200M when the same filename already exists in a stage (#1470)

Author: sfc-gh-yixie

DESCRIPTION.md

@@ -10,7 +10,8 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
 
 - v3.0.2(Unreleased)
 
-  - Fixed a memory leak in the logging module of the Cython extension
+  - Fixed a memory leak in the logging module of the Cython extension.
+  - Fixed a bug where the `put` command on AWS raised `AttributeError` when the file size was larger than 200M.
 
 - v3.0.1(February 28, 2023)
 

src/snowflake/connector/storage_client.py

@@ -185,8 +185,8 @@ def get_file_header(self, filename: str) -> FileHeader | None:
     def preprocess(self) -> None:
         meta = self.meta
         logger.debug(f"Preprocessing {meta.src_file_name}")
-
         if not meta.overwrite:
+            self.get_digest()  # self.get_file_header needs digest for multiparts upload when aws is used.
             self.get_file_header(meta.dst_file_name)  # Check if file exists on remote
             if meta.result_status == ResultStatus.UPLOADED:
                 # Skipped
@@ -202,7 +202,6 @@ def preprocess(self) -> None:
         if meta.require_compress:
             self.compress()
         self.get_digest()
-
         self.preprocessed = True
 
     def prepare_upload(self) -> None:

test/integ/test_large_put.py

@@ -56,10 +56,26 @@ def mocked_file_agent(*args, **kwargs):
                 "snowflake.connector.cursor.SnowflakeFileTransferAgent",
                 side_effect=mocked_file_agent,
             ):
+                # upload with auto compress = True
                 cnx.cursor().execute(
-                    f"put 'file://{files}' @%{db_parameters['name']}",
+                    f"put 'file://{files}' @%{db_parameters['name']} auto_compress=True",
                 )
                 assert mocked_file_agent.agent._multipart_threshold == 10000
+                cnx.cursor().execute(f"remove @%{db_parameters['name']}")
+
+                # upload with auto compress = False
+                cnx.cursor().execute(
+                    f"put 'file://{files}' @%{db_parameters['name']} auto_compress=False",
+                )
+                assert mocked_file_agent.agent._multipart_threshold == 10000
+
+                # Upload again. There was a bug when a large file is uploaded again while it already exists in a stage.
+                # Refer to preprocess(self) of storage_client.py.
+                # self.get_digest() needs to be called before self.get_file_header(meta.dst_file_name).
+                # SNOW-749141
+                cnx.cursor().execute(
+                    f"put 'file://{files}' @%{db_parameters['name']} auto_compress=False",
+                )  # do not add `overwrite=True` because overwrite will skip the code path to extract file header.
 
             c = cnx.cursor()
             try: