fix: address CodeRabbit v3 — allow UNION queries, document --api-key#2552
fix: address CodeRabbit v3 — allow UNION queries, document --api-key#2552tusharsharma20021114-rgb wants to merge 4 commits into
Conversation
…nnectors and Rust core
…ck caching, Cargo.toml lints, Python parameterized limits, mcp auth hardening
…_safety_check, update docs for --api-key auth
…CEPT; add --api-key to docs table
github-actions
Bot
added
documentation
WalkthroughChangesRust runtime hardening
Parameterized connector execution
Python policy and service controls
Workflow safety documentation
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related issues
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)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 Checkov (3.3.8).github/workflows/labeler.yamlTraceback (most recent call last): 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 (2)
core/wren-core/core/src/mdl/utils.rs (1)
182-186: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueEliminate unnecessary
Vecallocation.Since
modelsis immediately consumed via.into_iter()just below, you can leave it as aBTreeSetand skip the final.collect::<Vec<String>>().♻️ Proposed refactor
let models = required_fields .iter() .filter_map(|c| c.relation.as_ref().map(|r| r.table().to_string())) - .collect::<BTreeSet<_>>() // Collect into a BTreeSet to remove duplicates - .into_iter() // Convert BTreeSet back into an iterator - .collect::<Vec<String>>(); + .collect::<BTreeSet<_>>(); // Collect into a BTreeSet to remove duplicates🤖 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-core/core/src/mdl/utils.rs` around lines 182 - 186, Update the models collection in the surrounding relationship-processing function to retain the BTreeSet produced by the existing collect, removing the immediate into_iter and Vec<String> collection. Ensure subsequent consumers use the BTreeSet directly while preserving duplicate removal and iteration behavior.core/wren/src/wren/connector/postgres.py (1)
254-269: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winPreserve the original SQL query in error metadata.
When a query fails during execution, logging the original user-submitted query (rather than the internal
LIMIT-wrapped subquery) provides better visibility for debugging and matches thedry_runbehavior. The refactoring either overwrites the original SQL string or explicitly passes the wrapped string into the error metadata.
core/wren/src/wren/connector/postgres.py#L254-L269: Use a different variable name (executed, params = ...) so thesqlvariable retains the original query when passed tometadata={DIALECT_SQL: sql}in the except block.core/wren/src/wren/connector/canner.py#L262-L262: Changemetadata={DIALECT_SQL: executed}tometadata={DIALECT_SQL: sql}to retain the original query string.🤖 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 254 - 269, Preserve the original query in execution-error metadata: in core/wren/src/wren/connector/postgres.py lines 254-269, update the _apply_limit_param result variable to executed while keeping sql unchanged for metadata={DIALECT_SQL: sql}; in core/wren/src/wren/connector/canner.py line 262, change the metadata value from executed to 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.
Inline comments:
In `@core/wren/src/wren/connector/duckdb.py`:
- Around line 17-47: Update _init_duckdb_s3, _init_duckdb_minio, and
_init_duckdb_gcs so DuckDB CREATE SECRET and SET statements use safely quoted
literal values rather than bound parameters. Preserve the existing credential,
region, endpoint, and SSL values while escaping or quoting them through a safe
helper before embedding them in each SQL statement.
---
Nitpick comments:
In `@core/wren-core/core/src/mdl/utils.rs`:
- Around line 182-186: Update the models collection in the surrounding
relationship-processing function to retain the BTreeSet produced by the existing
collect, removing the immediate into_iter and Vec<String> collection. Ensure
subsequent consumers use the BTreeSet directly while preserving duplicate
removal and iteration behavior.
In `@core/wren/src/wren/connector/postgres.py`:
- Around line 254-269: Preserve the original query in execution-error metadata:
in core/wren/src/wren/connector/postgres.py lines 254-269, update the
_apply_limit_param result variable to executed while keeping sql unchanged for
metadata={DIALECT_SQL: sql}; in core/wren/src/wren/connector/canner.py line 262,
change the metadata value from executed to sql.
🪄 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: 4bf555da-da5f-4482-a526-5dad20d6fe0b
📒 Files selected for processing (52)
.github/workflows/labeler.yamlcore/wren-core-base/Cargo.tomlcore/wren-core-wasm/src/lib.rscore/wren-core/Cargo.tomlcore/wren-core/benchmarks/Cargo.tomlcore/wren-core/core/Cargo.tomlcore/wren-core/core/src/logical_plan/analyze/model_generation.rscore/wren-core/core/src/logical_plan/analyze/plan.rscore/wren-core/core/src/logical_plan/analyze/relation_chain.rscore/wren-core/core/src/logical_plan/unparser.rscore/wren-core/core/src/mdl/context.rscore/wren-core/core/src/mdl/dataset.rscore/wren-core/core/src/mdl/dialect/inner_dialect.rscore/wren-core/core/src/mdl/dialect/wren_dialect.rscore/wren-core/core/src/mdl/function/remote_function.rscore/wren-core/core/src/mdl/lineage.rscore/wren-core/core/src/mdl/mod.rscore/wren-core/core/src/mdl/type_planner.rscore/wren-core/core/src/mdl/utils.rscore/wren-core/sqllogictest/Cargo.tomlcore/wren-core/wren-example/Cargo.tomlcore/wren/src/wren/config.pycore/wren/src/wren/connector/athena.pycore/wren/src/wren/connector/base.pycore/wren/src/wren/connector/canner.pycore/wren/src/wren/connector/clickhouse.pycore/wren/src/wren/connector/databricks.pycore/wren/src/wren/connector/datafusion.pycore/wren/src/wren/connector/duckdb.pycore/wren/src/wren/connector/mssql.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/src/wren/engine.pycore/wren/src/wren/mcp_server.pycore/wren/src/wren/policy.pycore/wren/src/wren/serve_cli.pycore/wren/tests/conftest.pycore/wren/tests/unit/test_athena_connector.pycore/wren/tests/unit/test_athena_limit_pushdown.pycore/wren/tests/unit/test_canner_semicolon.pycore/wren/tests/unit/test_config.pycore/wren/tests/unit/test_connector_base.pycore/wren/tests/unit/test_duckdb_file_listing.pycore/wren/tests/unit/test_mysql_helpers.pycore/wren/tests/unit/test_postgres_semicolon_unlimited.pycore/wren/tests/unit/test_redshift_semicolon.pydocs/core/guides/mcp.mddocs/core/reference/cli.md
| connection.execute( | ||
| "CREATE SECRET wren_s3 (TYPE S3, KEY_ID ?, SECRET ?, REGION ?)", | ||
| [ | ||
| info.access_key.get_secret_value(), | ||
| info.secret_key.get_secret_value(), | ||
| info.region, | ||
| ], | ||
| ) | ||
|
|
||
|
|
||
| def _init_duckdb_minio(connection, info: MinioFileConnectionInfo): | ||
| connection.execute(f""" | ||
| CREATE SECRET wren_minio ( | ||
| TYPE S3, | ||
| KEY_ID '{_escape_sql(info.access_key.get_secret_value())}', | ||
| SECRET '{_escape_sql(info.secret_key.get_secret_value())}', | ||
| REGION 'ap-northeast-1' | ||
| )""") | ||
| connection.execute( | ||
| "CREATE SECRET wren_minio (TYPE S3, KEY_ID ?, SECRET ?, REGION 'ap-northeast-1')", | ||
| [ | ||
| info.access_key.get_secret_value(), | ||
| info.secret_key.get_secret_value(), | ||
| ], | ||
| ) | ||
| connection.execute("SET s3_endpoint=?", [info.endpoint]) | ||
| connection.execute("SET s3_url_style='path'") | ||
| connection.execute("SET s3_use_ssl=?", [info.ssl_enabled]) | ||
|
|
||
|
|
||
| def _init_duckdb_gcs(connection, info: GcsFileConnectionInfo): | ||
| connection.execute(f""" | ||
| CREATE SECRET wren_gcs ( | ||
| TYPE GCS, | ||
| KEY_ID '{_escape_sql(info.key_id.get_secret_value())}', | ||
| SECRET '{_escape_sql(info.secret_key.get_secret_value())}' | ||
| )""") | ||
| connection.execute( | ||
| "CREATE SECRET wren_gcs (TYPE GCS, KEY_ID ?, SECRET ?)", | ||
| [ | ||
| info.key_id.get_secret_value(), | ||
| info.secret_key.get_secret_value(), | ||
| ], | ||
| ) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
python -c "
import duckdb
conn = duckdb.connect()
try:
conn.execute('CREATE SECRET my_sec (TYPE S3, KEY_ID ?, SECRET ?)', ['test', 'test'])
print('CREATE SECRET parameterized: SUCCESS')
except Exception as e:
print('CREATE SECRET parameterized: ERROR:', str(e))
try:
conn.execute('SET s3_endpoint=?', ['test'])
print('SET parameterized: SUCCESS')
except Exception as e:
print('SET parameterized: ERROR:', str(e))
"Repository: Canner/WrenAI
Length of output: 271
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the target file and nearby context.
sed -n '1,220p' core/wren/src/wren/connector/duckdb.py
# Find DuckDB dependency/version declarations in the repo.
rg -n "duckdb|DuckDB" -g 'pyproject.toml' -g 'requirements*.txt' -g 'poetry.lock' -g 'Pipfile*' -g 'uv.lock' -g 'package.json' -g 'package-lock.json' -g 'pnpm-lock.yaml' -g 'yarn.lock' .
# Search for other CREATE SECRET / SET usages that might indicate the intended pattern.
rg -n "CREATE SECRET|SET s3_|s3_endpoint|s3_url_style|s3_use_ssl" core/wren/src .Repository: Canner/WrenAI
Length of output: 18956
🌐 Web query:
DuckDB CREATE SECRET parameter binding utility statements SET ? prepared parameters documentation
💡 Result:
In DuckDB, the CREATE SECRET statement does not support the use of prepared statement parameters (such as ?, $1, or named parameters) directly within the command [1]. Attempts to use this syntax will typically result in a parser error [1]. If you need to use dynamic values for secrets, the recommended approach is to use expressions instead of standard prepared parameters. Recent updates to DuckDB allow for the use of expressions within the CREATE SECRET statement [2][3]. For example, you can set a session variable and then use that variable within the secret definition [2][3]: SET VARIABLE my_secret_value = 'your_secret_here'; CREATE SECRET my_secret (TYPE s3, SECRET getvariable('my_secret_value')); Regarding the SET statement, DuckDB does support prepared parameters in combination with SET VARIABLE, allowing for dynamic assignment of session variables through prepared statements [4][5]. However, general configuration settings used with the SET command (e.g., SET memory_limit = '1GB') may behave differently depending on the specific implementation, and it is standard practice to use SET VARIABLE for dynamic data injection [4][5]. Top results: [1][4][5][2][3]
Citations:
- 1: Prepared statements fail for creating secrets duckdb/duckdb#13459
- 2: Expressions in create secret duckdb/duckdb#15801
- 3: https://github.com/duckdb/duckdb/blob/0da6e1e8/test/sql/secrets/create_secret_expression.test
- 4: https://github.com/duckdb/duckdb/blob/master/test/api/test_prepared_api.cpp
- 5: https://github.com/duckdb/duckdb/blob/0da6e1e8/test/api/test_prepared_api.cpp
🌐 Web query:
DuckDB SET statement prepared parameters s3_endpoint s3_use_ssl ? parameter binding documentation
💡 Result:
In DuckDB, the SET statement is used to modify configuration options [1][2]. While the s3_endpoint and s3_use_ssl settings were historically used to configure S3 access [3][4], they are now considered part of a legacy authentication scheme [3][4]. The recommended modern approach is to use the Secrets manager via the CREATE SECRET statement [3][5][6]. Regarding prepared parameters and the SET statement: DuckDB does support using prepared statements for SET operations [7]. You can use parameter binding syntax (such as?, $1, or $param) within a prepared SET statement to dynamically set configuration values, including legacy S3 settings if necessary [8][9][7]. For example, using positional parameters in a prepared SET statement: PREPARE set_s3_endpoint AS SET s3_endpoint = $1; EXECUTE set_s3_endpoint('my-endpoint.com'); However, for S3 configuration, you should prefer the modern secret-based approach [3][10]: CREATE OR REPLACE SECRET my_s3_secret ( TYPE s3, PROVIDER config, ENDPOINT $endpoint_param, USE_SSL $use_ssl_param,...); Note that the availability of parameter binding depends on the specific client driver being used; for instance, some clients like JDBC primarily support auto-incremented (?) parameters [8][9]. Top results: [3][8][7][2][5]
Citations:
- 1: https://duckdb.org/docs/current/configuration/overview
- 2: https://www.duckdb.org/docs/current/sql/statements/set
- 3: https://duckdb.org/docs/current/core_extensions/httpfs/s3api_legacy_authentication
- 4: https://duckdb.org/docs/current/core_extensions/httpfs/s3api_legacy_authentication.html
- 5: https://duckdb.org/docs/current/core%5Fextensions/httpfs/s3api
- 6: https://duckdb.org/docs/lts/core_extensions/httpfs/s3api.html
- 7: https://github.com/duckdb/duckdb/blob/0da6e1e8/test/api/test_prepared_api.cpp
- 8: https://duckdb.org/docs/lts/sql/query_syntax/prepared_statements
- 9: https://duckdb.org/docs/current/sql/query%5Fsyntax/prepared%5Fstatements
- 10: https://duckdb.org/docs/current/core_extensions/aws
🌐 Web query:
DuckDB Python connection.execute parameterized SET statement ? s3_endpoint s3_use_ssl
💡 Result:
In DuckDB, the SET statement is a DDL/configuration command, not a standard DML query (like SELECT or INSERT). Consequently, you cannot use parameterized inputs (e.g., ? or $1) directly with connection.execute("SET s3_endpoint = ?", [endpoint]) [1][2]. If you need to configure S3 settings programmatically, you should use f-strings or string formatting in Python to construct the SET statement, ensuring you sanitize the input to prevent injection if the values come from untrusted sources [1][2]. Recommended approach: con.execute(f"SET s3_endpoint = '{endpoint}'") con.execute(f"SET s3_use_ssl = {str(use_ssl).lower}") Important Note on S3 Configuration: While SET s3_endpoint and SET s3_use_ssl are available as legacy configuration methods, they are deprecated for managing secrets and sensitive credentials [3][4]. The modern, recommended way to configure S3 access is through DuckDB's Secret Manager [5][6]. This allows you to store configurations securely: CREATE OR REPLACE SECRET my_s3_secret ( TYPE S3, ENDPOINT 'your_endpoint_here', USE_SSL true, KEY_ID '...', SECRET '...'); Using the Secret Manager is safer as it avoids risks associated with leaking credentials through logs or accidental output [3][7].
Citations:
- 1: https://www.duckdb.org/docs/current/clients/python/dbapi
- 2: https://duckdb.org/docs/lts/sql/statements/set
- 3: https://duckdb.org/docs/current/core_extensions/httpfs/s3api_legacy_authentication
- 4: https://raw.githubusercontent.com/duckdb/duckdb-web/refs/heads/main/docs/current/core%5Fextensions/httpfs/s3api%5Flegacy%5Fauthentication.md
- 5: https://duckdb.org/docs/current/core%5Fextensions/httpfs/s3api
- 6: https://duckdb.org/docs/lts/core_extensions/httpfs/s3api.html
- 7: https://www.duckdb.org/docs/lts/core_extensions/httpfs/s3api_legacy_authentication
Inline literals for DuckDB secret/setup statements (core/wren/src/wren/connector/duckdb.py:17-37) CREATE SECRET and SET here don’t accept bound params, so S3/MinIO/GCS initialization will fail at runtime. Use safely quoted literals instead of execute(..., params).
🤖 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/duckdb.py` around lines 17 - 47, Update
_init_duckdb_s3, _init_duckdb_minio, and _init_duckdb_gcs so DuckDB CREATE
SECRET and SET statements use safely quoted literal values rather than bound
parameters. Preserve the existing credential, region, endpoint, and SSL values
while escaping or quoting them through a safe helper before embedding them in
each SQL statement.
1 participant
Summary
Fixes the 2 new CodeRabbit comments from the third review on PR #2551:
policy.py— Change whitelist from(exp.Select, exp.Explain)to(exp.Query, exp.Explain, exp.Subquery)to allowUNION,INTERSECT,EXCEPT, and parenthesized SELECT queries (BI tools compatibility).cli.md— Add--api-keyrow to thewren serve mcpoptions table.Verification
pytest: 987 passed, 65 failed (pre-existing), 59 skippedcargo check --all-targets --all-features: cleanSummary by CodeRabbit