fix(bigquery): push LIMIT into SQL and strip trailing semicolon#2465
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughBigQuery SQL submission now strips trailing semicolons, applies provided limits directly in SQL, and preserves dry-run configuration. Unit tests cover query, dry-run, helper behavior, and semicolons inside literals. ChangesBigQuery SQL preprocessing
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Sequence Diagram(s)sequenceDiagram
participant BigQueryConnector
participant BigQueryClient
participant QueryJob
BigQueryConnector->>BigQueryConnector: Strip trailing semicolon and append LIMIT
BigQueryConnector->>BigQueryClient: Submit SQL
BigQueryClient->>QueryJob: Create query job
BigQueryConnector->>QueryJob: Fetch results
Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
core/wren/tests/unit/test_bigquery_semicolon.py (1)
41-67: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDead
fake_bqsetup and missinguse_query_cacheassertion intest_dry_run_strips_trailing_semicolon.The first
patch.dict(line 48) andfake_bq(lines 43–45) are immediately overridden by the secondpatch.dict(lines 57–63) and have no effect on the test —dry_runis called inside the inner patch wheresys.modules["google.cloud.bigquery"]isfake_mod. This dead code is misleading for future maintainers.Additionally,
dry_run()setsuse_query_cache=Falsebut the test only assertsdry_run is True. Adding the missing assertion would fully verify the job config contract.♻️ Proposed cleanup
def test_dry_run_strips_trailing_semicolon() -> None: connector, client = _make_mock_connector() - fake_bq = type("BQ", (), { - "QueryJobConfig": staticmethod(lambda **kw: type("C", (), kw)()), - }) import sys from unittest.mock import patch - with patch.dict(sys.modules, {"google": MagicMock(), "google.cloud": MagicMock(), "google.cloud.bigquery": fake_bq}): - # Still import-side effect free: inject simple stub - import types - fake_mod = types.ModuleType("google.cloud.bigquery") - class QueryJobConfig: - def __init__(self, dry_run=False, use_query_cache=True): - self.dry_run = dry_run - self.use_query_cache = use_query_cache - fake_mod.QueryJobConfig = QueryJobConfig - with patch.dict(sys.modules, { - "google": types.ModuleType("google"), - "google.cloud": types.ModuleType("google.cloud"), - "google.cloud.bigquery": fake_mod, - "google.oauth2": types.ModuleType("google.oauth2"), - "google.oauth2.service_account": types.ModuleType("google.oauth2.service_account"), - }): - connector.dry_run("SELECT 1; \n") + import types + fake_mod = types.ModuleType("google.cloud.bigquery") + class QueryJobConfig: + def __init__(self, dry_run=False, use_query_cache=True): + self.dry_run = dry_run + self.use_query_cache = use_query_cache + fake_mod.QueryJobConfig = QueryJobConfig + with patch.dict(sys.modules, { + "google": types.ModuleType("google"), + "google.cloud": types.ModuleType("google.cloud"), + "google.cloud.bigquery": fake_mod, + "google.oauth2": types.ModuleType("google.oauth2"), + "google.oauth2.service_account": types.ModuleType("google.oauth2.service_account"), + }): + connector.dry_run("SELECT 1; \n") (sent,), kwargs = client.query.call_args assert sent == "SELECT 1" assert kwargs["job_config"].dry_run is True + assert kwargs["job_config"].use_query_cache is False🤖 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_bigquery_semicolon.py` around lines 41 - 67, Remove the unused fake_bq definition and the redundant outer patch.dict block from test_dry_run_strips_trailing_semicolon, keeping only the inner module stubs required by the test. Add an assertion that kwargs["job_config"].use_query_cache is False to verify the complete dry_run job configuration.
🤖 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.
Inline comments:
In `@core/wren/src/wren/connector/bigquery.py`:
- Around line 25-33: Update _apply_limit to avoid appending a second LIMIT when
the SQL already contains one: detect an existing LIMIT clause and return the
stripped query unchanged, or wrap the original query in a subquery before
applying the outer LIMIT. Preserve trailing-semicolon handling and integer limit
conversion.
---
Nitpick comments:
In `@core/wren/tests/unit/test_bigquery_semicolon.py`:
- Around line 41-67: Remove the unused fake_bq definition and the redundant
outer patch.dict block from test_dry_run_strips_trailing_semicolon, keeping only
the inner module stubs required by the test. Add an assertion that
kwargs["job_config"].use_query_cache is False to verify the complete dry_run job
configuration.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 0d75fb5f-70a9-4731-ba4a-19e5724f1942
📒 Files selected for processing (2)
core/wren/src/wren/connector/bigquery.pycore/wren/tests/unit/test_bigquery_semicolon.py
|
Thanks @coderabbitai — applied the cleanup in 82555c6: dropped the dead |
|
Tip For best results, initiate chat on the files or code changes.
🐇✨ |
Avoid second LIMIT on pre-limited SQL (CodeRabbit Canner#2465).
|
Addressed CodeRabbit: outer LIMIT already present → no second LIMIT appended. |
Avoid second LIMIT on pre-limited SQL (CodeRabbit Canner#2465).
96a4ba2 to
665fe09
Compare
|
#2478 fixed the building failure of |
Avoid second LIMIT on pre-limited SQL (CodeRabbit Canner#2465).
665fe09 to
bc0e9fb
Compare
|
@goldmedal Done — rebased onto current |
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 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.
Inline comments:
In `@core/wren/src/wren/connector/bigquery.py`:
- Around line 40-41: Update the limit-handling logic around _has_outer_limit so
a requested limit cannot exceed the caller’s smaller limit when the SQL already
contains an outer LIMIT. Apply the minimum of the existing and requested limits
or wrap the query with the requested limit, and add a regression test covering
query("... LIMIT 100", limit=5) returning at most five rows.
- Around line 10-22: Update _strip_trailing_semicolon and the _apply_limit flow
to recognize trailing line and block SQL comments, preserving those comments
while removing or relocating a terminal semicolon before appending LIMIT. Ensure
dry_run uses the same normalization so terminal semicolons are removed with
trailing comments present, and add regression tests covering both -- and /* ...
*/ comments.
- Around line 25-28: Replace the raw regex in _has_outer_limit with SQL-aware
tokenization or parsing that ignores comments, string literals, and nested
expressions while identifying only a trailing top-level LIMIT clause. Ensure
both SELECT 1 followed by a comment containing LIMIT and SELECT 1 LIMIT 10
followed by a comment are handled correctly, and add regression tests for these
cases.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: e2da6884-2234-425a-9eac-63177cbf6305
📒 Files selected for processing (2)
core/wren/src/wren/connector/bigquery.pycore/wren/tests/unit/test_bigquery_semicolon.py
🚧 Files skipped from review as they are similar to previous changes (1)
- core/wren/tests/unit/test_bigquery_semicolon.py
| _TRAILING_SEMICOLONS_RE = re.compile(r"[;\s]+\Z") | ||
|
|
||
|
|
||
| def _strip_trailing_semicolon(sql: str) -> str: | ||
| """Strip trailing ``;`` / whitespace for BigQuery jobs. | ||
|
|
||
| BigQuery job submission is more lenient than engines that subquery-wrap, | ||
| but a trailing semicolon still breaks when SQL is later composed, and | ||
| multi-statement batches with a terminal ``;`` are not the interface we | ||
| expose for single-statement GenBI queries. Only the terminating run is | ||
| stripped so ``SELECT ';'`` literals remain. | ||
| """ | ||
| return _TRAILING_SEMICOLONS_RE.sub("", sql) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
Handle trailing SQL comments before appending LIMIT.
For valid input such as SELECT 1; -- note, _apply_limit produces SELECT 1; -- note\nLIMIT 5; the semicolon remains before the appended clause, so the submitted SQL is invalid. dry_run also fails to strip this statement terminator. Preserve trailing comments while removing or relocating the terminal semicolon, and add -- and /* ... */ regression tests.
🤖 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/src/wren/connector/bigquery.py` around lines 10 - 22, Update
_strip_trailing_semicolon and the _apply_limit flow to recognize trailing line
and block SQL comments, preserving those comments while removing or relocating a
terminal semicolon before appending LIMIT. Ensure dry_run uses the same
normalization so terminal semicolons are removed with trailing comments present,
and add regression tests covering both -- and /* ... */ comments.
Avoid second LIMIT on pre-limited SQL (CodeRabbit Canner#2465).
bc0e9fb to
a8bcdad
Compare
|
This PR should also change to the shared |
|
@Bartok9, could you modify this PR? I think it should use the shared |
max_results only caps result pages; append LIMIT after stripping a terminal semicolon so the engine short-circuits expensive scans.
… dry_run Addresses CodeRabbit review: removes the overridden fake_bq/patch.dict block and verifies dry_run job config disables query cache.
Avoid second LIMIT on pre-limited SQL (CodeRabbit Canner#2465).
Drop fragile outer-LIMIT regex detection; subquery-wrap so nested limitations and comment-only LIMIT tokens cannot bypass max_results contract. Strip standalone query trailing semicolon; keep dry_run cleanup.
Drop the local _strip_trailing_semicolon duplicate and import the shared strip_trailing_semicolon from wren.connector.base, matching the other connectors (redshift/trino/mssql/clickhouse). Hoist test imports to module top to satisfy ruff PLC0415.
a8bcdad to
4670ae1
Compare
|
@goldmedal Done — switched to the shared |
2 participants
Summary
query(..., limit=N)now appendsLIMIT Nafter stripping a trailing semicolon instead of relying only onresult(max_results=...).dry_runalso strips trailing;for consistent single-statement validation.Motivation
google.cloud.bigquerymax_resultscontrols pagination of the results API; it does not push a limit into the job SQL, so expensive full-table scans still run when GenBI only needs a preview page. Other connectors already push limits into SQL (Athena append-LIMIT path, postgres subquery wrap). Match that pushdown and reuse the terminating-run semicolon strip soSELECT 1;cannot become multi-statement noise.Verification
core/wren/.venv/bin/python -m pytest tests/unit/test_bigquery_semicolon.py -q— 4 passedquery("SELECT 1;", limit=5)submitsSELECT 1\nLIMIT 5; dry_run sendsSELECT 1License
core/wren/**only (Apache-2.0).Summary by CodeRabbit
LIMITclause when needed.LIMITinjection logic, and dry-run SQL/config behavior.