diff --git a/pyathena/sqlalchemy/base.py b/pyathena/sqlalchemy/base.py index db0216bc3..ccaa54ceb 100644 --- a/pyathena/sqlalchemy/base.py +++ b/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 { diff --git a/tests/pyathena/sqlalchemy/test_base.py b/tests/pyathena/sqlalchemy/test_base.py index fbd40b342..b68a53a23 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,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