From 3aae3c72f924082312a3bbd7eb93a0bcfda74f0e Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 16 Oct 2026 11:21:31 +0900 Subject: [PATCH 1/3] Let DefaultDialect.do_ping() run the ping with the API cursor Since #1185, AthenaDialect.do_ping() ran SELECT 1 itself. Before SQLAlchemy 2.0.5, DefaultDialect.do_ping() is also where a failed ping is checked for a disconnect (is_disconnect() and the handle_error listeners), and it returns False so that the pool reconnects; the override made such a checkout raise instead. The override now passes DefaultDialect.do_ping() a connection whose cursor() opens the API cursor, so the ping keeps #1185's cursor and each SQLAlchemy version keeps its own error handling. Closes #1199 Co-Authored-By: Claude Opus 5.5 --- pyathena/sqlalchemy/base.py | 10 ++++-- tests/pyathena/sqlalchemy/test_base.py | 45 ++++++++++++++++++++++++++ 2 files changed, 52 insertions(+), 3 deletions(-) diff --git a/pyathena/sqlalchemy/base.py b/pyathena/sqlalchemy/base.py index db0216bc3..de8a4dd68 100644 --- a/pyathena/sqlalchemy/base.py +++ b/pyathena/sqlalchemy/base.py @@ -7,6 +7,7 @@ import re from collections.abc import Mapping, MutableMapping from re import Pattern +from types import SimpleNamespace from typing import ( TYPE_CHECKING, Any, @@ -544,9 +545,12 @@ 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. + api_connection = SimpleNamespace( + cursor=lambda: dbapi_connection._internal_cursor(Cursor) # type: ignore[attr-defined] + ) + return super().do_ping(cast("DBAPIConnection", api_connection)) def _column(self, name: str | None, type_: str, comment: str | None, partition: bool | None): return { diff --git a/tests/pyathena/sqlalchemy/test_base.py b/tests/pyathena/sqlalchemy/test_base.py index fbd40b342..84ec723be 100644 --- a/tests/pyathena/sqlalchemy/test_base.py +++ b/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,50 @@ 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. + + Before SQLAlchemy 2.0.5, it returns False when a failed ping is a disconnect, so + that the pool reconnects; the stand-in here returns False the same way. + """ + 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}, + ) + cursors = [] + + def do_ping(self, dbapi_connection): + cursors.append(dbapi_connection.cursor()) + return False + + monkeypatch.setattr(DefaultDialect, "do_ping", do_ping) + dbapi_connection = ( + AthenaAioAdaptedConnection(MagicMock(), conn) + if connection_class is AioConnection + else conn + ) + try: + assert dialect_class().do_ping(dbapi_connection) is False + [cursor] = cursors + 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 From 29d61f39769f98c67c182f43981289e43ac9af5f Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 16 Oct 2026 11:38:28 +0900 Subject: [PATCH 2/3] Show is_disconnect() the pinged connection's attributes Before SQLAlchemy 2.0.5, DefaultDialect.do_ping() passes the connection it pings to is_disconnect(). It received a namespace with only cursor(), so an is_disconnect() override that reads the connection, such as its schema_name or an adapter's driver_connection, raised AttributeError and the checkout failed instead of reconnecting. The ping now passes a small proxy whose cursor() opens the API cursor and whose other attributes are the pinged connection's. The regression test now fails a real ping: its stand-in for DefaultDialect.do_ping() asks is_disconnect() as SQLAlchemy before 2.0.5 does, so neither the earlier override nor the namespace pass it, and no AWS request is made. Co-Authored-By: Claude Opus 5.5 --- pyathena/sqlalchemy/base.py | 33 ++++++++++++++++--- tests/pyathena/sqlalchemy/test_base.py | 44 +++++++++++++++++++------- 2 files changed, 61 insertions(+), 16 deletions(-) diff --git a/pyathena/sqlalchemy/base.py b/pyathena/sqlalchemy/base.py index de8a4dd68..9474ea4b2 100644 --- a/pyathena/sqlalchemy/base.py +++ b/pyathena/sqlalchemy/base.py @@ -7,7 +7,6 @@ import re from collections.abc import Mapping, MutableMapping from re import Pattern -from types import SimpleNamespace from typing import ( TYPE_CHECKING, Any, @@ -100,6 +99,33 @@ } +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. + + Args: + connection: The pinged DBAPI connection, or the async adapter around one. + """ + + def __init__(self, connection: Any) -> None: + 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. @@ -547,10 +573,7 @@ def do_ping(self, dbapi_connection: DBAPIConnection) -> bool: # Using the plain cursor here performs the SELECT without the UNLOAD. # 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. - api_connection = SimpleNamespace( - cursor=lambda: dbapi_connection._internal_cursor(Cursor) # type: ignore[attr-defined] - ) - return super().do_ping(cast("DBAPIConnection", api_connection)) + return super().do_ping(cast("DBAPIConnection", _PingConnection(dbapi_connection))) def _column(self, name: str | None, type_: str, comment: str | None, partition: bool | None): return { diff --git a/tests/pyathena/sqlalchemy/test_base.py b/tests/pyathena/sqlalchemy/test_base.py index 84ec723be..b68a53a23 100644 --- a/tests/pyathena/sqlalchemy/test_base.py +++ b/tests/pyathena/sqlalchemy/test_base.py @@ -281,32 +281,54 @@ def test_do_ping_keeps_default_dialect_ping( ): """DefaultDialect.do_ping() runs the ping with the API cursor and decides its result. - Before SQLAlchemy 2.0.5, it returns False when a failed ping is a disconnect, so - that the pool reconnects; the stand-in here returns False the same way. + 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}, ) - cursors = [] + 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): - cursors.append(dbapi_connection.cursor()) - return False + 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) - dbapi_connection = ( - AthenaAioAdaptedConnection(MagicMock(), conn) - if connection_class is AioConnection - else conn - ) + monkeypatch.setattr(dialect_class, "is_disconnect", is_disconnect) try: assert dialect_class().do_ping(dbapi_connection) is False - [cursor] = cursors + [(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 From 35734209a2162fed99da8bb2f2421bbb82f208e6 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 16 Oct 2026 11:51:21 +0900 Subject: [PATCH 3/3] Document the ping proxy's constructor argument on __init__ The repository's docstring rules describe constructor arguments in the Args: section of __init__, not of the class. Co-Authored-By: Claude Opus 5.5 --- pyathena/sqlalchemy/base.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/pyathena/sqlalchemy/base.py b/pyathena/sqlalchemy/base.py index 9474ea4b2..ccaa54ceb 100644 --- a/pyathena/sqlalchemy/base.py +++ b/pyathena/sqlalchemy/base.py @@ -106,12 +106,14 @@ class _PingConnection: 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. - - Args: - connection: The pinged DBAPI connection, or the async adapter around one. """ 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: