Skip to content

Run the cursor_kwargs callbacks in the connection test instead of an S3FS query #1261

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
laughingman7743 merged 5 commits into from
Oct 18, 2026
Merged

Run the cursor_kwargs callbacks in the connection test instead of an S3FS query #1261

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
0c90e2a
Remove the S3FS cursor-level callback test
laughingman7743 Oct 18, 2026
5e1d56c
Check the callback that the connection's cursor_kwargs give a cursor
laughingman7743 Oct 18, 2026
cb9e49e
Run the callbacks given through cursor_kwargs for every cursor class
laughingman7743 Oct 18, 2026
26dbf10
Check that cursor_kwargs callbacks win over the connection's
laughingman7743 Oct 18, 2026
7b04c97
Keep a case of callbacks given only through cursor_kwargs
laughingman7743 Oct 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 0 additions & 20 deletions tests/pyathena/s3fs/test_cursor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):

@laughingman7743 laughingman7743 Oct 18, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Self-review of ea99b58e17b52252940951210ec445e2f167d9bc..b125bfa4c728f12b31c82ea6c118dedc149758f2.

Round 1 (behavior and test quality): CLEAN. The removed test passed the callback through the connection's cursor_kwargs, which Connection.cursor() merges into the constructor arguments for every cursor class (pyathena/connection.py:620), and S3FSCursor.execute() calls it at pyathena/s3fs/cursor.py:203. Removing that call fails test_connection.py::test_cursor_accepts_connection_callbacks[S3FSCursor] offline. contextlib, connect, and ENV stay in use by the remaining tests; just lint passed.

Round 2 (claims and operation): CLEAN. On ea99b58, START_CALLBACK_CURSORS includes S3FSCursor (tests/pyathena/test_connection.py:55), the cover asserts the started query ID and the polled state, and TestS3FSCursor::test_on_start_query_execution_execute keeps the execute() callback. The change saves one Athena query per run; no local live run, as the PR body states.

callback_query_id = None

Expand Down
20 changes: 16 additions & 4 deletions tests/pyathena/test_connection.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -240,11 +240,22 @@ def test_cursor_unknown_constructor_keyword(offline_connection, cursor_class):
offline_connection.cursor(cursor_class)


