APPS-32904: Allow reusing a connection from another connector via session and master token (#1818)

sfc-gh-klin · sfc-gh-mkeller · web-flow · commit 25f2547f11f1 · 2024-01-11T18:03:20.000Z
* Let connector take session_token and master_token to use an existing connection from another connector

* Integration tests added, skipping old driver

* Make heartbeat post request after token update

Co-authored-by: Mark Keller &lt;mark.keller@snowflake.com&gt;
diff --git a/DESCRIPTION.md b/DESCRIPTION.md
@@ -15,6 +15,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
   - Added support for Python 3.12
   - Make local testing more robust against implicit assumptions.
   - Fixed PyArrow Table type hinting
+  - Added support for connecting using an existing connection via the session and master token.
 
 - v3.6.0(December 09,2023)
 
diff --git a/src/snowflake/connector/connection.py b/src/snowflake/connector/connection.py
@@ -256,6 +256,18 @@ def _get_private_bytes_from_file(
         True,
         bool,
     ),  # Enable sending retryReason in response header for query-requests
+    "session_token": (
+        None,
+        (type(None), str),
+    ),  # session token from another connection, to be provided together with master token
+    "master_token": (
+        None,
+        (type(None), str),
+    ),  # master token from another connection, to be provided together with session token
+    "master_validity_in_seconds": (
+        None,
+        (type(None), int),
+    ),  # master token validity in seconds
 }
 
 APPLICATION_RE = re.compile(r"[\w\d_]+")
@@ -913,100 +925,123 @@ def __open_connection(self):
 
         # Setup authenticator
         auth = Auth(self.rest)
