Merge branch 'main' into SNOW-2062305-process-pool-batch-fetcher

Author: sfc-gh-mmishchenko

DESCRIPTION.md

@@ -11,9 +11,9 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
   - Bumped numpy dependency from <2.1.0 to <=2.2.4
   - Added Windows support for Python 3.13.
   - Add `bulk_upload_chunks` parameter to `write_pandas` function. Setting this parameter to True changes the behaviour of write_pandas function to first write all the data chunks to the local disk and then perform the wildcard upload of the chunks folder to the stage. In default behaviour the chunks are being saved, uploaded and deleted one by one.
+  - Added support for new authentication mechanism PAT with external session ID.
   - Added `client_fetch_use_mp` parameter that enables multiprocessed fetching of result batches.
 
-
 - v3.15.1(May 20, 2025)
   - Added basic arrow support for Interval types.
   - Fix `write_pandas` special characters usage in the location name.

prober/Dockerfile

@@ -43,7 +43,7 @@ ENV PATH="${PYENV_ROOT}/shims:${PYENV_ROOT}/bin:${PATH}"
 RUN git clone --depth=1 https://github.com/pyenv/pyenv.git ${PYENV_ROOT}
 
 # Build arguments for Python versions and Snowflake connector versions
-ARG MATRIX_VERSION='{"3.13.4": ["3.15.0", "3.13.2", "3.14.0", "3.12.3", "3.0.4", "3.12.1", "3.12.4", "3.11.0", "3.12.2", "3.6.0", "3.7.0"], "3.9.22": ["3.15.0", "3.13.2", "3.14.0", "3.12.3", "3.0.4", "3.12.1", "3.12.4", "3.11.0", "3.12.2", "3.6.0", "3.7.0"]}'
+ARG MATRIX_VERSION='{"3.13.4": ["3.15.0", "3.13.2", "3.14.0", "3.12.3", "3.12.1", "3.12.4", "3.11.0", "3.12.2", "3.6.0", "3.7.0"], "3.9.22": ["3.15.0", "3.13.2", "3.14.0", "3.12.3", "3.12.1", "3.12.4", "3.11.0", "3.12.2", "3.6.0", "3.7.0"]}'
 
 
 # Install Python versions from ARG MATRIX_VERSION

prober/testing_matrix.json

@@ -2,11 +2,11 @@
   "python-version": [
     {
       "version": "3.13.4",
-      "snowflake-connector-python": ["3.13.2", "3.14.0", "3.12.3", "3.0.4", "3.12.1", "3.12.4", "3.11.0", "3.12.2", "3.6.0", "3.7.0"]
+      "snowflake-connector-python": ["3.15.0" ,"3.13.2", "3.14.0", "3.12.3", "3.12.1", "3.12.4", "3.11.0", "3.12.2", "3.6.0", "3.7.0"]
     },
     {
       "version": "3.9.22",
-      "snowflake-connector-python": ["3.13.2", "3.14.0", "3.12.3", "3.0.4", "3.12.1", "3.12.4", "3.11.0", "3.12.2", "3.6.0", "3.7.0"]
+      "snowflake-connector-python": ["3.15.0", "3.13.2", "3.14.0", "3.12.3", "3.12.1", "3.12.4", "3.11.0", "3.12.2", "3.6.0", "3.7.0"]
     }
   ]
 }

src/snowflake/connector/auth/by_plugin.py

@@ -53,6 +53,7 @@ class AuthType(Enum):
     PAT = "PROGRAMMATIC_ACCESS_TOKEN'"
     NO_AUTH = "NO_AUTH"
     WORKLOAD_IDENTITY = "WORKLOAD_IDENTITY"
+    PAT_WITH_EXTERNAL_SESSION = "PAT_WITH_EXTERNAL_SESSION"
 
 
 class AuthByPlugin(ABC):

src/snowflake/connector/connection.py

@@ -107,6 +107,7 @@
     OAUTH_AUTHENTICATOR,
     OAUTH_AUTHORIZATION_CODE,
     OAUTH_CLIENT_CREDENTIALS,
+    PAT_WITH_EXTERNAL_SESSION,
     PROGRAMMATIC_ACCESS_TOKEN,
     REQUEST_ID,
     USR_PWD_MFA_AUTHENTICATOR,
@@ -363,6 +364,11 @@ def _get_private_bytes_from_file(
         True,
         bool,
     ),  # SNOW-XXXXX: remove the check_arrow_conversion_error_on_every_column flag
+    "external_session_id": (
+        None,
+        str,
+        # SNOW-2096721: External (Spark) session ID
+    ),
 }
 
 APPLICATION_RE = re.compile(r"[\w\d_]+")
@@ -1272,6 +1278,15 @@ def __open_connection(self):
                 )
             elif self._authenticator == PROGRAMMATIC_ACCESS_TOKEN:
                 self.auth_class = AuthByPAT(self._token)
