diff --git a/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts b/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts index ea9f849d16..516ce237fa 100644 --- a/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts @@ -69,8 +69,15 @@ export interface IbisSnowflakeConnectionInfo { } export interface IbisAthenaConnectionInfo { - aws_access_key_id: string; - aws_secret_access_key: string; + // AWS access key auth (optional if using OIDC) + aws_access_key_id?: string; + aws_secret_access_key?: string; + + // OIDC auth (optional if using access keys) + web_identity_token?: string; + role_arn?: string; + role_session_name?: string; + region_name: string; s3_staging_dir: string; schema_name: string; diff --git a/wren-ui/src/apollo/server/dataSource.ts b/wren-ui/src/apollo/server/dataSource.ts index 5346951dbb..c0fa99f106 100644 --- a/wren-ui/src/apollo/server/dataSource.ts +++ b/wren-ui/src/apollo/server/dataSource.ts @@ -73,22 +73,36 @@ interface IDataSourceConnectionInfo { const dataSource = { // Athena [DataSourceName.ATHENA]: { - sensitiveProps: ['awsSecretKey'], + sensitiveProps: ['awsSecretKey', 'webIdentityToken'], toIbisConnectionInfo(connectionInfo) { const decryptedConnectionInfo = decryptConnectionInfo( DataSourceName.ATHENA, connectionInfo, ); - const { awsAccessKey, awsRegion, awsSecretKey, s3StagingDir, schema } = + const { awsAccessKey, awsRegion, awsSecretKey, s3StagingDir, schema, webIdentityToken, roleArn, roleSessionName } = decryptedConnectionInfo as ATHENA_CONNECTION_INFO; - const res: IbisAthenaConnectionInfo = { - aws_access_key_id: awsAccessKey, - aws_secret_access_key: awsSecretKey, + + // Base fields shared by both authentication methods + const base = { region_name: awsRegion, s3_staging_dir: s3StagingDir, schema_name: schema, - }; - return res; + }; + // If OIDC fields are provided → send OIDC config + if (webIdentityToken && roleArn) { + return { + ...base, + web_identity_token: webIdentityToken, + role_arn: roleArn, + ...(roleSessionName && { role_session_name: roleSessionName }), + } satisfies IbisAthenaConnectionInfo; + } + // Otherwise → fallback to AWS access key authentication + return { + ...base, + aws_access_key_id: awsAccessKey, + aws_secret_access_key: awsSecretKey, + } satisfies IbisAthenaConnectionInfo; }, } as IDataSourceConnectionInfo< ATHENA_CONNECTION_INFO, diff --git a/wren-ui/src/apollo/server/repositories/projectRepository.ts b/wren-ui/src/apollo/server/repositories/projectRepository.ts index 830c5c2b2c..a56ae7ce2b 100644 --- a/wren-ui/src/apollo/server/repositories/projectRepository.ts +++ b/wren-ui/src/apollo/server/repositories/projectRepository.ts @@ -91,8 +91,11 @@ export interface ATHENA_CONNECTION_INFO { schema: string; s3StagingDir: string; awsRegion: string; - awsAccessKey: string; - awsSecretKey: string; + awsAccessKey?: string; + awsSecretKey?: string; + webIdentityToken?: string; + roleArn?: string; + roleSessionName?: string; } export interface REDSHIFT_PASSWORD_AUTH { diff --git a/wren-ui/src/components/pages/setup/dataSources/AthenaProperties.tsx b/wren-ui/src/components/pages/setup/dataSources/AthenaProperties.tsx index 0d8c257c34..7892b08019 100644 --- a/wren-ui/src/components/pages/setup/dataSources/AthenaProperties.tsx +++ b/wren-ui/src/components/pages/setup/dataSources/AthenaProperties.tsx @@ -1,14 +1,131 @@ -import { Form, Input } from 'antd'; -import { FORM_MODE } from '@/utils/enum'; +import { useEffect, useRef } from 'react'; +import { Form, Input, Radio } from 'antd'; +import { FORM_MODE, ATHENA_AUTH_METHOD } from '@/utils/enum'; import { ERROR_TEXTS } from '@/utils/error'; interface Props { mode?: FORM_MODE; } +function AthenaClassicFields() { + return ( + <> + + + + + + + + + ); +} + +function AthenaOIDCFields(props: { isEditMode: boolean }) { + const { isEditMode } = props; + + return ( + <> + + + + + + + + + + + + + ); +} + export default function AthenaProperties(props: Props) { const { mode } = props; const isEditMode = mode === FORM_MODE.EDIT; + + const form = Form.useFormInstance(); + + const initialTypeRef = useRef(null); + + const authType = Form.useWatch( + 'athenaAuthType', + form, + ) as ATHENA_AUTH_METHOD; + + // Set default auth type when creating + useEffect(() => { + if (!isEditMode) { + form.setFieldsValue({ + athenaAuthType: ATHENA_AUTH_METHOD.classic, + }); + } + }, [isEditMode, form]); + + // Preserve initial type on edit mode + useEffect(() => { + if (isEditMode && authType && initialTypeRef.current === null) { + initialTypeRef.current = authType; + } + }, [isEditMode, authType]); + + const getIsEditModeForComponent = (component: ATHENA_AUTH_METHOD) => { + if (!isEditMode) return false; + const initial = initialTypeRef.current || authType; + return initial === component; + }; + return ( <> + + {/* Common fields */} + The S3 path where Athena stores query results and metadata.
- You can find this in the Athena console under{' '} - Settings {'>'} Query result location. + Find this in Athena console under{' '} + Settings → Query result location. } rules={[ @@ -59,6 +179,7 @@ export default function AthenaProperties(props: Props) { >
+ - - - - - + + {/* Authentication method switch */} + + + + AWS credentials + + + OIDC (web identity token) + + + Instance Profile + + + + {/* Conditional auth fields */} + {authType === ATHENA_AUTH_METHOD.classic && } + + {authType === ATHENA_AUTH_METHOD.oidc && ( + + )} + + {authType === ATHENA_AUTH_METHOD.instance_profile && ( +
+ We will automatically detect AWS credentials from the Instance Profile role + assigned to this compute environment (EC2, ECS, EKS). +
+ )} ); } diff --git a/wren-ui/src/utils/enum/dataSources.ts b/wren-ui/src/utils/enum/dataSources.ts index c0ec244d52..650e2cfbf9 100644 --- a/wren-ui/src/utils/enum/dataSources.ts +++ b/wren-ui/src/utils/enum/dataSources.ts @@ -13,3 +13,9 @@ export enum DATA_SOURCES { ATHENA = 'ATHENA', REDSHIFT = 'REDSHIFT', } + +export enum ATHENA_AUTH_METHOD { + classic = 'classic', + oidc = 'oidc', + instance_profile = 'instance_profile', +} diff --git a/wren-ui/src/utils/error/dictionary.ts b/wren-ui/src/utils/error/dictionary.ts index 3c46e7483c..0c78dedfae 100644 --- a/wren-ui/src/utils/error/dictionary.ts +++ b/wren-ui/src/utils/error/dictionary.ts @@ -64,6 +64,12 @@ export const ERROR_TEXTS = { AWS_SECRET_KEY: { REQUIRED: 'Please input AWS secret access key.', }, + AWS_ROLE_ARN: { + REQUIRED: 'Please input AWS role ARN.', + }, + WEB_IDENTITY_TOKEN: { + REQUIRED: 'Please input web identity token.', + }, CLUSTER_IDENTIFIER: { REQUIRED: 'Please input cluster identifier.', },