SNOW-1789751: Pass GCS regional and virtual params (#2241)

Author: sfc-gh-pbulawa

DESCRIPTION.md

@@ -15,7 +15,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
   - Added handling of PAT provided in `password` field.
   - 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. See more: https://cloud.google.com/storage/docs/request-endpoints#xml-api
+  - 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.
 

src/snowflake/connector/file_transfer_agent.py

@@ -701,6 +701,7 @@ 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,6 +52,7 @@ 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.
 
@@ -85,6 +86,7 @@ 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
 
         if self.security_token:
             logger.debug(f"len(GCS_ACCESS_TOKEN): {len(self.security_token)}")
@@ -166,6 +168,8 @@ def generate_url_and_rest_args() -> (
                         if "region" not in self.stage_info
                         else self.stage_info["region"]
                     ),
+                    self.endpoint,
+                    self.use_virtual_endpoints,
                 )
                 access_token = self.security_token
             else:
@@ -204,6 +208,7 @@ def generate_url_and_rest_args() -> (
                         else self.stage_info["region"]
                     ),
                     self.endpoint,
+                    self.use_virtual_endpoints,
                 )
                 access_token = self.security_token
                 gcs_headers["Authorization"] = f"Bearer {access_token}"
@@ -368,6 +373,8 @@ def generate_url_and_authenticated_headers():
                         if "region" not in self.stage_info
                         else self.stage_info["region"]
                     ),
+                    self.endpoint,
+                    self.use_virtual_endpoints,
                 )
                 gcs_headers = {"Authorization": f"Bearer {self.security_token}"}
                 rest_args = {"headers": gcs_headers}