+            elif self._authenticator == PAT_WITH_EXTERNAL_SESSION:
+                # We don't need to do a POST to /v1/login-request to get session and master tokens at the startup
+                # time. PAT with external (Spark) session ID creates a new session when it encounters the unique
+                # (PAT, external session ID) combination for the first time and then onwards use the (PAT, external
+                # session id) as a key to identify and authenticate the session. So we bypass actual AuthN here.
+                self.auth_class = AuthNoAuth()
+                self._rest.set_pat_and_external_session(
+                    self._token, self._external_session_id
+                )
             elif self._authenticator == WORKLOAD_IDENTITY_AUTHENTICATOR:
                 self._check_experimental_authentication_flag()
                 # Standardize the provider enum.
@@ -1411,6 +1426,7 @@ def __config(self, **kwargs):
                 OAUTH_AUTHENTICATOR,
                 USR_PWD_MFA_AUTHENTICATOR,
                 WORKLOAD_IDENTITY_AUTHENTICATOR,
+                PAT_WITH_EXTERNAL_SESSION,
             ]:
                 self._authenticator = auth_tmp
 
@@ -1426,6 +1442,7 @@ def __config(self, **kwargs):
             NO_AUTH_AUTHENTICATOR,
             WORKLOAD_IDENTITY_AUTHENTICATOR,
             PROGRAMMATIC_ACCESS_TOKEN,
+            PAT_WITH_EXTERNAL_SESSION,
         }
 
         if not (self._master_token and self._session_token):
@@ -1474,6 +1491,7 @@ def __config(self, **kwargs):
                     KEY_PAIR_AUTHENTICATOR,
                     PROGRAMMATIC_ACCESS_TOKEN,
                     WORKLOAD_IDENTITY_AUTHENTICATOR,
+                    PAT_WITH_EXTERNAL_SESSION,
                 )
                 and not self._password
             ):

src/snowflake/connector/network.py

@@ -148,12 +148,12 @@
 
 HEADER_AUTHORIZATION_KEY = "Authorization"
 HEADER_SNOWFLAKE_TOKEN = 'Snowflake Token="{token}"'
+HEADER_EXTERNAL_SESSION_KEY = "X-Snowflake-External-Session-ID"
 
 REQUEST_ID = "requestId"
 REQUEST_GUID = "request_guid"
 SNOWFLAKE_HOST_SUFFIX = ".snowflakecomputing.com"
 
-
 SNOWFLAKE_CONNECTOR_VERSION = SNOWFLAKE_CONNECTOR_VERSION
 PYTHON_VERSION = PYTHON_VERSION
 OPERATING_SYSTEM = OPERATING_SYSTEM
@@ -166,6 +166,7 @@
 PYTHON_CONNECTOR_USER_AGENT = f"{CLIENT_NAME}/{SNOWFLAKE_CONNECTOR_VERSION} ({PLATFORM}) {IMPLEMENTATION}/{PYTHON_VERSION}"
 
 NO_TOKEN = "no-token"