@pytest.mark.parametrize(

@laughingman7743 laughingman7743 Oct 18, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Rebase onto fb22d08 after #1258 rewrote this test into an async test over 21 cursor classes, including the aio cursors; the patch series is unchanged, and the conflicts in this test were resolved by keeping master's async structure and adding the given_by parameter. Re-checked at 7b04c97 offline: 63 cases pass and tests/pyathena/test_connection.py passes (229); the mutations fail 3 (S3FS call), 24 (stored connection callback), 42 (connection wins), 12 (guard), and 42 (pop to get) cases; just lint passed.

Independent follow-up review of the rebase (relayed): Codex CLI 0.163.0, model gpt-6-astra, ChatGPT login, codex exec -s read-only, session 01a14c91-0029-75d3-b124-53270880a7ff, with the rebased diff and the range-diff against 911f8b4. Result: CLEAN. It found that the test keeps master's 21 classes, async execution, bounded future wait, and cleanup; that all seven aio constructors forward the callbacks; and that every other test of the file is unchanged from master.

"given_by", ["connection", "cursor_kwargs", "cursor_kwargs_over_connection"]

@laughingman7743 laughingman7743 Oct 18, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Independent follow-up review 3 (relayed): Codex CLI 0.163.0, model gpt-6-astra, ChatGPT login, codex exec -s read-only, session 01a14c7f-6fc4-77d1-83c5-0d90a279deb4. Static review of 8aeb6883..96408aa9. It confirmed that the repair rejects the connection's callbacks winning and invoking both sources. Result: FINDINGS (1).

P2: the repair put connection callbacks into every cursor_kwargs case, so calling the cursor's callback only when the connection also has one (pyathena/common.py:1270) passed.

Repaired in 911f8b4: the test has three sources, connection, cursor_kwargs alone, and cursor_kwargs_over_connection (42 cases). Mutations, offline: that guard fails 6 cases; the connection's callbacks winning fails 28; storing the connection's callback fails 12; pop to get fails 28 with TypeError; removing the S3FS call fails 3. just lint passed. Self-review of the repair: round 1 CLEAN; round 2: the PR body now lists the three sources, the 42 cases, and these mutation results.

@laughingman7743 laughingman7743 Oct 18, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Independent follow-up review 4 (relayed): Codex CLI 0.163.0, model gpt-6-astra, ChatGPT login, codex exec -s read-only, session 01a14c83-90bd-78e0-bf49-4fe2422356f2. Static review of 96408aa9..911f8b44 in a clean detached snapshot of 911f8b4. Result: CLEAN. Covered: the repair, fixture isolation, callback forwarding and execution for all 14 cursor classes, async completion, and cleanup. It found that all three sources assert delivery (on_poll receives the FAILED state; the six cursors that run it deliver the query ID to the start callback), and that the override case still rejects delivery to the connection's callbacks.

)
@pytest.mark.parametrize("cursor_class", CURSOR_CLASSES)
async def test_cursor_accepts_connection_callbacks(offline_connection, cursor_class):
started, polled = [], []
offline_connection.on_start_query_execution = started.append
offline_connection.on_poll = polled.append
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."""
started, polled, overridden = [], [], []
callbacks = {"on_start_query_execution": started.append, "on_poll": polled.append}
if given_by == "connection":

@laughingman7743 laughingman7743 Oct 18, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Independent follow-up review (relayed): Codex CLI 0.163.0, model gpt-6-astra, ChatGPT login, codex exec -s read-only, session 01a14c75-14f7-7f93-88bd-8ee554205975. Static review of b125bfa4..6349697b. It confirmed that the added check catches the pop to get mutation. Result: FINDINGS (1).

P2: RecordingCursor records the argument before the base constructor stores it, so storing connection.on_start_query_execution instead of the constructor argument (pyathena/common.py:259) passed; the other execution tests give the connection's callback or an execute() callback.

Repaired in 8aeb688: this test runs the callbacks given through the connection's attributes and through its cursor_kwargs for every cursor class, and the RecordingCursor check is reverted. Mutations, offline: that store fails the 6 cursor_kwargs cases of the cursors that run on_start_query_execution; pop to get fails the 14 cursor_kwargs cases; removing the S3FS call fails both S3FSCursor cases. 28 cases pass without a mutation; just lint passed.

Self-review of 6349697b..8aeb6883: round 1 CLEAN (the offline_connection fixture is function-scoped, so setting cursor_kwargs does not leak; on_poll is now also checked through cursor_kwargs); round 2: the PR title and body now describe the parametrized test, its renamed IDs, and the 28 cases.

for name, callback in callbacks.items():
setattr(offline_connection, name, callback)
else:
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):
async with _opened_cursor(offline_connection, cursor_class) as cursor:
Expand All @@ -260,6 +271,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 == []

@laughingman7743 laughingman7743 Oct 18, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Independent follow-up review 2 (relayed): Codex CLI 0.163.0, model gpt-6-astra, ChatGPT login, codex exec -s read-only, session 01a14c79-d42b-7eb2-a2d2-bac88a609d35. Static review of 6349697b..8aeb6883. It confirmed that the parametrized test catches the stored-callback and pop to get mutations. Result: FINDINGS (2).

1 (P2): the cursor_kwargs case assigns them after the connection is built, so a constructor that dropped only on_start_query_execution from cursor_kwargs (after pyathena/connection.py:365) would pass. Not repaired: the constructor stores cursor_kwargs as one dict for every key, which TestConnection::test_cursor_arraysize and test_cursor_kwargs_override_connection_defaults check through _connection(cursor_kwargs=...); a key-specific drop is not a path the code has.

2 (P2): with one source per case, the connection's callback winning over cursor_kwargs passed. Repaired in 96408aa: the cursor_kwargs case also sets other callbacks on the connection and asserts that they do not run; letting the connection's callbacks win fails the 14 cursor_kwargs cases, and storing the connection's callback still fails 6 (offline). just lint passed.

Self-review of the repair: round 1 CLEAN; round 2: the PR body's description of the test still holds (two sources per class, 28 cases).



# aio/test_connection.py tests the native asyncio cursors.
Expand Down
Loading