SNOW-1789751: Add GCP regional and virtual endpoints support (#2233)

Author: sfc-gh-pbulawa

src/snowflake/connector/azure_storage_client.py

@@ -64,7 +64,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:
         super().__init__(

src/snowflake/connector/connection.py

@@ -312,6 +312,10 @@ def _get_private_bytes_from_file(
         None,
         (type(None), int),
     ),  # SNOW-1817982: limit iobound TPE sizes when executing PUT/GET
+    "gcs_use_virtual_endpoints": (
+        False,
+        bool,
+    ),  # use https://{bucket}.storage.googleapis.com instead of https://storage.googleapis.com/{bucket}
     "unsafe_file_write": (
         False,
         bool,
@@ -395,6 +399,7 @@ class SnowflakeConnection:
           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.
+        gcs_use_virtual_endpoints: When true, the virtual endpoint url is used, see: https://cloud.google.com/storage/docs/request-endpoints#xml-api
     """
 
     OCSP_ENV_LOCK = Lock()
@@ -783,6 +788,14 @@ def unsafe_file_write(self) -> bool:
     def unsafe_file_write(self, value: bool) -> None:
         self._unsafe_file_write = value
 
+    @property
+    def gcs_use_virtual_endpoints(self) -> bool:
+        return self._gcs_use_virtual_endpoints
+
+    @gcs_use_virtual_endpoints.setter
+    def gcs_use_virtual_endpoints(self, value: bool) -> None:
+        self._gcs_use_virtual_endpoints = value
+
     def connect(self, **kwargs) -> None:
         """Establishes connection to Snowflake."""
         logger.debug("connect")

src/snowflake/connector/cursor.py

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

src/snowflake/connector/file_transfer_agent.py

@@ -356,6 +356,7 @@ def __init__(
         use_s3_regional_url: bool = False,
         iobound_tpe_limit: int | None = None,
         unsafe_file_write: bool = False,
+        gcs_use_virtual_endpoints: bool = False,
     ) -> None:
         self._cursor = cursor
         self._command = command
@@ -388,6 +389,7 @@ def __init__(
         self._credentials: StorageCredential | None = None
         self._iobound_tpe_limit = iobound_tpe_limit
         self._unsafe_file_write = unsafe_file_write
+        self._gcs_use_virtual_endpoints = gcs_use_virtual_endpoints
 
     def execute(self) -> None:
         self._parse_command()
@@ -683,7 +685,6 @@ 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:
@@ -703,7 +704,6 @@ 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,
             )
         raise Exception(f"{self._stage_location_type} is an unknown stage type")

src/snowflake/connector/gcs_storage_client.py

@@ -36,13 +36,15 @@
 GCS_FILE_HEADER_DIGEST = "gcs-file-header-digest"
 GCS_FILE_HEADER_CONTENT_LENGTH = "gcs-file-header-content-length"
 GCS_FILE_HEADER_ENCRYPTION_METADATA = "gcs-file-header-encryption-metadata"
+GCS_REGION_ME_CENTRAL_2 = "me-central2"
 CONTENT_CHUNK_SIZE = 10 * kilobyte
 ACCESS_TOKEN = "GCS_ACCESS_TOKEN"
 
 
 class GcsLocation(NamedTuple):
     bucket_name: str
     path: str
+    endpoint: str = "https://storage.googleapis.com"
 
 
 class SnowflakeGCSRestClient(SnowflakeStorageClient):
@@ -53,7 +55,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.
@@ -79,6 +80,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"]
+        )
 
         if self.security_token:
             logger.debug(f"len(GCS_ACCESS_TOKEN): {len(self.security_token)}")
@@ -91,7 +101,7 @@ def _has_expired_token(self, response: requests.Response) -> bool:
 
     def _has_expired_presigned_url(self, response: requests.Response) -> 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
 
@@ -152,7 +162,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:
@@ -182,7 +199,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}"
@@ -339,7 +364,14 @@ 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}
@@ -383,7 +415,13 @@ def generate_url_and_authenticated_headers():
             return None
 
     @staticmethod
-    def extract_bucket_name_and_path(stage_location: str) -> GcsLocation:
+    def get_location(
+        stage_location: str,
+        use_regional_url: str = False,
+        region: str = None,
+        endpoint: str = None,
+        use_virtual_endpoints: bool = False,
+    ) -> GcsLocation:
         container_name = stage_location
         path = ""
 
@@ -393,13 +431,40 @@ def extract_bucket_name_and_path(stage_location: str) -> GcsLocation:
             path = stage_location[stage_location.index("/") + 1 :]
             if path and not path.endswith("/"):
                 path += "/"
-
-        return GcsLocation(bucket_name=container_name, path=path)
+        if endpoint:
+            if endpoint.endswith("/"):
+                endpoint = endpoint[:-1]
+            return GcsLocation(bucket_name=container_name, path=path, endpoint=endpoint)
+        elif use_virtual_endpoints:
+            return GcsLocation(
+                bucket_name=container_name,
+                path=path,
+                endpoint=f"https://{container_name}.storage.googleapis.com",
+            )
+        elif use_regional_url:
+            return GcsLocation(
+                bucket_name=container_name,
+                path=path,
+                endpoint=f"https://storage.{region.lower()}.rep.googleapis.com",
+            )
+        else:
+            return GcsLocation(bucket_name=container_name, path=path)
 
     @staticmethod
-    def generate_file_url(stage_location: str, filename: str) -> str:
-        gcs_location = SnowflakeGCSRestClient.extract_bucket_name_and_path(
-            stage_location
+    def generate_file_url(
+        stage_location: str,
+        filename: str,
+        use_regional_url: str = False,
+        region: str = None,
+        endpoint: str = None,
+        use_virtual_endpoints: bool = False,
+    ) -> str:
+        gcs_location = SnowflakeGCSRestClient.get_location(
+            stage_location, use_regional_url, region, endpoint
         )
         full_file_path = f"{gcs_location.path}{filename}"
-        return f"https://storage.googleapis.com/{gcs_location.bucket_name}/{quote(full_file_path)}"
+
+        if use_virtual_endpoints:
+            return f"{gcs_location.endpoint}/{quote(full_file_path)}"
+        else:
+            return f"{gcs_location.endpoint}/{gcs_location.bucket_name}/{quote(full_file_path)}"

src/snowflake/connector/s3_storage_client.py

@@ -86,7 +86,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/test_connection.py

@@ -1369,6 +1369,34 @@ def test_server_session_keep_alive(conn_cnx):
     mock_delete_session.assert_called_once()
 
 
+@pytest.mark.skipolddriver
+@pytest.mark.parametrize(
+    "value",
+    [
+        True,
+        False,
+    ],
+)
+def test_gcs_use_virtual_endpoints(conn_cnx, value):
+    with mock.patch(
+        "snowflake.connector.network.SnowflakeRestful.fetch",
+        return_value={"data": {"token": None, "masterToken": None}, "success": True},
+    ):
+        with snowflake.connector.connect(
+            user="test-user",
+            password="test-password",
+            host="test-host",
+            port="443",
+            account="test-account",
+            gcs_use_virtual_endpoints=value,
+        ) as cnx:
+            assert cnx
+            cnx.commit = cnx.rollback = (
+                lambda: None
+            )  # Skip tear down, there's only a mocked rest api
+            assert cnx.gcs_use_virtual_endpoints == value
+
+
 @pytest.mark.skipolddriver
 def test_ocsp_mode_disable_ocsp_checks(
     conn_cnx, is_public_test, is_local_dev_setup, caplog