refactor(connector): centralize LIMIT coercion - #2624
Open
Bartok9 wants to merge 6 commits into
Open
Conversation
Closes: n/a — inventory fix for LIMIT interpolation safety across connectors.
Contributor
|
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:
WalkthroughAdds shared ChangesLimit coercion
Estimated code review effort: 3 (Moderate) | ~25 minutes Possibly related PRs
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}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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/base.py`:
- Around line 23-35: Apply coerce_limit consistently to every connector utility
that interpolates limit into SQL, including wren.connector.oracle’s ROWNUM
condition and connectors with conditional string coercion; normalize limit
before formatting while preserving None handling. Add coerce_limit to the shared
connector import/export surface so all call sites use the same non-negative
validation.
- Around line 23-32: Update the coerce_limit parameter annotation to include
string inputs alongside int and None, matching its documented and tested support
for numeric strings; leave the existing coercion and validation behavior
unchanged.
🪄 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 Plus
Run ID: e6081ff5-99b9-49e3-ba3f-90249db7e76e
📒 Files selected for processing (2)
core/wren/src/wren/connector/base.pycore/wren/tests/unit/test_coerce_limit.py
Comment on lines
+23
to
+35
| def coerce_limit(limit: int | None) -> int | None: | ||
| """Validate and coerce a user-supplied ``limit`` to a non-negative ``int``. | ||
|
|
||
| Connectors train-plan LIMIT by interpolating the value into SQL. ``int()`` | ||
| rejects injection strings like ``"5 OR 1=1"``; negatives are rejected so | ||
| engines never see ``LIMIT -1`` (undefined / dialect-dependent). | ||
| """ | ||
| if limit is None: | ||
| return None | ||
| coerced = int(limit) | ||
| if coerced < 0: | ||
| raise ValueError(f"limit must be non-negative, got {coerced}") | ||
| return coerced |
Contributor
There was a problem hiding this comment.
🔒 Security & Privacy | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
rg -n -C 2 'coerce_limit|strip_trailing_semicolon|__all__' core/wren/src/wren
rg -n -C 3 '\bLIMIT\b|\blimit\b' core/wren/src/wren/connectorRepository: Canner/WrenAI
Length of output: 50369
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== connector files =="
git ls-files 'core/wren/src/wren/connector/*.py' | sort
echo
echo "== limited query locations importing util/using limit =="
python3 - <<'PY'
from pathlib import Path
root = Path("core/wren/src/wren/connector")
for p in sorted(root.glob("*.py")):
text = p.read_text()
lines = text.splitlines()
has_query_sig = any("def query(self, sql: str, limit: int | None = None)" in line for line in lines)
if not has_query_sig:
continue
print(f"\n--- {p} ---")
for i,l in enumerate(lines, start=1):
if i in (1,23,24,25) or "def query" in l or "def _apply_limit" in l or "LIMIT" in l or "ROWNUM" in l:
print(f"{i}: {l}")
PY
echo
echo "== coerce_limit usage/ref =="
rg -n 'from wren\.connector\.base import|coerce_limit|strip_trailing_semicolon|_apply_limit|_coerce_limit' core/wren/src/wren/connectorRepository: Canner/WrenAI
Length of output: 15604
Apply coerce_limit to every connector limit interpolation.
core/wren/src/wren/connector/base.py defines the helper, but connector utilities such as wren.connector.oracle still format limits directly (for example ROWNUM <= {limit}), and several connectors only coerce string-looking values instead of applying the shared non-negative check. Add coerce_limit to the shared import/export surface and normalize limit at each connector call site before SQL interpolation.
🤖 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/base.py` around lines 23 - 35, Apply
coerce_limit consistently to every connector utility that interpolates limit
into SQL, including wren.connector.oracle’s ROWNUM condition and connectors with
conditional string coercion; normalize limit before formatting while preserving
None handling. Add coerce_limit to the shared connector import/export surface so
all call sites use the same non-negative validation.
This was referenced Jul 30, 2026
This was referenced Aug 3, 2026
Collaborator
|
Keeping this one open — it is the right shape, and the rest of the batch (#2625, #2626, #2627, #2634, #2635, #2636, #2637) has been closed in favour of folding into it. To land, please extend this PR to:
Please also rebase onto current |
Contributor
Author
|
Thanks @goldmedal — clear direction, and agreed on consolidating here. I'll:
Working on the follow-up push now. |
Use one shared coerce_limit across interpolating connectors with strict semantics (reject bool, non-integral values, negatives). Drop mysql's private _coerce_limit. Invalid limits raise a consistent ValueError.
Bartok9
changed the title
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 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/athena.py`:
- Line 301: Update the shared coerce_limit() helper to reject any numeric value
that would be truncated by int(), including Decimal("1.5") and Fraction(3, 2),
while preserving valid integral limits and existing validation behavior. Add
shared-helper tests covering both cases. The calls in
core/wren/src/wren/connector/athena.py:301, bigquery.py:54, canner.py:245,
clickhouse.py:395, datafusion.py:32, duckdb.py:76, and mysql.py:79 require no
direct changes because they will inherit the corrected helper behavior.
In `@core/wren/src/wren/connector/base.py`:
- Around line 53-59: Update the limit validation around the coercion logic in
the base connector so non-string numeric values are compared directly with
coerced, avoiding float conversion and ensuring Decimal("1.5") and Fraction(3,
2) are rejected instead of truncated. Preserve the documented ValueError for
oversized integers rather than allowing OverflowError to escape, and add
regression tests covering both oversized integers and fractional
Decimal/Fraction inputs.
In `@core/wren/src/wren/connector/datafusion.py`:
- Around line 31-38: Update DataFusionConnector.query to normalize the input
with strip_trailing_semicolon before branching on limit, then use that
normalized SQL for both the limited wrapper and the direct SessionContext.query
path. Preserve the existing LIMIT behavior while ensuring queries such as SELECT
1; execute consistently when no limit is provided.
In `@core/wren/src/wren/connector/duckdb.py`:
- Around line 75-77: Move the coerce_limit(limit) call in DuckDBConnector.query
below the method’s docstring so the string literal remains the first statement
and is assigned to query.__doc__; preserve the existing limit coercion behavior
after the documentation block.
🪄 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 Plus
Run ID: bc073e11-0cb6-460f-8258-0e8448f77ef4
📒 Files selected for processing (14)
core/wren/src/wren/connector/athena.pycore/wren/src/wren/connector/base.pycore/wren/src/wren/connector/bigquery.pycore/wren/src/wren/connector/canner.pycore/wren/src/wren/connector/clickhouse.pycore/wren/src/wren/connector/datafusion.pycore/wren/src/wren/connector/duckdb.pycore/wren/src/wren/connector/mysql.pycore/wren/src/wren/connector/oracle.pycore/wren/src/wren/connector/postgres.pycore/wren/src/wren/connector/redshift.pycore/wren/src/wren/connector/snowflake.pycore/wren/src/wren/connector/trino.pycore/wren/tests/unit/test_coerce_limit.py
…ce_limit CI: ruff format wanted one-line LIMIT SQL in datafusion/redshift; unit collection broke after removing mysql._coerce_limit — point helpers tests at base.coerce_limit.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/tests/unit/test_mysql_helpers.py`:
- Around line 44-68: Update coerce_limit to validate non-string numeric inputs
for exact integrality and reject negative values before calling int(), so
Decimal("1.5"), Fraction(3, 2), and Decimal("-0.5") raise ValueError while
existing None, integer, and numeric-string behavior remains unchanged. Add
shared-helper regression tests covering fractional Decimal/Fraction values and
negative fractional values.
🪄 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 Plus
Run ID: fc48f565-2880-4687-baae-b565983e60de
📒 Files selected for processing (3)
core/wren/src/wren/connector/datafusion.pycore/wren/src/wren/connector/redshift.pycore/wren/tests/unit/test_mysql_helpers.py
🚧 Files skipped from review as they are similar to previous changes (2)
- core/wren/src/wren/connector/redshift.py
- core/wren/src/wren/connector/datafusion.py
Reject Decimal/Fraction truncation without float() overflow side paths. Keep DuckDB query docstring first; strip DataFusion SQL on both limit paths.
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
core/wren/tests/unit/test_coerce_limit.py (1)
66-69: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRename the test to match its assertion.
test_oversized_int_is_value_error_not_overflowsuggests thatcoerce_limit(huge)raisesValueError, but Line 69 verifies that the integer is preserved. Rename the test and update the comment to describe the no-float-overflow behavior.Proposed clarification
-def test_oversized_int_is_value_error_not_overflow() -> None: +def test_preserves_oversized_int_without_float_overflow() -> None: huge = 10**400 - # int stays int; ensure path still returns or rejects cleanly as ValueError only on bad types + # Preserve arbitrary-size integers without converting through float. assert coerce_limit(huge) == huge🤖 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_coerce_limit.py` around lines 66 - 69, Rename test_oversized_int_is_value_error_not_overflow to reflect that coerce_limit preserves oversized integers, and update its comment to describe avoiding float overflow while retaining the integer value. Keep the existing assertion unchanged.
🤖 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_coerce_limit.py`:
- Around line 66-69: Rename test_oversized_int_is_value_error_not_overflow to
reflect that coerce_limit preserves oversized integers, and update its comment
to describe avoiding float overflow while retaining the integer value. Keep the
existing assertion unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 2e87e125-ea8c-4853-8362-c5fe9e015abe
📒 Files selected for processing (4)
core/wren/src/wren/connector/base.pycore/wren/src/wren/connector/datafusion.pycore/wren/src/wren/connector/duckdb.pycore/wren/tests/unit/test_coerce_limit.py
🚧 Files skipped from review as they are similar to previous changes (3)
- core/wren/src/wren/connector/base.py
- core/wren/src/wren/connector/duckdb.py
- core/wren/src/wren/connector/datafusion.py
Keep CI lint green (I001 isort on connector imports) and pin Decimal("-0.5")
regression coverage for the shared helper / mysql helper suite.
CodeRabbit strip-on-both-paths change correctly strips trailing semicolons without LIMIT; update unit expectation. Rename oversized-int test to match preserve-not-raise behavior.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Centralize LIMIT handling on a single
coerce_limit()inconnector/base.pyand wire every interpolating connector through it.User-visible change
Invalid
limitvalues now raise a consistentValueErrorinstead of a driver-level error after SQL interpolation. This is a refactor of the coercion contract across connectors — not an injection fix (call paths that acceptsqlalready trust the caller equivalently forlimit).Semantics (strictest of the prior batch)
None→ unlimitedbool-0.5must not becomeLIMIT 0)ConnectorABC.queryaslimit: int | NoneConnectors wired
postgres,trino,clickhouse,oracle,canner,bigquery,duckdb,redshift,athena,snowflake,datafusion, andmysql(private_coerce_limitremoved).Verification
Notes
Consolidates the closed one-PR-per-connector batch (#2625–#2627, #2634–#2637) per maintainer direction.
Summary by CodeRabbit
Bug Fixes
Tests