SNOW-1258702:SnowSQL File Upload support for China (#1912)

Author: sfc-gh-yuwang

DESCRIPTION.md

@@ -8,6 +8,10 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
 
 # Release Notes
 
+- v3.9.1(TBD)
+
+  - Improved s3 acceleration logic when connecting to China endpoint.
+
 - v3.8.1(April 09, 2024)
 
   - Reverted the change "Updated `write_pandas` to skip TABLE IF NOT EXISTS in truncate mode." introduced in v3.8.0 (yanked) as it's a breaking change. `write_pandas` will be fixed in the future in a non-breaking way.

src/snowflake/connector/s3_storage_client.py

@@ -80,6 +80,7 @@ def __init__(
             )
         )
         self.use_s3_regional_url = use_s3_regional_url
+        self.location_type = stage_info.get("locationType")
 
         # if GS sends us an endpoint, it's likely for FIPS. Use it.
         self.endpoint: str | None = None
@@ -92,6 +93,9 @@ def __init__(
     def transfer_accelerate_config(
         self, use_accelerate_endpoint: bool | None = None
     ) -> bool:
+        # accelerate cannot be used in China and us government
+        if self.location_type and "S3China" in self.location_type:
+            return False
         # if self.endpoint has been set, e.g. by metadata, no more config is needed.
         if self.endpoint is not None:
             return self.endpoint.find("s3-accelerate.amazonaws.com") >= 0

test/unit/test_s3_util.py

@@ -448,3 +448,36 @@ def test_download_retry_exceeded_error():
                 match=r"GET with url .* failed for exceeding maximum retries",
             ):
                 rest_client.download_chunk(0)
+
+
+def test_accelerate_in_china_endpoint():
+    meta_info = {
+        "name": "data1.txt.gz",
+        "stage_location_type": "S3",
+        "no_sleeping_time": True,
+        "put_callback": None,
+        "put_callback_output_stream": None,
+        SHA256_DIGEST: "123456789abcdef",
+        "dst_file_name": "data1.txt.gz",
+        "src_file_name": "path/to/put_get_1.txt",
+        "overwrite": True,
+    }
+    meta = SnowflakeFileMeta(**meta_info)
+    creds = {"AWS_SECRET_KEY": "", "AWS_KEY_ID": "", "AWS_TOKEN": ""}
+    rest_client = SnowflakeS3RestClient(
+        meta,
+        StorageCredential(
+            creds,
+            MagicMock(autospec=SnowflakeConnection),
+            "GET file:/tmp/file.txt @~",
+        ),
+        {
+            "locationType": "S3China",
+            "location": "bucket/path",
+            "creds": creds,
+            "region": "test",
+            "endPoint": None,
+        },
+        8 * megabyte,
+    )
+    assert not rest_client.transfer_accelerate_config()