feat(wren-ui): Support Databricks connector#2059
Conversation
fredalai
changed the title
WalkthroughAdds Databricks support across client and server: new enums/types, adaptor and repository connection types, data source mappings and MDL type, UI form component and setup flow changes for Personal Access Token and Service Principal auth, plus related error messages and utilities. Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant UI as Setup UI
participant Hook as useSetupConnectionDataSource
participant Service as ProjectService
participant Adaptor as IbisAdaptor
Note over User,Adaptor: Databricks setup (PAT or Service Principal)
User->>UI: Select Databricks, fill form
UI->>Hook: transformFormToProperties(form)
Hook->>Hook: getDatabricksAuthentication() (filter placeholders)
Hook-->>Service: connectionInfo (databricksType + auth fields)
UI->>Service: saveConnection(connectionInfo)
Service->>Adaptor: toIbisConnectionInfo(connectionInfo)
Adaptor->>Adaptor: decrypt sensitive fields
Adaptor-->>Service: IbisDatabricksConnectionInfo (databricks_type, serverHostname, httpPath, auth)
Service-->>UI: success / saved
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes
Possibly related PRs
Suggested labels
Suggested reviewers
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
📜 Recent review detailsConfiguration used: CodeRabbit UI Review profile: CHILL Plan: Pro ⛔ Files ignored due to path filters (1)
📒 Files selected for processing (15)
🚧 Files skipped from review as they are similar to previous changes (8)
🧰 Additional context used🧠 Learnings (5)📓 Common learnings📚 Learning: 2025-08-13T05:17:30.180ZApplied to files:
📚 Learning: 2025-03-18T10:28:10.593ZApplied to files:
📚 Learning: 2025-03-18T10:28:26.608ZApplied to files:
📚 Learning: 2025-07-09T02:43:20.433ZApplied to files:
🧬 Code graph analysis (3)wren-ui/src/apollo/client/graphql/__types__.ts (1)
wren-ui/src/hooks/useSetupConnectionDataSource.tsx (1)
wren-ui/src/utils/dataSourceType.ts (1)
🔇 Additional comments (13)
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
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (1)
wren-ui/public/images/dataSource/databricks.svgis excluded by!**/*.svg
📒 Files selected for processing (15)
wren-ui/src/apollo/client/graphql/__types__.ts(2 hunks)wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts(3 hunks)wren-ui/src/apollo/server/dataSource.ts(3 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(3 hunks)wren-ui/src/apollo/server/schema.ts(1 hunks)wren-ui/src/apollo/server/services/projectService.ts(1 hunks)wren-ui/src/apollo/server/types/dataSource.ts(1 hunks)wren-ui/src/components/pages/setup/dataSources/DatabricksProperties.tsx(1 hunks)wren-ui/src/components/pages/setup/utils.tsx(1 hunks)wren-ui/src/hooks/useSetupConnectionDataSource.tsx(4 hunks)wren-ui/src/utils/dataSourceType.ts(4 hunks)wren-ui/src/utils/enum/dataSources.ts(2 hunks)wren-ui/src/utils/error/dictionary.ts(1 hunks)
🧰 Additional context used
🧠 Learnings (6)
📓 Common learnings
Learnt from: goldmedal
Repo: Canner/WrenAI PR: 1887
File: wren-launcher/commands/dbt/data_source.go:273-276
Timestamp: 2025-08-13T05:17:30.180Z
Learning: In the Wren AI codebase, there's an intentional distinction between dbt connection types and Wren AI data source types: dbt uses "sqlserver" as the connection type, while Wren AI expects "mssql" as the data source type. The conversion layer in convertConnectionToDataSource() correctly matches "sqlserver" from dbt profiles and the WrenMSSQLDataSource.GetType() method correctly returns "mssql" for Wren AI compatibility.
Learnt from: narsik123
Repo: Canner/WrenAI PR: 1606
File: wren-ui/src/apollo/server/dataSource.ts:135-150
Timestamp: 2025-04-28T20:09:06.106Z
Learning: The Oracle data source in WrenAI includes an SSL toggle in the UI (OracleProperties component) and should pass this setting to the connection info using the pattern `...(ssl && { kwargs: { ssl: true } })` for consistency with other host-based connections like MS SQL.
📚 Learning: 2025-08-13T05:17:30.180Z
Learnt from: goldmedal
Repo: Canner/WrenAI PR: 1887
File: wren-launcher/commands/dbt/data_source.go:273-276
Timestamp: 2025-08-13T05:17:30.180Z
Learning: In the Wren AI codebase, there's an intentional distinction between dbt connection types and Wren AI data source types: dbt uses "sqlserver" as the connection type, while Wren AI expects "mssql" as the data source type. The conversion layer in convertConnectionToDataSource() correctly matches "sqlserver" from dbt profiles and the WrenMSSQLDataSource.GetType() method correctly returns "mssql" for Wren AI compatibility.
Applied to files:
wren-ui/src/apollo/client/graphql/__types__.tswren-ui/src/apollo/server/schema.tswren-ui/src/apollo/server/repositories/projectRepository.tswren-ui/src/apollo/server/types/dataSource.tswren-ui/src/apollo/server/adaptors/ibisAdaptor.tswren-ui/src/apollo/server/mdl/type.tswren-ui/src/utils/enum/dataSources.tswren-ui/src/components/pages/setup/utils.tsxwren-ui/src/apollo/server/dataSource.tswren-ui/src/utils/dataSourceType.tswren-ui/src/apollo/server/mdl/mdlBuilder.tswren-ui/src/hooks/useSetupConnectionDataSource.tsx
📚 Learning: 2025-03-18T10:28:10.593Z
Learnt from: andreashimin
Repo: Canner/WrenAI PR: 1414
File: wren-ui/src/apollo/server/utils/error.ts:0-0
Timestamp: 2025-03-18T10:28:10.593Z
Learning: The typo in error code enum names ("IDENTIED_AS_GENERAL" and "IDENTIED_AS_MISLEADING_QUERY" instead of "IDENTIFIED_AS_GENERAL" and "IDENTIFIED_AS_MISLEADING_QUERY") in wren-ui/src/apollo/server/utils/error.ts was recognized but intentionally not fixed as it was considered out of scope for PR #1414.
Applied to files:
wren-ui/src/apollo/client/graphql/__types__.tswren-ui/src/apollo/server/schema.tswren-ui/src/apollo/server/types/dataSource.tswren-ui/src/apollo/server/mdl/type.tswren-ui/src/utils/enum/dataSources.tswren-ui/src/apollo/server/services/projectService.tswren-ui/src/apollo/server/mdl/mdlBuilder.ts
📚 Learning: 2025-03-18T10:28:26.608Z
Learnt from: andreashimin
Repo: Canner/WrenAI PR: 1414
File: wren-ui/src/apollo/server/utils/error.ts:0-0
Timestamp: 2025-03-18T10:28:26.608Z
Learning: The enum names `IDENTIED_AS_GENERAL` and `IDENTIED_AS_MISLEADING_QUERY` in GeneralErrorCodes enum in wren-ui/src/apollo/server/utils/error.ts have typos (missing "FI" in "IDENTIFIED") but fixing them is not in scope for the current PR.
Applied to files:
wren-ui/src/apollo/server/mdl/type.ts
📚 Learning: 2025-07-09T02:43:20.433Z
Learnt from: goldmedal
Repo: Canner/WrenAI PR: 1827
File: wren-launcher/commands/dbt/data_source.go:50-53
Timestamp: 2025-07-09T02:43:20.433Z
Learning: In wren-launcher/commands/dbt/data_source.go, when encountering unsupported database types in convertConnectionToDataSource function, the user prefers to log a warning and continue processing instead of failing. The function should return nil, nil to allow processing to continue.
Applied to files:
wren-ui/src/apollo/server/mdl/type.ts
📚 Learning: 2025-07-09T02:43:07.560Z
Learnt from: goldmedal
Repo: Canner/WrenAI PR: 1827
File: wren-launcher/commands/dbt/data_source.go:154-156
Timestamp: 2025-07-09T02:43:07.560Z
Learning: In the WrenAI codebase, the MapType methods for data sources should return the original type name for unknown types, rather than a default fallback like "string". This maintains consistency across all data source implementations.
Applied to files:
wren-ui/src/apollo/server/mdl/mdlBuilder.ts
🧬 Code graph analysis (6)
wren-ui/src/apollo/client/graphql/__types__.ts (1)
wren-ui/src/utils/enum/dataSources.ts (1)
DatabricksConnectionType(3-3)
wren-ui/src/components/pages/setup/dataSources/DatabricksProperties.tsx (2)
wren-ui/src/utils/validator/hostValidator.ts (1)
hostValidator(3-13)wren-ui/src/utils/error/dictionary.ts (1)
ERROR_TEXTS(1-192)
wren-ui/src/components/pages/setup/utils.tsx (1)
wren-ui/src/utils/dataSourceType.ts (1)
getDataSourceConfig(108-114)
wren-ui/src/apollo/server/dataSource.ts (2)
wren-ui/src/apollo/server/repositories/projectRepository.ts (3)
DATABRICKS_CONNECTION_INFO(143-145)DATABRICKS_PERSONAL_ACCESS_TOKEN_AUTH(123-128)DATABRICKS_SERVICE_PRINCIPAL_AUTH(130-137)wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts (1)
IbisDatabricksConnectionInfo(135-137)
wren-ui/src/utils/dataSourceType.ts (1)
wren-ui/src/components/pages/setup/dataSources/DatabricksProperties.tsx (1)
DatabricksProperties(133-209)
wren-ui/src/hooks/useSetupConnectionDataSource.tsx (1)
wren-ui/migrations/20240531085916_transfer_model_properties.js (1)
dataSourceType(11-11)
🔇 Additional comments (15)
wren-ui/src/apollo/server/services/projectService.ts (1)
37-38: LGTM! Sensitive fields correctly masked.The addition of
accessTokenandclientSecretto the sensitive properties set ensures that Databricks authentication credentials are properly filtered out when exposing connection info, consistent with the existing security pattern for other data sources.wren-ui/src/utils/enum/dataSources.ts (1)
1-4: LGTM! Databricks enum additions are consistent.The export of
DatabricksConnectionTypeasDATABRICKS_AUTH_METHODfollows the same pattern asREDSHIFT_AUTH_METHOD, and the DATABRICKS enum member is correctly added to the DATA_SOURCES enum.Also applies to: 18-18
wren-ui/src/apollo/server/mdl/mdlBuilder.ts (1)
517-518: LGTM! Databricks data source mapping is correct.The case statement correctly maps
DataSourceName.DATABRICKStoWrenEngineDataSourceType.DATABRICKS, consistent with all other data source mappings in the switch statement.wren-ui/src/apollo/server/mdl/type.ts (1)
99-99: LGTM! Enum member correctly added.The DATABRICKS member is appropriately added to the WrenEngineDataSourceType enum, positioned before DATAFUSION which is noted as not supported by the UI.
wren-ui/src/apollo/server/schema.ts (1)
72-72: LGTM! GraphQL schema additions are correct.The DATABRICKS member is properly added to the DataSourceName enum, and the new DatabricksConnectionType enum follows the established pattern from RedshiftConnectionType. The snake_case enum values (service_principal, token) are consistent with the codebase convention.
Also applies to: 80-83
wren-ui/src/apollo/server/types/dataSource.ts (1)
13-13: LGTM! Enum member addition is correct.The DATABRICKS member is properly added to the DataSourceName enum, maintaining consistency with the GraphQL schema definition.
wren-ui/src/apollo/client/graphql/__types__.ts (1)
333-333: LGTM! Client types correctly mirror server schema.The DataSourceName enum and DatabricksConnectionType enum on the client side correctly mirror the GraphQL schema definitions from schema.ts, ensuring type consistency across the client-server boundary.
Also applies to: 344-347
wren-ui/src/hooks/useSetupConnectionDataSource.tsx (2)
134-145: Placeholder handling for Databricks authentication fields is correctly implemented.Verification confirms the placeholder logic properly scopes PASSWORD_PLACEHOLDER to only sensitive authentication fields. The hook applies placeholders only to
clientSecret(for service_principal auth) andaccessToken(for token auth), while non-sensitive fields likeserverHostname,httpPath,clientId, andazureTenantIdretain their original values without placeholder treatment. The UI component appropriately usesInput.Passwordfor sensitive fields and regularInputcomponents for others.
83-87: The review comment contains an inaccurate technical claim about pattern consistency.The Databricks authentication implementation is functionally correct and handles all required fields properly, but it does not mirror the Snowflake authentication pattern at the client level:
Pattern Differences:
Aspect Snowflake Databricks Client-side handling Explicitly sets unused field to nullOnly returns the used field Server-side handling Extracts both passwordandprivateKeyConditionally extracts only relevant field Result Both credentials sent, server normalizes Only used credential sent Verification Results:
✓ Service Principal:
clientSecretcorrectly filters outPASSWORD_PLACEHOLDER
✓ Token:accessTokencorrectly filters outPASSWORD_PLACEHOLDER
✓ All required fields handled:serverHostname,httpPath,clientId,clientSecret,azureTenantId,accessTokenare all properly processed at server-side
✓ Server-side filtering intoIbisConnectionInfoensures only relevant credentials are usedThe implementation is correct and actually more selective than Snowflake (server only uses the intended credential), but the client-side approach differs: Snowflake explicitly clears the unused field to
null, while Databricks omits it from the return object.wren-ui/src/apollo/server/repositories/projectRepository.ts (1)
12-15: LGTM! Databricks type definitions follow established patterns.The type structure mirrors the existing Redshift implementation with multiple authentication methods. The interfaces are well-defined with appropriate required and optional fields.
Also applies to: 123-159
wren-ui/src/components/pages/setup/utils.tsx (1)
115-119: LGTM! Databricks data source option configured correctly.The configuration follows the established pattern and includes the guide URL specified in the PR objectives.
wren-ui/src/utils/error/dictionary.ts (1)
79-90: LGTM! Error messages complete for Databricks fields.All required validation messages for both authentication methods are properly defined with clear, consistent wording.
wren-ui/src/utils/dataSourceType.ts (1)
13-13: LGTM! Databricks properly integrated into data source mappings.The image, name, and properties mappings are consistent with other data sources in the codebase.
Also applies to: 39-40, 70-71, 101-102
wren-ui/src/apollo/server/dataSource.ts (1)
442-492: LGTM! Databricks data source mapping implemented correctly.The implementation properly handles:
- Decryption of sensitive fields (
accessToken,clientSecret)- Both authentication methods with appropriate field mappings
- Error handling for invalid authentication types
- Consistent pattern with the existing Redshift implementation
The mix of snake_case (
databricks_type) and camelCase (other fields) matches the expected Ibis adaptor interface.wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts (1)
110-137: LGTM! Databricks adaptor types and mappings are well-defined.The connection type enum, authentication interfaces, and data source mapping are properly implemented. The field naming convention (snake_case for
databricks_type, camelCase for other fields) is intentional and defines the contract with the Ibis server API.Also applies to: 150-150, 164-164
fredalai
force-pushed
the
feature/databricks
branch
from
b9f498a to
dfecdb2
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
wren-ui/src/components/pages/setup/dataSources/DatabricksProperties.tsx (1)
95-100: Remove unused helper function.The
getIsEditModeForComponenthelper is defined but never used in the component. The shared fields (serverHostname and httpPath) are correctly disabled usingisEditModedirectly, which prevents the bypass issue mentioned in previous reviews.- const getIsEditModeForComponent = (componentType: DATABRICKS_AUTH_METHOD) => { - if (!isEditMode) return false; - - const initialType = initialDatabricksTypeRef.current || databricksType; - return initialType === componentType; - }; -
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (1)
wren-ui/public/images/dataSource/databricks.svgis excluded by!**/*.svg
📒 Files selected for processing (15)
wren-ui/src/apollo/client/graphql/__types__.ts(2 hunks)wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts(3 hunks)wren-ui/src/apollo/server/dataSource.ts(3 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(3 hunks)wren-ui/src/apollo/server/schema.ts(1 hunks)wren-ui/src/apollo/server/services/projectService.ts(1 hunks)wren-ui/src/apollo/server/types/dataSource.ts(1 hunks)wren-ui/src/components/pages/setup/dataSources/DatabricksProperties.tsx(1 hunks)wren-ui/src/components/pages/setup/utils.tsx(1 hunks)wren-ui/src/hooks/useSetupConnectionDataSource.tsx(4 hunks)wren-ui/src/utils/dataSourceType.ts(4 hunks)wren-ui/src/utils/enum/dataSources.ts(2 hunks)wren-ui/src/utils/error/dictionary.ts(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (7)
- wren-ui/src/components/pages/setup/utils.tsx
- wren-ui/src/utils/error/dictionary.ts
- wren-ui/src/utils/enum/dataSources.ts
- wren-ui/src/apollo/server/mdl/mdlBuilder.ts
- wren-ui/src/apollo/server/schema.ts
- wren-ui/src/apollo/server/services/projectService.ts
- wren-ui/src/apollo/server/types/dataSource.ts
🧰 Additional context used
🧠 Learnings (6)
📓 Common learnings
Learnt from: goldmedal
Repo: Canner/WrenAI PR: 1887
File: wren-launcher/commands/dbt/data_source.go:273-276
Timestamp: 2025-08-13T05:17:30.180Z
Learning: In the Wren AI codebase, there's an intentional distinction between dbt connection types and Wren AI data source types: dbt uses "sqlserver" as the connection type, while Wren AI expects "mssql" as the data source type. The conversion layer in convertConnectionToDataSource() correctly matches "sqlserver" from dbt profiles and the WrenMSSQLDataSource.GetType() method correctly returns "mssql" for Wren AI compatibility.
📚 Learning: 2025-03-18T10:28:10.593Z
Learnt from: andreashimin
Repo: Canner/WrenAI PR: 1414
File: wren-ui/src/apollo/server/utils/error.ts:0-0
Timestamp: 2025-03-18T10:28:10.593Z
Learning: The typo in error code enum names ("IDENTIED_AS_GENERAL" and "IDENTIED_AS_MISLEADING_QUERY" instead of "IDENTIFIED_AS_GENERAL" and "IDENTIFIED_AS_MISLEADING_QUERY") in wren-ui/src/apollo/server/utils/error.ts was recognized but intentionally not fixed as it was considered out of scope for PR #1414.
Applied to files:
wren-ui/src/apollo/client/graphql/__types__.tswren-ui/src/apollo/server/adaptors/ibisAdaptor.tswren-ui/src/apollo/server/mdl/type.ts
📚 Learning: 2025-08-13T05:17:30.180Z
Learnt from: goldmedal
Repo: Canner/WrenAI PR: 1887
File: wren-launcher/commands/dbt/data_source.go:273-276
Timestamp: 2025-08-13T05:17:30.180Z
Learning: In the Wren AI codebase, there's an intentional distinction between dbt connection types and Wren AI data source types: dbt uses "sqlserver" as the connection type, while Wren AI expects "mssql" as the data source type. The conversion layer in convertConnectionToDataSource() correctly matches "sqlserver" from dbt profiles and the WrenMSSQLDataSource.GetType() method correctly returns "mssql" for Wren AI compatibility.
Applied to files:
wren-ui/src/apollo/client/graphql/__types__.tswren-ui/src/hooks/useSetupConnectionDataSource.tsxwren-ui/src/apollo/server/dataSource.tswren-ui/src/apollo/server/adaptors/ibisAdaptor.tswren-ui/src/apollo/server/repositories/projectRepository.tswren-ui/src/utils/dataSourceType.tswren-ui/src/apollo/server/mdl/type.ts
📚 Learning: 2025-03-18T10:28:26.608Z
Learnt from: andreashimin
Repo: Canner/WrenAI PR: 1414
File: wren-ui/src/apollo/server/utils/error.ts:0-0
Timestamp: 2025-03-18T10:28:26.608Z
Learning: The enum names `IDENTIED_AS_GENERAL` and `IDENTIED_AS_MISLEADING_QUERY` in GeneralErrorCodes enum in wren-ui/src/apollo/server/utils/error.ts have typos (missing "FI" in "IDENTIFIED") but fixing them is not in scope for the current PR.
Applied to files:
wren-ui/src/apollo/client/graphql/__types__.tswren-ui/src/apollo/server/mdl/type.ts
📚 Learning: 2025-04-28T20:09:06.106Z
Learnt from: narsik123
Repo: Canner/WrenAI PR: 1606
File: wren-ui/src/apollo/server/dataSource.ts:135-150
Timestamp: 2025-04-28T20:09:06.106Z
Learning: The Oracle data source in WrenAI includes an SSL toggle in the UI (OracleProperties component) and should pass this setting to the connection info using the pattern `...(ssl && { kwargs: { ssl: true } })` for consistency with other host-based connections like MS SQL.
Applied to files:
wren-ui/src/hooks/useSetupConnectionDataSource.tsx
📚 Learning: 2025-07-09T02:43:20.433Z
Learnt from: goldmedal
Repo: Canner/WrenAI PR: 1827
File: wren-launcher/commands/dbt/data_source.go:50-53
Timestamp: 2025-07-09T02:43:20.433Z
Learning: In wren-launcher/commands/dbt/data_source.go, when encountering unsupported database types in convertConnectionToDataSource function, the user prefers to log a warning and continue processing instead of failing. The function should return nil, nil to allow processing to continue.
Applied to files:
wren-ui/src/apollo/server/mdl/type.ts
🧬 Code graph analysis (5)
wren-ui/src/apollo/client/graphql/__types__.ts (1)
wren-ui/src/utils/enum/dataSources.ts (1)
DatabricksConnectionType(3-3)
wren-ui/src/hooks/useSetupConnectionDataSource.tsx (1)
wren-ui/migrations/20240531085916_transfer_model_properties.js (1)
dataSourceType(11-11)
wren-ui/src/apollo/server/dataSource.ts (2)
wren-ui/src/apollo/server/repositories/projectRepository.ts (3)
DATABRICKS_CONNECTION_INFO(143-145)DATABRICKS_PERSONAL_ACCESS_TOKEN_AUTH(123-128)DATABRICKS_SERVICE_PRINCIPAL_AUTH(130-137)wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts (1)
IbisDatabricksConnectionInfo(135-137)
wren-ui/src/utils/dataSourceType.ts (1)
wren-ui/src/components/pages/setup/dataSources/DatabricksProperties.tsx (1)
DatabricksProperties(65-167)
wren-ui/src/components/pages/setup/dataSources/DatabricksProperties.tsx (2)
wren-ui/src/utils/error/dictionary.ts (1)
ERROR_TEXTS(1-192)wren-ui/src/utils/validator/hostValidator.ts (1)
hostValidator(3-13)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Analyze (go)
🔇 Additional comments (26)
wren-ui/src/apollo/server/mdl/type.ts (1)
99-100: LGTM!The DATABRICKS enum member is correctly added to WrenEngineDataSourceType, following the established naming convention and placed appropriately before DATAFUSION.
wren-ui/src/apollo/client/graphql/__types__.ts (2)
333-333: LGTM!The DATABRICKS member is correctly added to the DataSourceName enum, following the established naming convention.
344-347: LGTM!The DatabricksConnectionType enum is correctly defined with both authentication methods (token for PAT and service_principal), following the same pattern as the existing RedshiftConnectionType enum.
wren-ui/src/utils/dataSourceType.ts (4)
13-13: LGTM!The import statement correctly brings in the DatabricksProperties component, consistent with other data source property imports in this file.
39-40: LGTM!The data source image mapping for Databricks is correctly implemented, following the established pattern and naming convention.
70-71: LGTM!The data source name mapping returns 'Databricks' with appropriate capitalization, consistent with other data source name mappings.
101-102: LGTM!The data source properties mapping correctly returns the DatabricksProperties component, following the pattern established for other data sources.
wren-ui/src/components/pages/setup/dataSources/DatabricksProperties.tsx (4)
11-27: LGTM!The DatabricksTokenFields component correctly implements the Personal Access Token authentication form with appropriate validation and placeholder text.
29-63: LGTM!The DatabricksServicePrincipalFields component correctly implements the Service Principal authentication form. The azureTenantId field is appropriately optional, consistent with the PR description.
77-93: LGTM!The useEffect hooks correctly handle initialization of databricksType to 'token' in create mode and track the initial authentication method in edit mode.
102-167: LGTM!The main form implementation correctly handles both authentication methods with proper validation and conditional field rendering. The serverHostname and httpPath fields are appropriately disabled in edit mode to prevent modification as specified in the PR requirements.
wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts (4)
110-129: LGTM!The Databricks connection types and interfaces are well-defined, following the established pattern used for Redshift authentication. The enum, interfaces, and union type correctly model both Personal Access Token and Service Principal authentication methods.
135-137: LGTM!The IbisDatabricksConnectionInfo union type correctly combines both authentication methods, following the pattern established for IbisRedshiftConnectionInfo.
150-150: LGTM!DATABRICKS is correctly added to the SupportedDataSource enum, maintaining alphabetical ordering with other data sources.
164-164: LGTM!The dataSourceUrlMap entry correctly maps DATABRICKS to the 'databricks' API endpoint string, following the established pattern.
wren-ui/src/apollo/server/repositories/projectRepository.ts (4)
12-15: LGTM!The import statement is correctly updated to include IbisDatabricksConnectionType from the Ibis adaptor, following the same pattern as IbisRedshiftConnectionType.
123-137: LGTM!The Databricks authentication interfaces are well-defined with appropriate fields for both Personal Access Token and Service Principal authentication. The interfaces correctly use the IbisDatabricksConnectionType discriminator and match the corresponding Ibis adaptor interfaces.
143-145: LGTM!The DATABRICKS_CONNECTION_INFO union type correctly combines both authentication interfaces, following the pattern established for other data sources with multiple authentication methods.
158-159: LGTM!The WREN_AI_CONNECTION_INFO union is correctly extended to include DATABRICKS_CONNECTION_INFO, maintaining consistency with the addition of other data source connection types.
wren-ui/src/apollo/server/dataSource.ts (3)
11-12: LGTM!The import statements correctly bring in the Databricks connection types and info from the Ibis adaptor, following the pattern established for other data sources.
29-31: LGTM!The repository imports correctly include all three Databricks connection info types needed for the data source mapping.
442-492: LGTM!The Databricks data source implementation is comprehensive and correct:
- Sensitive properties (accessToken, clientSecret) are properly identified for encryption
- The toIbisConnectionInfo function correctly handles both authentication methods with proper field mapping
- The databricks_type discriminator is correctly passed through for both branches
- Error handling appropriately throws for invalid authentication types
- The field naming convention (databricks_type in snake_case, other fields in camelCase) matches the example payloads in the PR description
The implementation follows the established pattern for multi-authentication data sources (similar to Redshift) and correctly integrates with the repository and Ibis adaptor layers.
wren-ui/src/hooks/useSetupConnectionDataSource.tsx (4)
3-7: LGTM - Import follows established pattern.The addition of
DATABRICKS_AUTH_METHODfollows the same pattern as the existingREDSHIFT_AUTH_METHODimport.
83-87: Good pattern following Snowflake implementation.The structure correctly delegates authentication-specific field handling to the helper function while preserving all other properties.
134-145: Correct placeholder handling for UI.The conditional logic properly provides password placeholders for the appropriate authentication field based on the selected Databricks authentication type.
174-174: DATABRICKS_AUTH_METHOD enum values are correct and consistent throughout the codebase.Verification confirms that
DATABRICKS_AUTH_METHOD.service_principalequals"service_principal"and the token variant equals"token", matching the connectionInfo payload structure. The enum is imported from the GraphQL type definition and consistently used across form components, hooks, and server-side data transformations.
| function getDatabricksAuthentication(properties: Record<string, any>) { | ||
| if (properties?.databricksType === DATABRICKS_AUTH_METHOD.service_principal) { | ||
| return { | ||
| clientSecret: | ||
| properties?.clientSecret === PASSWORD_PLACEHOLDER | ||
| ? undefined | ||
| : properties?.clientSecret, | ||
| }; | ||
| } | ||
|
|
||
| return { | ||
| accessToken: | ||
| properties?.accessToken === PASSWORD_PLACEHOLDER | ||
| ? undefined | ||
| : properties?.accessToken, | ||
| }; | ||
| } |
There was a problem hiding this comment.
Missing explicit nullification of unused authentication field.
Unlike the getSnowflakeAuthentication helper (lines 156-171), this function doesn't explicitly set the unused authentication field to null. This creates a potential issue when switching authentication methods:
Scenario: User switches from Service Principal to Token
- Form properties may still contain
clientSecretfrom previous configuration getDatabricksAuthenticationreturns only{ accessToken: 'token456' }- Final spread
{ ...properties, ...getDatabricksAuthentication(properties) }preserves oldclientSecret - Server receives both
clientSecretandaccessTokeninstead of justaccessToken
This could cause:
- Security issues (old credentials sent unnecessarily)
- Server-side confusion about which authentication method to use
- Inconsistency with the established Snowflake pattern
Apply this diff to match the Snowflake pattern:
function getDatabricksAuthentication(properties: Record<string, any>) {
if (properties?.databricksType === DATABRICKS_AUTH_METHOD.service_principal) {
return {
clientSecret:
properties?.clientSecret === PASSWORD_PLACEHOLDER
? undefined
: properties?.clientSecret,
+ accessToken: null,
};
}
return {
accessToken:
properties?.accessToken === PASSWORD_PLACEHOLDER
? undefined
: properties?.accessToken,
+ clientSecret: null,
};
}📝 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.
| function getDatabricksAuthentication(properties: Record<string, any>) { | |
| if (properties?.databricksType === DATABRICKS_AUTH_METHOD.service_principal) { | |
| return { | |
| clientSecret: | |
| properties?.clientSecret === PASSWORD_PLACEHOLDER | |
| ? undefined | |
| : properties?.clientSecret, | |
| }; | |
| } | |
| return { | |
| accessToken: | |
| properties?.accessToken === PASSWORD_PLACEHOLDER | |
| ? undefined | |
| : properties?.accessToken, | |
| }; | |
| } | |
| function getDatabricksAuthentication(properties: Record<string, any>) { | |
| if (properties?.databricksType === DATABRICKS_AUTH_METHOD.service_principal) { | |
| return { | |
| clientSecret: | |
| properties?.clientSecret === PASSWORD_PLACEHOLDER | |
| ? undefined | |
| : properties?.clientSecret, | |
| accessToken: null, | |
| }; | |
| } | |
| return { | |
| accessToken: | |
| properties?.accessToken === PASSWORD_PLACEHOLDER | |
| ? undefined | |
| : properties?.accessToken, | |
| clientSecret: null, | |
| }; | |
| } |
🤖 Prompt for AI Agents
In wren-ui/src/hooks/useSetupConnectionDataSource.tsx around lines 173 to 189,
the databricks auth helper returns only the selected credential field and does
not explicitly null out the unused field; update getDatabricksAuthentication to
mirror the Snowflake helper by returning both fields with the unused one set to
null (i.e., when databricksType is service_principal return clientSecret (or
undefined if placeholder) and accessToken: null; otherwise return accessToken
(or undefined if placeholder) and clientSecret: null) so that spreading
properties into the final payload will not preserve old credentials.
dfecdb2 to
f187c20
Compare
fredalai
force-pushed
the
feature/databricks
branch
from
f187c20 to
919a188
Compare
2 participants
This pull request introduces support for the Databricks data source.
Description
Support Databricks connector
https://docs.getwren.ai/oss/guide/connect/databricksServer hostnameandHTTP pathfields cannot be edited.Databricks Data Source:
Engine API information
Development version: 0.22.0
Ref UI
PAT
Service Principal
Summary by CodeRabbit
New Features
Bug Fixes / Improvements
✏️ Tip: You can customize this high-level summary in your review settings.