diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 8a032e96ab..6642e5ad2e 100644 --- a/CONTRIBUTING.md +++ b/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: diff --git a/wren-ui/src/apollo/client/graphql/__types__.ts b/wren-ui/src/apollo/client/graphql/__types__.ts index 18454a756e..388dfa0aff 100644 --- a/wren-ui/src/apollo/client/graphql/__types__.ts +++ b/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; }; diff --git a/wren-ui/src/apollo/client/graphql/sql.generated.ts b/wren-ui/src/apollo/client/graphql/sql.generated.ts index 24fe293d0b..4228cdb601 100644 --- a/wren-ui/src/apollo/client/graphql/sql.generated.ts +++ b/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; export type GenerateQuestionMutationResult = Apollo.MutationResult; -export type GenerateQuestionMutationOptions = Apollo.BaseMutationOptions; \ No newline at end of file +export type GenerateQuestionMutationOptions = Apollo.BaseMutationOptions; +export const ModelSubstituteDocument = gql` + mutation ModelSubstitute($data: ModelSubstituteInput!) { + modelSubstitute(data: $data) +} + `; +export type ModelSubstituteMutationFn = Apollo.MutationFunction; + +/** + * __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) { + const options = {...defaultOptions, ...baseOptions} + return Apollo.useMutation(ModelSubstituteDocument, options); + } +export type ModelSubstituteMutationHookResult = ReturnType; +export type ModelSubstituteMutationResult = Apollo.MutationResult; +export type ModelSubstituteMutationOptions = Apollo.BaseMutationOptions; \ No newline at end of file diff --git a/wren-ui/src/apollo/client/graphql/sql.ts b/wren-ui/src/apollo/client/graphql/sql.ts index 7843e77eaf..e09219cabb 100644 --- a/wren-ui/src/apollo/client/graphql/sql.ts +++ b/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) + } +`; diff --git a/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts b/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts index a36af9d8ea..eb18765c85 100644 --- a/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts +++ b/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.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; @@ -140,6 +146,16 @@ export interface IIbisAdaptor { mdl: Manifest, parameters: Record, ) => Promise; + modelSubstitute: ( + sql: DialectSQL, + options: { + dataSource: DataSourceName; + connectionInfo: WREN_AI_CONNECTION_INFO; + mdl: Manifest; + catalog?: string; + schema?: string; + }, + ) => Promise; } 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 { + 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; + } } diff --git a/wren-ui/src/apollo/server/adaptors/tests/ibisAdaptor.test.ts b/wren-ui/src/apollo/server/adaptors/tests/ibisAdaptor.test.ts index 5f3d30a6b2..6acb5ddf50 100644 --- a/wren-ui/src/apollo/server/adaptors/tests/ibisAdaptor.test.ts +++ b/wren-ui/src/apollo/server/adaptors/tests/ibisAdaptor.test.ts @@ -8,6 +8,7 @@ import { } from '../ibisAdaptor'; import { DataSourceName } from '../../types'; import { Manifest } from '../../mdl/type'; +import { DialectSQL } from '../../models/adaptor'; import { BIG_QUERY_CONNECTION_INFO, CLICK_HOUSE_CONNECTION_INFO, @@ -257,16 +258,18 @@ describe('IbisAdaptor', () => { mockTrinoConnectionInfo, ); - const { username, host, password, port, schemas, ssl } = - mockTrinoConnectionInfo; + const { username, host, password, port, schemas } = mockTrinoConnectionInfo; const schemasArray = schemas.split(','); const [catalog, schema] = schemasArray[0].split('.'); const expectConnectionInfo = { - connectionUrl: `trino://${username}:${password}@${host}:${port}/${catalog}/${schema}`, + catalog, + host: `https://${host}`, + password, + port, + schema, + user: username, }; - if (ssl) expectConnectionInfo.connectionUrl += '&SSL=true'; - expect(result).toEqual([]); expect(mockedAxios.post).toHaveBeenCalledWith( `${ibisServerEndpoint}/v2/connector/trino/metadata/constraints`, @@ -395,7 +398,7 @@ describe('IbisAdaptor', () => { expect(result).toEqual({ valid: true, message: null }); expect(mockedAxios.post).toHaveBeenCalledWith( - `${ibisServerEndpoint}/v2/connector/postgres/validate/column_is_valid`, + `${ibisServerEndpoint}/v3/connector/postgres/validate/column_is_valid`, { connectionInfo: { connectionUrl: postgresConnectionUrl }, manifestStr: Buffer.from(JSON.stringify(mockManifest)).toString( @@ -427,7 +430,7 @@ describe('IbisAdaptor', () => { expect(result).toEqual({ valid: false, message: 'Error' }); expect(mockedAxios.post).toHaveBeenCalledWith( - `${ibisServerEndpoint}/v2/connector/postgres/validate/column_is_valid`, + `${ibisServerEndpoint}/v3/connector/postgres/validate/column_is_valid`, { connectionInfo: { connectionUrl: postgresConnectionUrl }, manifestStr: Buffer.from(JSON.stringify(mockManifest)).toString( @@ -470,7 +473,7 @@ describe('IbisAdaptor', () => { }); it('should throw an exception with correlationId and processTime when query fails', async () => { - mockedAxios.post.mockRejectedValue({ + const mockError = { response: { data: 'Error message', headers: { @@ -478,7 +481,8 @@ describe('IbisAdaptor', () => { 'x-process-time': '1s', }, }, - }); + }; + mockedAxios.post.mockRejectedValue(mockError); mockedEncryptor.prototype.decrypt.mockReturnValue( JSON.stringify({ password: mockPostgresConnectionInfo.password }), ); @@ -526,7 +530,7 @@ describe('IbisAdaptor', () => { }); it('should throw an exception with correlationId and processTime when dry run fails', async () => { - mockedAxios.post.mockRejectedValue({ + const mockError = { response: { data: 'Error message', headers: { @@ -534,7 +538,8 @@ describe('IbisAdaptor', () => { 'x-process-time': '1s', }, }, - }); + }; + mockedAxios.post.mockRejectedValue(mockError); mockedEncryptor.prototype.decrypt.mockReturnValue( JSON.stringify({ password: mockPostgresConnectionInfo.password }), ); @@ -555,4 +560,177 @@ describe('IbisAdaptor', () => { }, }); }); + + it('should successfully substitute SQL with model', async () => { + const mockResponse = { data: 'SELECT * FROM substituted_table' }; + mockedAxios.post.mockResolvedValue(mockResponse); + mockedEncryptor.prototype.decrypt.mockReturnValue( + JSON.stringify({ password: mockPostgresConnectionInfo.password }), + ); + + const result = await ibisAdaptor.modelSubstitute( + 'SELECT * FROM test_table' as DialectSQL, + { + dataSource: DataSourceName.POSTGRES, + connectionInfo: mockPostgresConnectionInfo, + mdl: mockManifest, + }, + ); + + expect(result).toEqual('SELECT * FROM substituted_table'); + expect(mockedAxios.post).toHaveBeenCalledWith( + `${ibisServerEndpoint}/v3/connector/postgres/model-substitute`, + { + sql: 'SELECT * FROM test_table', + connectionInfo: { connectionUrl: postgresConnectionUrl }, + manifestStr: Buffer.from(JSON.stringify(mockManifest)).toString( + 'base64', + ), + }, + { + headers: { + 'X-User-CATALOG': undefined, + 'X-User-SCHEMA': undefined, + }, + }, + ); + }); + + it('should handle error when model substitution fails with MODEL_NOT_FOUND', async () => { + const mockError = { + response: { + data: 'Model not found: test_table', + headers: { + 'x-correlation-id': '123', + 'x-process-time': '1s', + }, + }, + }; + mockedAxios.post.mockRejectedValue(mockError); + mockedEncryptor.prototype.decrypt.mockReturnValue( + JSON.stringify({ password: mockPostgresConnectionInfo.password }), + ); + + await expect( + ibisAdaptor.modelSubstitute('SELECT * FROM test_table' as DialectSQL, { + dataSource: DataSourceName.POSTGRES, + connectionInfo: mockPostgresConnectionInfo, + mdl: mockManifest, + }), + ).rejects.toMatchObject({ + message: + 'Model not found: test_table. Try to add catalog and schema in front of your table. eg: my_database.public.test_table', + extensions: { + other: { + correlationId: '123', + processTime: '1s', + }, + }, + }); + }); + + it('should handle error when model substitution fails with PARSING_EXCEPTION', async () => { + const mockError = { + response: { + data: 'sql.parser.ParsingException: Invalid SQL syntax', + headers: { + 'x-correlation-id': '123', + 'x-process-time': '1s', + }, + }, + }; + mockedAxios.post.mockRejectedValue(mockError); + mockedEncryptor.prototype.decrypt.mockReturnValue( + JSON.stringify({ password: mockPostgresConnectionInfo.password }), + ); + + await expect( + ibisAdaptor.modelSubstitute('SELECT * FROM test_table' as DialectSQL, { + dataSource: DataSourceName.POSTGRES, + connectionInfo: mockPostgresConnectionInfo, + mdl: mockManifest, + }), + ).rejects.toMatchObject({ + message: + 'sql.parser.ParsingException: Invalid SQL syntax. Please check your selected column and make sure its quoted for columns with non-alphanumeric characters.', + extensions: { + other: { + correlationId: '123', + processTime: '1s', + }, + }, + }); + }); + + it('should handle error when model substitution fails with generic error', async () => { + const mockError = { + response: { + data: 'Generic error occurred', + headers: { + 'x-correlation-id': '123', + 'x-process-time': '1s', + }, + }, + }; + mockedAxios.post.mockRejectedValue(mockError); + mockedEncryptor.prototype.decrypt.mockReturnValue( + JSON.stringify({ password: mockPostgresConnectionInfo.password }), + ); + + await expect( + ibisAdaptor.modelSubstitute('SELECT * FROM test_table' as DialectSQL, { + dataSource: DataSourceName.POSTGRES, + connectionInfo: mockPostgresConnectionInfo, + mdl: mockManifest, + }), + ).rejects.toMatchObject({ + message: 'Generic error occurred', + extensions: { + other: { + correlationId: '123', + processTime: '1s', + }, + }, + }); + }); + + it('should include catalog and schema in headers when provided', async () => { + const mockResponse = { data: 'SELECT * FROM substituted_table' }; + mockedAxios.post.mockResolvedValue(mockResponse); + mockedEncryptor.prototype.decrypt.mockReturnValue( + JSON.stringify({ password: mockPostgresConnectionInfo.password }), + ); + + const catalog = 'my_catalog'; + const schema = 'my_schema'; + + const result = await ibisAdaptor.modelSubstitute( + 'SELECT * FROM test_table' as DialectSQL, + { + dataSource: DataSourceName.POSTGRES, + connectionInfo: mockPostgresConnectionInfo, + mdl: mockManifest, + catalog, + schema, + }, + ); + + expect(result).toEqual('SELECT * FROM substituted_table'); + expect(mockedAxios.post).toHaveBeenCalledWith( + `${ibisServerEndpoint}/v3/connector/postgres/model-substitute`, + { + sql: 'SELECT * FROM test_table', + connectionInfo: { connectionUrl: postgresConnectionUrl }, + manifestStr: Buffer.from(JSON.stringify(mockManifest)).toString( + 'base64', + ), + }, + { + headers: { + 'X-User-CATALOG': catalog, + 'X-User-SCHEMA': schema, + }, + }, + ); + }); }); diff --git a/wren-ui/src/apollo/server/models/adaptor.ts b/wren-ui/src/apollo/server/models/adaptor.ts index 0af534fb53..6305e2c98b 100644 --- a/wren-ui/src/apollo/server/models/adaptor.ts +++ b/wren-ui/src/apollo/server/models/adaptor.ts @@ -2,6 +2,12 @@ import * as Errors from '@server/utils/error'; import { Manifest } from '@server/mdl/type'; import { ThreadResponse } from '../repositories'; +// Add branded types for SQL strings +type Brand = T & { __brand: B }; + +export type DialectSQL = Brand; +export type WrenSQL = Brand; + export interface WrenAIError { code: Errors.GeneralErrorCodes; message: string; diff --git a/wren-ui/src/apollo/server/resolvers.ts b/wren-ui/src/apollo/server/resolvers.ts index fc01d77279..039469e39f 100644 --- a/wren-ui/src/apollo/server/resolvers.ts +++ b/wren-ui/src/apollo/server/resolvers.ts @@ -8,6 +8,7 @@ import { DashboardResolver } from './resolvers/dashboardResolver'; import { SqlPairResolver } from './resolvers/sqlPairResolver'; import { InstructionResolver } from './resolvers/instructionResolver'; import { convertColumnType } from '@server/utils'; +import { DialectSQLScalar } from './scalars'; const projectResolver = new ProjectResolver(); const modelResolver = new ModelResolver(); @@ -19,6 +20,7 @@ const sqlPairResolver = new SqlPairResolver(); const instructionResolver = new InstructionResolver(); const resolvers = { JSON: GraphQLJSON, + DialectSQL: DialectSQLScalar, Query: { listDataSourceTables: projectResolver.listDataSourceTables, autoGenerateRelation: projectResolver.autoGenerateRelation, @@ -162,6 +164,7 @@ const resolvers = { updateSqlPair: sqlPairResolver.updateSqlPair, deleteSqlPair: sqlPairResolver.deleteSqlPair, generateQuestion: sqlPairResolver.generateQuestion, + modelSubstitute: sqlPairResolver.modelSubstitute, // Instructions createInstruction: instructionResolver.createInstruction, updateInstruction: instructionResolver.updateInstruction, diff --git a/wren-ui/src/apollo/server/resolvers/sqlPairResolver.ts b/wren-ui/src/apollo/server/resolvers/sqlPairResolver.ts index e5a78a7d5a..d9bb4e9e24 100644 --- a/wren-ui/src/apollo/server/resolvers/sqlPairResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/sqlPairResolver.ts @@ -2,6 +2,9 @@ import { IContext } from '@server/types/context'; import { SqlPair } from '@server/repositories'; import * as Errors from '@server/utils/error'; import { TelemetryEvent, TrackTelemetry } from '@server/telemetry/telemetry'; +import { DialectSQL, WrenSQL } from '@server/models/adaptor'; +import { format } from 'sql-formatter'; + export class SqlPairResolver { constructor() { this.getProjectSqlPairs = this.getProjectSqlPairs.bind(this); @@ -9,6 +12,7 @@ export class SqlPairResolver { this.updateSqlPair = this.updateSqlPair.bind(this); this.deleteSqlPair = this.deleteSqlPair.bind(this); this.generateQuestion = this.generateQuestion.bind(this); + this.modelSubstitute = this.modelSubstitute.bind(this); } public async getProjectSqlPairs( @@ -85,6 +89,31 @@ export class SqlPairResolver { return questions[0]; } + public async modelSubstitute( + _root: unknown, + arg: { + data: { + sql: DialectSQL; + }; + }, + ctx: IContext, + ): Promise { + const project = await ctx.projectService.getCurrentProject(); + const lastDeployment = await ctx.deployService.getLastDeployment( + project.id, + ); + const manifest = lastDeployment.manifest; + + const wrenSQL = await ctx.sqlPairService.modelSubstitute( + arg.data.sql as DialectSQL, + { + project, + manifest, + }, + ); + return format(wrenSQL, { language: 'postgresql' }) as WrenSQL; + } + private async validateSql(sql: string, ctx: IContext) { const project = await ctx.projectService.getCurrentProject(); const lastDeployment = await ctx.deployService.getLastDeployment( diff --git a/wren-ui/src/apollo/server/scalars.ts b/wren-ui/src/apollo/server/scalars.ts new file mode 100644 index 0000000000..d137671730 --- /dev/null +++ b/wren-ui/src/apollo/server/scalars.ts @@ -0,0 +1,25 @@ +import { GraphQLScalarType } from 'graphql'; +import { DialectSQL } from '@server/models/adaptor'; + +export const DialectSQLScalar = new GraphQLScalarType({ + name: 'DialectSQL', + description: 'A string representing a SQL query in a specific dialect', + serialize(value: unknown): string { + if (typeof value !== 'string') { + throw new Error('DialectSQL must be a string'); + } + return value; + }, + parseValue(value: unknown): DialectSQL { + if (typeof value !== 'string') { + throw new Error('DialectSQL must be a string'); + } + return value as DialectSQL; + }, + parseLiteral(ast: any): DialectSQL { + if (ast.kind !== 'StringValue') { + throw new Error('DialectSQL must be a string'); + } + return ast.value as DialectSQL; + }, +}); diff --git a/wren-ui/src/apollo/server/schema.ts b/wren-ui/src/apollo/server/schema.ts index 2a7835d4f2..83984e2a06 100644 --- a/wren-ui/src/apollo/server/schema.ts +++ b/wren-ui/src/apollo/server/schema.ts @@ -2,6 +2,7 @@ import { gql } from 'apollo-server-micro'; export const typeDefs = gql` scalar JSON + scalar DialectSQL enum DataSourceName { BIG_QUERY @@ -929,6 +930,10 @@ export const typeDefs = gql` sql: String! } + input ModelSubstituteInput { + sql: DialectSQL! + } + type Instruction { id: Int! projectId: Int! @@ -1137,6 +1142,7 @@ export const typeDefs = gql` ): SqlPair! deleteSqlPair(where: SqlPairWhereUniqueInput!): Boolean! generateQuestion(data: GenerateQuestionInput!): String! + modelSubstitute(data: ModelSubstituteInput!): String! # Instructions createInstruction(data: CreateInstructionInput!): Instruction! updateInstruction( diff --git a/wren-ui/src/apollo/server/services/modelService.ts b/wren-ui/src/apollo/server/services/modelService.ts index 5bb50403a8..5299ade75c 100644 --- a/wren-ui/src/apollo/server/services/modelService.ts +++ b/wren-ui/src/apollo/server/services/modelService.ts @@ -10,7 +10,7 @@ import { } from '@server/repositories'; import { getLogger, - parseJson, + safeParseJson, replaceAllowableSyntax, validateDisplayName, } from '@server/utils'; @@ -144,7 +144,7 @@ export class ModelService implements IModelService { } as CheckCalculatedFieldCanQueryData); logger.debug(`${logTitle} : checkCalculatedFieldCanQuery: ${canQuery}`); if (!canQuery) { - const parsedErrorMessage = parseJson(errorMessage); + const parsedErrorMessage = safeParseJson(errorMessage); throw Errors.create(Errors.GeneralErrorCodes.INVALID_CALCULATED_FIELD, { customMessage: parsedErrorMessage?.message || errorMessage, originalError: parsedErrorMessage || null, diff --git a/wren-ui/src/apollo/server/services/sqlPairService.ts b/wren-ui/src/apollo/server/services/sqlPairService.ts index 0ee60de7b3..fd9c8b186d 100644 --- a/wren-ui/src/apollo/server/services/sqlPairService.ts +++ b/wren-ui/src/apollo/server/services/sqlPairService.ts @@ -1,16 +1,22 @@ -import { IWrenAIAdaptor } from '../adaptors'; -import { - QuestionsResult, - QuestionsStatus, - SqlPairResult, - SqlPairStatus, - WrenAILanguage, -} from '../models/adaptor'; -import { ISqlPairRepository, SqlPair } from '../repositories/sqlPairRepository'; +import { SqlPair } from '@server/repositories'; +import { IWrenAIAdaptor } from '@server/adaptors/wrenAIAdaptor'; +import { ISqlPairRepository } from '@server/repositories/sqlPairRepository'; import { getLogger } from '@server/utils'; import { chunk } from 'lodash'; import * as Errors from '@server/utils/error'; import { Project } from '../repositories'; +import { IIbisAdaptor } from '../adaptors/ibisAdaptor'; +import { + DialectSQL, + WrenSQL, + WrenAILanguage, + SqlPairResult, + SqlPairStatus, + QuestionsResult, + QuestionsStatus, +} from '../models/adaptor'; +import { Manifest } from '@server/mdl/type'; +import { DataSourceName } from '@server/types'; const logger = getLogger('SqlPairService'); @@ -24,6 +30,12 @@ export interface EditSqlPair { question?: string; } +export interface ModelSubstituteOptions { + project: Project; + // if not given, will use the deployed manifest + manifest: Manifest; +} + export interface ISqlPairService { getProjectSqlPairs(projectId: number): Promise; createSqlPair(projectId: number, sqlPair: CreateSqlPair): Promise; @@ -38,21 +50,54 @@ export interface ISqlPairService { ): Promise; deleteSqlPair(projectId: number, sqlPairId: number): Promise; generateQuestions(project: Project, sqls: string[]): Promise; + modelSubstitute( + sql: DialectSQL, + options: ModelSubstituteOptions, + ): Promise; } export class SqlPairService implements ISqlPairService { private sqlPairRepository: ISqlPairRepository; private wrenAIAdaptor: IWrenAIAdaptor; + private ibisAdaptor: IIbisAdaptor; constructor({ sqlPairRepository, wrenAIAdaptor, + ibisAdaptor, }: { sqlPairRepository: ISqlPairRepository; wrenAIAdaptor: IWrenAIAdaptor; + ibisAdaptor: IIbisAdaptor; }) { this.sqlPairRepository = sqlPairRepository; this.wrenAIAdaptor = wrenAIAdaptor; + this.ibisAdaptor = ibisAdaptor; + } + + public async modelSubstitute( + sql: DialectSQL, + options: ModelSubstituteOptions, + ): Promise { + const { manifest: mdl, project } = options; + const { type: dataSource, connectionInfo } = project; + if (dataSource === DataSourceName.DUCKDB) { + // engine does not implement model substitute. + throw Errors.create(Errors.GeneralErrorCodes.IBIS_SERVER_ERROR, { + customMessage: 'DuckDB data source does not support model substitute.', + }); + } + // use the first model's table reference as default catalog and schema + const firstModel = mdl.models?.[0]; + const catalog = firstModel?.tableReference?.catalog; + const schema = firstModel?.tableReference?.schema; + return await this.ibisAdaptor.modelSubstitute(sql, { + dataSource, + connectionInfo, + mdl, + catalog, + schema, + }); } public async generateQuestions( diff --git a/wren-ui/src/apollo/server/utils/helper.ts b/wren-ui/src/apollo/server/utils/helper.ts index e8671fae31..e5b353b3a0 100644 --- a/wren-ui/src/apollo/server/utils/helper.ts +++ b/wren-ui/src/apollo/server/utils/helper.ts @@ -2,7 +2,7 @@ * @function * @description Retrieve json without error */ -export const parseJson = (data) => { +export const safeParseJson = (data) => { try { return JSON.parse(data); } catch (_e) { @@ -10,6 +10,17 @@ export const parseJson = (data) => { } }; +export const safeStringify = (data) => { + if (typeof data === 'string') { + return data; + } + try { + return JSON.stringify(data); + } catch (_e) { + return data; + } +}; + export const convertColumnType = (parent: { type: string }) => { return parent.type.includes('STRUCT') ? 'RECORD' : parent.type; }; diff --git a/wren-ui/src/common.ts b/wren-ui/src/common.ts index 8337b2b6c2..e06dd48783 100644 --- a/wren-ui/src/common.ts +++ b/wren-ui/src/common.ts @@ -138,6 +138,7 @@ export const initComponents = () => { const sqlPairService = new SqlPairService({ sqlPairRepository, wrenAIAdaptor, + ibisAdaptor, }); const instructionService = new InstructionService({ instructionRepository, diff --git a/wren-ui/src/components/editor/MarkdownEditor.tsx b/wren-ui/src/components/editor/MarkdownEditor.tsx index 8e6ee4e0c4..f5f42b5de6 100644 --- a/wren-ui/src/components/editor/MarkdownEditor.tsx +++ b/wren-ui/src/components/editor/MarkdownEditor.tsx @@ -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'; diff --git a/wren-ui/src/components/editor/SQLEditor.tsx b/wren-ui/src/components/editor/SQLEditor.tsx index 9cb629b0a7..4341ce4670 100644 --- a/wren-ui/src/components/editor/SQLEditor.tsx +++ b/wren-ui/src/components/editor/SQLEditor.tsx @@ -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(null); + const [focused, setFocused] = useState(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 ( - + + {!!toolbar && {toolbar}} + setFocused(true)} + onBlur={() => setFocused(false)} + name="sql_editor" + editorProps={{ $blockScrolling: true }} + enableLiveAutocompletion + enableBasicAutocompletion + showPrintMargin={false} + focus={autoFocus} + /> + ); } diff --git a/wren-ui/src/components/modals/AdjustReasoningStepsModal.tsx b/wren-ui/src/components/modals/AdjustReasoningStepsModal.tsx index 9f2fc65884..fea6d5a8ef 100644 --- a/wren-ui/src/components/modals/AdjustReasoningStepsModal.tsx +++ b/wren-ui/src/components/modals/AdjustReasoningStepsModal.tsx @@ -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, diff --git a/wren-ui/src/components/modals/AdjustSQLModal.tsx b/wren-ui/src/components/modals/AdjustSQLModal.tsx index 905152ac28..d0a8e77d1a 100644 --- a/wren-ui/src/components/modals/AdjustSQLModal.tsx +++ b/wren-ui/src/components/modals/AdjustSQLModal.tsx @@ -178,7 +178,7 @@ export default function AdjustSQLModal(props: Props) { }, ]} > - +
diff --git a/wren-ui/src/components/modals/FixSQLModal.tsx b/wren-ui/src/components/modals/FixSQLModal.tsx index c541880221..17a89a9579 100644 --- a/wren-ui/src/components/modals/FixSQLModal.tsx +++ b/wren-ui/src/components/modals/FixSQLModal.tsx @@ -98,7 +98,7 @@ export function FixSQLModal(props: Props) { }, ]} > - +
diff --git a/wren-ui/src/components/modals/ImportDataSourceSQLModal.tsx b/wren-ui/src/components/modals/ImportDataSourceSQLModal.tsx new file mode 100644 index 0000000000..bb1ef85015 --- /dev/null +++ b/wren-ui/src/components/modals/ImportDataSourceSQLModal.tsx @@ -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 ( + <> + + logo + {name} + + + ); +}; + +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 ( + reset()} + > +
+ + } + autoFocus + /> + +
+ {!!error && ( + } + /> + )} + + ); +} diff --git a/wren-ui/src/components/modals/QuestionSQLPairModal.tsx b/wren-ui/src/components/modals/QuestionSQLPairModal.tsx index 9fd5e5f884..9b146d2fb4 100644 --- a/wren-ui/src/components/modals/QuestionSQLPairModal.tsx +++ b/wren-ui/src/components/modals/QuestionSQLPairModal.tsx @@ -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 ( +
+ + + Wren SQL + + +
+ ); +}; + 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 ( - handleReset()} - footer={ -
-
- - + handleReset()} + footer={ +
+
- The SQL statement used here follows Wren SQL, which is - based on ANSI SQL and optimized for Wren AI.{` `} - + - Learn more about the syntax. - - -
-
- - -
-
- } - > - - - Question -
- Let AI create a matching question for this SQL statement. - -
+ Learn more about the syntax. + + +
+
+ +
- } - name="question" - required - rules={[ - { - validator: createSQLPairQuestionValidator( - ERROR_TEXTS.SQL_PAIR.QUESTION, - ), - }, - ]} - > - - - - - - -
- - Data preview (50 rows) - - - {showPreview && ( -
-
+ } + > + + + Question +
+ Let AI create a matching question for this SQL statement. + +
+
+ } + name="question" + required + rules={[ + { + validator: createSQLPairQuestionValidator( + ERROR_TEXTS.SQL_PAIR.QUESTION, + ), + }, + ]} + > + + + + + importDataSourceSQLModal.openModal({ + dataSource: dataSource.type, + }) + } + /> + ) + } + autoComplete + autoFocus + /> + + +
+ + Data preview (50 rows) + + + {showPreview && ( +
+ +
+ )} +
+ {!!error && ( + } + /> )} -
- {!!error && ( - } + + {dataSource.isSupportSubstitute && ( + { + form.setFieldsValue({ sql: convertedSql }); + }} /> )} - + ); } diff --git a/wren-ui/src/components/pages/setup/ConnectDataSource.tsx b/wren-ui/src/components/pages/setup/ConnectDataSource.tsx index 27d554185a..a95d73c48c 100644 --- a/wren-ui/src/components/pages/setup/ConnectDataSource.tsx +++ b/wren-ui/src/components/pages/setup/ConnectDataSource.tsx @@ -89,7 +89,7 @@ export default function ConnectDataSource(props: Props) { { - 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), ); }; diff --git a/wren-ui/src/hooks/useAutoComplete.tsx b/wren-ui/src/hooks/useAutoComplete.tsx new file mode 100644 index 0000000000..492e8e21bb --- /dev/null +++ b/wren-ui/src/hooks/useAutoComplete.tsx @@ -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 = (item: (Model | Field) & { parent?: Model }) => T; + +interface Props { + skip?: boolean; + includeColumns?: boolean; + convertor?: Convertor; +} + +const getDocHTML = (item: (Model | Field) & { parent?: Model }) => { + return [ + '
', + `${item.referenceName}`, + item.description + ? `
${item.description}
` + : null, + '
', + ] + .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; +export type Completer = ReturnType; + +export default function useAutoComplete(props: Props) { + const { includeColumns, skip } = props; + const { data } = useDiagramQuery({ skip }); + + // Defined convertor + const convertor = (props.convertor || convertCompleter) as Convertor; + + 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]); +} diff --git a/wren-ui/src/hooks/useMentions.tsx b/wren-ui/src/hooks/useMentions.tsx deleted file mode 100644 index bb47e27abe..0000000000 --- a/wren-ui/src/hooks/useMentions.tsx +++ /dev/null @@ -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; - -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 }; -} diff --git a/wren-ui/src/utils/dataSourceType.ts b/wren-ui/src/utils/dataSourceType.ts new file mode 100644 index 0000000000..d36f2e1b2c --- /dev/null +++ b/wren-ui/src/utils/dataSourceType.ts @@ -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) }; +}; diff --git a/wren-ui/src/utils/enum/dataSources.ts b/wren-ui/src/utils/enum/dataSources.ts index d5b395d406..f7822341b2 100644 --- a/wren-ui/src/utils/enum/dataSources.ts +++ b/wren-ui/src/utils/enum/dataSources.ts @@ -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', diff --git a/wren-ui/src/utils/error/dictionary.ts b/wren-ui/src/utils/error/dictionary.ts index 1c264e299b..383ff8f9d0 100644 --- a/wren-ui/src/utils/error/dictionary.ts +++ b/wren-ui/src/utils/error/dictionary.ts @@ -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.', + }, + }, };