Let DefaultDialect.do_ping() run the ping with the API cursor

pyathena/sqlalchemy/base.py

@@ -99,6 +99,35 @@
 }
 
 
+class _PingConnection:
+    """A DBAPI connection whose ``cursor()`` opens an API cursor.
+
+    ``AthenaDialect.do_ping()`` passes it to ``DefaultDialect.do_ping()``, so the ping
+    does not run through the connection's default cursor class. Other attributes are
+    those of the pinged connection, which ``is_disconnect()`` receives before
+    SQLAlchemy 2.0.5.
+    """
+
+    def __init__(self, connection: Any) -> None:
+        """Wrap the pinged connection.
+
+        Args:
+            connection: The pinged DBAPI connection, or the async adapter around one.
+        """
+        self._connection = connection
+
+    def cursor(self) -> Any:
+        """Open an API cursor on the pinged connection.
+
+        Returns:
+            The cursor, wrapped by the async adapter for an async connection.
+        """
+        return self._connection._internal_cursor(Cursor)
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._connection, name)
+
+
 class AthenaDialect(DefaultDialect):
     """SQLAlchemy dialect for Amazon Athena.
 
@@ -544,9 +573,9 @@ def do_ping(self, dbapi_connection: DBAPIConnection) -> bool:
         # When unload is enabled in the default cursor, SELECT 1 is wrapped in UNLOAD
         # which Athena rejects with MISSING_COLUMN_NAME because the column is unnamed.
         # Using the plain cursor here performs the SELECT without the UNLOAD.
-        with dbapi_connection._internal_cursor(Cursor) as cursor:  # type: ignore[attr-defined]
-            cursor.execute(self._dialect_specific_select_one)
-        return True
+        # The ping still runs in DefaultDialect.do_ping(), which before SQLAlchemy 2.0.5
+        # is also where a failed ping is checked for a disconnect.
+        return super().do_ping(cast("DBAPIConnection", _PingConnection(dbapi_connection)))
 
     def _column(self, name: str | None, type_: str, comment: str | None, partition: bool | None):
         return {

tests/pyathena/sqlalchemy/test_base.py

@@ -17,6 +17,7 @@
 import sqlalchemy
 from botocore.exceptions import ClientError
 from sqlalchemy import create_engine, engine_from_config, func, literal_column, select, text, types
+from sqlalchemy.engine.default import DefaultDialect
 from sqlalchemy.engine.url import make_url
 from sqlalchemy.exc import NoSuchTableError
 from sqlalchemy.sql import expression, type_coerce
@@ -268,6 +269,72 @@ def test_do_ping_uses_api_cursor_without_unload(
         finally:
             conn.close()
 
+    @pytest.mark.parametrize(
+        ("dialect_class", "connection_class", "cursor_class", "expected"),
+        [
+            (AthenaDialect, Connection, PandasCursor, Cursor),
+            (AthenaAioDialect, AioConnection, AioPandasCursor, AioCursor),
+        ],
+    )
+    def test_do_ping_keeps_default_dialect_ping(
+        self, monkeypatch, dialect_class, connection_class, cursor_class, expected
+    ):
+        """DefaultDialect.do_ping() runs the ping with the API cursor and decides its result.
+
+        The stand-in checks a failed ping as SQLAlchemy before 2.0.5 does: it asks
+        is_disconnect() with the pinged connection and the cursor, and returns False for
+        a disconnect so that the pool reconnects.
+        """
+        conn = connection_class(
+            region_name="us-west-2",
+            s3_staging_dir="s3://bucket/path/",
+            aws_access_key_id="access_key",
+            aws_secret_access_key="secret_key",
+            schema_name="pinged_schema",
+            cursor_class=cursor_class,
+            cursor_kwargs={"unload": True},
+        )
+        error = OperationalError("ExpiredTokenException")
+        checked = []
+
+        def execute(self, operation):
+            raise error
+
+        def is_disconnect(self, e, connection, cursor):
+            checked.append((e, connection.schema_name, cursor))
+            return True
+
+        def do_ping(self, dbapi_connection):
+            cursor = dbapi_connection.cursor()
+            try:
+                cursor.execute(self._dialect_specific_select_one)
+            except OperationalError as e:
+                if self.is_disconnect(e, dbapi_connection, cursor):
+                    return False
+                raise
+            finally:
+                cursor.close()
+            return True
+
+        if connection_class is AioConnection:
+            dbapi_connection = AthenaAioAdaptedConnection(MagicMock(), conn)
+            monkeypatch.setattr(AthenaAioAdaptedCursor, "execute", execute)
+        else:
+            dbapi_connection = conn
+            monkeypatch.setattr(Cursor, "execute", execute)
+        monkeypatch.setattr(DefaultDialect, "do_ping", do_ping)
+        monkeypatch.setattr(dialect_class, "is_disconnect", is_disconnect)
+        try:
+            assert dialect_class().do_ping(dbapi_connection) is False
+            [(e, schema_name, cursor)] = checked
+            assert e is error
+            assert schema_name == "pinged_schema"
+            if connection_class is AioConnection:
+                cursor = cursor._cursor
+            assert type(cursor) is expected
+        finally:
+            conn.close()
+
     @pytest.mark.parametrize("dialect_class", [AthenaDialect, AthenaAioDialect])
     def test_type_compiler(self, dialect_class):
         # SQLAlchemy 2.0 builds the type compiler from type_compiler_cls. A legacy