Add support for new authentication type - PAT with external session ID (#2355)

Author: sfc-gh-saroskar

DESCRIPTION.md

@@ -11,7 +11,7 @@ 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
 
 - v3.15.1(May 20, 2025)
   - Added basic arrow support for Interval types.

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,
@@ -362,6 +363,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_]+")
@@ -1265,6 +1271,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.
@@ -1404,6 +1419,7 @@ def __config(self, **kwargs):
                 OAUTH_AUTHENTICATOR,
                 USR_PWD_MFA_AUTHENTICATOR,
                 WORKLOAD_IDENTITY_AUTHENTICATOR,
+                PAT_WITH_EXTERNAL_SESSION,
             ]:
                 self._authenticator = auth_tmp
 
@@ -1419,6 +1435,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):
@@ -1467,6 +1484,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()
 

test/auth/authorization_parameters.py

@@ -216,3 +216,14 @@ def get_pat_connection_parameters(self) -> dict[str, str]:
         config["user"] = _get_env_variable("SNOWFLAKE_AUTH_TEST_BROWSER_USER")
 
         return config
+
+    def get_pat_with_external_session_connection_parameters(
+        self, external_session_id: str
+    ) -> dict[str, str]:
+        config = self.basic_config.copy()
+
+        config["authenticator"] = "PROGRAMMATIC_ACCESS_TOKEN_WITH_EXTERNAL_SESSION"
+        config["user"] = _get_env_variable("SNOWFLAKE_AUTH_TEST_BROWSER_USER")
+        config["external_session_id"] = external_session_id
+
+        return config

test/auth/authorization_test_helper.py

@@ -106,6 +106,33 @@ def connect_and_execute_simple_query(self):
             logger.error(e)
             return False
 
+    def connect_and_execute_set_session_state(self, key: str, value: str):
+        try:
+            logger.info("Trying to connect to Snowflake")
+            with snowflake.connector.connect(**self.configuration) as con:
+                result = con.cursor().execute(f"SET {key} = '{value}'")
+                logger.debug(result.fetchall())
+                logger.info("Successfully SET session variable")
+                return True
+        except Exception as e:
+            self.error_msg = e
+            logger.error(e)
+            return False
+
+    def connect_and_execute_check_session_state(self, key: str):
+        try:
+            logger.info("Trying to connect to Snowflake")
+            with snowflake.connector.connect(**self.configuration) as con:
+                result = con.cursor().execute(f"SELECT 1, ${key}")
+                value = result.fetchone()[1]
+                logger.debug(value)
+                logger.info("Successfully READ session variable")
+                return value
+        except Exception as e:
+            self.error_msg = e
+            logger.error(e)
+            return False
+
     def _provide_credentials(self, scenario: Scenario, login: str, password: str):
         try:
             webbrowser.register("xdg-open", None, webbrowser.GenericBrowser("xdg-open"))

test/auth/test_external_session_with_PAT.py

@@ -0,0 +1,63 @@
+import uuid
+from test.auth.authorization_parameters import (
+    AuthConnectionParameters,
+    get_pat_setup_command_variables,
+)
+
+import pytest
+from authorization_test_helper import AuthorizationTestHelper
+from test_pat import get_pat_token, remove_pat_token
+
+EXTERNAL_SESSION_ID = str(uuid.uuid4())
+SESSION_VAR_KEY = "PAT_WITH_EXTERNAL_SESSION_TEST_KEY"
+SESSION_VAR_VALUE = "PAT_WITH_EXTERNAL_SESSION_TEST_VALUE"
+
+
+@pytest.mark.auth
+def test_pat_with_external_session_authN_success() -> None:
+    pat_command_variables = get_pat_setup_command_variables()
+    connection_parameters = AuthConnectionParameters().get_pat_connection_parameters()
+    try:
+        pat_command_variables = get_pat_token(pat_command_variables)
+        connection_parameters["token"] = pat_command_variables["token"]
+        connection_parameters["external_session_id"] = EXTERNAL_SESSION_ID
+        connection_parameters["authenticator"] = "PAT_WITH_EXTERNAL_SESSION"
+        test_helper = AuthorizationTestHelper(connection_parameters)
+        test_helper.connect_and_execute_set_session_state(
+            SESSION_VAR_KEY, SESSION_VAR_VALUE
+        )
+        ret = test_helper.connect_and_execute_check_session_state(SESSION_VAR_KEY)
+        assert ret == SESSION_VAR_VALUE
+    finally:
+        remove_pat_token(pat_command_variables)
+    assert test_helper.get_error_msg() == "", "Error message should be empty"
+
+
+@pytest.mark.auth
+def test_pat_with_external_session_authN_fail() -> None:
+    pat_command_variables = get_pat_setup_command_variables()
+    try:
+        pat_command_variables = get_pat_token(pat_command_variables)
+        connection_parameters = (
+            AuthConnectionParameters().get_pat_connection_parameters()
+        )
+        connection_parameters["token"] = pat_command_variables["token"]
+        connection_parameters["external_session_id"] = EXTERNAL_SESSION_ID
+        connection_parameters["authenticator"] = "PAT_WITH_EXTERNAL_SESSION"
+        test_helper = AuthorizationTestHelper(connection_parameters)
+        test_helper.connect_and_execute_set_session_state(
+            SESSION_VAR_KEY, SESSION_VAR_VALUE
+        )
+        connection_parameters["external_session_id"] = str(
+            uuid.uuid4()
+        )  # User different external session
+        test_helper = AuthorizationTestHelper(connection_parameters)
+        ret = test_helper.connect_and_execute_check_session_state(SESSION_VAR_KEY)
+        assert ret != SESSION_VAR_VALUE
+    finally:
+        remove_pat_token(pat_command_variables)
+    print(test_helper.get_error_msg())
+    assert (
+        f"Session variable '${SESSION_VAR_KEY}' does not exist"
+        in test_helper.get_error_msg()
+    )