SNOW-2032699: Use GCS virtual url based on the stage response (#2274)

Author: sfc-gh-pbulawa

DESCRIPTION.md

@@ -16,14 +16,14 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
   - Added experimental support for OAuth authorization code and client credentials flows.
   - Improved error message for client-side query cancellations due to timeouts.
   - Added support of GCS regional endpoints.
-  - Added `gcs_use_virtual_endpoints` connection property that forces the usage of the virtual GCS usage. Thanks to this it should be possible to set up private DNS entry for the GCS endpoint. See more: https://cloud.google.com/storage/docs/request-endpoints#xml-api
   - Fixed a bug that caused driver to fail silently on `TO_DATE` arrow to python conversion when invalid date was followed by the correct one.
   - Added `check_arrow_conversion_error_on_every_column` connection property that can be set to `False` to restore previous behaviour in which driver will ignore errors until it occurs in the last column. This flag's purpose is to unblock workflows that may be impacted by the bugfix and will be removed in later releases.
   - Lower log levels from info to debug for some of the messages to make the output easier to follow.
   - Allow the connector to inherit a UUID4 generated upstream, provided in statement parameters (field: `requestId`), rather than automatically generate a UUID4 to use for the HTTP Request ID.
   - Improved logging in urllib3, boto3, botocore - assured data masking even after migration to the external owned library in the future.
   - Fix expired S3 credentials update and increment retry when expired credentials are found.
   - Added `client_fetch_threads` experimental parameter to better utilize threads for fetching query results.
+  - Added support of GCS virtual urls. See more: https://cloud.google.com/storage/docs/request-endpoints#xml-api
 
 - v3.14.0(March 03, 2025)
   - Bumped pyOpenSSL dependency upper boundary from <25.0.0 to <26.0.0.

src/snowflake/connector/connection.py

@@ -320,10 +320,6 @@ 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}
     "oauth_client_id": (
         None,
         (type(None), str),
@@ -439,7 +435,6 @@ 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
         check_arrow_conversion_error_on_every_column: When true, the error check after the conversion from arrow to python types will happen for every column in the row. This is a new behaviour which fixes the bug that caused the type errors to trigger silently when occurring at any place other than last column in a row. To revert the previous (faulty) behaviour, please set this flag to false.
     """
 
@@ -858,14 +853,6 @@ def oauth_security_features(self) -> _OAuthSecurityFeatures:
             refresh_token_enabled="refresh_token" in features,
         )
 
-    @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
-
     @property
     def check_arrow_conversion_error_on_every_column(self) -> bool:
         return self._check_arrow_conversion_error_on_every_column

src/snowflake/connector/cursor.py

@@ -1083,7 +1083,6 @@ 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

@@ -355,7 +355,6 @@ 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,7 +387,6 @@ 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()
@@ -704,7 +702,6 @@ def _create_file_transfer_client(
                 self._cursor._connection,
                 self._command,
                 unsafe_file_write=self._unsafe_file_write,
-                use_virtual_endpoints=self._gcs_use_virtual_endpoints,
             )
         raise Exception(f"{self._stage_location_type} is an unknown stage type")
 

src/snowflake/connector/gcs_storage_client.py

@@ -52,7 +52,6 @@ def __init__(
         cnx: SnowflakeConnection,
         command: str,
         unsafe_file_write: bool = False,
-        use_virtual_endpoints: bool = False,
     ) -> None:
         """Creates a client object with given stage credentials.
 
@@ -86,7 +85,9 @@ def __init__(
         self.endpoint: str | None = (
             None if "endPoint" not in stage_info else stage_info["endPoint"]
         )
-        self.use_virtual_endpoints: bool = use_virtual_endpoints
+        self.use_virtual_url: bool = (
+            "useVirtualUrl" in stage_info and stage_info["useVirtualUrl"]
+        )
 
         if self.security_token:
             logger.debug(f"len(GCS_ACCESS_TOKEN): {len(self.security_token)}")
@@ -169,7 +170,7 @@ def generate_url_and_rest_args() -> (
                         else self.stage_info["region"]
                     ),
                     self.endpoint,
-                    self.use_virtual_endpoints,
+                    self.use_virtual_url,
                 )
                 access_token = self.security_token
             else:
@@ -208,7 +209,7 @@ def generate_url_and_rest_args() -> (
                         else self.stage_info["region"]
                     ),
                     self.endpoint,
-                    self.use_virtual_endpoints,
+                    self.use_virtual_url,
                 )
                 access_token = self.security_token
                 gcs_headers["Authorization"] = f"Bearer {access_token}"
