Skip to content

Commit 8fdd89a

Browse files
douenergydouenergy
committed
support oidc connection
1 parent 4a05012 commit 8fdd89a

4 files changed

Lines changed: 173 additions & 31 deletions

File tree

‎wren-ui/src/apollo/client/graphql/__types__.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1246,6 +1246,12 @@ export enum RedshiftConnectionType {
12461246
redshift_iam = 'redshift_iam'
12471247
}
12481248

1249+
export enum ATHENA_AUTH_METHOD {
1250+
classic = 'classic',
1251+
oidc = 'oidc',
1252+
}
1253+
1254+
12491255
export type Relation = {
12501256
__typename?: 'Relation';
12511257
fromColumnId: Scalars['Int'];

‎wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -69,8 +69,15 @@ export interface IbisSnowflakeConnectionInfo {
6969
}
7070

7171
export interface IbisAthenaConnectionInfo {
72-
aws_access_key_id: string;
73-
aws_secret_access_key: string;
72+
// AWS access key auth (optional if using OIDC)
73+
aws_access_key_id?: string;
74+
aws_secret_access_key?: string;
75+
76+
// OIDC auth (optional if using access keys)
77+
web_identity_token?: string;
78+
role_arn?: string;
79+
role_session_name?: string;
80+
7481
region_name: string;
7582
s3_staging_dir: string;
7683
schema_name: string;

‎wren-ui/src/components/pages/setup/dataSources/AthenaProperties.tsx‎

Lines changed: 152 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,136 @@
1-
import { Form, Input } from 'antd';
1+
import { useEffect, useRef } from 'react';
2+
import { Form, Input, Radio } from 'antd';
23
import { FORM_MODE } from '@/utils/enum';
34
import { ERROR_TEXTS } from '@/utils/error';
45

6+
export enum ATHENA_AUTH_METHOD {
7+
classic = 'classic',
8+
oidc = 'oidc',
9+
}
10+
511
interface Props {
612
mode?: FORM_MODE;
713
}
814

15+
function AthenaClassicFields() {
16+
return (
17+
<>
18+
<Form.Item
19+
label="AWS access key ID"
20+
name="awsAccessKey"
21+
required
22+
rules={[
23+
{
24+
required: true,
25+
message: ERROR_TEXTS.CONNECTION.AWS_ACCESS_KEY.REQUIRED,
26+
},
27+
]}
28+
>
29+
<Input />
30+
</Form.Item>
31+
32+
<Form.Item
33+
label="AWS secret access key"
34+
name="awsSecretKey"
35+
required
36+
rules={[
37+
{
38+
required: true,
39+
message: ERROR_TEXTS.CONNECTION.AWS_SECRET_KEY.REQUIRED,
40+
},
41+
]}
42+
>
43+
<Input.Password />
44+
</Form.Item>
45+
</>
46+
);
47+
}
48+
49+
function AthenaOIDCFields(props: { isEditMode: boolean }) {
50+
const { isEditMode } = props;
51+
52+
return (
53+
<>
54+
<Form.Item
55+
label="Web identity token"
56+
name="webIdentityToken"
57+
required
58+
rules={[
59+
{
60+
required: true,
61+
message: ERROR_TEXTS.CONNECTION.WEB_IDENTITY_TOKEN.REQUIRED,
62+
},
63+
]}
64+
>
65+
<Input.TextArea
66+
rows={3}
67+
placeholder="Paste Google OIDC token here"
68+
/>
69+
</Form.Item>
70+
71+
<Form.Item
72+
label="AWS role ARN"
73+
name="roleArn"
74+
required
75+
rules={[
76+
{
77+
required: true,
78+
message: ERROR_TEXTS.CONNECTION.AWS_ROLE_ARN.REQUIRED,
79+
},
80+
]}
81+
>
82+
<Input
83+
placeholder="arn:aws:iam::<account-id>:role/<role-name>"
84+
disabled={isEditMode}
85+
/>
86+
</Form.Item>
87+
88+
<Form.Item
89+
label="Role session name"
90+
name="roleSessionName"
91+
extra="Optional session name used in STS AssumeRoleWithWebIdentity"
92+
>
93+
<Input placeholder="Optional session name" />
94+
</Form.Item>
95+
</>
96+
);
97+
}
98+
999
export default function AthenaProperties(props: Props) {
10100
const { mode } = props;
11101
const isEditMode = mode === FORM_MODE.EDIT;
102+
103+
const form = Form.useFormInstance();
104+
105+
const initialTypeRef = useRef<ATHENA_AUTH_METHOD | null>(null);
106+
107+
const authType = Form.useWatch(
108+
'athenaAuthType',
109+
form,
110+
) as ATHENA_AUTH_METHOD;
111+
112+
// Set default auth type when creating
113+
useEffect(() => {
114+
if (!isEditMode) {
115+
form.setFieldsValue({
116+
athenaAuthType: ATHENA_AUTH_METHOD.classic,
117+
});
118+
}
119+
}, [isEditMode, form]);
120+
121+
// Preserve initial type on edit mode
122+
useEffect(() => {
123+
if (isEditMode && authType && initialTypeRef.current === null) {
124+
initialTypeRef.current = authType;
125+
}
126+
}, [isEditMode, authType]);
127+
128+
const getIsEditModeForComponent = (component: ATHENA_AUTH_METHOD) => {
129+
if (!isEditMode) return false;
130+
const initial = initialTypeRef.current || authType;
131+
return initial === component;
132+
};
133+
12134
return (
13135
<>
14136
<Form.Item
@@ -24,10 +146,12 @@ export default function AthenaProperties(props: Props) {
24146
>
25147
<Input />
26148
</Form.Item>
149+
150+
{/* Common fields */}
27151
<Form.Item
28152
label="Database (schema)"
29153
name="schema"
30-
extra="The Athena database (also called schema) that contains the tables you want to query."
154+
extra="The Athena database (schema) that contains your tables."
31155
required
32156
rules={[
33157
{
@@ -38,6 +162,7 @@ export default function AthenaProperties(props: Props) {
38162
>
39163
<Input disabled={isEditMode} />
40164
</Form.Item>
165+
41166
<Form.Item
42167
label="S3 staging directory"
43168
name="s3StagingDir"
@@ -46,8 +171,8 @@ export default function AthenaProperties(props: Props) {
46171
<>
47172
The S3 path where Athena stores query results and metadata.
48173
<br />
49-
You can find this in the Athena console under{' '}
50-
<b>Settings {'>'} Query result location</b>.
174+
Find this in Athena console under{' '}
175+
<b>Settings → Query result location</b>.
51176
</>
52177
}
53178
rules={[
@@ -59,6 +184,7 @@ export default function AthenaProperties(props: Props) {
59184
>
60185
<Input placeholder="s3://bucket/path" />
61186
</Form.Item>
187+
62188
<Form.Item
63189
label="AWS region"
64190
name="awsRegion"
@@ -72,32 +198,29 @@ export default function AthenaProperties(props: Props) {
72198
>
73199
<Input placeholder="us-east-1" disabled={isEditMode} />
74200
</Form.Item>
75-
<Form.Item
76-
label="AWS access key ID"
77-
name="awsAccessKey"
78-
required
79-
rules={[
80-
{
81-
required: true,
82-
message: ERROR_TEXTS.CONNECTION.AWS_ACCESS_KEY.REQUIRED,
83-
},
84-
]}
85-
>
86-
<Input />
87-
</Form.Item>
88-
<Form.Item
89-
label="AWS secret access key"
90-
name="awsSecretKey"
91-
required
92-
rules={[
93-
{
94-
required: true,
95-
message: ERROR_TEXTS.CONNECTION.AWS_SECRET_KEY.REQUIRED,
96-
},
97-
]}
98-
>
99-
<Input.Password />
201+
202+
{/* Authentication method switch */}
203+
<Form.Item label="Authentication method" name="athenaAuthType">
204+
<Radio.Group buttonStyle="solid">
205+
<Radio.Button value={ATHENA_AUTH_METHOD.classic}>
206+
AWS credentials
207+
</Radio.Button>
208+
<Radio.Button value={ATHENA_AUTH_METHOD.oidc}>
209+
OIDC (web identity)
210+
</Radio.Button>
211+
</Radio.Group>
100212
</Form.Item>
213+
214+
{/* Conditional auth fields */}
215+
{authType === ATHENA_AUTH_METHOD.classic && <AthenaClassicFields />}
216+
217+
{authType === ATHENA_AUTH_METHOD.oidc && (
218+
<AthenaOIDCFields
219+
isEditMode={getIsEditModeForComponent(
220+
ATHENA_AUTH_METHOD.oidc,
221+
)}
222+
/>
223+
)}
101224
</>
102225
);
103226
}

‎wren-ui/src/utils/error/dictionary.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,12 @@ export const ERROR_TEXTS = {
6464
AWS_SECRET_KEY: {
6565
REQUIRED: 'Please input AWS secret access key.',
6666
},
67+
AWS_ROLE_ARN: {
68+
REQUIRED: 'Please input AWS role ARN.',
69+
},
70+
WEB_IDENTITY_TOKEN: {
71+
REQUIRED: 'Please input web identity token.',
72+
},
6773
CLUSTER_IDENTIFIER: {
6874
REQUIRED: 'Please input cluster identifier.',
6975
},

0 commit comments

Comments
 (0)