feat(wren-ui): support generate wren sql from dialect sql

CONTRIBUTING.md

@@ -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:

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

@@ -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;
 };

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

@@ -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>;

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

@@ -11,3 +11,9 @@ export const GENERATE_QUESTION = gql`
     generateQuestion(data: $data)
   }
 `;
+
+export const MODEL_SUBSTITUDE = gql`
+  mutation ModelSubstitute($data: ModelSubstituteInput!) {
+    modelSubstitute(data: $data)
+  }
+`;

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

@@ -18,6 +18,9 @@ import {
   toIbisConnectionInfo,
   toMultipleIbisConnectionInfos,
 } from '../dataSource';
+import { DialectSQL, WrenSQL } from '../models/adaptor';
+
+export type { WrenSQL };
 
 const logger = getLogger('IbisAdaptor');
 logger.level = 'debug';
@@ -90,6 +93,7 @@ const dataSourceUrlMap: Record<SupportedDataSource, string> = {
   [SupportedDataSource.CLICK_HOUSE]: 'clickhouse',
   [SupportedDataSource.TRINO]: 'trino',
 };
+
 export interface TableResponse {
   tables: CompactTable[];
 }
@@ -114,11 +118,13 @@ export interface IbisQueryOptions extends IbisBaseOptions {
 export interface IbisDryPlanOptions {
   dataSource: DataSourceName;
   mdl: Manifest;
+  // TODO: replace sql type with WrenSQL
   sql: string;
 }
 
 export interface IIbisAdaptor {
   query: (
+    // TODO: replace query type with WrenSQL
     query: string,
     options: IbisQueryOptions,
   ) => Promise<IbisQueryResponse>;
@@ -140,6 +146,16 @@ export interface IIbisAdaptor {
     mdl: Manifest,
     parameters: Record<string, any>,
   ) => Promise<ValidationResponse>;
+  modelSubstitute: (
+    sql: DialectSQL,
+    options: {
+      dataSource: DataSourceName;
+      connectionInfo: WREN_AI_CONNECTION_INFO;
+      mdl: Manifest;
+      catalog?: string;
+      schema?: string;
+    },
+  ) => Promise<WrenSQL>;
 }
 
 export interface IbisResponse {
@@ -162,6 +178,7 @@ enum IBIS_API_TYPE {
   METADATA = 'METADATA',
   VALIDATION = 'VALIDATION',
   ANALYSIS = 'ANALYSIS',
+  MODEL_SUBSTITUTE = 'MODEL_SUBSTITUTE',
 }
 
 export class IbisAdaptor implements IIbisAdaptor {
@@ -183,11 +200,8 @@ export class IbisAdaptor implements IIbisAdaptor {
       );
       return res.data;
     } catch (e) {
-      logger.debug(`Got error when dry plan with ibis: ${e.response.data}`);
-      throw Errors.create(Errors.GeneralErrorCodes.DRY_PLAN_ERROR, {
-        customMessage: e.response.data,
-        originalError: e,
-      });
+      logger.debug(`Dry plan error: ${e.response?.data || e.message}`);
+      this.throwError(e, 'Error during dry plan execution');
     }
   }
 
@@ -219,19 +233,8 @@ export class IbisAdaptor implements IIbisAdaptor {
         processTime: res.headers['x-process-time'],
       };
     } catch (e) {
-      logger.debug(
-        `Got error when querying ibis: ${e.response?.data || e.message}`,
-      );
-
-      throw Errors.create(Errors.GeneralErrorCodes.IBIS_SERVER_ERROR, {
-        customMessage:
-          e.response?.data || e.message || 'Error querying ibis server',
-        originalError: e,
-        other: {
-          correlationId: e.response?.headers['x-correlation-id'],
-          processTime: e.response?.headers['x-process-time'],
-        },
-      });
+      logger.debug(`Query error: ${e.response?.data || e.message}`);
+      this.throwError(e, 'Error querying ibis server');
     }
   }
 
@@ -259,15 +262,8 @@ export class IbisAdaptor implements IIbisAdaptor {
         processTime: response.headers['x-process-time'],
       };
     } catch (err) {
-      logger.info(`Got error when dry running ibis`);
-      throw Errors.create(Errors.GeneralErrorCodes.DRY_RUN_ERROR, {
-        customMessage: err.response?.data || err.message,
-        originalError: err,
-        other: {
-          correlationId: err.response?.headers['x-correlation-id'],
-          processTime: err.response?.headers['x-process-time'],
-        },
-      });
+      logger.debug(`Dry run error: ${err.response?.data || err.message}`);
+      this.throwError(err, 'Error during dry run execution');
     }
   }
 
@@ -310,16 +306,8 @@ export class IbisAdaptor implements IIbisAdaptor {
       );
       return await getTablesByConnectionInfo(ibisConnectionInfo);
     } catch (e) {
-      logger.debug(
-        `Got error when getting table: ${e.response?.data || e.message}`,
-      );
-      throw Errors.create(Errors.GeneralErrorCodes.IBIS_SERVER_ERROR, {
-        customMessage:
-          e.response?.data ||
-          e.message ||
-          'Error getting table from ibis server',
-        originalError: e,
-      });
+      logger.debug(`Get tables error: ${e.response?.data || e.message}`);
+      this.throwError(e, 'Error getting table from ibis server');
     }
   }
 