+NO_EXTERNAL_SESSION_ID = "no-external-session-id"
 
 STATUS_TO_EXCEPTION: dict[int, type[Error]] = {
     INTERNAL_SERVER_ERROR: InternalServerError,
@@ -189,6 +190,7 @@
 PROGRAMMATIC_ACCESS_TOKEN = "PROGRAMMATIC_ACCESS_TOKEN"
 NO_AUTH_AUTHENTICATOR = "NO_AUTH"
 WORKLOAD_IDENTITY_AUTHENTICATOR = "WORKLOAD_IDENTITY"
+PAT_WITH_EXTERNAL_SESSION = "PAT_WITH_EXTERNAL_SESSION"
 
 
 def is_retryable_http_code(code: int) -> bool:
@@ -313,6 +315,25 @@ def __call__(self, r: PreparedRequest) -> PreparedRequest:
         return r
 
 
+class PATWithExternalSessionAuth(AuthBase):
+    """Attaches HTTP Authorization headers for PAT with External Session."""
+
+    def __init__(self, token, external_session_id) -> None:
+        # setup any auth-related data here
+        self.token = token
+        self.external_session_id = external_session_id
+
+    def __call__(self, r: PreparedRequest) -> PreparedRequest:
+        """Modifies and returns the request."""
+        if HEADER_AUTHORIZATION_KEY in r.headers:
+            del r.headers[HEADER_AUTHORIZATION_KEY]
+        if self.token != NO_TOKEN:
+            r.headers[HEADER_AUTHORIZATION_KEY] = "Bearer " + self.token
+        if self.external_session_id != NO_EXTERNAL_SESSION_ID:
+            r.headers[HEADER_EXTERNAL_SESSION_KEY] = self.external_session_id
+        return r
+
+
 class SessionPool:
     def __init__(self, rest: SnowflakeRestful) -> None:
         # A stack of the idle sessions
@@ -404,6 +425,12 @@ def __init__(
     def token(self) -> str | None:
         return self._token if hasattr(self, "_token") else None
 
+    @property
+    def external_session_id(self) -> str | None:
+        return (
+            self._external_session_id if hasattr(self, "_external_session_id") else None
+        )
+
     @property
     def master_token(self) -> str | None:
         return self._master_token if hasattr(self, "_master_token") else None
@@ -513,6 +540,7 @@ def request(
                 headers,
                 json.dumps(body, cls=SnowflakeRestfulJsonEncoder),
                 token=self.token,
+                external_session_id=self.external_session_id,
                 _no_results=_no_results,
                 timeout=timeout,
                 _include_retry_params=_include_retry_params,
@@ -523,6 +551,7 @@ def request(
                 url,
                 headers,
                 token=self.token,
+                external_session_id=self.external_session_id,
                 timeout=timeout,
             )
 
@@ -542,6 +571,17 @@ def update_tokens(
             self._mfa_token = mfa_token
             self._master_validity_in_seconds = master_validity_in_seconds
 
+    def set_pat_and_external_session(
+        self,
+        personal_access_token,
+        external_session_id,
+    ) -> None:
+        """Updates session and master tokens and optionally temporary credential."""
+        with self._lock_token:
+            self._personal_access_token = personal_access_token
+            self._token = personal_access_token
+            self._external_session_id = external_session_id
+
     def _renew_session(self):
         """Renew a session and master token."""
         return self._token_request(REQUEST_TYPE_RENEW)
@@ -698,6 +738,7 @@ def _get_request(
         url: str,
         headers: dict[str, str],
         token: str = None,
+        external_session_id: str = None,
         timeout: int | None = None,
         is_fetch_query_status: bool = False,
     ) -> dict[str, Any]:
@@ -713,9 +754,13 @@ def _get_request(
             headers,
             timeout=timeout,
             token=token,
+            external_session_id=external_session_id,
             is_fetch_query_status=is_fetch_query_status,
         )
-        if ret.get("code") == SESSION_EXPIRED_GS_CODE:
+        if (
+            ret.get("code") == SESSION_EXPIRED_GS_CODE
+            and self._connection._authenticator != PAT_WITH_EXTERNAL_SESSION
+        ):
             try:
                 ret = self._renew_session()
             except ReauthenticationRequest as ex:
@@ -743,6 +788,7 @@ def _post_request(
         headers,
         body,
         token=None,
+        external_session_id: str | None = None,
         timeout: int | None = None,
         socket_timeout: int | None = None,
         _no_results: bool = False,
@@ -763,6 +809,7 @@ def _post_request(
             data=body,
             timeout=timeout,
             token=token,
+            external_session_id=external_session_id,
             no_retry=no_retry,
             _include_retry_params=_include_retry_params,
             socket_timeout=socket_timeout,
@@ -775,7 +822,10 @@ def _post_request(
 
         if ret.get("code") == MASTER_TOKEN_EXPIRED_GS_CODE:
             self._connection.expired = True
-        elif ret.get("code") == SESSION_EXPIRED_GS_CODE:
+        elif (
+            ret.get("code") == SESSION_EXPIRED_GS_CODE
+            and self._connection._authenticator != PAT_WITH_EXTERNAL_SESSION
+        ):
             try:
                 ret = self._renew_session()
             except ReauthenticationRequest as ex:
@@ -900,6 +950,7 @@ def _request_exec_wrapper(
         retry_ctx,
         no_retry: bool = False,
         token=NO_TOKEN,
+        external_session_id=NO_EXTERNAL_SESSION_ID,
         **kwargs,
     ):
         conn = self._connection
@@ -928,6 +979,7 @@ def _request_exec_wrapper(
                 headers=headers,
                 data=data,
                 token=token,
+                external_session_id=external_session_id,
                 raise_raw_http_failure=raise_raw_http_failure,
                 **kwargs,
             )
@@ -1061,6 +1113,7 @@ def _request_exec(
         headers,
         data,
         token,
+        external_session_id=None,
         catch_okta_unauthorized_error: bool = False,
         is_raw_text: bool = False,
         is_raw_binary: bool = False,
@@ -1088,6 +1141,11 @@ def _request_exec(
             # socket timeout is constant. You should be able to receive
             # the response within the time. If not, ConnectReadTimeout or
             # ReadTimeout is raised.
+            auth = (
+                PATWithExternalSessionAuth(token, external_session_id)
+                if (external_session_id is not None and token is not None)
+                else SnowflakeAuth(token)
+            )
             raw_ret = session.request(
                 method=method,
                 url=full_url,
@@ -1096,7 +1154,7 @@ def _request_exec(
                 timeout=socket_timeout,
                 verify=True,
                 stream=is_raw_binary,
-                auth=SnowflakeAuth(token),
+                auth=auth,
             )
             download_end_time = get_time_millis()