feat(wren-ui): Support Snowflake data source UI (#911)

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

@@ -131,7 +131,8 @@ export enum DataSourceName {
   MSSQL = 'MSSQL',
   MYSQL = 'MYSQL',
   POSTGRES = 'POSTGRES',
-  TRINO = 'TRINO'
+  TRINO = 'TRINO',
+  SNOWFLAKE = 'SNOWFLAKE',
 }
 
 export type DetailStep = {

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

@@ -52,12 +52,21 @@ export interface IbisTrinoConnectionInfo {
   password: string;
 }
 
+export interface IbisSnowflakeConnectionInfo {
+  user: string;
+  password: string;
+  account: string;
+  database: string;
+  schema: string;
+}
+
 export type IbisConnectionInfo =
   | UrlBasedConnectionInfo
   | HostBasedConnectionInfo
   | IbisPostgresConnectionInfo
   | IbisBigQueryConnectionInfo
-  | IbisTrinoConnectionInfo;
+  | IbisTrinoConnectionInfo
+  | IbisSnowflakeConnectionInfo;
 
 export enum SupportedDataSource {
   POSTGRES = 'POSTGRES',

wren-ui/src/apollo/server/adaptors/tests/ibisAdaptor.test.ts

@@ -15,6 +15,7 @@ import {
   MYSQL_CONNECTION_INFO,
   POSTGRES_CONNECTION_INFO,
   TRINO_CONNECTION_INFO,
+  SNOWFLAKE_CONNECTION_INFO,
 } from '../../repositories';
 import { snakeCase } from 'lodash';
 import { Encryptor } from '../../utils';
@@ -82,6 +83,14 @@ describe('IbisAdaptor', () => {
     username: 'my-username',
   };
 
+  const mockSnowflakeConnectionInfo: SNOWFLAKE_CONNECTION_INFO = {
+    user: 'my-user',
+    password: 'my-password',
+    account: 'my-account',
+    database: 'my-database',
+    schema: 'my-schema',
+  };
+
   const mockManifest: Manifest = {
     catalog: 'wrenai', // eg: "test-catalog"
     schema: 'wrenai', // eg: "test-schema"
@@ -264,6 +273,29 @@ describe('IbisAdaptor', () => {
     );
   });
 
+  it('should get snowflake constraints', async () => {
+    const mockResponse = { data: [] };
+    mockedAxios.post.mockResolvedValue(mockResponse);
+    // mock decrypt method in Encryptor to return the same password
+    mockedEncryptor.prototype.decrypt.mockReturnValue(
+      JSON.stringify({ password: mockSnowflakeConnectionInfo.password }),
+    );
+
+    const result = await ibisAdaptor.getConstraints(
+      DataSourceName.SNOWFLAKE,
+      mockSnowflakeConnectionInfo,
+    );
+    const expectConnectionInfo = Object.entries(
+      mockSnowflakeConnectionInfo,
+    ).reduce((acc, [key, value]) => ((acc[snakeCase(key)] = value), acc), {});
+
+    expect(result).toEqual([]);
+    expect(mockedAxios.post).toHaveBeenCalledWith(
+      `${ibisServerEndpoint}/v2/connector/snowflake/metadata/constraints`,
+      { connectionInfo: expectConnectionInfo },
+    );
+  });
+
   it('should get click house constraints', async () => {
     const mockResponse = { data: [] };
     mockedAxios.post.mockResolvedValue(mockResponse);

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

@@ -3,6 +3,7 @@ import {
   IbisPostgresConnectionInfo,
   HostBasedConnectionInfo,
   UrlBasedConnectionInfo,
+  IbisSnowflakeConnectionInfo,
 } from './adaptors/ibisAdaptor';
 import {
   BIG_QUERY_CONNECTION_INFO,
@@ -13,6 +14,7 @@ import {
   WREN_AI_CONNECTION_INFO,
   CLICK_HOUSE_CONNECTION_INFO,
   TRINO_CONNECTION_INFO,
+  SNOWFLAKE_CONNECTION_INFO,
 } from './repositories';
 import { DataSourceName } from './types';
 import { getConfig } from './config';
@@ -178,6 +180,23 @@ const dataSource = {
     },
   } as IDataSourceConnectionInfo<TRINO_CONNECTION_INFO, UrlBasedConnectionInfo>,
 
+  // Snowflake
+  [DataSourceName.SNOWFLAKE]: {
+    sensitiveProps: ['password'],
+    toIbisConnectionInfo(connectionInfo) {
+      const decryptedConnectionInfo = decryptConnectionInfo(
+        DataSourceName.SNOWFLAKE,
+        connectionInfo,
+      );
+      const { user, password, account, database, schema } =
+        decryptedConnectionInfo as SNOWFLAKE_CONNECTION_INFO;
+      return { user, password, account, database, schema };
+    },
+  } as IDataSourceConnectionInfo<
+    SNOWFLAKE_CONNECTION_INFO,
+    IbisSnowflakeConnectionInfo
+  >,
+
   // DuckDB
   [DataSourceName.DUCKDB]: {
     sensitiveProps: [],

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

@@ -60,6 +60,14 @@ export interface TRINO_CONNECTION_INFO {
   ssl: boolean;
 }
 
+export interface SNOWFLAKE_CONNECTION_INFO {
+  user: string;
+  password: string;
+  account: string;
+  database: string;
+  schema: string;
+}
+
 export interface DUCKDB_CONNECTION_INFO {
   initSql: string;
   extensions: Array<string>;
@@ -73,7 +81,8 @@ export type WREN_AI_CONNECTION_INFO =
   | DUCKDB_CONNECTION_INFO
   | MS_SQL_CONNECTION_INFO
   | CLICK_HOUSE_CONNECTION_INFO
-  | TRINO_CONNECTION_INFO;
+  | TRINO_CONNECTION_INFO
+  | SNOWFLAKE_CONNECTION_INFO;
 
 export interface Project {
   id: number; // ID

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

@@ -11,6 +11,7 @@ export const typeDefs = gql`
     MSSQL
     CLICK_HOUSE
     TRINO
+    SNOWFLAKE
   }
 
   enum ExpressionName {

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

@@ -6,6 +6,7 @@ export enum DataSourceName {
   MSSQL = 'MSSQL',
   CLICK_HOUSE = 'CLICK_HOUSE',
   TRINO = 'TRINO',
+  SNOWFLAKE = 'SNOWFLAKE',
 }
 
 export interface DataSource {

wren-ui/src/components/pages/setup/dataSources/SnowflakeProperties.tsx

@@ -0,0 +1,96 @@
+import { Form, Input } from 'antd';
+import { ERROR_TEXTS } from '@/utils/error';
+import { FORM_MODE } from '@/utils/enum';
+
+interface Props {
+  mode?: FORM_MODE;
+}
+
+export default function SnowflakeProperties(props: Props) {
+  const { mode } = props;
+  const isEditMode = mode === FORM_MODE.EDIT;
+  return (
+    <>
+      <Form.Item
+        label="Display name"
+        name="displayName"
+        required
+        rules={[
+          {
+            required: true,
+            message: ERROR_TEXTS.CONNECTION.DISPLAY_NAME.REQUIRED,
+          },
+        ]}
+      >
+        <Input />
+      </Form.Item>
+      <Form.Item
+        label="Username"
+        name="user"
+        rules={[
+          {
+            required: true,
+            message: ERROR_TEXTS.CONNECTION.USERNAME.REQUIRED,
+          },
+        ]}
+      >
+        <Input />
+      </Form.Item>
+      <Form.Item
+        label="Password"
+        name="password"
+        required
+        rules={[
+          {
+            required: true,
+            message: ERROR_TEXTS.CONNECTION.PASSWORD.REQUIRED,
+          },
+        ]}
+      >
+        <Input.Password placeholder="input password" />
+      </Form.Item>
+      <Form.Item
+        label="Account"
+        name="account"
+        required
+        rules={[
+          {
+            required: true,
+            message: ERROR_TEXTS.CONNECTION.ACCOUNT.REQUIRED,
+          },
+        ]}
+      >
+        <Input
+          placeholder="<snowflake_org_id>-<snowflake_user_id>"
+          disabled={isEditMode}
+        />
+      </Form.Item>
+      <Form.Item
+        label="Database name"
+        name="database"
+        required
+        rules={[
+          {
+            required: true,
+            message: ERROR_TEXTS.CONNECTION.DATABASE.REQUIRED,
+          },
+        ]}
+      >
+        <Input placeholder="Snowflake database name" disabled={isEditMode} />
+      </Form.Item>
+      <Form.Item
+        label="Schema"
+        name="schema"
+        required
+        rules={[
+          {
+            required: true,
+            message: ERROR_TEXTS.CONNECTION.SCHEMA.REQUIRED,
+          },
+        ]}
+      >
+        <Input />
+      </Form.Item>
+    </>
+  );
+}

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

@@ -14,6 +14,7 @@ import PostgreSQLProperties from './dataSources/PostgreSQLProperties';
 import SQLServerProperties from './dataSources/SQLServerProperties';
 import ClickHouseProperties from './dataSources/ClickHouseProperties';
 import TrinoProperties from './dataSources/TrinoProperties';
+import SnowflakeProperties from './dataSources/SnowflakeProperties';
 import { SampleDatasetName } from '@/apollo/client/graphql/__types__';
 import { ERROR_CODES } from '@/utils/errorHandler';
 
@@ -102,6 +103,12 @@ export const DATA_SOURCE_OPTIONS = {
     guide: 'https://docs.getwren.ai/oss/guide/connect/trino',
     disabled: false,
   },
+  [DATA_SOURCES.SNOWFLAKE]: {
+    label: 'Snowflake',
+    logo: '/images/dataSource/snowflake.svg',
+    guide: 'https://docs.getwren.ai/oss/guide/connect/snowflake',
+    disabled: false,
+  },
 } as { [key: string]: ButtonOption };
 
 export const DATA_SOURCE_FORM = {
@@ -112,6 +119,7 @@ export const DATA_SOURCE_FORM = {
   [DATA_SOURCES.MSSQL]: { component: SQLServerProperties },
   [DATA_SOURCES.CLICK_HOUSE]: { component: ClickHouseProperties },
   [DATA_SOURCES.TRINO]: { component: TrinoProperties },
+  [DATA_SOURCES.SNOWFLAKE]: { component: SnowflakeProperties },
 };
 
 export const TEMPLATE_OPTIONS = {
@@ -166,6 +174,10 @@ export const getDataSource = (dataSource: DATA_SOURCES) => {
         DATA_SOURCE_OPTIONS[DATA_SOURCES.TRINO],
         DATA_SOURCE_FORM[DATA_SOURCES.TRINO],
       ),
+      [DATA_SOURCES.SNOWFLAKE]: merge(
+        DATA_SOURCE_OPTIONS[DATA_SOURCES.SNOWFLAKE],
+        DATA_SOURCE_FORM[DATA_SOURCES.SNOWFLAKE],
+      ),
     }[dataSource] || defaultDataSource
   );
 };

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

@@ -6,4 +6,5 @@ export enum DATA_SOURCES {
   MSSQL = 'MSSQL',
   CLICK_HOUSE = 'CLICK_HOUSE',
   TRINO = 'TRINO',
+  SNOWFLAKE = 'SNOWFLAKE',
 }

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

@@ -46,6 +46,9 @@ export const ERROR_TEXTS = {
     CATALOG: {
       REQUIRED: 'Please input catalog name.',
     },
+    ACCOUNT: {
+      REQUIRED: 'Please input account.',
+    },
   },
   ADD_RELATION: {
     FROM_FIELD: {