From 0c90e2a22ac37f4110ab32141667d1db4523e227 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 09:37:30 +0900 Subject: [PATCH 1/5] Remove the S3FS cursor-level callback test test_connection.py::test_cursor_accepts_connection_callbacks now runs the connection's on_start_query_execution for every cursor class, S3FSCursor included, so this test repeats its contract with one more Athena query. The callback given to execute() stays tested by test_on_start_query_execution_execute. Refs #1246 Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/s3fs/test_cursor.py | 20 -------------------- 1 file changed, 20 deletions(-) diff --git a/tests/pyathena/s3fs/test_cursor.py b/tests/pyathena/s3fs/test_cursor.py index 12953a661..904f0ef04 100644 --- a/tests/pyathena/s3fs/test_cursor.py +++ b/tests/pyathena/s3fs/test_cursor.py @@ -252,26 +252,6 @@ def test_executemany(self, s3fs_cursor, empty_table): s3fs_cursor.execute(f"SELECT * FROM {empty_table}") assert sorted(s3fs_cursor.fetchall()) == rows - def test_on_start_query_execution(self): - callback_query_id = None - - def callback(query_id): - nonlocal callback_query_id - callback_query_id = query_id - - with ( - contextlib.closing( - connect( - schema_name=ENV.schema, - cursor_class=S3FSCursor, - cursor_kwargs={"on_start_query_execution": callback}, - ) - ) as conn, - conn.cursor() as cursor, - ): - cursor.execute("SELECT * FROM one_row") - assert callback_query_id == cursor.query_id - def test_on_start_query_execution_execute(self): callback_query_id = None From 5e1d56cc7a7851e1d92f7130fd359eb32b39eb94 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 09:41:22 +0900 Subject: [PATCH 2/5] Check the callback that the connection's cursor_kwargs give a cursor The removed S3FS test was the only one that gave on_start_query_execution through cursor_kwargs. The connection test of cursor_kwargs over the connection's defaults now includes it, so a callback left in the constructor arguments twice, or the connection's callback winning over the configured one, fails offline. Refs #1246 Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/test_connection.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/pyathena/test_connection.py b/tests/pyathena/test_connection.py index 94526da0f..5d598e890 100644 --- a/tests/pyathena/test_connection.py +++ b/tests/pyathena/test_connection.py @@ -340,16 +340,23 @@ def test_cursor_arguments_override_cursor_kwargs(self, key, configured, explicit assert conn.cursor(RecordingCursor).kwargs[key] is configured def test_cursor_kwargs_override_connection_defaults(self): + connection_callback, configured_callback = MagicMock(), MagicMock() conn = _connection( schema_name="connection", kill_on_interrupt=True, - cursor_kwargs={"schema_name": "configured", "kill_on_interrupt": False}, + on_start_query_execution=connection_callback, + cursor_kwargs={ + "schema_name": "configured", + "kill_on_interrupt": False, + "on_start_query_execution": configured_callback, + }, ) cursor = conn.cursor(RecordingCursor) assert cursor.kwargs["schema_name"] == "configured" assert cursor.kwargs["kill_on_interrupt"] is False + assert cursor.kwargs["on_start_query_execution"] is configured_callback @pytest.mark.parametrize("cursor_class", PAGED_CURSORS + UNCAPPED_CURSORS) def test_cursor_arraysize(self, cursor_class): From cb9e49e96751e8a9f7f6f0c5fe73dee7b3217f47 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 09:46:32 +0900 Subject: [PATCH 3/5] Run the callbacks given through cursor_kwargs for every cursor class test_cursor_accepts_connection_callbacks gives the callbacks through the connection's attributes and through its cursor_kwargs, so a cursor that stores the connection's callback instead of its own argument fails. This replaces the constructor-argument check added to test_cursor_kwargs_override_connection_defaults. Refs #1246 Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/test_connection.py | 20 +++++++++----------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/tests/pyathena/test_connection.py b/tests/pyathena/test_connection.py index 5d598e890..e6eb55b04 100644 --- a/tests/pyathena/test_connection.py +++ b/tests/pyathena/test_connection.py @@ -240,11 +240,16 @@ def test_cursor_unknown_constructor_keyword(offline_connection, cursor_class): offline_connection.cursor(cursor_class) +@pytest.mark.parametrize("given_by", ["connection", "cursor_kwargs"]) @pytest.mark.parametrize("cursor_class", CURSOR_CLASSES) -async def test_cursor_accepts_connection_callbacks(offline_connection, cursor_class): +async def test_cursor_accepts_connection_callbacks(offline_connection, cursor_class, given_by): started, polled = [], [] - offline_connection.on_start_query_execution = started.append - offline_connection.on_poll = polled.append + callbacks = {"on_start_query_execution": started.append, "on_poll": polled.append} + if given_by == "connection": + for name, callback in callbacks.items(): + setattr(offline_connection, name, callback) + else: + offline_connection.cursor_kwargs = callbacks with answered_requests(offline_connection.client): async with _opened_cursor(offline_connection, cursor_class) as cursor: @@ -340,23 +345,16 @@ def test_cursor_arguments_override_cursor_kwargs(self, key, configured, explicit assert conn.cursor(RecordingCursor).kwargs[key] is configured def test_cursor_kwargs_override_connection_defaults(self): - connection_callback, configured_callback = MagicMock(), MagicMock() conn = _connection( schema_name="connection", kill_on_interrupt=True, - on_start_query_execution=connection_callback, - cursor_kwargs={ - "schema_name": "configured", - "kill_on_interrupt": False, - "on_start_query_execution": configured_callback, - }, + cursor_kwargs={"schema_name": "configured", "kill_on_interrupt": False}, ) cursor = conn.cursor(RecordingCursor) assert cursor.kwargs["schema_name"] == "configured" assert cursor.kwargs["kill_on_interrupt"] is False - assert cursor.kwargs["on_start_query_execution"] is configured_callback @pytest.mark.parametrize("cursor_class", PAGED_CURSORS + UNCAPPED_CURSORS) def test_cursor_arraysize(self, cursor_class): From 26dbf10318ad695db9ff49c66efc86485ebf27ff Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 09:52:49 +0900 Subject: [PATCH 4/5] Check that cursor_kwargs callbacks win over the connection's The cursor_kwargs case of test_cursor_accepts_connection_callbacks also sets other callbacks on the connection and checks that they do not run. Refs #1246 Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/test_connection.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/pyathena/test_connection.py b/tests/pyathena/test_connection.py index e6eb55b04..633813368 100644 --- a/tests/pyathena/test_connection.py +++ b/tests/pyathena/test_connection.py @@ -243,12 +243,15 @@ def test_cursor_unknown_constructor_keyword(offline_connection, cursor_class): @pytest.mark.parametrize("given_by", ["connection", "cursor_kwargs"]) @pytest.mark.parametrize("cursor_class", CURSOR_CLASSES) async def test_cursor_accepts_connection_callbacks(offline_connection, cursor_class, given_by): - started, polled = [], [] + """The connection's callbacks, or the ones its cursor_kwargs give over them, run.""" + started, polled, overridden = [], [], [] callbacks = {"on_start_query_execution": started.append, "on_poll": polled.append} if given_by == "connection": for name, callback in callbacks.items(): setattr(offline_connection, name, callback) else: + for name in callbacks: + setattr(offline_connection, name, overridden.append) offline_connection.cursor_kwargs = callbacks with answered_requests(offline_connection.client): @@ -265,6 +268,7 @@ async def test_cursor_accepts_connection_callbacks(offline_connection, cursor_cl assert [execution.state for execution in polled] == ["FAILED"] # The Async and Spark cursors do not run on_start_query_execution. assert started == (["query-id"] if cursor_class in START_CALLBACK_CURSORS else []) + assert overridden == [] # aio/test_connection.py tests the native asyncio cursors. From 7b04c97bf4e096864b684139789aac47e850a0f5 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sun, 18 Oct 2026 09:56:50 +0900 Subject: [PATCH 5/5] Keep a case of callbacks given only through cursor_kwargs test_cursor_accepts_connection_callbacks gives the callbacks through the connection, through cursor_kwargs alone, and through cursor_kwargs over other callbacks on the connection. Refs #1246 Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/test_connection.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/pyathena/test_connection.py b/tests/pyathena/test_connection.py index 633813368..70328bbae 100644 --- a/tests/pyathena/test_connection.py +++ b/tests/pyathena/test_connection.py @@ -240,7 +240,9 @@ def test_cursor_unknown_constructor_keyword(offline_connection, cursor_class): offline_connection.cursor(cursor_class) -@pytest.mark.parametrize("given_by", ["connection", "cursor_kwargs"]) +@pytest.mark.parametrize( + "given_by", ["connection", "cursor_kwargs", "cursor_kwargs_over_connection"] +) @pytest.mark.parametrize("cursor_class", CURSOR_CLASSES) async def test_cursor_accepts_connection_callbacks(offline_connection, cursor_class, given_by): """The connection's callbacks, or the ones its cursor_kwargs give over them, run.""" @@ -250,8 +252,9 @@ async def test_cursor_accepts_connection_callbacks(offline_connection, cursor_cl for name, callback in callbacks.items(): setattr(offline_connection, name, callback) else: - for name in callbacks: - setattr(offline_connection, name, overridden.append) + if given_by == "cursor_kwargs_over_connection": + for name in callbacks: + setattr(offline_connection, name, overridden.append) offline_connection.cursor_kwargs = callbacks with answered_requests(offline_connection.client):