Implemented the functionality of poll_pre_ping which executes a query to see if the connection is alive. If not it'll recycle the session

Author: msrathore-db

src/databricks/sqlalchemy/base.py

@@ -336,6 +336,31 @@ def do_rollback(self, dbapi_connection):
         # Databricks SQL Does not support transactions
         pass
 
+    def do_ping(self, dbapi_connection):
+        """Test if a database connection is alive.
+
+        This method is called by SQLAlchemy when pool_pre_ping=True to verify
+        connections are still valid before using them from the pool.
+
+        Args:
+            dbapi_connection: A raw DBAPI connection (from databricks-sql-connector)
+
+        Returns:
+            True if the connection is alive, False otherwise.
+        """
+        try:
+            cursor = dbapi_connection.cursor()
+            try:
+                cursor.execute("SELECT VERSION()")
+                cursor.fetchone()
+                return True
+            finally:
+                cursor.close()
+        except Exception:
+            # Any exception means the connection is dead
+            # SQLAlchemy will discard it and create a new one
+            return False
+
     @reflection.cache
     def has_table(
         self, connection, table_name, schema=None, catalog=None, **kwargs

tests/test_local/e2e/test_basic.py

@@ -541,3 +541,58 @@ def test_table_comment_reflection(self, inspector: Inspector, table: Table):
     def test_column_comment(self, inspector: Inspector, table: Table):
         result = inspector.get_columns(table.name)[0].get("comment")
         assert result == "column comment"
+
+
+def test_pool_pre_ping_with_closed_connection(connection_details):
+    """Test that pool_pre_ping detects closed connections and creates new ones.
+
+    This test verifies that when a connection is manually closed (simulating
+    session expiration), pool_pre_ping detects it and automatically creates
+    a new connection without raising an error to the user.
+    """
+    conn_string, connect_args = version_agnostic_connect_arguments(connection_details)
+
+    # Create engine with pool_pre_ping enabled
+    engine = create_engine(
+        conn_string,
+        connect_args=connect_args,
+        pool_pre_ping=True,
+        pool_size=1,
+        max_overflow=0
+    )
+
+    # Step 1: Create connection and get session ID
+    with engine.connect() as conn:
+        result = conn.execute(text("SELECT VERSION()")).scalar()
+        assert result is not None
+
+        # Get session ID of first connection
+        raw_conn = conn.connection.dbapi_connection
+        session_id_1 = raw_conn.get_session_id_hex()
+        assert session_id_1 is not None
+
+    # Step 2: Manually close the connection to simulate expiration
+    pooled_conn = engine.pool._pool.queue[0]
+    pooled_conn.driver_connection.close()
+
+    # Verify connection is closed
+    assert not pooled_conn.driver_connection.open
+
+    # Step 3: Try to use the closed connection - pool_pre_ping should detect and recycle
+    with engine.connect() as conn:
+        result = conn.execute(text("SELECT VERSION()")).scalar()
+        assert result is not None
+
+        # Get session ID of new connection
+        raw_conn = conn.connection.dbapi_connection
+        session_id_2 = raw_conn.get_session_id_hex()
+        assert session_id_2 is not None
+
+        # Verify a NEW connection was created (different session ID)
+        assert session_id_1 != session_id_2, (
+            "pool_pre_ping should have detected the closed connection "
+            "and created a new one with a different session ID"
+        )
+
+    # Cleanup
+    engine.dispose()

tests/test_local/test_ping.py

@@ -0,0 +1,47 @@
+"""Unit tests for do_ping() method in DatabricksDialect."""
+import pytest
+from unittest.mock import Mock
+from databricks.sqlalchemy import DatabricksDialect
+
+
+class TestDoPing:
+    """Test the do_ping() method for connection health checks."""
+
+    @pytest.fixture
+    def dialect(self):
+        """Create a DatabricksDialect instance."""
+        return DatabricksDialect()
+
+    def test_do_ping_success(self, dialect):
+        """Test do_ping returns True when connection is alive."""
+        mock_connection = Mock()
+        mock_cursor = Mock()
+        mock_connection.cursor.return_value = mock_cursor
+
+        result = dialect.do_ping(mock_connection)
+
+        assert result is True
+        mock_cursor.execute.assert_called_once_with("SELECT VERSION()")
+        mock_cursor.fetchone.assert_called_once()
+        mock_cursor.close.assert_called_once()
+
+    def test_do_ping_failure_cursor_creation(self, dialect):
+        """Test do_ping returns False when cursor creation fails."""
+        mock_connection = Mock()
+        mock_connection.cursor.side_effect = Exception("Connection closed")
+
+        result = dialect.do_ping(mock_connection)
+
+        assert result is False
+
+    def test_do_ping_failure_execute_and_cursor_closes(self, dialect):
+        """Test do_ping returns False on execute error and cursor is closed."""
+        mock_connection = Mock()
+        mock_cursor = Mock()
+        mock_connection.cursor.return_value = mock_cursor
+        mock_cursor.execute.side_effect = Exception("Query failed")
+
+        result = dialect.do_ping(mock_connection)
+
+        assert result is False
+        mock_cursor.close.assert_called_once()