docs: add support for Aurora PostgreSQL and MySQL#2320
Conversation
github-actions
Bot
added
documentation
WalkthroughThis PR updates skill metadata and documentation to reflect Aurora PostgreSQL and Aurora MySQL as supported data sources. Version numbers are incremented for three skills, and their documentation is expanded with Aurora-specific guidance, connection examples, and datasource lists. ChangesSkill Metadata and Documentation Updates
Estimated code review effort🎯 1 (Trivial) | ⏱️ ~3 minutes Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (3 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}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
core/wren/src/wren/model/data_source.py (1)
107-139:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winMissing timeout configuration for Aurora PostgreSQL.
The
get_connection_info()method applies defaultconnect_timeoutandstatement_timeoutconfiguration forDataSource.postgres(lines 108-118), butDataSource.aurora_postgresis not handled. SinceAuroraPostgresConnectionInfoinherits fromPostgresConnectionInfoand supports the same connection parameters, Aurora PostgreSQL connections should receive the same timeout defaults.🛠️ Proposed fix to add Aurora PostgreSQL timeout handling
match self: - case DataSource.postgres: + case DataSource.postgres | DataSource.aurora_postgres: kwargs = info.kwargs if info.kwargs else {}🤖 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/model/data_source.py` around lines 107 - 139, get_connection_info currently applies connect_timeout and statement_timeout only for DataSource.postgres; add the same handling for DataSource.aurora_postgres by adding a case for DataSource.aurora_postgres that mirrors the Postgres branch: ensure info.kwargs exists, set kwargs["connect_timeout"] to 120 if missing, compose/append the -c statement_timeout=...s into kwargs["options"] using headers.get(X_WREN_DB_STATEMENT_TIMEOUT, 180), and assign back to info.kwargs; reference the match on self and the DataSource.postgres branch to copy behavior.
🧹 Nitpick comments (1)
skills/index.json (1)
72-85: 💤 Low valueConsider adding tags for other documented datasources.
The tag list includes Aurora variants but omits several datasources that are documented in
wren-usage/SKILL.mdline 219 (clickhouse, trino, mssql, databricks, redshift, spark, athena, oracle). Adding these tags would improve skill discoverability when users search by datasource.🤖 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 `@skills/index.json` around lines 72 - 85, The tags array under the "tags" key in skills/index.json is missing several documented datasource tags; update that "tags" list (the array shown in the diff) to include clickhouse, trino, mssql, databricks, redshift, spark, athena, and oracle so the skill matches the datasources listed in wren-usage/SKILL.md and improves discoverability.
🤖 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.
Outside diff comments:
In `@core/wren/src/wren/model/data_source.py`:
- Around line 107-139: get_connection_info currently applies connect_timeout and
statement_timeout only for DataSource.postgres; add the same handling for
DataSource.aurora_postgres by adding a case for DataSource.aurora_postgres that
mirrors the Postgres branch: ensure info.kwargs exists, set
kwargs["connect_timeout"] to 120 if missing, compose/append the -c
statement_timeout=...s into kwargs["options"] using
headers.get(X_WREN_DB_STATEMENT_TIMEOUT, 180), and assign back to info.kwargs;
reference the match on self and the DataSource.postgres branch to copy behavior.
---
Nitpick comments:
In `@skills/index.json`:
- Around line 72-85: The tags array under the "tags" key in skills/index.json is
missing several documented datasource tags; update that "tags" list (the array
shown in the diff) to include clickhouse, trino, mssql, databricks, redshift,
spark, athena, and oracle so the skill matches the datasources listed in
wren-usage/SKILL.md and improves discoverability.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 20fa4f74-992c-4633-8744-af8afcecdfe2
📒 Files selected for processing (11)
core/wren-core-base/manifest-macro/src/lib.rscore/wren-core/core/src/mdl/dialect/inner_dialect.rscore/wren/src/wren/connector/factory.pycore/wren/src/wren/model/__init__.pycore/wren/src/wren/model/data_source.pycore/wren/src/wren/model/field_registry.pyskills/index.jsonskills/versions.jsonskills/wren-generate-mdl/SKILL.mdskills/wren-onboarding/SKILL.mdskills/wren-usage/SKILL.md
kartikshastrakar
left a comment
kartikshastrakar
left a comment
There was a problem hiding this comment.
check
Commented in CodeRabbit Change Stack
kartikshastrakar
left a comment
kartikshastrakar
left a comment
There was a problem hiding this comment.
check
Commented in CodeRabbit Change Stack
kartikshastrakar
force-pushed
the
feat/add-clickhouse-aurora-support
branch
from
c4c2e67 to
35e7ded
Compare
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/model/data_source.py`:
- Around line 88-89: The get_connection_info logic only applies Postgres timeout
defaults to DataSource.postgres, so update the conditional to also include
DataSource.aurora_postgres (e.g., change checks like "if datasource ==
DataSource.postgres" to "if datasource in (DataSource.postgres,
DataSource.aurora_postgres)" or add an equivalent branch) so that
connect_timeout and the "x-wren-db-statement_timeout" option are injected for
Aurora Postgres as well; reference the DataSource enum members postgres and
aurora_postgres and the get_connection_info function to locate and adjust the
code.
🪄 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: e642a4be-2859-44f6-83cf-2991da317217
📒 Files selected for processing (11)
core/wren-core-base/manifest-macro/src/lib.rscore/wren-core/core/src/mdl/dialect/inner_dialect.rscore/wren/src/wren/connector/factory.pycore/wren/src/wren/model/__init__.pycore/wren/src/wren/model/data_source.pycore/wren/src/wren/model/field_registry.pyskills/index.jsonskills/versions.jsonskills/wren-generate-mdl/SKILL.mdskills/wren-onboarding/SKILL.mdskills/wren-usage/SKILL.md
✅ Files skipped from review due to trivial changes (3)
- skills/versions.json
- skills/wren-usage/SKILL.md
- skills/wren-onboarding/SKILL.md
| aurora_postgres = auto() | ||
| aurora_mysql = auto() |
There was a problem hiding this comment.
Apply Postgres timeout defaults to Aurora Postgres too.
Line 108 only matches DataSource.postgres, so Aurora Postgres skips connect_timeout and x-wren-db-statement_timeout option injection in get_connection_info.
Proposed fix
- case DataSource.postgres:
+ case DataSource.postgres | DataSource.aurora_postgres:
kwargs = info.kwargs if info.kwargs else {}
if "connect_timeout" not in kwargs:
kwargs["connect_timeout"] = 120
options = kwargs.get("options", "")
if "statement_timeout" not in options:
if options:
options += " "
options += f"-c statement_timeout={headers.get(X_WREN_DB_STATEMENT_TIMEOUT, 180)}s"
kwargs["options"] = options
info.kwargs = kwargs🤖 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/model/data_source.py` around lines 88 - 89, The
get_connection_info logic only applies Postgres timeout defaults to
DataSource.postgres, so update the conditional to also include
DataSource.aurora_postgres (e.g., change checks like "if datasource ==
DataSource.postgres" to "if datasource in (DataSource.postgres,
DataSource.aurora_postgres)" or add an equivalent branch) so that
connect_timeout and the "x-wren-db-statement_timeout" option are injected for
Aurora Postgres as well; reference the DataSource enum members postgres and
aurora_postgres and the get_connection_info function to locate and adjust the
code.
|
@kartikshastrakar, Thanks for working on this. However, I don't see anything different between the official PostgreSQL and Aurora ones. To reduce maintenance effort, if the behavior of the connection and syntax are the same, I think the Aurora user can just use the Postgres and MySQL connectors to connect to Aurora. No need to add a special data source in Wren AI. |
kartikshastrakar
force-pushed
the
feat/add-clickhouse-aurora-support
branch
from
35e7ded to
0b218d0
Compare
kartikshastrakar
left a comment
kartikshastrakar
left a comment
There was a problem hiding this comment.
@goldmedal That completely makes sense! I agree that adding separate data sources creates unnecessary maintenance overhead since the standard ones work out of the box for Aurora.
I have reverted all the code changes from the core platform (wren-core, wren-core-base, wren-core-py, etc.). The codebase is now completely clean of any duplicate data sources.
Instead, I just updated the AI Agent Skills (wren-onboarding, wren-generate-mdl, wren-usage). This will teach the AI agents that if a user wants to connect to AWS Aurora, it should seamlessly route them through the standard postgres or mysql connectors.
I force-pushed these documentation updates. Let me know if these tweaks look good to merge!
goldmedal
changed the title
goldmedal
changed the title
goldmedal
left a comment
goldmedal
left a comment
There was a problem hiding this comment.
Thanks @kartikshastrakar, look great 👍
|
@goldmedal Thank you |
3 participants
Summary by CodeRabbit
Release Notes
New Features
Documentation