feat(wren-ui): Import from data source SQL (#1553)

2026-09-01 15:34:04 +08:00 · 2025-04-18 13:58:16 +08:00
parent 2b252ae2b7
commit b2f29f4555
18 changed files with 693 additions and 287 deletions
@@ -148,7 +148,8 @@ If you prefer to learn by example, you can refer to this Trino [issue](https://g
   - Implement the data source form template in this file

 3. Set up the data source template:
-   - Navigate to `wren-ui/src/components/pages/setup/utils` > `DATA_SOURCE_FORM`
+   - Navigate to `wren-ui/src/utils/dataSourceType.ts`
+   - Add new data source image, name, properties
   - Update the necessary files to include the new data source template settings

 4. Update the data source list:
@@ -11,6 +11,7 @@ export type Scalars = {
  Boolean: boolean;
  Int: number;
  Float: number;
+  DialectSQL: any;
  JSON: any;
 };

@@ -551,6 +552,10 @@ export type ModelInfo = {
  sourceTableName: Scalars['String'];
 };

+export type ModelSubstituteInput = {
+  sql: Scalars['DialectSQL'];
+};
+
 export type ModelSyncResponse = {
  __typename?: 'ModelSyncResponse';
  status: SyncStatus;
@@ -592,6 +597,7 @@ export type Mutation = {
  generateThreadResponseAnswer: ThreadResponse;
  generateThreadResponseBreakdown: ThreadResponse;
  generateThreadResponseChart: ThreadResponse;
+  modelSubstitute: Scalars['String'];
  previewBreakdownData: Scalars['JSON'];
  previewData: Scalars['JSON'];
  previewItemSQL: Scalars['JSON'];
@@ -773,6 +779,11 @@ export type MutationGenerateThreadResponseChartArgs = {
 };


+export type MutationModelSubstituteArgs = {
+  data: ModelSubstituteInput;
+};
+
+
 export type MutationPreviewBreakdownDataArgs = {
  where: PreviewDataInput;
 };
@@ -17,6 +17,13 @@ export type GenerateQuestionMutationVariables = Types.Exact<{

 export type GenerateQuestionMutation = { __typename?: 'Mutation', generateQuestion: string };

+export type ModelSubstituteMutationVariables = Types.Exact<{
+  data: Types.ModelSubstituteInput;
+}>;
+
+
+export type ModelSubstituteMutation = { __typename?: 'Mutation', modelSubstitute: string };
+

 export const PreviewSqlDocument = gql`
    mutation PreviewSQL($data: PreviewSQLDataInput!) {
@@ -79,4 +86,35 @@ export function useGenerateQuestionMutation(baseOptions?: Apollo.MutationHookOpt
      }
 export type GenerateQuestionMutationHookResult = ReturnType<typeof useGenerateQuestionMutation>;
 export type GenerateQuestionMutationResult = Apollo.MutationResult<GenerateQuestionMutation>;
-export type GenerateQuestionMutationOptions = Apollo.BaseMutationOptions<GenerateQuestionMutation, GenerateQuestionMutationVariables>;
+export type GenerateQuestionMutationOptions = Apollo.BaseMutationOptions<GenerateQuestionMutation, GenerateQuestionMutationVariables>;
+export const ModelSubstituteDocument = gql`
+    mutation ModelSubstitute($data: ModelSubstituteInput!) {
+  modelSubstitute(data: $data)
+}
+    `;
+export type ModelSubstituteMutationFn = Apollo.MutationFunction<ModelSubstituteMutation, ModelSubstituteMutationVariables>;
+
+/**
+ * __useModelSubstituteMutation__
+ *
+ * To run a mutation, you first call `useModelSubstituteMutation` within a React component and pass it any options that fit your needs.
+ * When your component renders, `useModelSubstituteMutation` returns a tuple that includes:
+ * - A mutate function that you can call at any time to execute the mutation
+ * - An object with fields that represent the current status of the mutation's execution
+ *
+ * @param baseOptions options that will be passed into the mutation, supported options are listed on: https://www.apollographql.com/docs/react/api/react-hooks/#options-2;
+ *
+ * @example
+ * const [modelSubstituteMutation, { data, loading, error }] = useModelSubstituteMutation({
+ *   variables: {
+ *      data: // value for 'data'
+ *   },
+ * });
+ */
+export function useModelSubstituteMutation(baseOptions?: Apollo.MutationHookOptions<ModelSubstituteMutation, ModelSubstituteMutationVariables>) {
+        const options = {...defaultOptions, ...baseOptions}
+        return Apollo.useMutation<ModelSubstituteMutation, ModelSubstituteMutationVariables>(ModelSubstituteDocument, options);
+      }
+export type ModelSubstituteMutationHookResult = ReturnType<typeof useModelSubstituteMutation>;
+export type ModelSubstituteMutationResult = Apollo.MutationResult<ModelSubstituteMutation>;
+export type ModelSubstituteMutationOptions = Apollo.BaseMutationOptions<ModelSubstituteMutation, ModelSubstituteMutationVariables>;
@@ -11,3 +11,9 @@ export const GENERATE_QUESTION = gql`
    generateQuestion(data: $data)
  }
 `;
+
+export const MODEL_SUBSTITUDE = gql`
+  mutation ModelSubstitute($data: ModelSubstituteInput!) {
+    modelSubstitute(data: $data)
+  }
+`;
@@ -5,7 +5,7 @@ import { useState, useContext, useRef } from 'react';
 import ReadOutlined from '@ant-design/icons/ReadOutlined';
 import EditOutlined from '@ant-design/icons/EditOutlined';
 import { nextTick } from '@/utils/time';
-import { Mention } from '@/hooks/useMentions';
+import { Mention } from '@/hooks/useAutoComplete';
 import { FormItemInputContext } from 'antd/lib/form/context';
 import MarkdownBlock from './MarkdownBlock';

@@ -1,37 +1,130 @@
-import { useState } from 'react';
+import clsx from 'clsx';
+import { useState, useContext, useRef, useEffect } from 'react';
+import styled from 'styled-components';
 import AceEditor from '@/components/editor/AceEditor';
+import { FormItemInputContext } from 'antd/lib/form/context';
+import useAutoComplete from '@/hooks/useAutoComplete';
+
+const Wrapper = styled.div`
+  transition: all 0.3s cubic-bezier(0.645, 0.045, 0.355, 1);
+
+  &:hover {
+    border-color: var(--geekblue-5) !important;
+  }
+
+  &.adm-markdown-editor-error {
+    border-color: var(--red-5) !important;
+
+    .adm-markdown-editor-length {
+      color: var(--red-5) !important;
+    }
+  }
+  &:not(.adm-markdown-editor-error).adm-markdown-editor-focused {
+    border-color: var(--geekblue-5) !important;
+    box-shadow: 0 0 0 2px rgba(47, 84, 235, 0.2);
+  }
+
+  &.adm-markdown-editor-focused.adm-markdown-editor-error {
+    borer-color: var(--red-4) !important;
+    box-shadow: 0 0 0 2px rgba(255, 77, 79, 0.2);
+  }
+`;
+
+const Toolbar = styled.div`
+  color: var(--gray-8);
+  background-color: var(--gray-3);
+  border-bottom: 1px solid var(--gray-5);
+  height: 32px;
+  padding: 4px 8px;
+  border-radius: 4px 4px 0px 0px;
+`;

 interface Props {
  value?: string;
  onChange?: (value: string) => void;
  autoFocus?: boolean;
+  autoComplete?: boolean;
+  toolbar?: React.ReactNode;
 }

+const getLangTools = () => {
+  const { ace } = window as any;
+  return ace ? ace.require('ace/ext/language_tools') : null;
+};
+
 export default function SQLEditor(props: Props) {
-  const { value, onChange, autoFocus } = props;
+  const { value, onChange, autoFocus, autoComplete, toolbar } = props;
+  const $wrapper = useRef<HTMLDivElement>(null);
+  const [focused, setFocused] = useState<boolean>(false);
+
+  const formItemContext = useContext(FormItemInputContext);
+  const { status } = formItemContext;
+
+  const completers = useAutoComplete({
+    includeColumns: true,
+    skip: !autoComplete,
+  });
+
+  const resetCompleters = () => {
+    // clear custom completer
+    const langTools = getLangTools();
+    langTools?.setCompleters([
+      langTools.keyWordCompleter,
+      langTools.snippetCompleter,
+      langTools.textCompleter,
+    ]);
+  };
+
+  useEffect(() => {
+    resetCompleters();
+    if (!autoComplete || completers.length === 0) return;
+
+    const langTools = getLangTools();
+    const customCompleter = {
+      getCompletions: (_editor, _session, _pos, _prefix, callback) => {
+        callback(null, completers);
+      },
+    };
+    langTools?.addCompleter(customCompleter);
+
+    return () => resetCompleters();
+  }, [focused, autoComplete, completers]);

  const [sql, setSql] = useState(value || '');

  const change = (sql) => {
    setSql(sql);
-    onChange && onChange(sql);
+    onChange?.(sql);
  };

  return (
-    <AceEditor
-      mode="sql"
-      width="100%"
-      height="300px"
-      fontSize={14}
-      theme="tomorrow"
-      value={value || sql}
-      onChange={change}
-      name="sql_editor"
-      editorProps={{ $blockScrolling: true }}
-      enableLiveAutocompletion
-      enableBasicAutocompletion
-      showPrintMargin={false}
-      focus={autoFocus}
-    />
+    <Wrapper
+      ref={$wrapper}
+      className={clsx(
+        'border border-gray-5 rounded overflow-hidden',
+        status ? `adm-markdown-editor-${status}` : '',
+        focused ? 'adm-markdown-editor-focused' : '',
+      )}
+      tabIndex={-1}
+    >
+      {!!toolbar && <Toolbar>{toolbar}</Toolbar>}
+      <AceEditor
+        mode="sql"
+        width="100%"
+        height="300px"
+        fontSize={14}
+        theme="tomorrow"
+        value={value || sql}
+        onChange={change}
+        onFocus={() => setFocused(true)}
+        onBlur={() => setFocused(false)}
+        name="sql_editor"
+        editorProps={{ $blockScrolling: true }}
+        enableLiveAutocompletion
+        enableBasicAutocompletion
+        showPrintMargin={false}
+        focus={autoFocus}
+      />
+    </Wrapper>
  );
 }
@@ -4,7 +4,7 @@ import styled from 'styled-components';
 import { Form, Modal, Select, Tag } from 'antd';
 import QuestionCircleOutlined from '@ant-design/icons/QuestionCircleOutlined';
 import { ERROR_TEXTS } from '@/utils/error';
-import useMentions from '@/hooks/useMentions';
+import useAutoComplete, { convertMention } from '@/hooks/useAutoComplete';
 import { ModalAction } from '@/hooks/useModalAction';
 import MarkdownEditor from '@/components/editor/MarkdownEditor';
 import { useListModelsQuery } from '@/apollo/client/graphql/model.generated';
@@ -36,7 +36,11 @@ export default function AdjustReasoningStepsModal(props: Props) {
  const { visible, defaultValue, loading, onSubmit, onClose } = props;
  const [form] = Form.useForm();

-  const { mentions } = useMentions({ includeColumns: true, skip: !visible });
+  const mentions = useAutoComplete({
+    convertor: convertMention,
+    includeColumns: true,
+    skip: !visible,
+  });
  const listModelsResult = useListModelsQuery({ skip: !visible });
  const modelNameMap = keyBy(
    listModelsResult.data?.listModels,
@@ -178,7 +178,7 @@ export default function AdjustSQLModal(props: Props) {
            },
          ]}
        >
-          <SQLEditor autoFocus />
+          <SQLEditor autoComplete autoFocus />
        </Form.Item>
      </Form>
      <div className="my-3">
@@ -98,7 +98,7 @@ export function FixSQLModal(props: Props) {
            },
          ]}
        >
-          <SQLEditor autoFocus />
+          <SQLEditor autoComplete autoFocus />
        </Form.Item>
      </Form>
      <div className="my-3">
@@ -0,0 +1,120 @@
+import { useMemo } from 'react';
+import { Modal, Form, Alert } from 'antd';
+import { ModalAction } from '@/hooks/useModalAction';
+import { getDataSourceImage, getDataSourceName } from '@/utils/dataSourceType';
+import { DATA_SOURCES } from '@/utils/enum';
+import { ERROR_TEXTS } from '@/utils/error';
+import { parseGraphQLError } from '@/utils/errorHandler';
+import SQLEditor from '@/components/editor/SQLEditor';
+import ErrorCollapse from '@/components/ErrorCollapse';
+import { useModelSubstituteMutation } from '@/apollo/client/graphql/sql.generated';
+import { DataSource, DataSourceName } from '@/apollo/client/graphql/__types__';
+
+type Props = ModalAction<{ dataSource: DATA_SOURCES }> & {
+  loading?: boolean;
+};
+
+const Toolbar = (props) => {
+  const { dataSource } = props;
+  if (!dataSource) return null;
+  const logo = getDataSourceImage(dataSource);
+  const name = getDataSourceName(dataSource);
+  return (
+    <>
+      <span className="d-flex align-center gx-2">
+        <img src={logo} alt="logo" width="20" height="20" />
+        {name}
+      </span>
+    </>
+  );
+};
+
+export const isSupportSubstitute = (dataSource: DataSource) => {
+  // DuckDB not supported, sample dataset as well
+  return (
+    !dataSource?.sampleDataset && dataSource?.type !== DataSourceName.DUCKDB
+  );
+};
+
+export default function ImportDataSourceSQLModal(props: Props) {
+  const { visible, defaultValue, loading, onSubmit, onClose } = props;
+  const name = getDataSourceName(defaultValue?.dataSource) || 'data source';
+
+  const [substituteDialectSQL, modelSubstitudeResult] =
+    useModelSubstituteMutation();
+  const error = useMemo(
+    () =>
+      modelSubstitudeResult.error
+        ? {
+            ...parseGraphQLError(modelSubstitudeResult.error),
+            shortMessage: `Invalid ${name} SQL syntax`,
+          }
+        : null,
+    [modelSubstitudeResult.error],
+  );
+
+  const [form] = Form.useForm();
+
+  const reset = () => {
+    form.resetFields();
+    modelSubstitudeResult.reset();
+  };
+
+  const submit = async () => {
+    form
+      .validateFields()
+      .then(async (values) => {
+        const response = await substituteDialectSQL({
+          variables: { data: { sql: values.dialectSql } },
+        });
+        await onSubmit(response.data?.modelSubstitute);
+        onClose();
+      })
+      .catch(console.error);
+  };
+
+  return (
+    <Modal
+      title={`Import from ${name} SQL`}
+      centered
+      closable
+      confirmLoading={loading}
+      destroyOnClose
+      maskClosable={false}
+      onCancel={onClose}
+      onOk={submit}
+      okText="Convert"
+      visible={visible}
+      width={600}
+      cancelButtonProps={{ disabled: loading }}
+      okButtonProps={{ disabled: loading }}
+      afterClose={() => reset()}
+    >
+      <Form form={form} layout="vertical">
+        <Form.Item
+          name="dialectSql"
+          label="SQL statement"
+          rules={[
+            {
+              required: true,
+              message: ERROR_TEXTS.IMPORT_DATA_SOURCE_SQL.SQL.REQUIRED,
+            },
+          ]}
+        >
+          <SQLEditor
+            toolbar={<Toolbar dataSource={defaultValue?.dataSource} />}
+            autoFocus
+          />
+        </Form.Item>
+      </Form>
+      {!!error && (
+        <Alert
+          showIcon
+          type="error"
+          message={error.shortMessage}
+          description={<ErrorCollapse message={error.message} />}
+        />
+      )}
+    </Modal>
+  );
+}
@@ -1,16 +1,23 @@
-import { useEffect, useState } from 'react';
+import { useEffect, useMemo, useState } from 'react';
 import styled from 'styled-components';
 import { Alert, Button, Form, Input, Modal, Typography } from 'antd';
+import { Logo } from '@/components/Logo';
 import InfoCircleOutlined from '@ant-design/icons/InfoCircleOutlined';
+import SelectOutlined from '@ant-design/icons/SelectOutlined';
 import { ERROR_TEXTS } from '@/utils/error';
 import { FORM_MODE } from '@/utils/enum';
-import { ModalAction } from '@/hooks/useModalAction';
+import { getDataSourceName } from '@/utils/dataSourceType';
+import useModalAction, { ModalAction } from '@/hooks/useModalAction';
 import SQLEditor from '@/components/editor/SQLEditor';
 import { parseGraphQLError } from '@/utils/errorHandler';
 import { createSQLPairQuestionValidator } from '@/utils/validator';
 import ErrorCollapse from '@/components/ErrorCollapse';
 import PreviewData from '@/components/dataPreview/PreviewData';
+import ImportDataSourceSQLModal, {
+  isSupportSubstitute,
+} from '@/components/modals/ImportDataSourceSQLModal';
 import { usePreviewSqlMutation } from '@/apollo/client/graphql/sql.generated';
+import { useGetSettingsQuery } from '@/apollo/client/graphql/settings.generated';
 import { useGenerateQuestionMutation } from '@/apollo/client/graphql/sql.generated';
 import { SqlPair } from '@/apollo/client/graphql/__types__';

@@ -27,6 +34,23 @@ const StyledForm = styled(Form)`
  }
 `;

+const Toolbar = (props: { dataSource: string; onClick: () => void }) => {
+  const { dataSource, onClick } = props;
+  const name = getDataSourceName(dataSource);
+  return (
+    <div className="d-flex justify-space-between align-center px-1">
+      <span className="d-flex align-center gx-2">
+        <Logo size={16} />
+        Wren SQL
+      </span>
+      <Button className="px-0" type="link" size="small" onClick={onClick}>
+        <SelectOutlined />
+        Import from {name} SQL
+      </Button>
+    </div>
+  );
+};
+
 export default function QuestionSQLPairModal(props: Props) {
  const {
    defaultValue,
@@ -40,6 +64,17 @@ export default function QuestionSQLPairModal(props: Props) {

  // pass payload?.isCreateMode to prevent formMode from being set to Update when passing defaultValue, for the 'Add a SQL pair from an existing answer' scenario use.
  const isCreateMode = formMode === FORM_MODE.CREATE || payload?.isCreateMode;
+  const importDataSourceSQLModal = useModalAction();
+
+  const { data: settingsResult } = useGetSettingsQuery();
+  const settings = settingsResult?.settings;
+  const dataSource = useMemo(
+    () => ({
+      isSupportSubstitute: isSupportSubstitute(settings?.dataSource),
+      type: settings?.dataSource?.type,
+    }),
+    [settings?.dataSource],
+  );

  const [form] = Form.useForm();
  const [error, setError] =
@@ -151,133 +186,159 @@ export default function QuestionSQLPairModal(props: Props) {
  const disabled = !sqlValue;

  return (
-    <Modal
-      title={`${isCreateMode ? 'Add' : 'Update'} question-SQL pair`}
-      centered
-      closable
-      confirmLoading={confirmLoading}
-      destroyOnClose
-      maskClosable={false}
-      onCancel={onClose}
-      visible={visible}
-      width={640}
-      cancelButtonProps={{ disabled: confirmLoading }}
-      okButtonProps={{ disabled: previewSqlResult.loading }}
-      afterClose={() => handleReset()}
-      footer={
-        <div className="d-flex justify-space-between align-center">
-          <div
-            className="text-sm ml-2 d-flex justify-space-between align-center"
-            style={{ width: 300 }}
-          >
-            <InfoCircleOutlined className="mr-2 text-sm gray-7" />
-            <Typography.Text
-              type="secondary"
-              className="text-sm gray-7 text-left"
-            >
-              The SQL statement used here follows <b>Wren SQL</b>, which is
-              based on ANSI SQL and optimized for Wren AI.{` `}
-              <Typography.Link
-                type="secondary"
-                href="https://docs.getwren.ai/oss/guide/home/wren_sql"
-                target="_blank"
-                rel="noopener noreferrer"
-              >
-                Learn more about the syntax.
-              </Typography.Link>
-            </Typography.Text>
-          </div>
-          <div>
-            <Button onClick={onClose}>Cancel</Button>
-            <Button
-              type="primary"
-              onClick={onSubmitButton}
-              loading={confirmLoading}
-            >
-              Submit
-            </Button>
-          </div>
-        </div>
-      }
-    >
-      <StyledForm form={form} preserve={false} layout="vertical">
-        <Form.Item
-          className="adm-question-form-item"
-          label={
+    <>
+      <Modal
+        title={`${isCreateMode ? 'Add' : 'Update'} question-SQL pair`}
+        centered
+        closable
+        confirmLoading={confirmLoading}
+        destroyOnClose
+        maskClosable={false}
+        onCancel={onClose}
+        visible={visible}
+        width={640}
+        cancelButtonProps={{ disabled: confirmLoading }}
+        okButtonProps={{ disabled: previewSqlResult.loading }}
+        afterClose={() => handleReset()}
+        footer={
+          <div className="d-flex justify-space-between align-center">
            <div
-              className="d-flex justify-space-between"
-              style={{ width: '100%' }}
+              className="text-sm ml-2 d-flex justify-space-between align-center"
+              style={{ width: 300 }}
            >
-              <span>Question</span>
-              <div className="gray-8 text-sm">
-                Let AI create a matching question for this SQL statement.
-                <Button
-                  className="ml-2"
-                  size="small"
-                  loading={generatingQuestion}
-                  onClick={onGenerateQuestion}
-                  disabled={disabled}
+              <InfoCircleOutlined className="mr-2 text-sm gray-7" />
+              <Typography.Text
+                type="secondary"
+                className="text-sm gray-7 text-left"
+              >
+                The SQL statement used here follows <b>Wren SQL</b>, which is
+                based on ANSI SQL and optimized for Wren AI.{` `}
+                <Typography.Link
+                  type="secondary"
+                  href="https://docs.getwren.ai/oss/guide/home/wren_sql"
+                  target="_blank"
+                  rel="noopener noreferrer"
                >
-                  <span className="text-sm">Generate question</span>
-                </Button>
-              </div>
+                  Learn more about the syntax.
+                </Typography.Link>
+              </Typography.Text>
+            </div>
+            <div>
+              <Button onClick={onClose}>Cancel</Button>
+              <Button
+                type="primary"
+                onClick={onSubmitButton}
+                loading={confirmLoading}
+              >
+                Submit
+              </Button>
            </div>
-          }
-          name="question"
-          required
-          rules={[
-            {
-              validator: createSQLPairQuestionValidator(
-                ERROR_TEXTS.SQL_PAIR.QUESTION,
-              ),
-            },
-          ]}
-        >
-          <Input />
-        </Form.Item>
-        <Form.Item
-          label="SQL statement"
-          name="sql"
-          required
-          rules={[
-            {
-              required: true,
-              message: ERROR_TEXTS.SQL_PAIR.SQL.REQUIRED,
-            },
-          ]}
-        >
-          <SQLEditor autoFocus />
-        </Form.Item>
-      </StyledForm>
-      <div className="my-3">
-        <Typography.Text className="d-block gray-7 mb-2">
-          Data preview (50 rows)
-        </Typography.Text>
-        <Button
-          onClick={onPreviewData}
-          loading={previewing}
-          disabled={disabled}
-        >
-          Preview data
-        </Button>
-        {showPreview && (
-          <div className="my-3">
-            <PreviewData
-              loading={previewing}
-              previewData={previewSqlResult?.data?.previewSql}
-              copyable={false}
-            />
          </div>
+        }
+      >
+        <StyledForm form={form} preserve={false} layout="vertical">
+          <Form.Item
+            className="adm-question-form-item"
+            label={
+              <div
+                className="d-flex justify-space-between"
+                style={{ width: '100%' }}
+              >
+                <span>Question</span>
+                <div className="gray-8 text-sm">
+                  Let AI create a matching question for this SQL statement.
+                  <Button
+                    className="ml-2"
+                    size="small"
+                    loading={generatingQuestion}
+                    onClick={onGenerateQuestion}
+                    disabled={disabled}
+                  >
+                    <span className="text-sm">Generate question</span>
+                  </Button>
+                </div>
+              </div>
+            }
+            name="question"
+            required
+            rules={[
+              {
+                validator: createSQLPairQuestionValidator(
+                  ERROR_TEXTS.SQL_PAIR.QUESTION,
+                ),
+              },
+            ]}
+          >
+            <Input />
+          </Form.Item>
+          <Form.Item
+            label="SQL statement"
+            name="sql"
+            required
+            rules={[
+              {
+                required: true,
+                message: ERROR_TEXTS.SQL_PAIR.SQL.REQUIRED,
+              },
+            ]}
+          >
+            <SQLEditor
+              toolbar={
+                dataSource.isSupportSubstitute && (
+                  <Toolbar
+                    dataSource={dataSource.type}
+                    onClick={() =>
+                      importDataSourceSQLModal.openModal({
+                        dataSource: dataSource.type,
+                      })
+                    }
+                  />
+                )
+              }
+              autoComplete
+              autoFocus
+            />
+          </Form.Item>
+        </StyledForm>
+        <div className="my-3">
+          <Typography.Text className="d-block gray-7 mb-2">
+            Data preview (50 rows)
+          </Typography.Text>
+          <Button
+            onClick={onPreviewData}
+            loading={previewing}
+            disabled={disabled}
+          >
+            Preview data
+          </Button>
+          {showPreview && (
+            <div className="my-3">
+              <PreviewData
+                loading={previewing}
+                previewData={previewSqlResult?.data?.previewSql}
+                copyable={false}
+              />
+            </div>
+          )}
+        </div>
+        {!!error && (
+          <Alert
+            showIcon
+            type="error"
+            message={error.shortMessage}
+            description={<ErrorCollapse message={error.message} />}
+          />
        )}
-      </div>
-      {!!error && (
-        <Alert
-          showIcon
-          type="error"
-          message={error.shortMessage}
-          description={<ErrorCollapse message={error.message} />}
+      </Modal>
+      {dataSource.isSupportSubstitute && (
+        <ImportDataSourceSQLModal
+          {...importDataSourceSQLModal.state}
+          onClose={importDataSourceSQLModal.closeModal}
+          onSubmit={async (convertedSql: string) => {
+            form.setFieldsValue({ sql: convertedSql });
+          }}
        />
      )}
-    </Modal>
+    </>
  );
 }
@@ -89,7 +89,7 @@ export default function ConnectDataSource(props: Props) {
        <Alert
          message={connectError.shortMessage}
          description={
-            dataSource === DATA_SOURCES.PG_SQL
+            dataSource === DATA_SOURCES.POSTGRES
              ? getPostgresErrorMessage(connectError)
              : connectError.message
          }
@@ -7,16 +7,12 @@ import Starter from './Starter';
 import ConnectDataSource from './ConnectDataSource';
 import SelectModels from './SelectModels';
 import DefineRelations from './DefineRelations';
-import BigQueryProperties from './dataSources/BigQueryProperties';
-import DuckDBProperties from './dataSources/DuckDBProperties';
-import MySQLProperties from './dataSources/MySQLProperties';
-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';
+import {
+  getDataSourceConfig,
+  getDataSourceFormComponent,
+} from '@/utils/dataSourceType';

 type SetupStep = {
  step: number;
@@ -62,66 +58,47 @@ export const SETUP_STEPS = {

 export const DATA_SOURCE_OPTIONS = {
  [DATA_SOURCES.BIG_QUERY]: {
-    label: 'BigQuery',
-    logo: '/images/dataSource/bigQuery.svg',
+    ...getDataSourceConfig(DATA_SOURCES.BIG_QUERY),
    guide: 'https://docs.getwren.ai/oss/guide/connect/bigquery',
    disabled: false,
  },
  [DATA_SOURCES.DUCKDB]: {
-    label: 'DuckDB',
-    logo: '/images/dataSource/duckDb.svg',
+    ...getDataSourceConfig(DATA_SOURCES.DUCKDB),
    guide: 'https://docs.getwren.ai/oss/guide/connect/duckdb',
    disabled: false,
  },
-  [DATA_SOURCES.PG_SQL]: {
-    label: 'PostgreSQL',
-    logo: '/images/dataSource/postgreSql.svg',
+  [DATA_SOURCES.POSTGRES]: {
+    ...getDataSourceConfig(DATA_SOURCES.POSTGRES),
    guide: 'https://docs.getwren.ai/oss/guide/connect/postgresql',
    disabled: false,
  },
  [DATA_SOURCES.MYSQL]: {
-    label: 'MySQL',
-    logo: '/images/dataSource/mysql.svg',
+    ...getDataSourceConfig(DATA_SOURCES.MYSQL),
    guide: 'https://docs.getwren.ai/oss/guide/connect/mysql',
    disabled: false,
  },
  [DATA_SOURCES.MSSQL]: {
-    label: 'SQL Server',
-    logo: '/images/dataSource/sqlserver.svg',
+    ...getDataSourceConfig(DATA_SOURCES.MSSQL),
    guide: 'https://docs.getwren.ai/oss/guide/connect/sqlserver',
    disabled: false,
  },
  [DATA_SOURCES.CLICK_HOUSE]: {
-    label: 'ClickHouse',
-    logo: '/images/dataSource/clickhouse.svg',
+    ...getDataSourceConfig(DATA_SOURCES.CLICK_HOUSE),
    guide: 'https://docs.getwren.ai/oss/guide/connect/clickhouse',
    disabled: false,
  },
  [DATA_SOURCES.TRINO]: {
-    label: 'Trino',
-    logo: '/images/dataSource/trino.svg',
+    ...getDataSourceConfig(DATA_SOURCES.TRINO),
    guide: 'https://docs.getwren.ai/oss/guide/connect/trino',
    disabled: false,
  },
  [DATA_SOURCES.SNOWFLAKE]: {
-    label: 'Snowflake',
-    logo: '/images/dataSource/snowflake.svg',
+    ...getDataSourceConfig(DATA_SOURCES.SNOWFLAKE),
    guide: 'https://docs.getwren.ai/oss/guide/connect/snowflake',
    disabled: false,
  },
 } as { [key: string]: ButtonOption };

-export const DATA_SOURCE_FORM = {
-  [DATA_SOURCES.BIG_QUERY]: { component: BigQueryProperties },
-  [DATA_SOURCES.DUCKDB]: { component: DuckDBProperties },
-  [DATA_SOURCES.PG_SQL]: { component: PostgreSQLProperties },
-  [DATA_SOURCES.MYSQL]: { component: MySQLProperties },
-  [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 = {
  [SampleDatasetName.ECOMMERCE]: {
    label: 'E-commerce',
@@ -136,49 +113,13 @@ export const TEMPLATE_OPTIONS = {
 };

 export const getDataSources = () => {
-  return Object.keys(DATA_SOURCE_OPTIONS).map((key) => ({
-    ...DATA_SOURCE_OPTIONS[key],
-    value: key,
-  })) as ButtonOption[];
+  return Object.values(DATA_SOURCE_OPTIONS) as ButtonOption[];
 };

 export const getDataSource = (dataSource: DATA_SOURCES) => {
-  const defaultDataSource = merge(
-    DATA_SOURCE_OPTIONS[DATA_SOURCES.BIG_QUERY],
-    DATA_SOURCE_FORM[DATA_SOURCES.BIG_QUERY],
-  );
-  return (
-    {
-      [DATA_SOURCES.BIG_QUERY]: defaultDataSource,
-      [DATA_SOURCES.DUCKDB]: merge(
-        DATA_SOURCE_OPTIONS[DATA_SOURCES.DUCKDB],
-        DATA_SOURCE_FORM[DATA_SOURCES.DUCKDB],
-      ),
-      [DATA_SOURCES.PG_SQL]: merge(
-        DATA_SOURCE_OPTIONS[DATA_SOURCES.PG_SQL],
-        DATA_SOURCE_FORM[DATA_SOURCES.PG_SQL],
-      ),
-      [DATA_SOURCES.MYSQL]: merge(
-        DATA_SOURCE_OPTIONS[DATA_SOURCES.MYSQL],
-        DATA_SOURCE_FORM[DATA_SOURCES.MYSQL],
-      ),
-      [DATA_SOURCES.MSSQL]: merge(
-        DATA_SOURCE_OPTIONS[DATA_SOURCES.MSSQL],
-        DATA_SOURCE_FORM[DATA_SOURCES.MSSQL],
-      ),
-      [DATA_SOURCES.CLICK_HOUSE]: merge(
-        DATA_SOURCE_OPTIONS[DATA_SOURCES.CLICK_HOUSE],
-        DATA_SOURCE_FORM[DATA_SOURCES.CLICK_HOUSE],
-      ),
-      [DATA_SOURCES.TRINO]: merge(
-        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
+  return merge(
+    DATA_SOURCE_OPTIONS[dataSource],
+    getDataSourceFormComponent(dataSource),
  );
 };

@@ -0,0 +1,96 @@
+import { useMemo } from 'react';
+import { capitalize } from 'lodash';
+import { useDiagramQuery } from '@/apollo/client/graphql/diagram.generated';
+import { getNodeTypeIcon } from '@/utils/nodeType';
+import {
+  DiagramModel,
+  DiagramView,
+  DiagramModelField,
+  DiagramViewField,
+} from '@/apollo/client/graphql/__types__';
+
+type Model = DiagramModel | DiagramView;
+type Field = DiagramModelField | DiagramViewField;
+type Convertor<T> = (item: (Model | Field) & { parent?: Model }) => T;
+
+interface Props<T> {
+  skip?: boolean;
+  includeColumns?: boolean;
+  convertor?: Convertor<T>;
+}
+
+const getDocHTML = (item: (Model | Field) & { parent?: Model }) => {
+  return [
+    '<div style="max-width: 380px;">',
+    `<b style="display: block;color: var(--gray-8); padding: 0 4px 4px;">${item.referenceName}</b>`,
+    item.description
+      ? `<div style="color: var(--gray-7); padding: 4px 4px 0; border-top: 1px var(--gray-4) solid;">${item.description}</div>`
+      : null,
+    '</div>',
+  ]
+    .filter(Boolean)
+    .join('');
+};
+
+const shouldQuoteIdentifier = (word: string) => {
+  return /[^a-z0-9_]/.test(word) || /^\d/.test(word);
+};
+
+// For mention usage
+export const convertMention = (item: (Model | Field) & { parent?: Model }) => {
+  return {
+    id: `${item.id}-${item.referenceName}`,
+    label: item.displayName,
+    value: item.referenceName,
+    nodeType: capitalize(item.nodeType),
+    meta: item.parent ? `${item.displayName}.${item.displayName}` : undefined,
+    icon: getNodeTypeIcon(
+      { nodeType: item.nodeType, type: (item as Field).type },
+      { className: 'gray-8 mr-2' },
+    ),
+  };
+};
+
+// For ace completer usage
+export const convertCompleter = (
+  item: (Model | Field) & { parent?: Model },
+) => {
+  return {
+    caption: item.parent
+      ? `${item.parent.displayName}.${item.displayName}`
+      : item.displayName,
+    value: shouldQuoteIdentifier(item.referenceName)
+      ? `"${item.referenceName}"`
+      : item.referenceName,
+    meta: item.nodeType.toLowerCase(),
+    // Higher score for models, views
+    score: item.parent ? 1 : 10,
+    docHTML: getDocHTML(item),
+  };
+};
+
+export type Mention = ReturnType<typeof convertMention>;
+export type Completer = ReturnType<typeof convertCompleter>;
+
+export default function useAutoComplete<T = Completer>(props: Props<T>) {
+  const { includeColumns, skip } = props;
+  const { data } = useDiagramQuery({ skip });
+
+  // Defined convertor
+  const convertor = (props.convertor || convertCompleter) as Convertor<T>;
+
+  return useMemo(() => {
+    const models = data?.diagram.models || [];
+    const views = data?.diagram.views || [];
+
+    return [...models, ...views].reduce((result, item) => {
+      result.push(convertor(item));
+      if (includeColumns) {
+        item.fields.forEach((field) => {
+          result.push(convertor({ ...field, parent: item }));
+        });
+      }
+      return result;
+    }, [] as T[]);
+  }, [data?.diagram, includeColumns]);
+}
@@ -1,62 +0,0 @@
-import { useMemo } from 'react';
-import { capitalize } from 'lodash';
-import { useDiagramQuery } from '@/apollo/client/graphql/diagram.generated';
-import { getNodeTypeIcon } from '@/utils/nodeType';
-import {
-  DiagramModel,
-  DiagramView,
-  DiagramModelField,
-  DiagramViewField,
-} from '@/apollo/client/graphql/__types__';
-
-type Model = DiagramModel | DiagramView;
-type Field = DiagramModelField | DiagramViewField;
-
-interface Props {
-  skip?: boolean;
-  includeColumns?: boolean;
-}
-
-const convertMention = (item: (Model | Field) & { meta?: string }) => {
-  return {
-    id: `${item.id}-${item.referenceName}`,
-    label: item.displayName,
-    value: item.referenceName,
-    nodeType: capitalize(item.nodeType),
-    meta: item.meta,
-    icon: getNodeTypeIcon(
-      { nodeType: item.nodeType, type: (item as Field).type },
-      { className: 'gray-8 mr-2' },
-    ),
-  };
-};
-
-export type Mention = ReturnType<typeof convertMention>;
-
-export default function useMentions(props: Props) {
-  const { includeColumns, skip } = props;
-  const { data } = useDiagramQuery({ skip });
-
-  // handle mentions data
-  const mentions = useMemo(() => {
-    const models = data?.diagram.models || [];
-    const views = data?.diagram.views || [];
-
-    return [...models, ...views].reduce((result, item) => {
-      result.push(convertMention(item));
-      if (includeColumns) {
-        item.fields.forEach((field) => {
-          result.push(
-            convertMention({
-              ...field,
-              meta: `${item.displayName}.${field.displayName}`,
-            }),
-          );
-        });
-      }
-      return result;
-    }, [] as Mention[]);
-  }, [data?.diagram, includeColumns]);
-
-  return { mentions };
-}
@@ -0,0 +1,92 @@
+import { DATA_SOURCES } from '@/utils/enum';
+import BigQueryProperties from '@/components/pages/setup/dataSources/BigQueryProperties';
+import DuckDBProperties from '@/components/pages/setup/dataSources/DuckDBProperties';
+import MySQLProperties from '@/components/pages/setup/dataSources/MySQLProperties';
+import PostgreSQLProperties from '@/components/pages/setup/dataSources/PostgreSQLProperties';
+import SQLServerProperties from '@/components/pages/setup/dataSources/SQLServerProperties';
+import ClickHouseProperties from '@/components/pages/setup/dataSources/ClickHouseProperties';
+import TrinoProperties from '@/components/pages/setup/dataSources/TrinoProperties';
+import SnowflakeProperties from '@/components/pages/setup/dataSources/SnowflakeProperties';
+
+export const getDataSourceImage = (dataSource: DATA_SOURCES | string) => {
+  switch (dataSource) {
+    case DATA_SOURCES.BIG_QUERY:
+      return '/images/dataSource/bigQuery.svg';
+    case DATA_SOURCES.POSTGRES:
+      return '/images/dataSource/postgreSql.svg';
+    case DATA_SOURCES.MYSQL:
+      return '/images/dataSource/mysql.svg';
+    case DATA_SOURCES.MSSQL:
+      return '/images/dataSource/sqlserver.svg';
+    case DATA_SOURCES.CLICK_HOUSE:
+      return '/images/dataSource/clickhouse.svg';
+    case DATA_SOURCES.DUCKDB:
+      return '/images/dataSource/duckdb.svg';
+    case DATA_SOURCES.TRINO:
+      return '/images/dataSource/trino.svg';
+    case DATA_SOURCES.SNOWFLAKE:
+      return '/images/dataSource/snowflake.svg';
+    default:
+      return null;
+  }
+};
+
+export const getDataSourceName = (dataSource: DATA_SOURCES | string) => {
+  switch (dataSource) {
+    case DATA_SOURCES.BIG_QUERY:
+      return 'BigQuery';
+    case DATA_SOURCES.POSTGRES:
+      return 'PostgreSQL';
+    case DATA_SOURCES.MYSQL:
+      return 'MySQL';
+    case DATA_SOURCES.MSSQL:
+      return 'SQL Server';
+    case DATA_SOURCES.CLICK_HOUSE:
+      return 'ClickHouse';
+    case DATA_SOURCES.DUCKDB:
+      return 'DuckDB';
+    case DATA_SOURCES.TRINO:
+      return 'Trino';
+    case DATA_SOURCES.SNOWFLAKE:
+      return 'Snowflake';
+    default:
+      return '';
+  }
+};
+
+export const getDataSourceProperties = (dataSource: DATA_SOURCES | string) => {
+  switch (dataSource) {
+    case DATA_SOURCES.BIG_QUERY:
+      return BigQueryProperties;
+    case DATA_SOURCES.POSTGRES:
+      return PostgreSQLProperties;
+    case DATA_SOURCES.MYSQL:
+      return MySQLProperties;
+    case DATA_SOURCES.MSSQL:
+      return SQLServerProperties;
+    case DATA_SOURCES.CLICK_HOUSE:
+      return ClickHouseProperties;
+    case DATA_SOURCES.DUCKDB:
+      return DuckDBProperties;
+    case DATA_SOURCES.TRINO:
+      return TrinoProperties;
+    case DATA_SOURCES.SNOWFLAKE:
+      return SnowflakeProperties;
+    default:
+      return null;
+  }
+};
+
+export const getDataSourceConfig = (dataSource: DATA_SOURCES | string) => {
+  return {
+    label: getDataSourceName(dataSource),
+    logo: getDataSourceImage(dataSource),
+    value: DATA_SOURCES[dataSource],
+  };
+};
+
+export const getDataSourceFormComponent = (
+  dataSource: DATA_SOURCES | string,
+) => {
+  return { component: getDataSourceProperties(dataSource) || (() => null) };
+};
@@ -1,7 +1,7 @@
 export enum DATA_SOURCES {
  BIG_QUERY = 'BIG_QUERY',
  DUCKDB = 'DUCKDB',
-  PG_SQL = 'POSTGRES',
+  POSTGRES = 'POSTGRES',
  MYSQL = 'MYSQL',
  MSSQL = 'MSSQL',
  CLICK_HOUSE = 'CLICK_HOUSE',
@@ -133,4 +133,9 @@ export const ERROR_TEXTS = {
      MAX_LENGTH: 'Reasoning steps must be 3000 characters or fewer.',
    },
  },
+  IMPORT_DATA_SOURCE_SQL: {
+    SQL: {
+      REQUIRED: 'Please input SQL statement.',
+    },
+  },
 };