fix(athena): push LIMIT into SQL and strip trailing semicolon on wrap - #2457
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
Walkthrough
ChangesAthena LIMIT Pushdown
Estimated code review effort: 2 (Simple) | ~10 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant AthenaConnector
participant AthenaCursor
Client->>AthenaConnector: query(sql, limit)
AthenaConnector->>AthenaConnector: strip trailing semicolon
AthenaConnector->>AthenaConnector: wrap SQL with LIMIT
AthenaConnector->>AthenaCursor: execute(executed)
AthenaCursor-->>AthenaConnector: result table
AthenaConnector-->>Client: return Arrow table
Possibly related PRs
Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://cdn.github.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
core/wren/tests/unit/test_athena_limit_pushdown.py (1)
18-38: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd an error-path test for the
DIALECT_SQLmetadata change.The PR changes
WrenErrormetadata fromDIALECT_SQL: sqltoDIALECT_SQL: executed(line 339 inathena.py), but no test verifies this. Whencursor.executeraises, the resultingWrenErrorshould carry the composed subquery-wrapped SQL, not the original input.♻️ Suggested error-path test
def test_query_without_limit_runs_original_sql(): connector = AthenaConnector.__new__(AthenaConnector) connector.connection = MagicMock() cursor = MagicMock() cursor.close = MagicMock() with patch( "wren.connector.athena._build_athena_arrow_table", return_value=MagicMock() ), patch("wren.connector.athena.contextlib.closing", side_effect=lambda c: _CM(c)): connector.connection.cursor.return_value = cursor connector.query("SELECT 1") cursor.execute.assert_called_once_with("SELECT 1") +def test_query_error_metadata_reflects_executed_sql(): + from wren.errors import WrenError + + connector = AthenaConnector.__new__(AthenaConnector) + connector.connection = MagicMock() + cursor = MagicMock() + cursor.close = MagicMock() + cursor.execute.side_effect = RuntimeError("boom") + connector.connection.cursor.return_value = cursor + + with patch("wren.connector.athena.contextlib.closing", side_effect=lambda c: _CM(c)): + with pytest.raises(WrenError) as exc_info: + connector.query("SELECT 1;", limit=3) + + assert exc_info.value.metadata["DIALECT_SQL"] == ( + "SELECT * FROM (SELECT 1) AS _wren_sub LIMIT 3" + )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@core/wren/tests/unit/test_athena_limit_pushdown.py` around lines 18 - 38, The Athena query path needs an error-path unit test to verify the `DIALECT_SQL` metadata now uses the executed SQL. Extend the existing `test_query_pushes_limit_into_sql` style in `test_athena_limit_pushdown.py` by making `cursor.execute` raise and asserting `AthenaConnector.query` raises a `WrenError` whose metadata comes from the wrapped subquery SQL produced inside `AthenaConnector.query`, not the original input SQL. Use the existing `AthenaConnector`, `cursor.execute`, and `WrenError` symbols to locate the code path and ensure the test checks the composed SQL string is attached under `DIALECT_SQL`.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@core/wren/tests/unit/test_athena_limit_pushdown.py`:
- Around line 18-38: The Athena query path needs an error-path unit test to
verify the `DIALECT_SQL` metadata now uses the executed SQL. Extend the existing
`test_query_pushes_limit_into_sql` style in `test_athena_limit_pushdown.py` by
making `cursor.execute` raise and asserting `AthenaConnector.query` raises a
`WrenError` whose metadata comes from the wrapped subquery SQL produced inside
`AthenaConnector.query`, not the original input SQL. Use the existing
`AthenaConnector`, `cursor.execute`, and `WrenError` symbols to locate the code
path and ensure the test checks the composed SQL string is attached under
`DIALECT_SQL`.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 3b846a6e-f353-4929-b13d-26209585e7f1
📒 Files selected for processing (2)
core/wren/src/wren/connector/athena.pycore/wren/tests/unit/test_athena_limit_pushdown.py
|
Nice cleanup — pushing LIMIT server-side brings Athena in line with the other pushdown connectors (postgres/trino/redshift/clickhouse), and dropping the One thing worth fixing before merge: Single-line wrap is broken by a trailing everything after This is exactly the case flagged on the sibling Snowflake PR (#2456), which fixed it by putting the inner SQL on its own line and adding a regression test. Applying the same guard here keeps the two PRs consistent: executed = (
"SELECT * FROM (\n"
f"{_strip_trailing_semicolon(sql)}\n"
f") AS _wren_sub LIMIT {int(limit)}"
)plus a small test mirroring For fairness: the single-line form is shared by several existing connectors, so this fragility is pre-existing across the family — but within this Snowflake+Athena series it's an avoidable inconsistency, and a trailing comment in generated SQL isn't far-fetched. Non-blocking, but a thought for a follow-up: |
query(limit=N) previously executed the full Athena statement and sliced the Arrow table in Python, so Presto/Trino engines scanned and returned the whole result even for page/sample calls. Wrap user SQL as a subquery with LIMIT and strip a trailing semicolon so composition stays valid.
query(limit=N) now wraps SQL with LIMIT rather than post-slicing the Arrow table. Align the preexisting mock test with that contract.
goldmedal review on Canner#2457 — same newline guard as Snowflake Canner#2456 so query("SELECT 1 -- pick", limit=n) does not comment out the closing wrap.
a51f834 to
969298c
Compare
goldmedal
left a comment
goldmedal
left a comment
There was a problem hiding this comment.
Review
Verdict: sound change, safe to merge after fixing one misleading comment. The core swap — full result + table.slice(0, limit) → wrapping in SELECT * FROM (<sql>) AS _wren_sub LIMIT N — is correct, well-tested, and consistent with the existing pushdown helpers in postgres/trino/clickhouse/duckdb/redshift/canner/datafusion/oracle/databricks. No public API / MDL / schema surface is touched, and the fix: prefix is valid for release-please. Athena is AWS-only and can't be testcontainerized, so unit-only coverage is the right call.
Findings (severity order)
1. [medium — accuracy] The new inline comment references a Snowflake guard that doesn't exist.
The wrap block comments ... same guard as Snowflake #2456, but snowflake.py still does client-side arrow_table.slice(0, limit) and has no _strip_trailing_semicolon helper and no subquery-LIMIT wrapping — there is no analogous guard there. Additionally, Athena is the only connector using a multiline wrap ((\n{sql}\n)); every sibling is single-line, so the shapes wouldn't match even if Snowflake did push down. The cross-reference is incorrect and will send the next maintainer on a dead-end. Please drop it and keep only the real justification (see #2).
2. [low — style/consistency] Reviewer handle + PR numbers baked into permanent source.
No sibling connector embeds PR numbers or reviewer names in its wrap comment. These references go stale immediately and read as review-thread residue. The substantive reason, though, is genuinely good and worth keeping — suggest replacing the whole parenthetical with something like:
Multiline wrap so a trailing
-- line commentin the inner SQL is terminated by the newline instead of swallowing the closing) AS _wren_sub LIMIT n. (Single-line sibling connectors don't guard this.)
3. [low — tests] test_athena_limit_pushdown.py over-mocks and duplicates existing coverage.
It bypasses __init__ via __new__, hand-rolls a _CM context manager, and patches contextlib.closing — brittle against refactors. And test_query_pushes_limit_into_sql asserts the same wrapped SQL that test_connector_query_returns_arrow_table_and_respects_limit already asserts in test_athena_connector.py using that file's existing _make_cursor harness. The one net-new case worth keeping — test_query_limit_survives_trailing_line_comment — could live in the existing file and reuse that harness instead of standing up a parallel mocking style.
4. [nit — pre-existing, cross-connector] ; followed by a trailing comment isn't stripped.
_TRAILING_SEMICOLONS_RE = r"[;\s]+\Z" only anchors at the end, so SELECT 1; -- foo keeps the interior ; and the wrap becomes invalid Trino SQL. Shared by every sibling helper and not realistic planner output — noting for completeness, not asking for a change here.
Things done right (no action)
LIMIT {int(limit)}casts to int, closing the injection vector — safer than siblings that interpolate{limit}raw.if limit is not Nonecorrectly handleslimit=0(empty table) rather than the buggyif limit:.- Error metadata switched to the executed SQL, so failures show what actually ran.
- No ORDER-BY determinism regression: the old
slice(0, N)on an unordered result was equally non-deterministic; this is a preview-row cap, not a semantic limit.
Drop the incorrect Snowflake Canner#2456 cross-reference (snowflake.py has no subquery-LIMIT wrap / semicolon guard) and the review-thread residue (PR numbers, reviewer handle). Keep only the substantive justification: multiline wrap so a trailing line comment can't swallow the closing wrap.
|
Thanks @goldmedal — addressed the review. Fixed the misleading comment (findings #1 + #2). Dropped the incorrect # Multiline wrap so a trailing `-- line comment` in the inner SQL
# is terminated by the newline instead of swallowing the closing
# `) AS _wren_sub LIMIT n`. (Single-line sibling connectors don't guard this.)On the non-blocking items:
Pushed as 895460b. |
2 participants
Summary
query(..., limit=N)into Athena SQL viaSELECT * FROM (<sql>) AS _wren_sub LIMIT N.;-terminated user SQL cannot break composition.Motivation
Previously the connector ran the full statement and used
table.slice(0, limit). Athena (Presto/Trino-flavoured) can stop early when LIMIT is in the plan; client-side slicing paid full scan + transfer cost. Same pattern as redshift/postgres/clickhouse/snowflake pushdown helpers.Verification
PYTHONPATH=src pytest tests/unit/test_athena_limit_pushdown.py→ 3 passedSELECT 1;+limit=3isSELECT * FROM (SELECT 1) AS _wren_sub LIMIT 3limit, original SQL is executed unchanged.License
core/wren/**only (Apache-2.0).Summary by CodeRabbit
Bug Fixes
limitinto the generated SQL via subquery wrapping and appendingLIMIT <n>(instead of slicing results in Python).--comments are handled safely.Tests
LIMITpushdown and that returned Arrow tables respect the requested row count.