feat(wren-ui): Athena OIDC auth

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;

wren-ui/src/apollo/server/dataSource.ts

@@ -73,22 +73,36 @@ interface IDataSourceConnectionInfo<C, I> {
 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,

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 {

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 (
+    <>
+      <Form.Item
+        label="AWS access key ID"
+        name="awsAccessKey"
+        required
+        rules={[
+          {
+            required: true,
+            message: ERROR_TEXTS.CONNECTION.AWS_ACCESS_KEY.REQUIRED,
+          },
+        ]}
+      >
+        <Input />
+      </Form.Item>
+
+      <Form.Item
+        label="AWS secret access key"
+        name="awsSecretKey"
+        required
+        rules={[
+          {
+            required: true,
+            message: ERROR_TEXTS.CONNECTION.AWS_SECRET_KEY.REQUIRED,
+          },
+        ]}
+      >
+        <Input.Password />
+      </Form.Item>
+    </>
+  );
+}
+
+function AthenaOIDCFields(props: { isEditMode: boolean }) {
+  const { isEditMode } = props;
+
+  return (
+    <>
+      <Form.Item
+        label="Web identity token"
+        name="webIdentityToken"
+        required
+        rules={[
+          {
+            required: true,
+            message: ERROR_TEXTS.CONNECTION.WEB_IDENTITY_TOKEN.REQUIRED,
+          },
+        ]}
+      >
+        <Input.TextArea
+          rows={3}
+          placeholder="OAuth 2.0 access token or OpenID Connect ID token"
+        />
+      </Form.Item>
+
+      <Form.Item
+        label="AWS role ARN"
+        name="roleArn"
+        required
+        rules={[
+          {
+            required: true,
+            message: ERROR_TEXTS.CONNECTION.AWS_ROLE_ARN.REQUIRED,
+          },
+        ]}
+      >
+        <Input
+          placeholder="arn:aws:iam::<account-id>:role/<role-name>"
+          disabled={isEditMode}
+        />
+      </Form.Item>
+
+      <Form.Item
+        label="Role session name"
+        name="roleSessionName"
+        extra="Optional session name used in AWS STS assume role operation."
+      >
+        <Input placeholder="session name" />
+      </Form.Item>
+    </>
+  );
+}
+
 export default function AthenaProperties(props: Props) {
   const { mode } = props;
   const isEditMode = mode === FORM_MODE.EDIT;
+
+  const form = Form.useFormInstance();
+
+  const initialTypeRef = useRef<ATHENA_AUTH_METHOD | null>(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 (
     <>
       <Form.Item
@@ -24,10 +141,12 @@ export default function AthenaProperties(props: Props) {
       >
         <Input />
       </Form.Item>
+
+      {/* Common fields */}
       <Form.Item
         label="Database (schema)"
         name="schema"
-        extra="The Athena database (also called schema) that contains the tables you want to query."
+        extra="The Athena database (schema) that contains your tables."
         required
         rules={[
           {
@@ -38,6 +157,7 @@ export default function AthenaProperties(props: Props) {
       >
         <Input disabled={isEditMode} />
       </Form.Item>
+
       <Form.Item
         label="S3 staging directory"
         name="s3StagingDir"
@@ -46,8 +166,8 @@ export default function AthenaProperties(props: Props) {
           <>
             The S3 path where Athena stores query results and metadata.
             <br />
-            You can find this in the Athena console under{' '}
-            <b>Settings {'>'} Query result location</b>.
+            Find this in Athena console under{' '}
+            <b>Settings → Query result location</b>.
           </>
         }
         rules={[
@@ -59,6 +179,7 @@ export default function AthenaProperties(props: Props) {
       >
         <Input placeholder="s3://bucket/path" />
       </Form.Item>
+
       <Form.Item
         label="AWS region"
         name="awsRegion"
@@ -72,32 +193,39 @@ export default function AthenaProperties(props: Props) {
       >
         <Input placeholder="us-east-1" disabled={isEditMode} />
       </Form.Item>
-      <Form.Item
-        label="AWS access key ID"
-        name="awsAccessKey"
-        required
-        rules={[
-          {
-            required: true,
-            message: ERROR_TEXTS.CONNECTION.AWS_ACCESS_KEY.REQUIRED,
-          },
-        ]}
-      >
-        <Input />
-      </Form.Item>
-      <Form.Item
-        label="AWS secret access key"
-        name="awsSecretKey"
-        required
-        rules={[
-          {
-            required: true,
-            message: ERROR_TEXTS.CONNECTION.AWS_SECRET_KEY.REQUIRED,
-          },
-        ]}
-      >
-        <Input.Password />
+
+      {/* Authentication method switch */}
+      <Form.Item label="Authentication method" name="athenaAuthType">
+        <Radio.Group buttonStyle="solid">
+          <Radio.Button value={ATHENA_AUTH_METHOD.classic}>
+            AWS credentials
+          </Radio.Button>
+          <Radio.Button value={ATHENA_AUTH_METHOD.oidc}>
+            OIDC (web identity token)
+          </Radio.Button>
+          <Radio.Button value={ATHENA_AUTH_METHOD.instance_profile}>
+            Instance Profile
+          </Radio.Button>
+        </Radio.Group>
       </Form.Item>
+
+      {/* Conditional auth fields */}
+      {authType === ATHENA_AUTH_METHOD.classic && <AthenaClassicFields />}
+
+      {authType === ATHENA_AUTH_METHOD.oidc && (
+        <AthenaOIDCFields
+          isEditMode={getIsEditModeForComponent(
+            ATHENA_AUTH_METHOD.oidc,
+          )}
+        />
+      )}
+
+      {authType === ATHENA_AUTH_METHOD.instance_profile && (
+        <div className="gray-8" style={{ fontStyle: 'italic'}}>
+          We will automatically detect AWS credentials from the Instance Profile role
+          assigned to this compute environment (EC2, ECS, EKS).
+        </div>
+      )}
     </>
   );
 }

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',
+}

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.',
     },