feat(wren-ui): Support Athena connector#1754
Conversation
WalkthroughThis change adds support for the Athena data source throughout the application. It introduces new types, schema entries, configuration options, UI components for Athena, and updates related enums and mappings. Additionally, it modifies dashboard schedule types to make certain fields optional and nullable, and updates error messages for Athena-specific fields. Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant UI
participant Server
participant Athena
User->>UI: Selects Athena as data source
UI->>UI: Renders AthenaProperties form
User->>UI: Fills AWS credentials and config
UI->>Server: Submits Athena connection info
Server->>Athena: Uses credentials to connect
Athena-->>Server: Returns connection status
Server-->>UI: Responds with result
UI-->>User: Displays connection status
Possibly related PRs
Suggested reviewers
Poem
📜 Recent review detailsConfiguration used: CodeRabbit UI ⛔ Files ignored due to path filters (1)
📒 Files selected for processing (15)
✅ Files skipped from review due to trivial changes (1)
🚧 Files skipped from review as they are similar to previous changes (14)
⏰ Context from checks skipped due to timeout of 90000ms (3)
✨ Finishing Touches
🧪 Generate Unit Tests
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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
115ba58 to
2b13d70
Compare
There was a problem hiding this comment.
Actionable comments posted: 6
🧹 Nitpick comments (4)
wren-ui/src/apollo/server/schema.ts (1)
47-58: Ensure consistent enum ordering for DataSourceNameThe
ATHENAmember is inserted at the top of this GraphQL enum, whereas TypeScript enums place it elsewhere. For predictable ordering in introspection-driven UIs, consider sorting the values alphabetically or aligning the order with your TypeScript definitions.wren-ui/src/pages/home/dashboard.tsx (1)
12-14: Circular-dependency / tree-shaking risk – export theScheduletype from a dedicated file
CacheSettingsDrawernow re-exportsSchedule, forcing consumers of a type to import the whole drawer bundle.
Consider extractingScheduleto its own module (e.g.types.ts) or to the GraphQL-generated types to avoid unnecessary bundle weight and potential circular imports.wren-ui/src/utils/dataSourceType.ts (1)
11-12: Import path will break the chunk if the component file is code-split elsewhereIf
AthenaPropertiesis heavy (SDK validation, etc.), consider lazy loading it with dynamicimport()as done for some other data-sources to keep initial bundle small.wren-ui/src/utils/error/dictionary.ts (1)
52-63: Great – dedicated error texts for AWS fieldsClear, actionable messages improve UX. Consider adding simple format validation (e.g. S3 URI must start with
s3://, region regex) in the form component for even earlier feedback.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (1)
wren-ui/public/images/dataSource/athena.svgis excluded by!**/*.svg
📒 Files selected for processing (15)
wren-ui/src/apollo/client/graphql/__types__.ts(3 hunks)wren-ui/src/apollo/client/graphql/dashboard.generated.ts(1 hunks)wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts(3 hunks)wren-ui/src/apollo/server/dataSource.ts(2 hunks)wren-ui/src/apollo/server/mdl/mdlBuilder.ts(1 hunks)wren-ui/src/apollo/server/mdl/type.ts(1 hunks)wren-ui/src/apollo/server/repositories/projectRepository.ts(1 hunks)wren-ui/src/apollo/server/schema.ts(1 hunks)wren-ui/src/apollo/server/types/dataSource.ts(1 hunks)wren-ui/src/components/pages/setup/dataSources/AthenaProperties.tsx(1 hunks)wren-ui/src/components/pages/setup/utils.tsx(1 hunks)wren-ui/src/pages/home/dashboard.tsx(2 hunks)wren-ui/src/utils/dataSourceType.ts(4 hunks)wren-ui/src/utils/enum/dataSources.ts(1 hunks)wren-ui/src/utils/error/dictionary.ts(1 hunks)
🔇 Additional comments (13)
wren-ui/src/utils/enum/dataSources.ts (1)
11-11: Add Athena to supported data sourcesThe new enum entry is correct and aligns with the PR objectives.
wren-ui/src/apollo/server/mdl/type.ts (1)
86-90: Add Athena to WrenEngineDataSourceTypeThe new enum entry correctly integrates Athena as a supported engine type.
wren-ui/src/apollo/server/mdl/mdlBuilder.ts (1)
494-497: Map Athena in buildDataSourceThe new switch-case correctly maps
DataSourceName.ATHENAtoWrenEngineDataSourceType.ATHENA, ensuring Athena manifests are built properly.wren-ui/src/components/pages/setup/utils.tsx (1)
105-109: 👍 Athena option correctly wiredEntry looks consistent with existing options (guide URL, disabled flag). No issues spotted.
wren-ui/src/utils/dataSourceType.ts (3)
60-62: Display name alignment looks good“Athena” follows existing naming conventions.
87-88: Properties component registeredRegistration completes the triad (image, label, form component). Nice.
33-35: ```shell
#!/bin/bashSearch the entire repository for athena.svg (case-insensitive)
find . -type f -iname athena.svg
</details> <details> <summary>wren-ui/src/apollo/server/repositories/projectRepository.ts (1)</summary> `96-97`: **Union order fine; don’t forget downstream narrowing** Wherever `connectionInfo` is type-narrowed (e.g. adaptor layer) add a case for `ATHENA`, otherwise TypeScript’s exhaustiveness checks will fail silently if you used `as unknown as`. </details> <details> <summary>wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts (2)</summary> `69-75`: **Athena interface is clear and consistent with ibis-py expectations** The field names follow ibis’ parameter conventions (`aws_access_key_id`, `s3_staging_dir`, etc.), so no mapping friction is expected on the Python side. --- `86-99`: **Enum / URL map updated correctly** `SupportedDataSource.ATHENA` and its URL mapping were added without disturbing existing entries—good housekeeping. </details> <details> <summary>wren-ui/src/apollo/client/graphql/__types__.ts (2)</summary> `293-300`: **`frequency`, `hour`, `minute` turned into `Maybe` – tighten caller checks** The schedule fields are now nullable/optional. Any consumer code that previously did: ```ts const { frequency, hour, minute } = dashboard.schedule;without null-guards will now break at runtime when
scheduleor any of its children come backnull.
Please grep the UI for direct property access and add safe-navigation (?.) or fallback defaults where needed.
315-323: NewATHENAentry added toDataSourceNameenum – review exhaustivenessEvery
switch/ifthat branches onDataSourceName(logo mapping, validator, connection form, etc.) must handle theATHENAcase, otherwise the compiler won’t complain (string enum) but users will see empty UI / “unsupported” errors.wren-ui/src/apollo/client/graphql/dashboard.generated.ts (1)
59-59: GeneratedDashboardQueryreflects nullable schedule – audit UI for null-checks
dashboard.scheduleand its nested fields are now optional in the generated hook result.
Before merging, runtsc --noEmit(or your CI) and manually test the dashboard page with no schedule configured to be sure there are no uncontrolledundefinedaccesses or rendering glitches.
| CLICK_HOUSE = 'CLICK_HOUSE', | ||
| TRINO = 'TRINO', | ||
| SNOWFLAKE = 'SNOWFLAKE', | ||
| ATHENA = 'ATHENA', |
There was a problem hiding this comment.
Missing Athena properties in DataSourceProperties union
You’ve added ATHENA to DataSourceName but haven’t defined or included an AthenaDataSourceProperties interface in the DataSourceProperties type. Please create an interface for Athena-specific fields (s3_staging_dir, aws_access_key_id, aws_secret_access_key, region_name, schema_name) and add it to the union to maintain type safety.
🤖 Prompt for AI Agents
In wren-ui/src/apollo/server/types/dataSource.ts at line 11, you added ATHENA to
DataSourceName but did not define AthenaDataSourceProperties or include it in
the DataSourceProperties union. Define an interface AthenaDataSourceProperties
with fields s3_staging_dir, aws_access_key_id, aws_secret_access_key,
region_name, and schema_name, then add this interface to the
DataSourceProperties union type to ensure type safety for Athena data sources.
| <DashboardHeader | ||
| isSupportCached={isSupportCached} | ||
| schedule={data?.dashboard?.schedule} | ||
| schedule={data?.dashboard?.schedule as Schedule} |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Avoid unchecked as Schedule; rely on proper typing instead
The explicit cast suppresses the compiler even when data?.dashboard?.schedule is:
undefined | null(now allowed by the schema change);- structurally incompatible after future edits.
Prefer a safe-guard:
- schedule={data?.dashboard?.schedule as Schedule}
+ schedule={data?.dashboard?.schedule ?? undefined}and declare the prop in DashboardHeader as schedule?: Schedule to keep strict-null-checks useful.
📝 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.
| schedule={data?.dashboard?.schedule as Schedule} | |
| - schedule={data?.dashboard?.schedule as Schedule} | |
| + schedule={data?.dashboard?.schedule ?? undefined} |
🤖 Prompt for AI Agents
In wren-ui/src/pages/home/dashboard.tsx at line 112, remove the explicit cast to
Schedule on data?.dashboard?.schedule to avoid unchecked type assertion.
Instead, rely on proper typing by allowing the schedule prop in DashboardHeader
to be optional (schedule?: Schedule). This preserves strict-null-checks and
safely handles cases where schedule may be undefined or null.
| export interface ATHENA_CONNECTION_INFO { | ||
| database: string; | ||
| schema: string; | ||
| s3StagingDir: string; | ||
| awsRegion: string; | ||
| awsAccessKey: string; | ||
| awsSecretKey: string; | ||
| } |
There was a problem hiding this comment.
Interface added, but remember to mask secrets when logging
awsAccessKey / awsSecretKey are now part of connectionInfo. Ensure any logging/serialization that surfaces connectionInfo strips these fields; other connectors already have masking helpers – please extend them to Athena.
🤖 Prompt for AI Agents
In wren-ui/src/apollo/server/repositories/projectRepository.ts around lines 86
to 93, the new ATHENA_CONNECTION_INFO interface includes sensitive fields
awsAccessKey and awsSecretKey. To fix this, identify all places where
connectionInfo is logged or serialized and update the code to mask or remove
these secret fields before output. Reuse or extend existing masking helper
functions used by other connectors to ensure consistent handling of sensitive
data for Athena connection info.
| label="AWS access key ID" | ||
| name="awsAccessKey" | ||
| required | ||
| rules={[ | ||
| { | ||
| required: true, | ||
| message: ERROR_TEXTS.CONNECTION.AWS_ACCESS_KEY.REQUIRED, | ||
| }, | ||
| ]} | ||
| > | ||
| <Input /> | ||
| </Form.Item> |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Mask the AWS access-key field & treat it as sensitive
Even though the access-key ID is less sensitive than the secret key, showing it in plain text invites shoulder-surfing and screenshots in bug reports.
- >
- <Input />
+ >
+ <Input.Password />If you adopt this, remember to
- add
"awsAccessKey"tosensitivePropsfor Athena indataSource.ts, and - re-encrypt existing records (if any).
📝 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.
| label="AWS access key ID" | |
| name="awsAccessKey" | |
| required | |
| rules={[ | |
| { | |
| required: true, | |
| message: ERROR_TEXTS.CONNECTION.AWS_ACCESS_KEY.REQUIRED, | |
| }, | |
| ]} | |
| > | |
| <Input /> | |
| </Form.Item> | |
| label="AWS access key ID" | |
| name="awsAccessKey" | |
| required | |
| rules={[ | |
| { | |
| required: true, | |
| message: ERROR_TEXTS.CONNECTION.AWS_ACCESS_KEY.REQUIRED, | |
| }, | |
| ]} | |
| > | |
| <Input.Password /> | |
| </Form.Item> |
🤖 Prompt for AI Agents
In wren-ui/src/components/pages/setup/dataSources/AthenaProperties.tsx around
lines 76 to 87, the AWS access key ID input field is currently shown in plain
text, which risks exposure through shoulder-surfing or screenshots. Change the
Input component to a password type input to mask the value. Additionally, add
"awsAccessKey" to the sensitiveProps list for Athena in dataSource.ts and ensure
any existing stored access keys are re-encrypted to maintain security.
| const dataSource = { | ||
| // Athena | ||
| [DataSourceName.ATHENA]: { | ||
| sensitiveProps: ['awsSecretKey'], | ||
| toIbisConnectionInfo(connectionInfo) { | ||
| const decryptedConnectionInfo = decryptConnectionInfo( | ||
| DataSourceName.ATHENA, | ||
| connectionInfo, | ||
| ); | ||
| const { awsAccessKey, awsRegion, awsSecretKey, s3StagingDir, schema } = | ||
| decryptedConnectionInfo as ATHENA_CONNECTION_INFO; | ||
| const res: IbisAthenaConnectionInfo = { | ||
| aws_access_key_id: awsAccessKey, | ||
| aws_secret_access_key: awsSecretKey, | ||
| region_name: awsRegion, | ||
| s3_staging_dir: s3StagingDir, | ||
| schema_name: schema, | ||
| }; | ||
| return res; | ||
| }, | ||
| } as IDataSourceConnectionInfo< | ||
| ATHENA_CONNECTION_INFO, | ||
| IbisAthenaConnectionInfo | ||
| >, | ||
|
|
There was a problem hiding this comment.
Encrypt the AWS access-key ID as well
Currently only awsSecretKey is marked sensitive. Storing the access-key ID in clear text leaks half of the credential pair and complicates key rotation audits.
- sensitiveProps: ['awsSecretKey'],
+ // access-key ID is not strictly secret, but encrypting both halves of
+ // the credential pair avoids accidental leaks in logs / DB exports.
+ sensitiveProps: ['awsSecretKey', 'awsAccessKey'],Encryption / decryption will keep functioning because both helpers iterate over sensitiveProps.
📝 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.
| const dataSource = { | |
| // Athena | |
| [DataSourceName.ATHENA]: { | |
| sensitiveProps: ['awsSecretKey'], | |
| toIbisConnectionInfo(connectionInfo) { | |
| const decryptedConnectionInfo = decryptConnectionInfo( | |
| DataSourceName.ATHENA, | |
| connectionInfo, | |
| ); | |
| const { awsAccessKey, awsRegion, awsSecretKey, s3StagingDir, schema } = | |
| decryptedConnectionInfo as ATHENA_CONNECTION_INFO; | |
| const res: IbisAthenaConnectionInfo = { | |
| aws_access_key_id: awsAccessKey, | |
| aws_secret_access_key: awsSecretKey, | |
| region_name: awsRegion, | |
| s3_staging_dir: s3StagingDir, | |
| schema_name: schema, | |
| }; | |
| return res; | |
| }, | |
| } as IDataSourceConnectionInfo< | |
| ATHENA_CONNECTION_INFO, | |
| IbisAthenaConnectionInfo | |
| >, | |
| const dataSource = { | |
| // Athena | |
| [DataSourceName.ATHENA]: { | |
| // access-key ID is not strictly secret, but encrypting both halves of | |
| // the credential pair avoids accidental leaks in logs / DB exports. | |
| sensitiveProps: ['awsSecretKey', 'awsAccessKey'], | |
| toIbisConnectionInfo(connectionInfo) { | |
| const decryptedConnectionInfo = decryptConnectionInfo( | |
| DataSourceName.ATHENA, | |
| connectionInfo, | |
| ); | |
| const { awsAccessKey, awsRegion, awsSecretKey, s3StagingDir, schema } = | |
| decryptedConnectionInfo as ATHENA_CONNECTION_INFO; | |
| const res: IbisAthenaConnectionInfo = { | |
| aws_access_key_id: awsAccessKey, | |
| aws_secret_access_key: awsSecretKey, | |
| region_name: awsRegion, | |
| s3_staging_dir: s3StagingDir, | |
| schema_name: schema, | |
| }; | |
| return res; | |
| }, | |
| } as IDataSourceConnectionInfo< | |
| ATHENA_CONNECTION_INFO, | |
| IbisAthenaConnectionInfo | |
| >, |
🤖 Prompt for AI Agents
In wren-ui/src/apollo/server/dataSource.ts between lines 68 and 92, the
awsAccessKey is not included in the sensitiveProps array, causing it to be
stored in clear text. To fix this, add 'awsAccessKey' to the sensitiveProps
array alongside 'awsSecretKey' so that both keys are encrypted and decrypted
properly, ensuring full credential security and easier key rotation audits.
| schedule?: Maybe<DashboardSchedule>; | ||
| }; |
There was a problem hiding this comment.
🛠️ Refactor suggestion
DetailedDashboard.schedule can now be null – update component props
Components such as <DashboardHeader> and any helpers that expect schedule! need to accept undefined / null. Prop-drilling a possibly-null object into components typed as non-null will silently widen to any and defeat type-safety.
🤖 Prompt for AI Agents
In wren-ui/src/apollo/client/graphql/__types__.ts at lines 388-389, the
DetailedDashboard.schedule field is now nullable, but components like
DashboardHeader still expect a non-null schedule prop. Update the component
props and any helper functions to accept schedule as possibly undefined or null,
ensuring type definitions reflect this change to maintain type safety and
prevent implicit any types.
2b13d70 to
8bc18e8
Compare
This pull request introduces support for the Athena data source.
Description
Support Athena (Trino SQL) connector
DatabaseandAWS regionfields cannot be edited.Athena Data Source:
Logo:
Engine API information
Ref UI