-        if self.auth_class is not None:
-            if type(
-                self.auth_class
-            ) not in FIRST_PARTY_AUTHENTICATORS and not issubclass(
-                type(self.auth_class), AuthByKeyPair
-            ):
-                raise TypeError("auth_class must be a child class of AuthByKeyPair")
-                # TODO: add telemetry for custom auth
-            self.auth_class = self.auth_class
-        elif self._authenticator == DEFAULT_AUTHENTICATOR:
-            self.auth_class = AuthByDefault(
-                password=self._password,
-                timeout=self._login_timeout,
-                backoff_generator=self._backoff_generator,
-            )
-        elif self._authenticator == EXTERNAL_BROWSER_AUTHENTICATOR:
-            self._session_parameters[PARAMETER_CLIENT_STORE_TEMPORARY_CREDENTIAL] = (
-                self._client_store_temporary_credential if IS_LINUX else True
-            )
-            auth.read_temporary_credentials(
-                self.host,
-                self.user,
-                self._session_parameters,
+
+        if self._session_token and self._master_token:
+            auth._rest.update_tokens(
+                self._session_token,
+                self._master_token,
+                self._master_validity_in_seconds,
             )
-            # Depending on whether self._rest.id_token is available we do different
-            #  auth_instance
-            if self._rest.id_token is None:
-                self.auth_class = AuthByWebBrowser(
-                    application=self.application,
-                    protocol=self._protocol,
-                    host=self.host,
-                    port=self.port,
-                    timeout=self._login_timeout,
-                    backoff_generator=self._backoff_generator,
+            heartbeat_ret = auth._rest._heartbeat()
+            logger.debug(heartbeat_ret)
+            if not heartbeat_ret or not heartbeat_ret.get("success"):
+                Error.errorhandler_wrapper(
+                    self,
+                    None,
+                    ProgrammingError,
+                    {
+                        "msg": "Session and master tokens invalid",
+                        "errno": ER_INVALID_VALUE,
+                    },
                 )
             else:
-                self.auth_class = AuthByIdToken(
-                    id_token=self._rest.id_token,
-                    application=self.application,
-                    protocol=self._protocol,
-                    host=self.host,
-                    port=self.port,
+                logger.debug("Session and master token validation successful.")
+
+        else:
+            if self.auth_class is not None:
+                if type(
+                    self.auth_class
+                ) not in FIRST_PARTY_AUTHENTICATORS and not issubclass(
+                    type(self.auth_class), AuthByKeyPair
+                ):
+                    raise TypeError("auth_class must be a child class of AuthByKeyPair")
+                    # TODO: add telemetry for custom auth
+                self.auth_class = self.auth_class
+            elif self._authenticator == DEFAULT_AUTHENTICATOR:
+                self.auth_class = AuthByDefault(
+                    password=self._password,
                     timeout=self._login_timeout,
                     backoff_generator=self._backoff_generator,
                 )
-
-        elif self._authenticator == KEY_PAIR_AUTHENTICATOR:
-            private_key = self._private_key
-
-            if self._private_key_file:
-                private_key = _get_private_bytes_from_file(
-                    self._private_key_file,
-                    self._private_key_file_pwd,
-                )
-
-            self.auth_class = AuthByKeyPair(
-                private_key=private_key,
-                timeout=self._login_timeout,
-                backoff_generator=self._backoff_generator,
-            )
-        elif self._authenticator == OAUTH_AUTHENTICATOR:
-            self.auth_class = AuthByOAuth(
-                oauth_token=self._token,
-                timeout=self._login_timeout,
-                backoff_generator=self._backoff_generator,
-            )
-        elif self._authenticator == USR_PWD_MFA_AUTHENTICATOR:
-            self._session_parameters[PARAMETER_CLIENT_REQUEST_MFA_TOKEN] = (
-                self._client_request_mfa_token if IS_LINUX else True
-            )
-            if self._session_parameters[PARAMETER_CLIENT_REQUEST_MFA_TOKEN]:
+            elif self._authenticator == EXTERNAL_BROWSER_AUTHENTICATOR:
+                self._session_parameters[
+                    PARAMETER_CLIENT_STORE_TEMPORARY_CREDENTIAL
+                ] = (self._client_store_temporary_credential if IS_LINUX else True)
                 auth.read_temporary_credentials(
                     self.host,
                     self.user,
                     self._session_parameters,
                 )
-            self.auth_class = AuthByUsrPwdMfa(
-                password=self._password,
-                mfa_token=self.rest.mfa_token,
-                timeout=self._login_timeout,
-                backoff_generator=self._backoff_generator,
-            )
-        else:
-            # okta URL, e.g., https://<account>.okta.com/
-            self.auth_class = AuthByOkta(
-                application=self.application,
-                timeout=self._login_timeout,
-                backoff_generator=self._backoff_generator,
-            )
+                # Depending on whether self._rest.id_token is available we do different
+                #  auth_instance
+                if self._rest.id_token is None:
+                    self.auth_class = AuthByWebBrowser(
+                        application=self.application,
+                        protocol=self._protocol,
+                        host=self.host,
+                        port=self.port,
+                        timeout=self._login_timeout,
+                        backoff_generator=self._backoff_generator,
+                    )
+                else:
+                    self.auth_class = AuthByIdToken(
+                        id_token=self._rest.id_token,
+                        application=self.application,
+                        protocol=self._protocol,
+                        host=self.host,
+                        port=self.port,
+                        timeout=self._login_timeout,
+                        backoff_generator=self._backoff_generator,
+                    )
+
+            elif self._authenticator == KEY_PAIR_AUTHENTICATOR:
+                private_key = self._private_key
+
+                if self._private_key_file:
+                    private_key = _get_private_bytes_from_file(
+                        self._private_key_file,
+                        self._private_key_file_pwd,
+                    )
+
+                self.auth_class = AuthByKeyPair(
+                    private_key=private_key,
+                    timeout=self._login_timeout,
+                    backoff_generator=self._backoff_generator,
+                )
+            elif self._authenticator == OAUTH_AUTHENTICATOR:
+                self.auth_class = AuthByOAuth(
+                    oauth_token=self._token,
+                    timeout=self._login_timeout,
+                    backoff_generator=self._backoff_generator,
+                )
+            elif self._authenticator == USR_PWD_MFA_AUTHENTICATOR:
+                self._session_parameters[PARAMETER_CLIENT_REQUEST_MFA_TOKEN] = (
+                    self._client_request_mfa_token if IS_LINUX else True
+                )
+                if self._session_parameters[PARAMETER_CLIENT_REQUEST_MFA_TOKEN]:
+                    auth.read_temporary_credentials(
+                        self.host,
+                        self.user,
+                        self._session_parameters,
+                    )
+                self.auth_class = AuthByUsrPwdMfa(
+                    password=self._password,
+                    mfa_token=self.rest.mfa_token,
+                    timeout=self._login_timeout,
+                    backoff_generator=self._backoff_generator,
+                )
+            else:
+                # okta URL, e.g., https://<account>.okta.com/
+                self.auth_class = AuthByOkta(
+                    application=self.application,
+                    timeout=self._login_timeout,
+                    backoff_generator=self._backoff_generator,
+                )
 
-        self.authenticate_with_retry(self.auth_class)
+            self.authenticate_with_retry(self.auth_class)
 
-        self._password = None  # ensure password won't persist
-        self.auth_class.reset_secrets()
+            self._password = None  # ensure password won't persist
+            self.auth_class.reset_secrets()
 
         self.initialize_query_context_cache()
 
@@ -1115,34 +1150,35 @@ def __config(self, **kwargs):
             ]:
                 self._authenticator = auth_tmp
 
-        if not self.user and self._authenticator != OAUTH_AUTHENTICATOR:
-            # OAuth Authentication does not require a username
-            Error.errorhandler_wrapper(
-                self,
-                None,
-                ProgrammingError,
-                {"msg": "User is empty", "errno": ER_NO_USER},
-            )
+        if not (self._master_token and self._session_token):
+            if not self.user and self._authenticator != OAUTH_AUTHENTICATOR:
+                # OAuth Authentication does not require a username
+                Error.errorhandler_wrapper(
+                    self,
+                    None,
+                    ProgrammingError,
+                    {"msg": "User is empty", "errno": ER_NO_USER},
+                )
 
-        if self._private_key or self._private_key_file:
-            self._authenticator = KEY_PAIR_AUTHENTICATOR
+            if self._private_key or self._private_key_file:
+                self._authenticator = KEY_PAIR_AUTHENTICATOR
 
-        if (
-            self.auth_class is None
-            and self._authenticator
-            not in [
-                EXTERNAL_BROWSER_AUTHENTICATOR,
-                OAUTH_AUTHENTICATOR,
-                KEY_PAIR_AUTHENTICATOR,
-            ]
-            and not self._password
-        ):
-            Error.errorhandler_wrapper(
-                self,
-                None,
-                ProgrammingError,
-                {"msg": "Password is empty", "errno": ER_NO_PASSWORD},
-            )
+            if (
+                self.auth_class is None
+                and self._authenticator
+                not in [
+                    EXTERNAL_BROWSER_AUTHENTICATOR,
+                    OAUTH_AUTHENTICATOR,
+                    KEY_PAIR_AUTHENTICATOR,
+                ]
+                and not self._password
+            ):
+                Error.errorhandler_wrapper(
+                    self,
+                    None,
+                    ProgrammingError,
+                    {"msg": "Password is empty", "errno": ER_NO_PASSWORD},
+                )
 
         if not self._account:
             Error.errorhandler_wrapper(
diff --git a/src/snowflake/connector/network.py b/src/snowflake/connector/network.py
@@ -595,7 +595,7 @@ def _token_request(self, request_type):
                 },
             )
 
-    def _heartbeat(self) -> None:
+    def _heartbeat(self) -> Any | dict[Any, Any] | None:
         headers = {
             HTTP_HEADER_CONTENT_TYPE: CONTENT_TYPE_APPLICATION_JSON,
             HTTP_HEADER_ACCEPT: CONTENT_TYPE_APPLICATION_JSON,
@@ -614,6 +614,7 @@ def _heartbeat(self) -> None:
         )
         if not ret.get("success"):
             logger.error("Failed to heartbeat. code: %s, url: %s", ret.get("code"), url)
+        return ret
 
     def delete_session(self, retry: bool = False) -> None:
         """Deletes the session."""
diff --git a/test/integ/test_connection.py b/test/integ/test_connection.py
@@ -131,6 +131,60 @@ def test_with_config(db_parameters):
         cnx.close()
 
 
+@pytest.mark.skipolddriver
+def test_with_tokens(conn_cnx, db_parameters):
+    """Creates a connection using session and master token."""
+    try:
+        with conn_cnx(
+            timezone="UTC",
+        ) as initial_cnx:
+            assert initial_cnx, "invalid initial cnx"
+            master_token = initial_cnx.rest._master_token
+            session_token = initial_cnx.rest._token
+            with snowflake.connector.connect(
+                account=db_parameters["account"],
+                host=db_parameters["host"],
+                port=db_parameters["port"],
+                protocol=db_parameters["protocol"],
+                session_token=session_token,
+                master_token=master_token,
+            ) as token_cnx:
+                assert token_cnx, "invalid second cnx"
+    except Exception:
+        # This is my way of guaranteeing that we'll not expose the
+        # sensitive information that this test needs to handle.
+        # db_parameter contains passwords.
+        pytest.fail("something failed", pytrace=False)
+
+
+@pytest.mark.skipolddriver
+def test_with_tokens_expired(conn_cnx, db_parameters):
+    """Creates a connection using session and master token."""
+    try:
+        with conn_cnx(
+            timezone="UTC",
+        ) as initial_cnx:
+            assert initial_cnx, "invalid initial cnx"
+            master_token = initial_cnx._rest._master_token
+            session_token = initial_cnx._rest._token
+
+        with pytest.raises(ProgrammingError):
+            token_cnx = snowflake.connector.connect(
+                account=db_parameters["account"],
+                host=db_parameters["host"],
+                port=db_parameters["port"],
+                protocol=db_parameters["protocol"],
+                session_token=session_token,
+                master_token=master_token,
+            )
+            token_cnx.close()
+    except Exception:
+        # This is my way of guaranteeing that we'll not expose the
+        # sensitive information that this test needs to handle.
+        # db_parameter contains passwords.
+        pytest.fail("something failed", pytrace=False)
+
+
 def test_keep_alive_true(db_parameters):
     """Creates a connection with client_session_keep_alive parameter."""
     config = {

Original file line number	Diff line number	Diff line change
`@@ -595,7 +595,7 @@ def _token_request(self, request_type):`
`595`	`595`	`},`
`596`	`596`	`)`
`597`	`597`
`598`		`- def _heartbeat(self) -> None:`
	`598`	`+ def _heartbeat(self) -> Any \| dict[Any, Any] \| None:`
`599`	`599`	`headers = {`
`600`	`600`	`HTTP_HEADER_CONTENT_TYPE: CONTENT_TYPE_APPLICATION_JSON,`
`601`	`601`	`HTTP_HEADER_ACCEPT: CONTENT_TYPE_APPLICATION_JSON,`
`@@ -614,6 +614,7 @@ def _heartbeat(self) -> None:`
`614`	`614`	`)`
`615`	`615`	`if not ret.get("success"):`
`616`	`616`	`logger.error("Failed to heartbeat. code: %s, url: %s", ret.get("code"), url)`
	`617`	`+ return ret`
`617`	`618`
`618`	`619`	`def delete_session(self, retry: bool = False) -> None:`
`619`	`620`	`"""Deletes the session."""`