feat(wren): add strict query mode for SQL policy enforcement#1500
Open
goldmedal wants to merge 4 commits intoCanner:mainfrom
Open
feat(wren): add strict query mode for SQL policy enforcement#1500goldmedal wants to merge 4 commits intoCanner:mainfrom
goldmedal wants to merge 4 commits intoCanner:mainfrom
Conversation
Add configurable strict mode to scope all SQL access to MDL-defined models and block dangerous functions before queries reach the database. - New ~/.wren/config.json with strict_mode and denied_functions settings - Table validation: reject queries referencing tables not in MDL - Function deny list: block specified functions (case-insensitive) - Config file-only (no CLI args), WREN_HOME env var support - New error codes: MODEL_NOT_FOUND, BLOCKED_FUNCTION Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
github-actions
bot
added
the
documentation
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughAdds a new immutable Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant CLI as "CLI"
participant Config as "Config Loader\n(load_config)"
participant Engine as "WrenEngine"
participant Policy as "Policy Validator\n(validate_sql_policy)"
CLI->>Config: resolve WREN_HOME\ncall load_config(path)
Config-->>CLI: return WrenConfig or raise WrenError
CLI->>Engine: construct WrenEngine(..., config)
Engine->>Policy: validate_sql_policy(ast, model_names, config)
Policy-->>Engine: success or raise WrenError
Engine->>Engine: proceed with manifest rewrite / planning
Engine-->>CLI: return plan or raise WrenError
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
wren/src/wren/cli.py (1)
130-150:⚠️ Potential issue | 🟠 MajorConvert config/engine construction failures into CLI exits.
load_config()andWrenEngine(...)now run before the surrounding commandtryblocks, so a malformed or unreadableconfig.jsonwill bubble out as a Typer traceback instead of the existingError:+ exit code 1 behavior.Proposed fix
def _build_engine( datasource: str | None, mdl: str | None, connection_info: str | None, connection_file: str | None, *, conn_required: bool = True, ): from wren.config import load_config # noqa: PLC0415 from wren.engine import WrenEngine # noqa: PLC0415 from wren.model.data_source import DataSource # noqa: PLC0415 + from wren.model.error import WrenError # noqa: PLC0415 @@ - config = load_config(_WREN_HOME) - return WrenEngine( - manifest_str=manifest_str, - data_source=ds, - connection_info=conn_dict, - config=config, - ) + try: + config = load_config(_WREN_HOME) + return WrenEngine( + manifest_str=manifest_str, + data_source=ds, + connection_info=conn_dict, + config=config, + ) + except WrenError as e: + typer.echo(f"Error: {e}", err=True) + raise typer.Exit(1) from edef dry_plan( @@ ): """Plan SQL through MDL and print the expanded SQL (no DB required).""" from wren.config import load_config # noqa: PLC0415 from wren.engine import WrenEngine # noqa: PLC0415 from wren.model.data_source import DataSource # noqa: PLC0415 + from wren.model.error import WrenError # noqa: PLC0415 @@ - config = load_config(_WREN_HOME) + try: + config = load_config(_WREN_HOME) + except WrenError as e: + typer.echo(f"Error: {e}", err=True) + raise typer.Exit(1) from e with WrenEngine( manifest_str=manifest_str, data_source=ds, connection_info={}, config=config ) as engine:Also applies to: 290-312
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@wren/src/wren/cli.py` around lines 130 - 150, The config loading and WrenEngine construction can raise exceptions that currently escape the CLI and produce tracebacks; wrap the calls to load_config(_WREN_HOME) and the WrenEngine(...) constructor in the same kind of try/except handling used for DataSource so that any error is caught, printed via typer.echo with a clear "Error:" message, and then exit with raise typer.Exit(1); specifically update the block that creates config = load_config(_WREN_HOME) and the return WrenEngine(...) to handle and convert errors (e.g., FileNotFoundError, JSONDecodeError, ValueError or generic Exception) into a CLI-friendly error+typer.Exit(1), and apply the same wrapping to the analogous construction around lines 290-312.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@wren/src/wren/config.py`:
- Around line 54-62: Replace the current permissive coercions in config parsing
with explicit type checks: in the config loader (look for strict_mode,
denied_raw, denied_functions in wren/src/wren/config.py) ensure strict_mode is a
bool (raise WrenError with ErrorCode.GENERIC_USER_ERROR if not), ensure
denied_raw is a list (already checked) and then validate every element in
denied_raw is a string before lowercasing and turning into denied_functions
(raise WrenError if any element is non-string); then add regression unit tests
in wren/tests/unit/test_config.py that assert a WrenError is raised for
{"strict_mode": "false"} and for a config where denied_functions contains mixed
types (e.g., ["safe", 1, {"obj":true}]).
In `@wren/src/wren/engine.py`:
- Around line 178-179: The current fix flattens CTE names via
CTERewriter._collect_user_cte_names and passes a global set to
validate_sql_policy, which allows out-of-scope CTEs to hide real tables; change
the flow so validate_sql_policy receives scope-aware CTE visibility (e.g., have
CTERewriter._collect_user_cte_names return a scoped structure such as a mapping
of AST nodes/with-blocks to their CTE name sets or a stack-based scope resolver,
and update validate_sql_policy to consult that scoped map when checking table
references) and add a regression test using the provided nested-shadowing SQL
snippet to assert that an outer FROM secret_table is not considered safe by an
inner WITH secret_table. Ensure you update function signatures
(CTERewriter._collect_user_cte_names and validate_sql_policy) and all call sites
in engine.py (where user_cte_names is produced) to use the new scoped
representation.
In `@wren/src/wren/policy.py`:
- Around line 47-54: The current _check_tables loop only inspects exp.Table
nodes and skips nodes with empty name, letting table-valued functions (TVFs)
like unnest/read_csv/generate_series bypass strict-mode validation; update
_check_tables (the loop over ast.find_all(exp.Table)) to also detect TVFs by: 1)
scanning FROM-clause expressions for exp.Func or known TVF node types (e.g.,
exp.Unnest, exp.ReadCSV, exp.GenerateSeries) since exp.TableFunction doesn't
exist in sqlglot 30.1.0, 2) applying the same name/whitelist/model checks for
those TVF nodes (treat their function name or node type as the referenced table)
and 3) in strict mode raise WrenError when a TVF is used and not in user CTEs or
model_names; ensure you reference the same helper variables
(user_cte_names_lower, model_names_lower) and keep existing continue/raise
behavior.
---
Outside diff comments:
In `@wren/src/wren/cli.py`:
- Around line 130-150: The config loading and WrenEngine construction can raise
exceptions that currently escape the CLI and produce tracebacks; wrap the calls
to load_config(_WREN_HOME) and the WrenEngine(...) constructor in the same kind
of try/except handling used for DataSource so that any error is caught, printed
via typer.echo with a clear "Error:" message, and then exit with raise
typer.Exit(1); specifically update the block that creates config =
load_config(_WREN_HOME) and the return WrenEngine(...) to handle and convert
errors (e.g., FileNotFoundError, JSONDecodeError, ValueError or generic
Exception) into a CLI-friendly error+typer.Exit(1), and apply the same wrapping
to the analogous construction around lines 290-312.
🪄 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: defaults
Review profile: CHILL
Plan: Pro
Run ID: 35a6ce89-4d53-42e2-aa34-8dcc63d1a25f
📒 Files selected for processing (9)
wren/README.mdwren/src/wren/cli.pywren/src/wren/config.pywren/src/wren/engine.pywren/src/wren/model/error.pywren/src/wren/policy.pywren/tests/unit/test_config.pywren/tests/unit/test_engine.pywren/tests/unit/test_policy.py
…king - Validate config.json types strictly: reject non-bool strict_mode and non-string denied_functions entries instead of silently coercing - Wrap load_config/WrenEngine construction in try/except in CLI for friendly error output instead of raw tracebacks - Make CTE name visibility scope-aware: walk up the AST to find ancestor WITH clauses instead of using a flat global set, preventing nested CTEs from shadowing outer table references - Block table-valued functions (read_csv, generate_series, unnest) in strict mode — they previously bypassed validation via empty name or non-Table AST nodes - Remove user_cte_names parameter from validate_sql_policy (now internal) - Add 9 regression tests covering all review findings Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@wren/src/wren/cli.py`:
- Line 14: The WREN_HOME environment value isn't being expanded so values like
"~/ .wren" remain literal; update the _WREN_HOME initialization to expand the
environment variable with os.path.expanduser before converting to Path (i.e.,
get os.environ.get("WREN_HOME", Path.home() / ".wren"), pass the retrieved
string through os.path.expanduser or expanduser on Path, then wrap with Path) so
the _WREN_HOME variable correctly resolves leading "~"; reference the symbol
_WREN_HOME and use os.path.expanduser (or Path(...).expanduser()) to implement
the change.
🪄 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: defaults
Review profile: CHILL
Plan: Pro
Run ID: ecb2cac1-dc08-4211-ba3a-4b1660b231d2
📒 Files selected for processing (8)
wren/src/wren/cli.pywren/src/wren/config.pywren/src/wren/engine.pywren/src/wren/policy.pywren/tests/unit/test_config.pywren/tests/unit/test_cte_rewriter.pywren/tests/unit/test_engine.pywren/tests/unit/test_policy.py
✅ Files skipped from review due to trivial changes (2)
- wren/tests/unit/test_cte_rewriter.py
- wren/tests/unit/test_policy.py
🚧 Files skipped from review as they are similar to previous changes (3)
- wren/src/wren/config.py
- wren/tests/unit/test_engine.py
- wren/src/wren/policy.py
- Add .expanduser() to _WREN_HOME so WREN_HOME=~/.wren resolves correctly - Change error fallback guard from strict_mode-only to strict_mode or denied_functions, ensuring policy enforcement is never bypassed when only denied_functions is configured Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@wren/src/wren/cli.py`:
- Around line 315-319: The config loading block that calls
load_config(_WREN_HOME) only catches WrenError; update the exception handling to
also catch OSError (filesystem errors) so filesystem issues are reported the
same way as WrenError: catch OSError alongside WrenError around load_config,
call typer.echo with the error message and raise typer.Exit(1) from the caught
exception. Reference the load_config call, the WrenError exception class, and
typer.Exit to locate where to add the OSError handling.
- Around line 145-155: The current try/except only catches WrenError, but
load_config can raise other exceptions (e.g., PermissionError from
config_path.exists()), so broaden the exception handling around the load_config
call and WrenEngine creation: catch Exception (in addition to WrenError) when
calling load_config and constructing WrenEngine in cli.py, convert it to a clean
CLI error by calling typer.echo with the exception message and then raise
typer.Exit(1) from the caught exception; reference load_config, WrenError,
WrenEngine and typer.Exit to locate the block to modify.
🪄 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: defaults
Review profile: CHILL
Plan: Pro
Run ID: 4f0b34f0-2cbd-4693-8dbd-208d417d8888
📒 Files selected for processing (2)
wren/src/wren/cli.pywren/src/wren/engine.py
Filesystem errors (e.g. PermissionError) from load_config now produce a clean "Error:" message instead of a raw traceback. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
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
1 participant
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
SELECT * FROM pg_shadoware rejected withMODEL_NOT_FOUNDdenied_functionslist blocks dangerous functions (e.g.,pg_read_file,dblink) withBLOCKED_FUNCTIONerror~/.wren/config.jsononly (no CLI args), withWREN_HOMEenv var to override the config directoryConfig format
{ "strict_mode": true, "denied_functions": ["pg_read_file", "dblink", "lo_import"] }Files changed
wren/src/wren/config.pyWrenConfigdataclass +load_config()wren/src/wren/policy.pyvalidate_sql_policy()with table + function checkswren/src/wren/model/error.pyMODEL_NOT_FOUND,BLOCKED_FUNCTIONcodes +SQL_POLICY_CHECKphasewren/src/wren/engine.py_plan()wren/src/wren/cli.pyWREN_HOME, pass to enginewren/README.mdTest plan
🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Documentation
Tests