feat(wren-ui): Support sample data UI (#53)

* feat(wren-ui): Support sample data UI

* feat(wren-ui): Use SampleDatasetName enum instead of

* fixup! feat(wren-ui): Support sample data UI

wren-ui/src/apollo/client/graphql/dataSource.generated.ts

@@ -3,6 +3,13 @@ import * as Types from './__types__';
 import { gql } from '@apollo/client';
 import * as Apollo from '@apollo/client';
 const defaultOptions = {} as const;
+export type StartSampleDatasetMutationVariables = Types.Exact<{
+  data: Types.SampleDatasetInput;
+}>;
+
+
+export type StartSampleDatasetMutation = { __typename?: 'Mutation', startSampleDataset: any };
+
 export type ListDataSourceTablesQueryVariables = Types.Exact<{ [key: string]: never; }>;
 
 
@@ -35,6 +42,37 @@ export type SaveRelationsMutationVariables = Types.Exact<{
 export type SaveRelationsMutation = { __typename?: 'Mutation', saveRelations: any };
 
 
+export const StartSampleDatasetDocument = gql`
+    mutation StartSampleDataset($data: SampleDatasetInput!) {
+  startSampleDataset(data: $data)
+}
+    `;
+export type StartSampleDatasetMutationFn = Apollo.MutationFunction<StartSampleDatasetMutation, StartSampleDatasetMutationVariables>;
+
+/**
+ * __useStartSampleDatasetMutation__
+ *
+ * To run a mutation, you first call `useStartSampleDatasetMutation` within a React component and pass it any options that fit your needs.
+ * When your component renders, `useStartSampleDatasetMutation` 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 [startSampleDatasetMutation, { data, loading, error }] = useStartSampleDatasetMutation({
+ *   variables: {
+ *      data: // value for 'data'
+ *   },
+ * });
+ */
+export function useStartSampleDatasetMutation(baseOptions?: Apollo.MutationHookOptions<StartSampleDatasetMutation, StartSampleDatasetMutationVariables>) {
+        const options = {...defaultOptions, ...baseOptions}
+        return Apollo.useMutation<StartSampleDatasetMutation, StartSampleDatasetMutationVariables>(StartSampleDatasetDocument, options);
+      }
+export type StartSampleDatasetMutationHookResult = ReturnType<typeof useStartSampleDatasetMutation>;
+export type StartSampleDatasetMutationResult = Apollo.MutationResult<StartSampleDatasetMutation>;
+export type StartSampleDatasetMutationOptions = Apollo.BaseMutationOptions<StartSampleDatasetMutation, StartSampleDatasetMutationVariables>;
 export const ListDataSourceTablesDocument = gql`
     query ListDataSourceTables {
   listDataSourceTables {

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

@@ -1,5 +1,11 @@
 import { gql } from '@apollo/client';
 
+export const SET_SAMPLE_DATASETS = gql`
+  mutation StartSampleDataset($data: SampleDatasetInput!) {
+    startSampleDataset(data: $data)
+  }
+`;
+
 export const LIST_DATA_SOURCE_TABLES = gql`
   query ListDataSourceTables {
     listDataSourceTables {

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

@@ -1,10 +1,15 @@
 import Link from 'next/link';
 import Image from 'next/image';
+import { useState } from 'react';
 import { Button, Typography, Row, Col } from 'antd';
+import Icon from '@ant-design/icons';
 import styled from 'styled-components';
 import { ButtonOption, getDataSources, getTemplates } from './utils';
 import { makeIterable, IterableComponent } from '@/utils/iteration';
-import { DataSourceName } from '@/apollo/client/graphql/__types__';
+import {
+  DataSourceName,
+  SampleDatasetName,
+} from '@/apollo/client/graphql/__types__';
 
 const ButtonItem = styled(Button)`
   border: 1px var(--gray-4) solid;
@@ -15,6 +20,11 @@ const ButtonItem = styled(Button)`
   &:disabled {
     opacity: 0.5;
   }
+
+  // loading of button
+  .ant-btn-loading-icon .anticon {
+    font-size: 24px;
+  }
 `;
 
 const PlainImage = styled.div`
@@ -37,15 +47,32 @@ const CommingSoon = styled.div`
 
 const ButtonTemplate = (
   props: IterableComponent<
-    ButtonOption & { value: string; onSelect: (value: string) => void }
+    ButtonOption & {
+      value: string;
+      onSelect: (value: string) => void;
+      selectedTemplate?: string;
+    }
   >,
 ) => {
-  const { value, disabled, logo, label, onSelect } = props;
+  const {
+    value,
+    disabled,
+    submitting,
+    logo,
+    IconComponent,
+    label,
+    onSelect,
+    selectedTemplate,
+  } = props;
+
+  const loading = selectedTemplate === value;
+
   return (
     <Col span={6} key={value}>
       <ButtonItem
-        className="text-left px-4 py-2 bg-gray-2 gray-8 d-flex justify-space-between align-center"
-        disabled={disabled}
+        className={`text-left px-4 py-2 bg-gray-2 gray-8 d-flex align-center ${loading ? 'flex-start' : 'justify-space-between'}`}
+        disabled={disabled || submitting}
+        loading={loading}
         onClick={() => onSelect(value)}
       >
         <div className="d-flex align-center">
@@ -57,6 +84,12 @@ const ButtonTemplate = (
               width="40"
               height="40"
             />
+          ) : IconComponent ? (
+            <Icon
+              component={IconComponent}
+              className="mr-2 p-1"
+              style={{ width: 40, height: 40, fontSize: 32 }}
+            />
           ) : (
             <PlainImage className="mr-2" />
           )}
@@ -72,7 +105,10 @@ const DataSourceIterator = makeIterable(ButtonTemplate);
 const TemplatesIterator = makeIterable(ButtonTemplate);
 
 export default function Starter(props) {
-  const { onNext } = props;
+  const { onNext, submitting } = props;
+
+  const [template, setTemplate] = useState<SampleDatasetName>();
+
   const dataSources = getDataSources();
   const templates = getTemplates();
 
@@ -81,6 +117,7 @@ export default function Starter(props) {
   };
 
   const onSelectTemplate = (value: string) => {
+    setTemplate(value as SampleDatasetName);
     onNext && onNext({ template: value });
   };
 
@@ -101,7 +138,11 @@ export default function Starter(props) {
         .
       </Typography.Text>
       <Row className="mt-6" gutter={[16, 16]}>
-        <DataSourceIterator data={dataSources} onSelect={onSelectDataSource} />
+        <DataSourceIterator
+          data={dataSources}
+          onSelect={onSelectDataSource}
+          submitting={submitting}
+        />
       </Row>
 
       <div className="py-8" />
@@ -110,7 +151,12 @@ export default function Starter(props) {
         Play around with sample data
       </Typography.Title>
       <Row className="mt-6" gutter={[16, 16]}>
-        <TemplatesIterator data={templates} onSelect={onSelectTemplate} />
+        <TemplatesIterator
+          data={templates}
+          onSelect={onSelectTemplate}
+          submitting={submitting}
+          selectedTemplate={template}
+        />
       </Row>
 
       <div className="py-12" />

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

@@ -1,11 +1,16 @@
 import { merge } from 'lodash';
+import IconComponentProps from '@ant-design/icons';
+import CustomerServiceOutlined from '@ant-design/icons/CustomerServiceOutlined';
+import ShoppingCartOutlined from '@ant-design/icons/ShoppingCartOutlined';
+import TrophyOutlined from '@ant-design/icons/TrophyOutlined';
+import { SETUP, DATA_SOURCES } from '@/utils/enum';
 import Starter from './Starter';
 import ConnectDataSource from './ConnectDataSource';
 import SelectModels from './SelectModels';
 import DefineRelations from './DefineRelations';
-import { SETUP, DATA_SOURCES, DEMO_TEMPLATES } from '@/utils/enum';
 import BigQueryProperties from './dataSources/BigQueryProperties';
 import DuckDBProperties from './dataSources/DuckDBProperties';
+import { SampleDatasetName } from '@/apollo/client/graphql/__types__';
 
 type SetupStep = {
   step: number;
@@ -21,8 +26,10 @@ type SetupStep = {
 export type ButtonOption = {
   label: string;
   logo: string;
+  IconComponent?: typeof IconComponentProps;
   guide: string;
   disabled: boolean;
+  submitting?: boolean;
 };
 
 export const SETUP_STEPS = {
@@ -73,13 +80,17 @@ export const DATA_SOURCE_FORM = {
 };
 
 export const TEMPLATE_OPTIONS = {
-  [DEMO_TEMPLATES.CRM]: {
-    label: 'CRM',
-    logo: '',
-  },
-  [DEMO_TEMPLATES.ECORMERCE]: {
+  [SampleDatasetName.ECOMMERCE]: {
     label: 'E-commerce',
-    logo: '',
+    IconComponent: ShoppingCartOutlined,
+  },
+  [SampleDatasetName.MUSIC]: {
+    label: 'Music Store',
+    IconComponent: CustomerServiceOutlined,
+  },
+  [SampleDatasetName.NBA]: {
+    label: 'NBA',
+    IconComponent: TrophyOutlined,
   },
 };
 
@@ -111,6 +122,6 @@ export const getTemplates = () => {
     ...TEMPLATE_OPTIONS[key],
     value: key,
   })) as (Omit<ButtonOption, 'guide' | 'disabled'> & {
-    value: DEMO_TEMPLATES;
+    value: SampleDatasetName;
   })[];
 };

wren-ui/src/hooks/useSetupConnection.tsx

@@ -1,8 +1,14 @@
 import { useState, useEffect } from 'react';
-import { SETUP } from '@/utils/enum';
+import { Path, SETUP } from '@/utils/enum';
 import { useRouter } from 'next/router';
-import { useSaveDataSourceMutation } from '@/apollo/client/graphql/dataSource.generated';
-import { DataSourceName } from '@/apollo/client/graphql/__types__';
+import {
+  useSaveDataSourceMutation,
+  useStartSampleDatasetMutation,
+} from '@/apollo/client/graphql/dataSource.generated';
+import {
+  DataSourceName,
+  SampleDatasetName,
+} from '@/apollo/client/graphql/__types__';
 
 const transformProperties = (
   properties: Record<string, any>,
@@ -42,7 +48,13 @@ export default function useSetupConnection() {
   const [saveDataSourceMutation, { loading, error }] =
     useSaveDataSourceMutation({
       onError: (error) => console.error(error),
-      onCompleted: () => router.push('/setup/models'),
+      onCompleted: () => router.push(Path.OnboardingModels),
+    });
+
+  const [startSampleDatasetMutation, { loading: startSampleDatasetLoading }] =
+    useStartSampleDatasetMutation({
+      onError: (error) => console.error(error),
+      onCompleted: () => router.push(Path.Home),
     });
 
   useEffect(() => {
@@ -60,6 +72,14 @@ export default function useSetupConnection() {
     });
   };
 
+  const submitTemplate = async (template: SampleDatasetName) => {
+    await startSampleDatasetMutation({
+      variables: {
+        data: { name: template },
+      },
+    });
+  };
+
   const onBack = () => {
     if (stepKey === SETUP.CREATE_DATA_SOURCE) {
       setStepKey(SETUP.STARTER);
@@ -68,6 +88,7 @@ export default function useSetupConnection() {
 
   const onNext = (data?: {
     dataSource?: DataSourceName;
+    template?: SampleDatasetName;
     properties?: JSON;
   }) => {
     if (stepKey === SETUP.STARTER) {
@@ -75,7 +96,7 @@ export default function useSetupConnection() {
         setDataSource(data?.dataSource);
         setStepKey(SETUP.CREATE_DATA_SOURCE);
       } else {
-        // TODO: implement template chosen
+        submitTemplate(data.template);
       }
     } else if (stepKey === SETUP.CREATE_DATA_SOURCE) {
       submitDataSource(data.properties);
@@ -87,7 +108,7 @@ export default function useSetupConnection() {
     dataSource,
     onBack,
     onNext,
-    submitting: loading,
+    submitting: loading || startSampleDatasetLoading,
     connectErrorMessage,
   };
 }

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

@@ -5,8 +5,3 @@ export enum SETUP {
   RECOMMEND_RELATIONS = 'recommendRelations',
   DEFINE_RELATIONS = 'defineRelations',
 }
-
-export enum DEMO_TEMPLATES {
-  CRM = 'CRM',
-  ECORMERCE = 'Ecommerce',
-}