SNOW-901157 new expired feature (#1810)

Author: sfc-gh-mkeller

DESCRIPTION.md

@@ -14,6 +14,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
   - Added support for Vector types
   - Changed urllib3 version pin to only affect Python versions < 3.10.
   - Support for `private_key_file` and `private_key_file_pwd` connection parameters
+  - Added a new flag `expired` to `SnowflakeConnection` class, that keeps track of whether the connection's master token has expired.
 
 - v3.5.0(November 13,2023)
 

src/snowflake/connector/connection.py

@@ -412,6 +412,7 @@ def __init__(
         self.__set_error_attributes()
         self.connect(**kwargs)
         self._telemetry = TelemetryClient(self._rest)
+        self.expired = False
 
         # get the imported modules from sys.modules
         self._log_telemetry_imported_packages()

src/snowflake/connector/network.py

@@ -743,7 +743,9 @@ def _post_request(
             )
         )
 
-        if ret.get("code") == SESSION_EXPIRED_GS_CODE:
+        if ret.get("code") == MASTER_TOKEN_EXPIRED_GS_CODE:
+            self._connection.expired = True
+        elif ret.get("code") == SESSION_EXPIRED_GS_CODE:
             try:
                 ret = self._renew_session()
             except ReauthenticationRequest as ex:

test/unit/test_connection.py

@@ -426,3 +426,41 @@ def test_private_key_file_reading(tmp_path: Path):
             )
     assert m.call_count == 1
     assert m.call_args_list[0].kwargs["private_key"] == pkb
+
+
+def test_expired_detection():
+    with mock.patch(
+        "snowflake.connector.network.SnowflakeRestful._post_request",
+        return_value={
+            "data": {
+                "masterToken": "some master token",
+                "token": "some token",
+                "validityInSeconds": 3600,
+                "masterValidityInSeconds": 14400,
+                "displayUserName": "TEST_USER",
+                "serverVersion": "7.42.0",
+            },
+            "code": None,
+            "message": None,
+            "success": True,
+        },
+    ):
+        conn = fake_connector()
+    assert not conn.expired
+    with conn.cursor() as cur:
+        with mock.patch(
+            "snowflake.connector.network.SnowflakeRestful.fetch",
+            return_value={
+                "data": {
+                    "errorCode": "390114",
+                    "reAuthnMethods": ["USERNAME_PASSWORD"],
+                },
+                "code": "390114",
+                "message": "Authentication token has expired.  The user must authenticate again.",
+                "success": False,
+                "headers": None,
+            },
+        ):
+            with pytest.raises(ProgrammingError):
+                cur.execute("select 1;")
+    assert conn.expired