SNOW-1657037: Pass Azure SAS Token via params

src/snowflake/connector/azure_storage_client.py

@@ -12,6 +12,7 @@
 from random import choice
 from string import hexdigits
 from typing import TYPE_CHECKING, Any, NamedTuple
+from urllib.parse import parse_qsl
 
 from .compat import quote
 from .constants import FileHeader, ResultStatus
@@ -95,12 +96,9 @@ def generate_authenticated_url_and_rest_args() -> tuple[bytes, dict[str, Any]]:
             sas_token = self.credentials.creds["AZURE_SAS_TOKEN"]
             if sas_token and sas_token.startswith("?"):
                 sas_token = sas_token[1:]
-            if "?" in url:
-                _url = url + "&" + sas_token
-            else:
-                _url = url + "?" + sas_token
+            params = {k: v for k, v in parse_qsl(sas_token)}
             headers["Date"] = timestamp
-            rest_args = {"headers": headers}
+            rest_args = {"headers": headers, "params": params}
             if data:
                 rest_args["data"] = data
             return _url, rest_args

src/snowflake/connector/network.py

@@ -451,6 +451,7 @@ def request(
         method: str = "post",
         client: str = "sfsql",
         timeout: int | None = None,
+        params: dict | None = None,
         _no_results: bool = False,
         _include_retry_params: bool = False,
         _no_retry: bool = False,
@@ -493,6 +494,7 @@ def request(
                 headers,
                 json.dumps(body),
                 token=self.token,
+                params=params,
                 _no_results=_no_results,
                 timeout=timeout,
                 _include_retry_params=_include_retry_params,
@@ -503,6 +505,7 @@ def request(
                 url,
                 headers,
                 token=self.token,
+                params=params,
                 timeout=timeout,
             )
 
@@ -678,6 +681,7 @@ def _get_request(
         url: str,
         headers: dict[str, str],
         token: str = None,
+        params: dict | None = None,
         timeout: int | None = None,
         is_fetch_query_status: bool = False,
     ) -> dict[str, Any]:
@@ -693,6 +697,7 @@ def _get_request(
             headers,
             timeout=timeout,
             token=token,
+            params=params,
             is_fetch_query_status=is_fetch_query_status,
         )
         if ret.get("code") == SESSION_EXPIRED_GS_CODE:
@@ -712,6 +717,7 @@ def _get_request(
                     url,
                     headers,
                     token=self.token,
+                    params=params,
                     is_fetch_query_status=is_fetch_query_status,
                 )
 
@@ -723,6 +729,7 @@ def _post_request(
         headers,
         body,
         token=None,
+        params: dict | None = None,
         timeout: int | None = None,
         socket_timeout: int | None = None,
         _no_results: bool = False,
@@ -743,6 +750,7 @@ def _post_request(
             data=body,
             timeout=timeout,
             token=token,
+            params=params,
             no_retry=no_retry,
             _include_retry_params=_include_retry_params,
             socket_timeout=socket_timeout,
@@ -769,7 +777,7 @@ def _post_request(
             )
             if ret.get("success"):
                 return self._post_request(
-                    url, headers, body, token=self.token, timeout=timeout
+                    url, headers, body, token=self.token, params=params, timeout=timeout
                 )
 
         if isinstance(ret.get("data"), dict) and ret["data"].get("queryId"):
@@ -789,6 +797,7 @@ def _post_request(
                 result_url,
                 headers,
                 token=self.token,
+                params=params,
                 timeout=timeout,
                 is_fetch_query_status=bool(
                     re.match(r"^/queries/.+/result$", result_url)
@@ -880,6 +889,7 @@ def _request_exec_wrapper(
         retry_ctx,
         no_retry: bool = False,
         token=NO_TOKEN,
+        params: dict | None = None,
         **kwargs,
     ):
         conn = self._connection
@@ -908,6 +918,7 @@ def _request_exec_wrapper(
                 headers=headers,
                 data=data,
                 token=token,
+                params=params,
                 raise_raw_http_failure=raise_raw_http_failure,
                 **kwargs,
             )
@@ -1041,6 +1052,7 @@ def _request_exec(
         headers,
         data,
         token,
+        params,
         catch_okta_unauthorized_error: bool = False,
         is_raw_text: bool = False,
         is_raw_binary: bool = False,
@@ -1077,6 +1089,7 @@ def _request_exec(
                 verify=True,
                 stream=is_raw_binary,
                 auth=SnowflakeAuth(token),
+                params=params,
             )
             download_end_time = get_time_millis()