@@ -340,17 +328,8 @@ export class IbisAdaptor implements IIbisAdaptor {
       );
       return res.data;
     } catch (e) {
-      logger.debug(
-        `Got error when getting constraint: ${e.response?.data || e.message}`,
-      );
-
-      throw Errors.create(Errors.GeneralErrorCodes.IBIS_SERVER_ERROR, {
-        customMessage:
-          e.response?.data ||
-          e.message ||
-          'Error getting constraint from ibis server',
-        originalError: e,
-      });
+      logger.debug(`Get constraints error: ${e.response?.data || e.message}`);
+      this.throwError(e, 'Error getting constraint from ibis server');
     }
   }
 
@@ -375,12 +354,54 @@ export class IbisAdaptor implements IIbisAdaptor {
         body,
       );
       return { valid: true, message: null };
+    } catch (e) {
+      logger.debug(`Validation error: ${e.response?.data || e.message}`);
+      return { valid: false, message: e.response?.data || e.message };
+    }
+  }
+
+  public async modelSubstitute(
+    sql: DialectSQL,
+    options: {
+      dataSource: DataSourceName;
+      connectionInfo: WREN_AI_CONNECTION_INFO;
+      mdl: Manifest;
+      catalog?: string;
+      schema?: string;
+    },
+  ): Promise<WrenSQL> {
+    const { dataSource, mdl, catalog, schema } = options;
+    let connectionInfo = options.connectionInfo;
+    connectionInfo = this.updateConnectionInfo(connectionInfo);
+    const headers = {
+      'X-User-CATALOG': catalog,
+      'X-User-SCHEMA': schema,
+    };
+    const ibisConnectionInfo = toIbisConnectionInfo(dataSource, connectionInfo);
+    const body = {
+      sql,
+      connectionInfo: ibisConnectionInfo,
+      manifestStr: Buffer.from(JSON.stringify(mdl)).toString('base64'),
+    };
+    try {
+      logger.debug(`Running model substitution with ibis`);
+      const res = await axios.post(
+        `${this.ibisServerEndpoint}/${this.getIbisApiVersion(IBIS_API_TYPE.MODEL_SUBSTITUTE)}/connector/${dataSourceUrlMap[dataSource]}/model-substitute`,
+        body,
+        {
+          headers,
+        },
+      );
+      return res.data as WrenSQL;
     } catch (e) {
       logger.debug(
-        `Got error when validating connection: ${e.response?.data || e.message}`,
+        `Model substitution error: ${e.response?.data || e.message}`,
+      );
+      this.throwError(
+        e,
+        'Error running model substitution with ibis server',
+        this.modelSubstituteErrorMessageBuilder,
       );
-
-      return { valid: false, message: e.response?.data || e.message };
     }
   }
 
@@ -437,8 +458,55 @@ export class IbisAdaptor implements IIbisAdaptor {
       IBIS_API_TYPE.DRY_RUN,
       IBIS_API_TYPE.DRY_PLAN,
       IBIS_API_TYPE.VALIDATION,
+      IBIS_API_TYPE.MODEL_SUBSTITUTE,
     ].includes(apiType);
     if (useV3) logger.debug('Using ibis v3 api');
     return useV3 ? 'v3' : 'v2';
   }
+
+  private throwError(
+    e: any,
+    defaultMessage: string,
+    errorMessageBuilder?: CallableFunction,
+  ) {
+    const customMessage =
+      e.response?.data?.message ||
+      e.response?.data ||
+      e.message ||
+      defaultMessage;
+    throw Errors.create(Errors.GeneralErrorCodes.IBIS_SERVER_ERROR, {
+      customMessage: errorMessageBuilder
+        ? errorMessageBuilder(customMessage)
+        : customMessage,
+      originalError: e,
+      other: {
+        correlationId: e.response?.headers['x-correlation-id'],
+        processTime: e.response?.headers['x-process-time'],
+      },
+    });
+  }
+
+  private modelSubstituteErrorMessageBuilder(message: string) {
+    const ModelSubstituteErrorEnum = {
+      MODEL_NOT_FOUND: () => {
+        return message.includes('Model not found');
+      },
+      PARSING_EXCEPTION: () => {
+        return message.includes('sql.parser.ParsingException');
+      },
+    };
+    if (ModelSubstituteErrorEnum.MODEL_NOT_FOUND()) {
+      const modelName = message.split(': ')[1];
+      return (
+        message +
+        `. Try to add catalog and schema in front of your table. eg: my_database.public.${modelName}`
+      );
+    } else if (ModelSubstituteErrorEnum.PARSING_EXCEPTION()) {
+      return (
+        message +
+        '. Please check your selected column and make sure its quoted for columns with non-alphanumeric characters.'
+      );
+    }
+    return message;
+  }
 }