[ASYNC] Apply #2233 to async code

Author: sfc-gh-pczajka

src/snowflake/connector/aio/_azure_storage_client.py

@@ -49,7 +49,6 @@ def __init__(
         credentials: StorageCredential | None,
         chunk_size: int,
         stage_info: dict[str, Any],
-        use_s3_regional_url: bool = False,
         unsafe_file_write: bool = False,
     ) -> None:
         SnowflakeAzureRestClientSync.__init__(

src/snowflake/connector/aio/_cursor.py

@@ -663,6 +663,7 @@ async def execute(
                     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,
+                    gcs_use_virtual_endpoints=self._connection.gcs_use_virtual_endpoints,
                 )
                 await sf_file_transfer_agent.execute()
                 data = sf_file_transfer_agent.result()

src/snowflake/connector/aio/_file_transfer_agent.py

@@ -63,6 +63,7 @@ def __init__(
         source_from_stream: IO[bytes] | None = None,
         use_s3_regional_url: bool = False,
         unsafe_file_write: bool = False,
+        gcs_use_virtual_endpoints: bool = False,
     ) -> None:
         super().__init__(
             cursor=cursor,
@@ -82,6 +83,7 @@ def __init__(
             source_from_stream=source_from_stream,
             use_s3_regional_url=use_s3_regional_url,
             unsafe_file_write=unsafe_file_write,
+            gcs_use_virtual_endpoints=gcs_use_virtual_endpoints,
         )
 
     async def execute(self) -> None:
@@ -281,7 +283,6 @@ async def _create_file_transfer_client(
                 self._credentials,
                 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:
@@ -303,7 +304,6 @@ async def _create_file_transfer_client(
                 self._stage_info,
                 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:

src/snowflake/connector/aio/_gcs_storage_client.py

@@ -27,6 +27,7 @@
     GCS_METADATA_ENCRYPTIONDATAPROP,
     GCS_METADATA_MATDESC_KEY,
     GCS_METADATA_SFC_DIGEST,
+    GCS_REGION_ME_CENTRAL_2,
 )
 
 
@@ -38,7 +39,6 @@ def __init__(
         stage_info: dict[str, Any],
         cnx: SnowflakeConnection,
         command: str,
-        use_s3_regional_url: bool = False,
         unsafe_file_write: bool = False,
     ) -> None:
         """Creates a client object with given stage credentials.
@@ -65,6 +65,15 @@ def __init__(
         # presigned_url in meta is for downloading
         self.presigned_url: str = meta.presigned_url or stage_info.get("presignedUrl")
         self.security_token = credentials.creds.get("GCS_ACCESS_TOKEN")
+        self.use_regional_url = (
+            "region" in stage_info
+            and stage_info["region"].lower() == GCS_REGION_ME_CENTRAL_2
+            or "useRegionalUrl" in stage_info
+            and stage_info["useRegionalUrl"]
+        )
+        self.endpoint: str | None = (
+            None if "endPoint" not in stage_info else stage_info["endPoint"]
+        )
 
     async def _has_expired_token(self, response: aiohttp.ClientResponse) -> bool:
         return self.security_token and response.status == 401
@@ -73,7 +82,7 @@ async def _has_expired_presigned_url(
         self, response: aiohttp.ClientResponse
     ) -> bool:
         # Presigned urls can be generated for any xml-api operation
-        # offered by GCS. Hence the error codes expected are similar
+        # offered by GCS. Hence, the error codes expected are similar
         # to xml api.
         # https://cloud.google.com/storage/docs/xml-api/reference-status
 
@@ -132,7 +141,14 @@ def generate_url_and_rest_args() -> (
         ):
             if not self.presigned_url:
                 upload_url = self.generate_file_url(
-                    self.stage_info["location"], meta.dst_file_name.lstrip("/")
+                    self.stage_info["location"],
+                    meta.dst_file_name.lstrip("/"),
+                    self.use_regional_url,
+                    (
+                        None
+                        if "region" not in self.stage_info
+                        else self.stage_info["region"]
+                    ),
                 )
                 access_token = self.security_token
             else:
@@ -162,7 +178,15 @@ def generate_url_and_rest_args() -> (
             gcs_headers = {}
             if not self.presigned_url:
                 download_url = self.generate_file_url(
-                    self.stage_info["location"], meta.src_file_name.lstrip("/")
+                    self.stage_info["location"],
+                    meta.src_file_name.lstrip("/"),
+                    self.use_regional_url,
+                    (
+                        None
+                        if "region" not in self.stage_info
+                        else self.stage_info["region"]
+                    ),
+                    self.endpoint,
                 )
                 access_token = self.security_token
                 gcs_headers["Authorization"] = f"Bearer {access_token}"
@@ -279,7 +303,14 @@ async def get_file_header(self, filename: str) -> FileHeader | None:
 
             def generate_url_and_authenticated_headers():
                 url = self.generate_file_url(
-                    self.stage_info["location"], filename.lstrip("/")
+                    self.stage_info["location"],
+                    filename.lstrip("/"),
+                    self.use_regional_url,
+                    (
+                        None
+                        if "region" not in self.stage_info
+                        else self.stage_info["region"]
+                    ),
                 )
                 gcs_headers = {"Authorization": f"Bearer {self.security_token}"}
                 rest_args = {"headers": gcs_headers}

src/snowflake/connector/aio/_s3_storage_client.py

@@ -74,7 +74,13 @@ def __init__(
                 self.stage_info["location"]
             )
         )
-        self.use_s3_regional_url = use_s3_regional_url
+        self.use_s3_regional_url = (
+            use_s3_regional_url
+            or "useS3RegionalUrl" in stage_info
+            and stage_info["useS3RegionalUrl"]
+            or "useRegionalUrl" in stage_info
+            and stage_info["useRegionalUrl"]
+        )
         self.location_type = stage_info.get("locationType")
 
         # if GS sends us an endpoint, it's likely for FIPS. Use it.

test/integ/aio/test_connection_async.py

@@ -1686,3 +1686,34 @@ async def test_no_auth_connection_negative_case():
         await conn.execute_string("select 1")
 
     await conn.close()
+
+
+@pytest.mark.skipolddriver
+@pytest.mark.parametrize(
+    "value",
+    [
+        True,
+        False,
+    ],
+)
+async def test_gcs_use_virtual_endpoints(value):
+    with mock.patch(
+        "snowflake.connector.aio._network.SnowflakeRestful.fetch",
+        return_value={"data": {"token": None, "masterToken": None}, "success": True},
+    ):
+        cnx = snowflake.connector.aio.SnowflakeConnection(
+            user="test-user",
+            password="test-password",
+            host="test-host",
+            port="443",
+            account="test-account",
+            gcs_use_virtual_endpoints=value,
+        )
+        try:
+            await cnx.connect()
+            cnx.commit = cnx.rollback = (
+                lambda: None
+            )  # Skip tear down, there's only a mocked rest api
+            assert cnx.gcs_use_virtual_endpoints == value
+        finally:
+            await cnx.close()

test/unit/aio/test_gcs_client_async.py

@@ -330,7 +330,7 @@ async def test_get_file_header_none_with_presigned_url(tmp_path):
     )
     storage_credentials = Mock()
     storage_credentials.creds = {}
-    stage_info = Mock()
+    stage_info: dict[str, any] = dict()
     connection = Mock()
     client = SnowflakeGCSRestClient(
         meta, storage_credentials, stage_info, connection, ""
@@ -339,3 +339,102 @@ async def test_get_file_header_none_with_presigned_url(tmp_path):
         await client._update_presigned_url()
     file_header = await client.get_file_header(meta.name)
     assert file_header is None
+
+
+@pytest.mark.parametrize(
+    "region,return_url,use_regional_url,endpoint,gcs_use_virtual_endpoints",
+    [
+        (
+            "US-CENTRAL1",
+            "https://storage.us-central1.rep.googleapis.com",
+            True,
+            None,
+            False,
+        ),
+        (
+            "ME-CENTRAL2",
+            "https://storage.me-central2.rep.googleapis.com",
+            True,
+            None,
+            False,
+        ),
+        ("US-CENTRAL1", "https://storage.googleapis.com", False, None, False),
+        ("US-CENTRAL1", "https://storage.googleapis.com", False, None, False),
+        ("US-CENTRAL1", "https://location.storage.googleapis.com", False, None, True),
+        ("US-CENTRAL1", "https://location.storage.googleapis.com", True, None, True),
+        (
+            "US-CENTRAL1",
+            "https://overriddenurl.com",
+            False,
+            "https://overriddenurl.com",
+            False,
+        ),
+        (
+            "US-CENTRAL1",
+            "https://overriddenurl.com",
+            True,
+            "https://overriddenurl.com",
+            False,
+        ),
+        (
+            "US-CENTRAL1",
+            "https://overriddenurl.com",
+            True,
+            "https://overriddenurl.com",
+            True,
+        ),
+        (
+            "US-CENTRAL1",
+            "https://overriddenurl.com",
+            False,
+            "https://overriddenurl.com",
+            False,
+        ),
+        (
+            "US-CENTRAL1",
+            "https://overriddenurl.com",
+            False,
+            "https://overriddenurl.com",
+            True,
+        ),
+    ],
+)
+def test_url(region, return_url, use_regional_url, endpoint, gcs_use_virtual_endpoints):
+    gcs_location = SnowflakeGCSRestClient.get_location(
+        stage_location="location",
+        use_regional_url=use_regional_url,
+        region=region,
+        endpoint=endpoint,
+        use_virtual_endpoints=gcs_use_virtual_endpoints,
+    )
+    assert gcs_location.endpoint == return_url
+
+
+@pytest.mark.parametrize(
+    "region,use_regional_url,return_value",
+    [
+        ("ME-CENTRAL2", False, True),
+        ("ME-CENTRAL2", True, True),
+        ("US-CENTRAL1", False, False),
+        ("US-CENTRAL1", True, True),
+    ],
+)
+def test_use_regional_url(region, use_regional_url, return_value):
+    meta = SnowflakeFileMeta(
+        name="path/some_file",
+        src_file_name="path/some_file",
+        stage_location_type="GCS",
+        presigned_url="www.example.com",
+    )
+    storage_credentials = Mock()
+    storage_credentials.creds = {}
+    stage_info: dict[str, any] = dict()
+    stage_info["region"] = region
+    stage_info["useRegionalUrl"] = use_regional_url
+    connection = Mock()
+
+    client = SnowflakeGCSRestClient(
+        meta, storage_credentials, stage_info, connection, ""
+    )
+
+    assert client.use_regional_url == return_value

test/unit/aio/test_s3_util_async.py

@@ -29,14 +29,11 @@
         SnowflakeFileMeta,
         StorageCredential,
     )
-    from snowflake.connector.s3_storage_client import ERRORNO_WSAECONNABORTED
     from snowflake.connector.vendored.requests import HTTPError
 except ImportError:
     # Compatibility for olddriver tests
     from requests import HTTPError
 
-    from snowflake.connector.s3_util import ERRORNO_WSAECONNABORTED  # NOQA
-
     SnowflakeFileMeta = dict
     SnowflakeS3RestClient = None
     RequestExceedMaxRetryError = None
@@ -500,3 +497,46 @@ async def test_accelerate_in_china_endpoint():
         8 * megabyte,
     )
     assert not await rest_client.transfer_accelerate_config()
+
+
+@pytest.mark.parametrize(
+    "use_s3_regional_url,stage_info_flags,expected",
+    [
+        (False, {}, False),
+        (True, {}, True),
+        (False, {"useS3RegionalUrl": True}, True),
+        (False, {"useRegionalUrl": True}, True),
+        (True, {"useS3RegionalUrl": False}, True),
+        (False, {"useS3RegionalUrl": True, "useRegionalUrl": False}, True),
+        (False, {"useS3RegionalUrl": False, "useRegionalUrl": True}, True),
+        (False, {"useS3RegionalUrl": False, "useRegionalUrl": False}, False),
+    ],
+)
+def test_s3_regional_url_logic_async(use_s3_regional_url, stage_info_flags, expected):
+    """Tests that the async S3 storage client correctly handles regional URL flags from stage_info."""
+    if SnowflakeS3RestClient is None:
+        pytest.skip("S3 storage client not available")
+
+    meta = SnowflakeFileMeta(
+        name="path/some_file",
+        src_file_name="path/some_file",
+        stage_location_type="S3",
+    )
+    storage_credentials = StorageCredential({}, mock.Mock(), "test")
+
+    stage_info = {
+        "region": "us-west-2",
+        "location": "test-bucket",
+        "endPoint": None,
+    }
+    stage_info.update(stage_info_flags)
+
+    client = SnowflakeS3RestClient(
+        meta=meta,
+        credentials=storage_credentials,
+        stage_info=stage_info,
+        chunk_size=1024,
+        use_s3_regional_url=use_s3_regional_url,
+    )
+
+    assert client.use_s3_regional_url == expected