Added functionality for is_disconnect()

Author: msrathore-db

src/databricks/sqlalchemy/base.py

@@ -342,6 +342,10 @@ def do_ping(self, dbapi_connection):
         This method is called by SQLAlchemy when pool_pre_ping=True to verify
         connections are still valid before using them from the pool.
 
+        This implementation improves upon SQLAlchemy's default do_ping() by
+        wrapping the cursor creation in a try block, which properly handles
+        cases where the connection is closed and cursor() itself raises an exception.
+
         Args:
             dbapi_connection: A raw DBAPI connection (from databricks-sql-connector)
 
@@ -361,6 +365,39 @@ def do_ping(self, dbapi_connection):
             # SQLAlchemy will discard it and create a new one
             return False
 
+    def is_disconnect(self, e, connection, cursor):
+        """Determine if an exception indicates the connection was lost.
+
+        This method is called by SQLAlchemy after exceptions occur during query
+        execution to determine if the error was due to a lost connection. If this
+        returns True, SQLAlchemy will invalidate the connection and create a new
+        one for the next operation.
+
+        This is complementary to do_ping():
+        - do_ping() is proactive: checks connection health BEFORE queries
+        - is_disconnect() is reactive: classifies errors AFTER they occur
+
+        Args:
+            e: The exception that was raised
+            connection: The connection that raised the exception (may be None)
+            cursor: The cursor that raised the exception (may be None)
+
+        Returns:
+            True if the error indicates a disconnect, False otherwise
+        """
+        from databricks.sql.exc import InterfaceError, DatabaseError
+
+        # InterfaceError: Client-side errors (e.g., connection already closed)
+        if isinstance(e, InterfaceError):
+            return True
+
+        # DatabaseError: Server-side errors with invalid handle indicate session expired
+        if isinstance(e, DatabaseError):
+            error_msg = str(e).lower()
+            return "invalid" in error_msg and "handle" in error_msg
+
+        return False
+
     @reflection.cache
     def has_table(
         self, connection, table_name, schema=None, catalog=None, **kwargs

tests/test_local/e2e/test_basic.py

@@ -596,3 +596,73 @@ def test_pool_pre_ping_with_closed_connection(connection_details):
 
     # Cleanup
     engine.dispose()
+
+
+def test_is_disconnect_handles_runtime_errors(db_engine):
+    """Test that is_disconnect() properly classifies disconnect errors during query execution.
+
+    This tests the reactive error handling (complementary to pool_pre_ping's proactive checking).
+    When a connection fails DURING a query, is_disconnect() should recognize the error
+    and tell SQLAlchemy to invalidate the connection.
+    """
+    from sqlalchemy import create_engine, text
+    from sqlalchemy.exc import DBAPIError
+
+    engine = create_engine(
+        db_engine.url,
+        pool_pre_ping=False,  # Disabled - we want to test is_disconnect, not do_ping
+        pool_size=1,
+        max_overflow=0,
+    )
+
+    # Step 1: Execute a successful query
+    with engine.connect() as conn:
+        result = conn.execute(text("SELECT VERSION()")).scalar()
+        assert result is not None
+
+        # Get session ID of working 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 server-side session expiration
+    pooled_conn = engine.pool._pool.queue[0]
+    pooled_conn.driver_connection.close()
+
+    # Step 3: Try to execute query on closed connection
+    # This should:
+    # 1. Fail with an exception
+    # 2. is_disconnect() gets called by SQLAlchemy
+    # 3. Returns True (recognizes it as disconnect error)
+    # 4. SQLAlchemy invalidates the connection
+    # 5. Next operation gets a fresh connection
+
+    # First query will fail because connection is closed
+    try:
+        with engine.connect() as conn:
+            conn.execute(text("SELECT VERSION()")).scalar()
+        # If we get here without exception, the connection wasn't actually closed
+        pytest.skip("Connection wasn't properly closed - cannot test is_disconnect")
+    except DBAPIError as e:
+        # Expected - connection was closed
+        # is_disconnect() should have been called and returned True
+        # This causes SQLAlchemy to invalidate the connection
+        assert "closed" in str(e).lower() or "invalid" in str(e).lower()
+
+    # Step 4: Next query should work because is_disconnect() invalidated the bad connection
+    with engine.connect() as conn:
+        result = conn.execute(text("SELECT VERSION()")).scalar()
+        assert result is not None
+
+        # Verify we got a NEW connection
+        raw_conn = conn.connection.dbapi_connection
+        session_id_2 = raw_conn.get_session_id_hex()
+        assert session_id_2 is not None
+
+        # Different session ID proves connection was invalidated and recreated
+        assert session_id_1 != session_id_2, (
+            "is_disconnect() should have invalidated the bad connection, "
+            "causing SQLAlchemy to create a new one with different session ID"
+        )
+
+    engine.dispose()

tests/test_local/test_is_disconnect.py

@@ -0,0 +1,45 @@
+"""Tests for DatabricksDialect.is_disconnect() method."""
+import pytest
+from databricks.sqlalchemy import DatabricksDialect
+from databricks.sql.exc import InterfaceError, DatabaseError, OperationalError
+
+
+class TestIsDisconnect:
+    @pytest.fixture
+    def dialect(self):
+        return DatabricksDialect()
+
+    def test_interface_error_is_disconnect(self, dialect):
+        """InterfaceError (client-side) is always a disconnect."""
+        error = InterfaceError("Cannot create cursor from closed connection")
+        assert dialect.is_disconnect(error, None, None) is True
+
+    def test_database_error_with_invalid_handle(self, dialect):
+        """DatabaseError with 'invalid handle' is a disconnect."""
+        test_cases = [
+            DatabaseError("Invalid SessionHandle"),
+            DatabaseError("[Errno INVALID_HANDLE] Session does not exist"),
+            DatabaseError("INVALID HANDLE"),
+            DatabaseError("invalid handle"),
+        ]
+        for error in test_cases:
+            assert dialect.is_disconnect(error, None, None) is True
+
+    def test_database_error_without_invalid_handle(self, dialect):
+        """DatabaseError without 'invalid handle' is not a disconnect."""
+        test_cases = [
+            DatabaseError("Syntax error in SQL"),
+            DatabaseError("Table not found"),
+            DatabaseError("Permission denied"),
+        ]
+        for error in test_cases:
+            assert dialect.is_disconnect(error, None, None) is False
+
+    def test_other_errors_not_disconnect(self, dialect):
+        """Other exception types are not disconnects."""
+        test_cases = [
+            OperationalError("Timeout waiting for query"),
+            Exception("Some random error"),
+        ]
+        for error in test_cases:
+            assert dialect.is_disconnect(error, None, None) is False