SNOW-1848371 adding connection.is_valid to perform connection validation on TCP/IP and Session levels (#2117)

Co-authored-by: Mark Keller <[email protected]>

Author: sfc-gh-dszmolka

DESCRIPTION.md

@@ -13,6 +13,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
   - Adding support for the new PAT authentication method.
   - Updated README.md to include instructions on how to verify package signatures using `cosign`.
   - Updated the log level for cursor's chunk rowcount from INFO to DEBUG
+  - Added a feature to verify if the connection is still good enough to send queries over.
 
 - v3.12.4(December 3,2024)
   - Fixed a bug where multipart uploads to Azure would be missing their MD5 hashes.

src/snowflake/connector/connection.py

@@ -1677,7 +1677,7 @@ def _log_telemetry(self, telemetry_data) -> None:
             self._telemetry.try_add_log_to_batch(telemetry_data)
 
     def _add_heartbeat(self) -> None:
-        """Add an hourly heartbeat query in order to keep connection alive."""
+        """Add a periodic heartbeat query in order to keep connection alive."""
         if not self.heartbeat_thread:
             self._validate_client_session_keep_alive_heartbeat_frequency()
             heartbeat_wref = weakref.WeakMethod(self._heartbeat_tick)
@@ -1703,7 +1703,7 @@ def _cancel_heartbeat(self) -> None:
             logger.debug("stopped heartbeat")
 
     def _heartbeat_tick(self) -> None:
-        """Execute a hearbeat if connection isn't closed yet."""
+        """Execute a heartbeat if connection isn't closed yet."""
         if not self.is_closed():
             logger.debug("heartbeating!")
             self.rest._heartbeat()
@@ -1990,3 +1990,21 @@ def _log_telemetry_imported_packages(self) -> None:
                     connection=self,
                 )
             )
+
+    def is_valid(self) -> bool:
+        """This function tries to answer the question: Is this connection still good for sending queries?
+        Attempts to validate the connections both on the TCP/IP and Session levels."""
+        logger.debug("validating connection and session")
+        if self.is_closed():
+            logger.debug("connection is already closed and not valid")
+            return False
+
+        try:
+            logger.debug("trying to heartbeat into the session to validate")
+            hb_result = self.rest._heartbeat()
+            session_valid = hb_result.get("success")
+            logger.debug("session still valid? %s", session_valid)
+            return bool(session_valid)
+        except Exception as e:
+            logger.debug("session could not be validated due to exception: %s", e)
+            return False

test/integ/test_connection.py

@@ -1462,3 +1462,12 @@ def test_disable_telemetry(conn_cnx, caplog):
                 cur.execute("select 1").fetchall()
             assert not conn.telemetry_enabled
     assert "POST /telemetry/send" not in caplog.text
+
+
+@pytest.mark.skipolddriver
+def test_is_valid(conn_cnx):
+    """Tests whether connection and session validation happens."""
+    with conn_cnx() as conn:
+        assert conn
+        assert conn.is_valid() is True
+    assert conn.is_valid() is False