Skip to content

Commit 1b1a5b9

Browse files
fredalaifredalai
authored
feat(wren-ui): Support Athena (Trino) connector (#1754)
1 parent b3b3e6d commit 1b1a5b9

File tree

16 files changed

+204
-15
lines changed

16 files changed

+204
-15
lines changed

‎wren-ui/public/images/dataSource/athena.svg

Lines changed: 15 additions & 0 deletions
Loading

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

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -294,9 +294,9 @@ export type DashboardSchedule = {
294294
__typename?: 'DashboardSchedule';
295295
cron?: Maybe<Scalars['String']>;
296296
day?: Maybe<CacheScheduleDayEnum>;
297-
frequency: ScheduleFrequencyEnum;
298-
hour: Scalars['Int'];
299-
minute: Scalars['Int'];
297+
frequency?: Maybe<ScheduleFrequencyEnum>;
298+
hour?: Maybe<Scalars['Int']>;
299+
minute?: Maybe<Scalars['Int']>;
300300
timezone?: Maybe<Scalars['String']>;
301301
};
302302

@@ -313,12 +313,13 @@ export type DataSourceInput = {
313313
};
314314

315315
export enum DataSourceName {
316+
ATHENA = 'ATHENA',
316317
BIG_QUERY = 'BIG_QUERY',
317318
CLICK_HOUSE = 'CLICK_HOUSE',
318319
DUCKDB = 'DUCKDB',
319320
MSSQL = 'MSSQL',
320-
ORACLE = 'ORACLE',
321321
MYSQL = 'MYSQL',
322+
ORACLE = 'ORACLE',
322323
POSTGRES = 'POSTGRES',
323324
SNOWFLAKE = 'SNOWFLAKE',
324325
TRINO = 'TRINO'
@@ -384,7 +385,7 @@ export type DetailedDashboard = {
384385
items: Array<DashboardItem>;
385386
name: Scalars['String'];
386387
nextScheduledAt?: Maybe<Scalars['String']>;
387-
schedule: DashboardSchedule;
388+
schedule?: Maybe<DashboardSchedule>;
388389
};
389390

390391
export type DetailedModel = {

‎wren-ui/src/apollo/client/graphql/dashboard.generated.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ export type SetDashboardScheduleMutation = { __typename?: 'Mutation', setDashboa
5656
export type DashboardQueryVariables = Types.Exact<{ [key: string]: never; }>;
5757

5858

59-
export type DashboardQuery = { __typename?: 'Query', dashboard: { __typename?: 'DetailedDashboard', id: number, name: string, description?: string | null, cacheEnabled: boolean, nextScheduledAt?: string | null, schedule: { __typename?: 'DashboardSchedule', frequency: Types.ScheduleFrequencyEnum, hour: number, minute: number, day?: Types.CacheScheduleDayEnum | null, timezone?: string | null, cron?: string | null }, items: Array<{ __typename?: 'DashboardItem', id: number, dashboardId: number, type: Types.DashboardItemType, displayName?: string | null, layout: { __typename?: 'DashboardItemLayout', x: number, y: number, w: number, h: number }, detail: { __typename?: 'DashboardItemDetail', sql: string, chartSchema?: any | null } }> } };
59+
export type DashboardQuery = { __typename?: 'Query', dashboard: { __typename?: 'DetailedDashboard', id: number, name: string, description?: string | null, cacheEnabled: boolean, nextScheduledAt?: string | null, schedule?: { __typename?: 'DashboardSchedule', frequency?: Types.ScheduleFrequencyEnum | null, hour?: number | null, minute?: number | null, day?: Types.CacheScheduleDayEnum | null, timezone?: string | null, cron?: string | null } | null, items: Array<{ __typename?: 'DashboardItem', id: number, dashboardId: number, type: Types.DashboardItemType, displayName?: string | null, layout: { __typename?: 'DashboardItemLayout', x: number, y: number, w: number, h: number }, detail: { __typename?: 'DashboardItemDetail', sql: string, chartSchema?: any | null } }> } };
6060

6161
export const CommonDashboardItemFragmentDoc = gql`
6262
fragment CommonDashboardItem on DashboardItem {

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

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -66,13 +66,13 @@ export interface IbisSnowflakeConnectionInfo {
6666
schema: string;
6767
}
6868

69-
export type IbisConnectionInfo =
70-
| UrlBasedConnectionInfo
71-
| HostBasedConnectionInfo
72-
| IbisPostgresConnectionInfo
73-
| IbisBigQueryConnectionInfo
74-
| IbisTrinoConnectionInfo
75-
| IbisSnowflakeConnectionInfo;
69+
export interface IbisAthenaConnectionInfo {
70+
aws_access_key_id: string;
71+
aws_secret_access_key: string;
72+
region_name: string;
73+
s3_staging_dir: string;
74+
schema_name: string;
75+
}
7676

7777
export enum SupportedDataSource {
7878
POSTGRES = 'POSTGRES',
@@ -83,6 +83,7 @@ export enum SupportedDataSource {
8383
MSSQL = 'MSSQL',
8484
CLICK_HOUSE = 'CLICK_HOUSE',
8585
TRINO = 'TRINO',
86+
ATHENA = 'ATHENA',
8687
}
8788

8889
const dataSourceUrlMap: Record<SupportedDataSource, string> = {
@@ -94,6 +95,7 @@ const dataSourceUrlMap: Record<SupportedDataSource, string> = {
9495
[SupportedDataSource.MSSQL]: 'mssql',
9596
[SupportedDataSource.CLICK_HOUSE]: 'clickhouse',
9697
[SupportedDataSource.TRINO]: 'trino',
98+
[SupportedDataSource.ATHENA]: 'athena',
9799
};
98100

99101
export interface TableResponse {

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

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,10 @@ import {
55
UrlBasedConnectionInfo,
66
IbisSnowflakeConnectionInfo,
77
IbisTrinoConnectionInfo,
8+
IbisAthenaConnectionInfo,
89
} from './adaptors/ibisAdaptor';
910
import {
11+
ATHENA_CONNECTION_INFO,
1012
BIG_QUERY_CONNECTION_INFO,
1113
DUCKDB_CONNECTION_INFO,
1214
MYSQL_CONNECTION_INFO,
@@ -64,6 +66,30 @@ interface IDataSourceConnectionInfo<C, I> {
6466
}
6567

6668
const dataSource = {
69+
// Athena
70+
[DataSourceName.ATHENA]: {
71+
sensitiveProps: ['awsSecretKey'],
72+
toIbisConnectionInfo(connectionInfo) {
73+
const decryptedConnectionInfo = decryptConnectionInfo(
74+
DataSourceName.ATHENA,
75+
connectionInfo,
76+
);
77+
const { awsAccessKey, awsRegion, awsSecretKey, s3StagingDir, schema } =
78+
decryptedConnectionInfo as ATHENA_CONNECTION_INFO;
79+
const res: IbisAthenaConnectionInfo = {
80+
aws_access_key_id: awsAccessKey,
81+
aws_secret_access_key: awsSecretKey,
82+
region_name: awsRegion,
83+
s3_staging_dir: s3StagingDir,
84+
schema_name: schema,
85+
};
86+
return res;
87+
},
88+
} as IDataSourceConnectionInfo<
89+
ATHENA_CONNECTION_INFO,
90+
IbisAthenaConnectionInfo
91+
>,
92+
6793
// BigQuery
6894
[DataSourceName.BIG_QUERY]: {
6995
sensitiveProps: ['credentials'],

‎wren-ui/src/apollo/server/mdl/mdlBuilder.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -492,6 +492,8 @@ export class MDLBuilder implements IMDLBuilder {
492492
return;
493493
}
494494
switch (type) {
495+
case DataSourceName.ATHENA:
496+
return WrenEngineDataSourceType.ATHENA;
495497
case DataSourceName.BIG_QUERY:
496498
return WrenEngineDataSourceType.BIGQUERY;
497499
case DataSourceName.DUCKDB:

‎wren-ui/src/apollo/server/mdl/type.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,7 @@ export interface TableReference {
8484
}
8585

8686
export enum WrenEngineDataSourceType {
87+
ATHENA = 'ATHENA',
8788
BIGQUERY = 'BIGQUERY',
8889
CANNER = 'CANNER',
8990
CLICKHOUSE = 'CLICKHOUSE',

‎wren-ui/src/apollo/server/repositories/projectRepository.ts

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,17 @@ export interface DUCKDB_CONNECTION_INFO {
8383
configurations: Record<string, any>;
8484
}
8585

86+
export interface ATHENA_CONNECTION_INFO {
87+
database: string;
88+
schema: string;
89+
s3StagingDir: string;
90+
awsRegion: string;
91+
awsAccessKey: string;
92+
awsSecretKey: string;
93+
}
94+
8695
export type WREN_AI_CONNECTION_INFO =
96+
| ATHENA_CONNECTION_INFO
8797
| BIG_QUERY_CONNECTION_INFO
8898
| POSTGRES_CONNECTION_INFO
8999
| MYSQL_CONNECTION_INFO

‎wren-ui/src/apollo/server/schema.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ export const typeDefs = gql`
4545
}
4646
4747
enum DataSourceName {
48+
ATHENA
4849
BIG_QUERY
4950
DUCKDB
5051
POSTGRES

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ export enum DataSourceName {
88
CLICK_HOUSE = 'CLICK_HOUSE',
99
TRINO = 'TRINO',
1010
SNOWFLAKE = 'SNOWFLAKE',
11+
ATHENA = 'ATHENA',
1112
}
1213

1314
export interface DataSource {

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

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
import { Form, Input } from 'antd';
2+
import { FORM_MODE } from '@/utils/enum';
3+
import { ERROR_TEXTS } from '@/utils/error';
4+
5+
interface Props {
6+
mode?: FORM_MODE;
7+
}
8+
9+
export default function AthenaProperties(props: Props) {
10+
const { mode } = props;
11+
const isEditMode = mode === FORM_MODE.EDIT;
12+
return (
13+
<>
14+
<Form.Item
15+
label="Display name"
16+
name="displayName"
17+
required
18+
rules={[
19+
{
20+
required: true,
21+
message: ERROR_TEXTS.CONNECTION.DISPLAY_NAME.REQUIRED,
22+
},
23+
]}
24+
>
25+
<Input />
26+
</Form.Item>
27+
<Form.Item
28+
label="Database (schema)"
29+
name="schema"
30+
extra="The Athena database (also called schema) that contains the tables you want to query."
31+
required
32+
rules={[
33+
{
34+
required: true,
35+
message: ERROR_TEXTS.CONNECTION.DATABASE.REQUIRED,
36+
},
37+
]}
38+
>
39+
<Input disabled={isEditMode} />
40+
</Form.Item>
41+
<Form.Item
42+
label="S3 staging directory"
43+
name="s3StagingDir"
44+
required
45+
extra={
46+
<>
47+
The S3 path where Athena stores query results and metadata.
48+
<br />
49+
You can find this in the Athena console under{' '}
50+
<b>Settings {'>'} Query result location</b>.
51+
</>
52+
}
53+
rules={[
54+
{
55+
required: true,
56+
message: ERROR_TEXTS.CONNECTION.S3_STAGING_DIR.REQUIRED,
57+
},
58+
]}
59+
>
60+
<Input placeholder="s3://bucket/path" />
61+
</Form.Item>
62+
<Form.Item
63+
label="AWS region"
64+
name="awsRegion"
65+
required
66+
rules={[
67+
{
68+
required: true,
69+
message: ERROR_TEXTS.CONNECTION.AWS_REGION.REQUIRED,
70+
},
71+
]}
72+
>
73+
<Input placeholder="us-east-1" disabled={isEditMode} />
74+
</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 />
100+
</Form.Item>
101+
</>
102+
);
103+
}

‎wren-ui/src/components/pages/setup/utils.tsx

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,11 @@ export const DATA_SOURCE_OPTIONS = {
102102
guide: 'https://docs.getwren.ai/oss/guide/connect/snowflake',
103103
disabled: false,
104104
},
105+
[DATA_SOURCES.ATHENA]: {
106+
...getDataSourceConfig(DATA_SOURCES.ATHENA),
107+
guide: 'https://docs.getwren.ai/oss/guide/connect/athena',
108+
disabled: false,
109+
},
105110
} as { [key: string]: ButtonOption };
106111

107112
export const TEMPLATE_OPTIONS = {

‎wren-ui/src/pages/home/dashboard.tsx

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,9 @@ import { LoadingWrapper } from '@/components/PageLoading';
99
import DashboardGrid from '@/components/pages/home/dashboardGrid';
1010
import EmptyDashboard from '@/components/pages/home/dashboardGrid/EmptyDashboard';
1111
import DashboardHeader from '@/components/pages/home/dashboardGrid/DashboardHeader';
12-
import CacheSettingsDrawer from '@/components/pages/home/dashboardGrid/CacheSettingsDrawer';
12+
import CacheSettingsDrawer, {
13+
Schedule,
14+
} from '@/components/pages/home/dashboardGrid/CacheSettingsDrawer';
1315
import {
1416
useDashboardQuery,
1517
useDeleteDashboardItemMutation,
@@ -107,7 +109,7 @@ export default function Dashboard() {
107109
<EmptyDashboard show={dashboardItems.length === 0}>
108110
<DashboardHeader
109111
isSupportCached={isSupportCached}
110-
schedule={data?.dashboard?.schedule}
112+
schedule={data?.dashboard?.schedule as Schedule}
111113
nextScheduleTime={data?.dashboard?.nextScheduledAt}
112114
onCacheSettings={() => {
113115
cacheSettingsDrawer.openDrawer({

‎wren-ui/src/utils/dataSourceType.ts

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import SQLServerProperties from '@/components/pages/setup/dataSources/SQLServerP
88
import ClickHouseProperties from '@/components/pages/setup/dataSources/ClickHouseProperties';
99
import TrinoProperties from '@/components/pages/setup/dataSources/TrinoProperties';
1010
import SnowflakeProperties from '@/components/pages/setup/dataSources/SnowflakeProperties';
11+
import AthenaProperties from '@/components/pages/setup/dataSources/AthenaProperties';
1112

1213
export const getDataSourceImage = (dataSource: DATA_SOURCES | string) => {
1314
switch (dataSource) {
@@ -29,6 +30,8 @@ export const getDataSourceImage = (dataSource: DATA_SOURCES | string) => {
2930
return '/images/dataSource/trino.svg';
3031
case DATA_SOURCES.SNOWFLAKE:
3132
return '/images/dataSource/snowflake.svg';
33+
case DATA_SOURCES.ATHENA:
34+
return '/images/dataSource/athena.svg';
3235
default:
3336
return null;
3437
}
@@ -54,6 +57,8 @@ export const getDataSourceName = (dataSource: DATA_SOURCES | string) => {
5457
return 'Trino';
5558
case DATA_SOURCES.SNOWFLAKE:
5659
return 'Snowflake';
60+
case DATA_SOURCES.ATHENA:
61+
return 'Athena (Trino)';
5762
default:
5863
return '';
5964
}
@@ -79,6 +84,8 @@ export const getDataSourceProperties = (dataSource: DATA_SOURCES | string) => {
7984
return TrinoProperties;
8085
case DATA_SOURCES.SNOWFLAKE:
8186
return SnowflakeProperties;
87+
case DATA_SOURCES.ATHENA:
88+
return AthenaProperties;
8289
default:
8390
return null;
8491
}

‎wren-ui/src/utils/enum/dataSources.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,4 +8,5 @@ export enum DATA_SOURCES {
88
CLICK_HOUSE = 'CLICK_HOUSE',
99
TRINO = 'TRINO',
1010
SNOWFLAKE = 'SNOWFLAKE',
11+
ATHENA = 'ATHENA',
1112
}

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,18 @@ export const ERROR_TEXTS = {
4949
ACCOUNT: {
5050
REQUIRED: 'Please input account.',
5151
},
52+
S3_STAGING_DIR: {
53+
REQUIRED: 'Please input S3 staging directory.',
54+
},
55+
AWS_REGION: {
56+
REQUIRED: 'Please input AWS region.',
57+
},
58+
AWS_ACCESS_KEY: {
59+
REQUIRED: 'Please input AWS access key ID.',
60+
},
61+
AWS_SECRET_KEY: {
62+
REQUIRED: 'Please input AWS secret access key.',
63+
},
5264
},
5365
ADD_RELATION: {
5466
FROM_FIELD: {

0 commit comments

Comments
 (0)