Repository navigation
fix: run ping with plain cursor - #1185
Conversation
|
I hope that I did this correctly, I ran into the issue while I was trying to set up a database in superset and it wouldn't allow me to change credentials since the ping wasn't going through properly. Let me know if there's anything else I need to do :) |
| def do_ping(self, dbapi_connection: DBAPIConnection) -> bool: | ||
| # The connection's default cursor would wrap SELECT 1 in UNLOAD when unload | ||
| # is enabled, which Athena rejects with MISSING_COLUMN_NAME because the | ||
| # column is unnamed, and would write a result file to S3 for every ping. |
There was a problem hiding this comment.
The first reason is accurate (I reproduced MISSING_COLUMN_NAME: line 1:1: Column name not specified at position 1 with the pandas driver and unload=True), but "would write a result file to S3 for every ping" is not a difference between the two cursors: the plain Cursor also writes a CSV result and its .metadata file to the staging directory for every SELECT 1. Could you drop that clause so the comment only states the UNLOAD failure?
There was a problem hiding this comment.
I updated the comment
| # is enabled, which Athena rejects with MISSING_COLUMN_NAME because the | ||
| # column is unnamed, and would write a result file to S3 for every ping. | ||
| cursor = dbapi_connection._internal_cursor(Cursor) # type: ignore[attr-defined] | ||
| try: |
There was a problem hiding this comment.
Optional nit: the reflection paths in this class open the same cursor with a context manager (with self._internal_cursor(raw_connection) as cursor:), and both Cursor and AthenaAioAdaptedCursor support it, so this could be with dbapi_connection._internal_cursor(Cursor) as cursor: followed by cursor.execute(...). The try/finally form matches SQLAlchemy's own do_ping, so keeping it is fine too.
There was a problem hiding this comment.
Changed to the same style as the other paths
|
|
||
| @pytest.mark.parametrize( | ||
| "engine", | ||
| [{"driver": driver, "unload": True} for driver in ("pandas", "arrow", "polars")], |
There was a problem hiding this comment.
do_ping no longer uses the connection's default cursor class at all, so parametrizing over pandas/arrow/polars runs the same code path three times. The unit test test_do_ping_uses_api_cursor_without_unload already checks the cursor selection for both dialects. Each case here runs four Athena queries, two of them UNLOAD queries that only serve to check out the pooled connection, and these tests run on every CI build. Could you keep only the pandas case?
There was a problem hiding this comment.
Makes sense, I changed it to only use pandas
| class TestAsyncSQLAlchemyAthena: | ||
| @pytest.mark.parametrize( | ||
| "async_engine", | ||
| [{"driver": driver, "unload": True} for driver in ("aiopandas", "aioarrow", "aiopolars")], |
There was a problem hiding this comment.
Same as the sync test: the ping path does not depend on the driver, so a single aiopandas case covers it and saves two Athena queries per CI run.
There was a problem hiding this comment.
Also makes sense, I changed it to only use aiopandas
|
Thanks for the fix and for the clear description! I reproduced the Our contribution guide asks for an issue before a pull request so that the bug report and the discussion are tracked separately from the implementation. Could you open an issue describing the problem (for example, Superset's connection test failing with |
b769e0d to
119cc70
Compare
|
Thanks for all the feedback! I updated the code accordingly and created the issue :) |
…I 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 (cherry picked from commit 3aae3c7) 3.x resolution, as in the #1185 backport: - pyathena/sqlalchemy/base.py: the namespace's cursor() opens the API cursor with dbapi_connection.cursor(Cursor), since 3.x has no Connection._internal_cursor() (#1102). - tests/pyathena/sqlalchemy/test_base.py: the test uses the 3.x adapter name AsyncAdapt_pyathena_connection, and only DefaultDialect is added to 3.x's imports; the other master tests and imports in the conflict hunks are not part of this change. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2 participants
This fixes the issue of running a ping when
unload=true is on which results in a
MISSING_COLUMN_NAME error.
WHAT
When unload=true is enabled for the cursor, a ping results in a query being run with UNLOAD but the select doesn't have a column name and Athena throws a MISSING_COLUMN_NAME error.
WHY
By using the plain cursor to run the ping, the UNLOAD isn't used even when enabled which allows for the SELECT 1 query to run without a column name but still checks that the connection is valid
TEST
A test was added to make sure that this works
Closes #1189