diff --git a/pyathena/sqlalchemy/base.py b/pyathena/sqlalchemy/base.py index 2a9b184d0..db0216bc3 100644 --- a/pyathena/sqlalchemy/base.py +++ b/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, diff --git a/tests/pyathena/aio/sqlalchemy/test_base.py b/tests/pyathena/aio/sqlalchemy/test_base.py index 77f463975..c6d1c21de 100644 --- a/tests/pyathena/aio/sqlalchemy/test_base.py +++ b/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", [ diff --git a/tests/pyathena/sqlalchemy/test_base.py b/tests/pyathena/sqlalchemy/test_base.py index f8ae551f7..fbd40b342 100644 --- a/tests/pyathena/sqlalchemy/test_base.py +++ b/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", [