fix(postgres): coerce LIMIT before SQL interpolation - #2625
Conversation
Reject negative and non-numeric limits so client values cannot break QUERY composition.
WalkthroughThe Postgres connector now coerces and validates query limits before SQL interpolation. Unit tests cover negative values, injection-like strings, and numeric string limits. ChangesPostgres limit validation
Estimated code review effort: 2 (Simple) | ~10 minutes 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://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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/src/wren/connector/postgres.py`:
- Around line 223-234: Update _coerce_limit so negative numeric inputs are
rejected before int() truncates them; alternatively reject non-integral limits
explicitly while preserving safe coercion for valid values. Add a regression
test verifying limit=-0.5 raises ValueError rather than producing LIMIT 0.
🪄 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: 4584742b-8b84-4390-9068-756c84dfca65
📒 Files selected for processing (2)
core/wren/src/wren/connector/postgres.pycore/wren/tests/unit/test_postgres_coerce_limit.py
| def _coerce_limit(limit: int | None) -> int | None: | ||
| """Validate and coerce a user-supplied ``limit`` to a non-negative ``int``. | ||
|
|
||
| ``int(limit)`` rejects strings like ``"5 OR 1=1"`` so the value can be | ||
| safely interpolated into SQL. Negative limits are also rejected. | ||
| """ | ||
| if limit is None: | ||
| return None | ||
| coerced = int(limit) | ||
| if coerced < 0: | ||
| raise ValueError(f"limit must be non-negative, got {coerced}") | ||
| return coerced |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Validate negativity before integer truncation.
int(-0.5) produces 0, so the current check allows a negative runtime value and emits LIMIT 0. Check the original numeric value before coercion, or reject non-integral limits explicitly; add a regression test for limit=-0.5.
🤖 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/postgres.py` around lines 223 - 234, Update
_coerce_limit so negative numeric inputs are rejected before int() truncates
them; alternatively reject non-integral limits explicitly while preserving safe
coercion for valid values. Add a regression test verifying limit=-0.5 raises
ValueError rather than producing LIMIT 0.
Unit CI lacks the postgres extra; match semicolon unlimited stub pattern.
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_postgres_coerce_limit.py`:
- Around line 17-24: Update _ensure_psycopg_stub() to expose the created errors
module as the psycopg module’s errors attribute and define a QueryCanceled
exception on it, matching the symbol caught by the PostgreSQL connector’s query
error handler.
🪄 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: 513c1419-04ce-4cf4-bc1e-8eca6d8c84e5
📒 Files selected for processing (1)
core/wren/tests/unit/test_postgres_coerce_limit.py
| def _ensure_psycopg_stub() -> None: | ||
| if "psycopg" in sys.modules: | ||
| return | ||
| mod = types.ModuleType("psycopg") | ||
| errors = types.ModuleType("psycopg.errors") | ||
| sys.modules["psycopg"] = mod | ||
| sys.modules["psycopg.errors"] = errors | ||
|
|
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== files =="
git ls-files | rg 'test_postgres_coerce_limit|postgres\.py$|settings|pytest|requirements|pyproject' || true
echo "== target test =="
if [ -f core/wren/tests/unit/test_postgres_coerce_limit.py ]; then
wc -l core/wren/tests/unit/test_postgres_coerce_limit.py
sed -n '1,90p' core/wren/tests/unit/test_postgres_coerce_limit.py | nl -ba
fi
echo "== postgres connector references =="
if [ -f core/wren/src/wren/connector/postgres.py ]; then
wc -l core/wren/src/wren/connector/postgres.py
rg -n "psycopg|QueryCanceled|coerce|limit" core/wren/src/wren/connector/postgres.py || true
echo "--- relevant context ---"
rg -n -C 4 "psycopg|QueryCanceled|execute|limit" core/wren/src/wren/connector/postgres.py || true
fi
echo "== imports/usages of coerce_limit tests or module =="
rg -n "test_postgres_coerce_limit|psycopg\.errors|QueryCanceled|_ensure_psycopg_stub|stub_psycopg" . || true
echo "== python semantic probe =="
python3 - <<'PY'
import types, sys
mod = types.ModuleType("psycopg")
errors = types.ModuleType("psycopg.errors")
sys.modules["psycopg"] = mod
sys.modules["psycopg.errors"] = errors
try:
errors.QueryCanceled
except AttributeError as e:
print("empty stub missing errors.QueryCanceled:", e)
try:
mod.QueryCanceled
except AttributeError as e:
pass
errors.QueryCanceled = Exception
mod.errors = errors
from psycopg.errors import QueryCanceled
print("fixed stub has errors.QueryCanceled:", QueryCanceled)
PYRepository: Canner/WrenAI
Length of output: 602
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== target test =="
cat -n core/wren/tests/unit/test_postgres_coerce_limit.py
echo "== postgres connector relevant symbols/context =="
wc -l core/wren/src/wren/connector/postgres.py
rg -n -C 5 "psycopg|QueryCanceled|execute|limit|get_limit|coerce" core/wren/src/wren/connector/postgres.py || true
echo "== imports/usages =="
rg -n "test_postgres_coerce_limit|psycopg\.errors|QueryCanceled|_ensure_psycopg_stub|stub_psycopg" core/wren/tests/unit/test_postgres_coerce_limit.py core/wren/tests/connectors/test_postgres.py core/wren/src/wren/connector/postgres.py || true
echo "== python semantic probe =="
python3 - <<'PY'
import types, sys
mod = types.ModuleType("psycopg")
errors = types.ModuleType("psycopg.errors")
sys.modules["psycopg"] = mod
sys.modules["psycopg.errors"] = errors
try:
errors.QueryCanceled
except AttributeError as e:
print("current_stub_missing_errors_QueryCanceled")
print(repr(e))
try:
mod.errors
except AttributeError:
print("current_stub_missing_mod_errors")
errors.QueryCanceled = Exception
mod.errors = errors
from psycopg.errors import QueryCanceled
print("fixed_import_has_errors_QueryCanceled")
print(QueryCanceled)
PYRepository: Canner/WrenAI
Length of output: 9269
Complete the psycopg stub’s error surface.
core/wren/src/wren/connector/postgres.py catches psycopg.errors.QueryCanceled, but _ensure_psycopg_stub() only registers psycopg.errors without exposing it on psycopg or providing QueryCanceled. Connector errors raised during query() can trigger AttributeError while evaluating the handler.
Proposed fix
mod = types.ModuleType("psycopg")
errors = types.ModuleType("psycopg.errors")
+ errors.QueryCanceled = type("QueryCanceled", (Exception,), {})
+ mod.errors = errors
sys.modules["psycopg"] = mod
sys.modules["psycopg.errors"] = errors📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def _ensure_psycopg_stub() -> None: | |
| if "psycopg" in sys.modules: | |
| return | |
| mod = types.ModuleType("psycopg") | |
| errors = types.ModuleType("psycopg.errors") | |
| sys.modules["psycopg"] = mod | |
| sys.modules["psycopg.errors"] = errors | |
| def _ensure_psycopg_stub() -> None: | |
| if "psycopg" in sys.modules: | |
| return | |
| mod = types.ModuleType("psycopg") | |
| errors = types.ModuleType("psycopg.errors") | |
| errors.QueryCanceled = type("QueryCanceled", (Exception,), {}) | |
| mod.errors = errors | |
| sys.modules["psycopg"] = mod | |
| sys.modules["psycopg.errors"] = errors |
🤖 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_postgres_coerce_limit.py` around lines 17 - 24,
Update _ensure_psycopg_stub() to expose the created errors module as the psycopg
module’s errors attribute and define a QueryCanceled exception on it, matching
the symbol caught by the PostgreSQL connector’s query error handler.
|
Closing in favour of a single consolidated change — thanks for the work, the underlying tidy-up is worth doing, just not as one PR per connector. Why this is being closed rather than reviewed:
On the
What remains is genuine but smaller: a negative limit currently surfaces as a driver-level error instead of a clear What we would take instead — a single PR,
#2624 is the natural home for that; it is being kept open with a note to that effect. On this PR specifically — |
|
Thanks @goldmedal — completely fair, and appreciated the detailed write-up. Agreed on consolidating into a single I'll reshape #2624 along those lines and close out the per-connector duplicates. Thanks again for the clear bar and for keeping #2624 open as the home. |
2 participants
Summary
Coerce
limitto a non-negative int before interpolating into Postgres SQL (same contract as MySQL).What failure does this repair?
PostgresConnector.querybuilt the clause as a raw f-stringLIMIT {limit}, so a non-numeric or negativelimitflowed straight into the SQL text. Example before this fix:_coerce_limitnow rejects negatives and non-numeric strings before interpolation, while still accepting numeric strings like"3".Motivation
Raw
LIMIT {limit}accepted injection-style strings and negatives.Verification
Duplicate check
Related open connector-limit PRs (parallel work across connectors, no overlap):
Summary by CodeRabbit
limitvalues before SQL generation.limitinputs are now rejected with an error.limitvalues provided as strings continue to be supported.