diff --git a/DESCRIPTION.md b/DESCRIPTION.md index b56e1dd5a..1803023e3 100644 --- a/DESCRIPTION.md +++ b/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. diff --git a/src/snowflake/connector/connection.py b/src/snowflake/connector/connection.py index c51c33c60..78cde8b01 100644 --- a/src/snowflake/connector/connection.py +++ b/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 diff --git a/test/integ/test_connection.py b/test/integ/test_connection.py index bec9de556..fe828a6ee 100644 --- a/test/integ/test_connection.py +++ b/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