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

DESCRIPTION.md

@@ -10,6 +10,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
 
 - v3.12.5(TBD)
   - Added a feature to limit the sizes of IO-bound ThreadPoolExecutors during PUT and GET commands.
+  - 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

@@ -1668,7 +1668,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)
@@ -1694,7 +1694,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()
@@ -1981,3 +1981,23 @@ 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
+        else:
+            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", str(session_valid))
+                return session_valid
+            except Exception as e:
+                logger.debug(
+                    "session could not be validated due to exception: %s", str(e)
+                )
+                return False

test/integ/test_connection.py

@@ -1462,3 +1462,13 @@ 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
+        conn.close()
+        assert conn.is_valid() is False