feat(CE): update model query types

Multiwoven · Dec 16, 2024 · 607f6eb · 607f6eb
1 parent 6f6a63d
commit 607f6eb
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 7 deletions.
diff --git a/ui/src/components/ModelTable/ModelTable.tsx b/ui/src/components/ModelTable/ModelTable.tsx
@@ -1,5 +1,5 @@
 import GenerateTable from '@/components/Table/Table';
-import { getAllModels, GetAllModelsResponse } from '@/services/models';
+import { AllDataModels, getAllModels, GetAllModelsResponse } from '@/services/models';
 import { addIconDataToArray, ConvertToTableData } from '@/utils';
 import NoModels from '@/views/Models/NoModels';
 import Loader from '@/components/Loader';
@@ -16,7 +16,7 @@ const ModelTable = ({ handleOnRowClick }: ModelTableProps): JSX.Element => {
 
   const { data } = useQueryWrapper<ApiResponse<GetAllModelsResponse[]>, Error>(
     ['models', activeWorkspaceId],
-    () => getAllModels({ type: 'data' }),
+    () => getAllModels({ type: AllDataModels }),
     {
       refetchOnMount: true,
       refetchOnWindowFocus: false,

diff --git a/ui/src/services/models.ts b/ui/src/services/models.ts
@@ -47,10 +47,10 @@ export type GetAllModelsResponse = {
   attributes: ModelAttributes;
 };
 
-export type ModelQueryType = 'data' | 'ai_ml' | 'raw_sql' | 'dbt' | 'soql' | 'table_selector';
+export const AllDataModels = 'raw_sql,dbt,soql,table_selector';
 
 export type GetAllModelsProps = {
-  type: ModelQueryType;
+  type: string;
 };
 
 export const getModelPreview = async (query: string, connector_id: string): Promise<any> => {
@@ -59,7 +59,7 @@ export const getModelPreview = async (query: string, connector_id: string): Prom
 };
 
 export const getAllModels = async ({
-  type = 'data',
+  type = AllDataModels,
 }: GetAllModelsProps): Promise<ApiResponse<GetAllModelsResponse[]>> =>
   multiwovenFetch<null, ApiResponse<GetAllModelsResponse[]>>({
     method: 'get',

diff --git a/ui/src/views/Models/ModelsList/ModelsList.tsx b/ui/src/views/Models/ModelsList/ModelsList.tsx
@@ -3,7 +3,7 @@ import TopBar from '@/components/TopBar';
 import { Box } from '@chakra-ui/react';
 import { FiPlus } from 'react-icons/fi';
 import { useQuery } from '@tanstack/react-query';
-import { getAllModels, GetAllModelsResponse } from '@/services/models';
+import { AllDataModels, getAllModels, GetAllModelsResponse } from '@/services/models';
 import Loader from '@/components/Loader';
 import NoModels from '@/views/Models/NoModels';
 import { useStore } from '@/stores';
@@ -22,7 +22,7 @@ const ModelsList = (): JSX.Element | null => {
 
   const { data, isLoading } = useQuery({
     queryKey: ['models', activeWorkspaceId, 'data'],
-    queryFn: () => getAllModels({ type: 'data' }),
+    queryFn: () => getAllModels({ type: AllDataModels }),
     refetchOnMount: true,
     refetchOnWindowFocus: false,
     enabled: activeWorkspaceId > 0,