fix: run ping with plain cursor

pyathena/sqlalchemy/base.py

@@ -63,6 +63,7 @@
         PoolProxiedConnection,
     )
     from sqlalchemy.engine.interfaces import (
+        DBAPIConnection,
         ReflectedForeignKeyConstraint,
         ReflectedIndex,
         ReflectedPrimaryKeyConstraint,
@@ -538,6 +539,15 @@ def _internal_cursor(raw_connection: PoolProxiedConnection) -> Any:
         """
         return raw_connection.driver_connection._internal_cursor(Cursor)  # type: ignore[union-attr]
 
+    @override
+    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
+
     def _column(self, name: str | None, type_: str, comment: str | None, partition: bool | None):
         return {
             "name": name,

tests/pyathena/aio/sqlalchemy/test_base.py

@@ -58,6 +58,17 @@ async def test_execute_keeps_fetch_error_when_close_fails(self, error, caplog):
 
 
 class TestAsyncSQLAlchemyAthena:
+    @pytest.mark.parametrize(
+        "async_engine",
+        [{"driver": "aiopandas", "unload": True}],
+        indirect=True,
+    )
+    async def test_ping_with_unload(self, async_engine):
+        _, conn = async_engine
+        assert await conn.run_sync(
+            lambda sync_conn: sync_conn.dialect.do_ping(sync_conn.connection.dbapi_connection)
+        )
+
     @pytest.mark.parametrize(
         "async_engine",
         [

tests/pyathena/sqlalchemy/test_base.py

@@ -28,7 +28,11 @@
 from pyathena.aio.connection import AioConnection
 from pyathena.aio.cursor import AioCursor
 from pyathena.aio.pandas.cursor import AioPandasCursor
-from pyathena.aio.sqlalchemy.base import AthenaAioAdaptedConnection, AthenaAioDialect
+from pyathena.aio.sqlalchemy.base import (
+    AthenaAioAdaptedConnection,
+    AthenaAioAdaptedCursor,
+    AthenaAioDialect,
+)
 from pyathena.connection import Connection
 from pyathena.cursor import Cursor
 from pyathena.error import DatabaseError, OperationalError
@@ -222,6 +226,48 @@ def test_internal_cursor_leaves_out_default_cursor_options(
         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_uses_api_cursor_without_unload(
+        self, monkeypatch, dialect_class, connection_class, cursor_class, expected
+    ):
+        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",
+            cursor_class=cursor_class,
+            cursor_kwargs={"unload": True},
+        )
+        executed = []
+        closed = []
+        monkeypatch.setattr(expected, "close", lambda self: closed.append(self))
+        if connection_class is AioConnection:
+            dbapi_connection = AthenaAioAdaptedConnection(MagicMock(), conn)
+            monkeypatch.setattr(
+                AthenaAioAdaptedCursor,
+                "execute",
+                lambda self, operation: executed.append((self._cursor, operation)),
+            )
+        else:
+            dbapi_connection = conn
+            monkeypatch.setattr(
+                Cursor, "execute", lambda self, operation: executed.append((self, operation))
+            )
+        try:
+            assert dialect_class().do_ping(dbapi_connection) is True
+            [(cursor, operation)] = executed
+            assert type(cursor) is expected
+            assert operation == "SELECT 1"
+            assert closed == [cursor]
+        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
@@ -961,6 +1007,23 @@ def test_native_array_results_across_cursors(self, engine):
             finally:
                 array_engine.dispose()
 
+    @pytest.mark.parametrize(
+        "engine",
+        [{"driver": "pandas", "unload": True}],
+        indirect=True,
+    )
+    def test_ping_with_unload(self, engine):
+        engine, conn = engine
+        assert engine.dialect.do_ping(conn.connection.dbapi_connection) is True
+        ping_engine = sqlalchemy.create_engine(engine.url, pool_pre_ping=True)
+        try:
+            # The second checkout of the pooled connection pings it.
+            for _ in range(2):
+                with ping_engine.connect() as ping_conn:
+                    assert ping_conn.execute(text("SELECT 1 AS one")).scalar() == 1
+        finally:
+            ping_engine.dispose()
+
     @pytest.mark.parametrize(
         "engine",
         [