@@ -374,7 +375,7 @@ def generate_url_and_authenticated_headers():
                         else self.stage_info["region"]
                     ),
                     self.endpoint,
-                    self.use_virtual_endpoints,
+                    self.use_virtual_url,
                 )
                 gcs_headers = {"Authorization": f"Bearer {self.security_token}"}
                 rest_args = {"headers": gcs_headers}
@@ -423,7 +424,7 @@ def get_location(
         use_regional_url: str = False,
         region: str = None,
         endpoint: str = None,
-        use_virtual_endpoints: bool = False,
+        use_virtual_url: bool = False,
     ) -> GcsLocation:
         container_name = stage_location
         path = ""
@@ -438,7 +439,7 @@ def get_location(
             if endpoint.endswith("/"):
                 endpoint = endpoint[:-1]
             return GcsLocation(bucket_name=container_name, path=path, endpoint=endpoint)
-        elif use_virtual_endpoints:
+        elif use_virtual_url:
             return GcsLocation(
                 bucket_name=container_name,
                 path=path,
@@ -460,14 +461,14 @@ def generate_file_url(
         use_regional_url: str = False,
         region: str = None,
         endpoint: str = None,
-        use_virtual_endpoints: bool = False,
+        use_virtual_url: bool = False,
     ) -> str:
         gcs_location = SnowflakeGCSRestClient.get_location(
             stage_location, use_regional_url, region, endpoint
         )
         full_file_path = f"{gcs_location.path}{filename}"
 
-        if use_virtual_endpoints:
+        if use_virtual_url:
             return f"{gcs_location.endpoint}/{quote(full_file_path)}"
         else:
             return f"{gcs_location.endpoint}/{gcs_location.bucket_name}/{quote(full_file_path)}"

test/integ/test_connection.py

@@ -1375,34 +1375,6 @@ 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, caplog):
     caplog.set_level(logging.DEBUG, "snowflake.connector.ocsp_snowflake")

test/unit/test_gcs_client.py

@@ -350,7 +350,7 @@ def test_get_file_header_none_with_presigned_url(tmp_path):
 
 
 @pytest.mark.parametrize(
-    "region,return_url,use_regional_url,endpoint,gcs_use_virtual_endpoints",
+    "region,return_url,use_regional_url,endpoint,use_virtual_url",
     [
         (
             "US-CENTRAL1",
@@ -407,13 +407,13 @@ def test_get_file_header_none_with_presigned_url(tmp_path):
         ),
     ],
 )
-def test_url(region, return_url, use_regional_url, endpoint, gcs_use_virtual_endpoints):
+def test_url(region, return_url, use_regional_url, endpoint, use_virtual_url):
     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,
+        use_virtual_url=use_virtual_url,
     )
     assert gcs_location.endpoint == return_url
 
@@ -446,3 +446,28 @@ def test_use_regional_url(region, use_regional_url, return_value):
     )
 
     assert client.use_regional_url == return_value
+
+
+@pytest.mark.parametrize(
+    "use_virtual_url,return_value",
+    [(False, False), (True, True), (None, False)],
+)
+def test_stage_info_use_virtual_url(use_virtual_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()
+    if use_virtual_url is not None:
+        stage_info["useVirtualUrl"] = use_virtual_url
+    connection = Mock()
+
+    client = SnowflakeGCSRestClient(
+        meta, storage_credentials, stage_info, connection, ""
+    )
+
+    assert client.use_virtual_url == return_value