From 267ce1eba50f2f5668ba38119b1fa63611cfc99c Mon Sep 17 00:00:00 2001 From: Andy Yen <38731840+onlyjackfrost@users.noreply.github.com> Date: Fri, 5 Jul 2024 15:46:38 +0800 Subject: [PATCH] Chore: some update to reflect the adjustment of Wren Engine v0.6.0 (#472) * no need to deploy wren engine anymore * fix typo * communicate with engine/ibis server with manifest. * remove unused method in wren engine adaptor * add comment at wren engine adaptor * use the limit option from ibis server instead of rewriting the sql ourself * prevent no successful deployment found * pass manifest when getting native sql --- .../src/apollo/server/adaptors/ibisAdaptor.ts | 18 ++++- .../server/adaptors/wrenEngineAdaptor.ts | 81 +++---------------- wren-ui/src/apollo/server/data/sample.ts | 4 +- .../apollo/server/resolvers/modelResolver.ts | 31 ++++--- wren-ui/src/apollo/server/schema.ts | 4 +- .../apollo/server/services/askingService.ts | 2 +- .../apollo/server/services/deployService.ts | 37 ++++----- .../apollo/server/services/queryService.ts | 42 +++++----- .../services/tests/deployService.test.ts | 9 +-- wren-ui/src/pages/api/graphql.ts | 1 - 10 files changed, 83 insertions(+), 146 deletions(-) diff --git a/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts b/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts index dd9e5c2e6..cb2dcf8af 100644 --- a/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts @@ -6,7 +6,11 @@ import { Manifest } from '@server/mdl/type'; import * as Errors from '@server/utils/error'; import { getConfig } from '@server/config'; import { toDockerHost } from '@server/utils'; -import { CompactTable, RecommendConstraint } from '@server/services'; +import { + CompactTable, + DEFAULT_PREVIEW_LIMIT, + RecommendConstraint, +} from '@server/services'; import { snakeCase } from 'lodash'; import { WREN_AI_CONNECTION_INFO } from '../repositories'; import { toIbisConnectionInfo } from '../dataSource'; @@ -88,18 +92,21 @@ export interface ValidationResponse { message: string | null; } -export interface IbisQueryOptions { +export interface IbisBaseOptions { dataSource: DataSourceName; connectionInfo: WREN_AI_CONNECTION_INFO; mdl: Manifest; } +export interface IbisQueryOptions extends IbisBaseOptions { + limit?: number; +} export interface IIbisAdaptor { query: ( query: string, options: IbisQueryOptions, ) => Promise; - dryRun: (query: string, options: IbisQueryOptions) => Promise; + dryRun: (query: string, options: IbisBaseOptions) => Promise; getTables: ( dataSource: DataSourceName, connectionInfo: WREN_AI_CONNECTION_INFO, @@ -148,6 +155,11 @@ export class IbisAdaptor implements IIbisAdaptor { const res = await axios.post( `${this.ibisServerEndpoint}/v2/ibis/${dataSourceUrlMap[dataSource]}/query`, body, + { + params: { + limit: options.limit || DEFAULT_PREVIEW_LIMIT, + }, + }, ); const response = res.data; return response; diff --git a/wren-ui/src/apollo/server/adaptors/wrenEngineAdaptor.ts b/wren-ui/src/apollo/server/adaptors/wrenEngineAdaptor.ts index 83a7da77a..8fb0a47e8 100644 --- a/wren-ui/src/apollo/server/adaptors/wrenEngineAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/wrenEngineAdaptor.ts @@ -2,38 +2,16 @@ import axios, { AxiosResponse } from 'axios'; import { Manifest } from '../mdl/type'; import { getLogger } from '@server/utils'; import * as Errors from '@server/utils/error'; -import { CompactTable } from '../services'; +import { CompactTable, DEFAULT_PREVIEW_LIMIT } from '../services'; const logger = getLogger('WrenEngineAdaptor'); logger.level = 'debug'; -const DEFAULT_PREVIEW_LIMIT = 500; - -export enum WrenEngineDeployStatusEnum { - SUCCESS = 'SUCCESS', - FAILED = 'FAILED', -} - export interface WrenEngineDeployStatusResponse { systemStatus: string; version: string; } -export interface DeployResponse { - status: WrenEngineDeployStatusEnum; - error?: string; -} - -interface DeployPayload { - manifest: Manifest; - version: string; -} - -export interface deployData { - manifest: Manifest; - hash: string; -} - export interface ColumnMetadata { name: string; type: string; @@ -89,18 +67,23 @@ export interface DryRunResponse { } export interface IWrenEngineAdaptor { - deploy(deployData: deployData): Promise; + // duckdb data source related prepareDuckDB(options: DuckDBPrepareOptions): Promise; - listTables(): Promise; - putSessionProps(props: Record): Promise; queryDuckdb(sql: string): Promise; + putSessionProps(props: Record): Promise; + + // metadata related, used to fetch metadata of duckdb + listTables(): Promise; + + // config wren engine patchConfig(config: Record): Promise; + + // query previewData( sql: string, + mdl: Manifest, limit?: number, - mdl?: Manifest, ): Promise; - describeStatement(sql: string): Promise; getNativeSQL(sql: string, options?: DryPlanOption): Promise; validateColumnIsValid( manifest: Manifest, @@ -182,33 +165,6 @@ export class WrenEngineAdaptor implements IWrenEngineAdaptor { return this.formatToCompactTable(response); } - public async deploy(deployData: deployData): Promise { - const { manifest, hash } = deployData; - const deployPayload = { manifest, version: hash } as DeployPayload; - - try { - // skip if the model has been deployed - const resp = await this.getDeployStatus(); - if (resp.version === hash) { - return { status: WrenEngineDeployStatusEnum.SUCCESS }; - } - - // start deploy to wren engine - await axios.post( - `${this.wrenEngineBaseEndpoint}/v1/mdl/deploy`, - deployPayload, - ); - logger.debug(`WrenEngine: Deploy wren engine success, hash: ${hash}`); - return { status: WrenEngineDeployStatusEnum.SUCCESS }; - } catch (err: any) { - logger.debug(`Got error when deploying to wren engine: ${err.message}`); - return { - status: WrenEngineDeployStatusEnum.FAILED, - error: `WrenEngine Error, deployment hash:${hash}: ${err.message}`, - }; - } - } - public async putSessionProps(props: Record) { const setSessionStatements = Object.entries(props) .map(([key, value]) => { @@ -270,8 +226,8 @@ export class WrenEngineAdaptor implements IWrenEngineAdaptor { public async previewData( sql: string, + manifest: Manifest, limit: number = DEFAULT_PREVIEW_LIMIT, - manifest?: Manifest, ): Promise { try { const url = new URL(this.previewUrlPath, this.wrenEngineBaseEndpoint); @@ -297,19 +253,6 @@ export class WrenEngineAdaptor implements IWrenEngineAdaptor { } } - public async describeStatement( - sql: string, - ): Promise { - try { - // preview data with limit 1 to get column metadata - const res = await this.previewData(sql, 1); - return { columns: res.columns }; - } catch (err: any) { - logger.debug(`Got error when describing statement: ${err.message}`); - throw err; - } - } - public async getNativeSQL( sql: string, options: DryPlanOption, diff --git a/wren-ui/src/apollo/server/data/sample.ts b/wren-ui/src/apollo/server/data/sample.ts index 04e3ac17d..b9c8653e3 100644 --- a/wren-ui/src/apollo/server/data/sample.ts +++ b/wren-ui/src/apollo/server/data/sample.ts @@ -657,7 +657,7 @@ export const sampleDatasets: Record = { }, { name: 'BlkHome', - properties: { description: 'number of bloacks by the home team.' }, + properties: { description: 'number of blocks by the home team.' }, }, { name: 'TovHome', @@ -748,7 +748,7 @@ export const sampleDatasets: Record = { }, { name: 'BlkAway', - properties: { description: 'number of bloacks by the away team.' }, + properties: { description: 'number of blocks by the away team.' }, }, { name: 'TovAway', diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 52b3e1ccc..a86383af5 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -32,8 +32,6 @@ export enum SyncStatusEnum { UNSYNCRONIZED = 'UNSYNCRONIZED', } -const PREVIEW_MAX_OUTPUT_ROW = 100; - export class ModelResolver { constructor() { // model & model column @@ -147,7 +145,7 @@ export class ModelResolver { const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); const currentHash = ctx.deployService.createMDLHash(manifest, id); const lastDeploy = await ctx.deployService.getLastDeployment(id); - const lastDeployHash = lastDeploy.hash; + const lastDeployHash = lastDeploy?.hash; const inProgressDeployment = await ctx.deployService.getInProgressDeployment(id); if (inProgressDeployment) { @@ -547,8 +545,7 @@ export class ModelResolver { // create view const project = await ctx.projectService.getCurrentProject(); - const deployment = await ctx.deployService.getLastDeployment(project.id); - const mdl = deployment.manifest; + const { manifest } = await ctx.deployService.getLastDeployment(project.id); // get sql statement of a response const response = await ctx.askingService.getResponse(responseId); @@ -563,9 +560,9 @@ export class ModelResolver { // describe columns const { columns } = await ctx.queryService.describeStatement(statement, { project, - limit: PREVIEW_MAX_OUTPUT_ROW, + limit: 1, modelingOnly: false, - mdl, + manifest, }); if (isEmpty(columns)) { @@ -634,9 +631,8 @@ export class ModelResolver { const data = (await ctx.queryService.preview(sql, { project, - limit: PREVIEW_MAX_OUTPUT_ROW, modelingOnly: false, - mdl: manifest, + manifest, })) as PreviewDataResponse; return data; @@ -653,8 +649,8 @@ export class ModelResolver { const data = (await ctx.queryService.preview(view.statement, { project, - limit: limit | PREVIEW_MAX_OUTPUT_ROW, - mdl: manifest, + limit, + manifest, modelingOnly: false, })) as PreviewDataResponse; return data; @@ -669,13 +665,12 @@ export class ModelResolver { const project = projectId ? await ctx.projectService.getProjectById(projectId) : await ctx.projectService.getCurrentProject(); - const deployment = await ctx.deployService.getLastDeployment(project.id); - const mdl = deployment.manifest; + const { manifest } = await ctx.deployService.getLastDeployment(project.id); const previewRes = await ctx.queryService.preview(sql, { project, - limit: limit || PREVIEW_MAX_OUTPUT_ROW, + limit: limit, modelingOnly: false, - mdl, + manifest, dryRun, }); return dryRun ? { dryRun: 'success' } : previewRes; @@ -693,6 +688,7 @@ export class ModelResolver { if (sampleDataset) { throw new Error(`Doesn't support Native SQL`); } + const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); // get sql statement of a response const response = await ctx.askingService.getResponse(responseId); @@ -704,7 +700,10 @@ export class ModelResolver { const steps = response.detail.steps; const sql = format(constructCteSql(steps)); - return await ctx.wrenEngineAdaptor.getNativeSQL(sql); + return await ctx.wrenEngineAdaptor.getNativeSQL(sql, { + manifest, + modelingOnly: false, + }); } public async updateViewMetadata( diff --git a/wren-ui/src/apollo/server/schema.ts b/wren-ui/src/apollo/server/schema.ts index aabec08a7..d8549fa20 100644 --- a/wren-ui/src/apollo/server/schema.ts +++ b/wren-ui/src/apollo/server/schema.ts @@ -316,8 +316,8 @@ export const typeDefs = gql` input PreviewViewDataInput { id: Int! - # It will return default 100 rows if not specified limit - # refer: PREVIEW_MAX_OUTPUT_ROW + # It will return default 500 rows if not specified limit + # refer: DEFAULT_PREVIEW_LIMIT limit: Int } diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index a9f40db5f..79d944bdc 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -473,7 +473,7 @@ export class AskingService implements IAskingService { const sql = format(constructCteSql(steps, stepIndex)); const data = (await this.queryService.preview(sql, { project, - mdl, + manifest: mdl, limit, })) as PreviewDataResponse; diff --git a/wren-ui/src/apollo/server/services/deployService.ts b/wren-ui/src/apollo/server/services/deployService.ts index 6b29fcb37..1ac08a0c1 100644 --- a/wren-ui/src/apollo/server/services/deployService.ts +++ b/wren-ui/src/apollo/server/services/deployService.ts @@ -2,10 +2,6 @@ import { IWrenAIAdaptor, WrenAIDeployStatusEnum, } from '../adaptors/wrenAIAdaptor'; -import { - IWrenEngineAdaptor, - WrenEngineDeployStatusEnum, -} from '../adaptors/wrenEngineAdaptor'; import { Deploy, DeployStatusEnum, @@ -43,23 +39,19 @@ export interface IDeployService { export class DeployService implements IDeployService { private wrenAIAdaptor: IWrenAIAdaptor; - private wrenEngineAdaptor: IWrenEngineAdaptor; private deployLogRepository: IDeployLogRepository; private telemetry: Telemetry; constructor({ wrenAIAdaptor, - wrenEngineAdaptor, deployLogRepository, telemetry, }: { wrenAIAdaptor: IWrenAIAdaptor; - wrenEngineAdaptor: IWrenEngineAdaptor; deployLogRepository: IDeployLogRepository; telemetry: Telemetry; }) { this.wrenAIAdaptor = wrenAIAdaptor; - this.wrenEngineAdaptor = wrenEngineAdaptor; this.deployLogRepository = deployLogRepository; this.telemetry = telemetry; } @@ -102,24 +94,25 @@ export class DeployService implements IDeployService { } as Deploy; const deploy = await this.deployLogRepository.createOne(deployData); - // deploy to wren-engine & AI-service - const [engineRes, aiRes] = await Promise.all([ - this.wrenEngineAdaptor.deploy({ manifest, hash }), - this.wrenAIAdaptor.deploy({ manifest, hash }), - ]); + // deploy to AI-service + const { status: aiStatus, error: aiError } = + await this.wrenAIAdaptor.deploy({ + manifest, + hash, + }); - // store deploy log + // update deploy status const status = - engineRes.status === WrenEngineDeployStatusEnum.SUCCESS && - aiRes.status === WrenAIDeployStatusEnum.SUCCESS + aiStatus === WrenAIDeployStatusEnum.SUCCESS ? DeployStatusEnum.SUCCESS : DeployStatusEnum.FAILED; - const error = engineRes.error || aiRes.error; - await this.deployLogRepository.updateOne(deploy.id, { status, error }); - if (status === DeployStatusEnum.SUCCESS) { - this.telemetry.send_event('deploy_model', { mdl: manifest }); - } - return { status, error }; + await this.deployLogRepository.updateOne(deploy.id, { + status, + error: aiError, + }); + this.telemetry.send_event('deploy_model', { mdl: manifest }); + + return { status, error: aiError }; } public createMDLHash(manifest: Manifest, projectId: number) { diff --git a/wren-ui/src/apollo/server/services/queryService.ts b/wren-ui/src/apollo/server/services/queryService.ts index 18540a6ac..9e9d69ee7 100644 --- a/wren-ui/src/apollo/server/services/queryService.ts +++ b/wren-ui/src/apollo/server/services/queryService.ts @@ -8,7 +8,6 @@ import { IbisPostgresConnectionInfo, IbisBigQueryConnectionInfo, ValidationRules, - IbisQueryOptions, } from '../adaptors/ibisAdaptor'; import { getLogger } from '@server/utils'; import { Project } from '../repositories'; @@ -16,6 +15,8 @@ import { Project } from '../repositories'; const logger = getLogger('QueryService'); logger.level = 'debug'; +export const DEFAULT_PREVIEW_LIMIT = 500; + export interface ColumnMetadata { name: string; type: string; @@ -33,10 +34,10 @@ export interface DescribeStatementResponse { export interface PreviewOptions { project: Project; modelingOnly?: boolean; - mdl: Manifest; + // if not given, will use the deployed manifest + manifest: Manifest; limit?: number; dryRun?: boolean; - // if not given, will use the deployed manifest } export interface SqlValidateOptions { @@ -93,42 +94,37 @@ export class QueryService implements IQueryService { sql: string, options: PreviewOptions, ): Promise { - const { project, mdl, limit, dryRun } = options; + const { project, manifest: mdl, limit, dryRun } = options; const { type: dataSource, connectionInfo } = project; if (this.useEngine(dataSource)) { if (dryRun) { logger.debug('Using wren engine to dry run'); await this.wrenEngineAdaptor.dryRun(sql, { manifest: mdl, - limit: limit, + limit, }); return true; } else { logger.debug('Using wren engine to preview'); - const data = await this.wrenEngineAdaptor.previewData(sql, limit, mdl); + const data = await this.wrenEngineAdaptor.previewData(sql, mdl, limit); return data as PreviewDataResponse; } } else { - logger.debug('Use ibis adaptor to preview'); - /** - * add alias to FROM clause to prevent ibis error - * ibis server does not have limit parameter, should handle it in sql - */ - const rewrittenSql = limit - ? `SELECT tmp.* FROM (${sql}) tmp LIMIT ${limit}` - : sql; - this.checkDataSourceIsSupported(dataSource); - const queryOptions = { - dataSource, - connectionInfo, - mdl, - dryRun, - } as IbisQueryOptions; + logger.debug('Use ibis adaptor to preview'); if (dryRun) { - await this.ibisAdaptor.dryRun(rewrittenSql, queryOptions); + await this.ibisAdaptor.dryRun(sql, { + dataSource, + connectionInfo, + mdl, + }); } else { - const data = await this.ibisAdaptor.query(rewrittenSql, queryOptions); + const data = await this.ibisAdaptor.query(sql, { + dataSource, + connectionInfo, + mdl, + limit, + }); return this.transformDataType(data); } } diff --git a/wren-ui/src/apollo/server/services/tests/deployService.test.ts b/wren-ui/src/apollo/server/services/tests/deployService.test.ts index dd783d080..b0a2a8e2d 100644 --- a/wren-ui/src/apollo/server/services/tests/deployService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/deployService.test.ts @@ -3,7 +3,7 @@ import { DeployStatusEnum } from '@server/repositories/deployLogRepository'; describe('DeployService', () => { let mockWrenAIAdaptor; - let mockWrenEngineAdaptor; + let mockDeployLogRepository; let deployService; let mockTelemetry; @@ -11,7 +11,6 @@ describe('DeployService', () => { beforeEach(() => { mockTelemetry = { send_event: jest.fn() }; mockWrenAIAdaptor = { deploy: jest.fn() }; - mockWrenEngineAdaptor = { deploy: jest.fn() }; mockDeployLogRepository = { findLastProjectDeployLog: jest.fn(), createOne: jest.fn(), @@ -21,7 +20,6 @@ describe('DeployService', () => { deployService = new DeployService({ telemetry: mockTelemetry, wrenAIAdaptor: mockWrenAIAdaptor, - wrenEngineAdaptor: mockWrenEngineAdaptor, deployLogRepository: mockDeployLogRepository, }); }); @@ -32,7 +30,6 @@ describe('DeployService', () => { mockDeployLogRepository.findLastProjectDeployLog.mockResolvedValue(null); mockWrenAIAdaptor.deploy.mockResolvedValue({ status: 'SUCCESS' }); - mockWrenEngineAdaptor.deploy.mockResolvedValue({ status: 'SUCCESS' }); mockDeployLogRepository.createOne.mockResolvedValue({ id: 123 }); const response = await deployService.deploy(manifest, projectId); @@ -44,7 +41,7 @@ describe('DeployService', () => { }); }); - it('should return failed status if either deployment fails', async () => { + it('should return failed status if ai-service deployment fails', async () => { const manifest = { key: 'value' }; const projectId = 1; @@ -53,7 +50,6 @@ describe('DeployService', () => { status: 'FAILED', error: 'AI error', }); - mockWrenEngineAdaptor.deploy.mockResolvedValue({ status: 'SUCCESS' }); mockDeployLogRepository.createOne.mockResolvedValue({ id: 123 }); const response = await deployService.deploy(manifest, projectId); @@ -74,7 +70,6 @@ describe('DeployService', () => { expect(response.status).toEqual(DeployStatusEnum.SUCCESS); expect(mockWrenAIAdaptor.deploy).not.toHaveBeenCalled(); - expect(mockWrenEngineAdaptor.deploy).not.toHaveBeenCalled(); }); // Add more tests here to cover other scenarios and error handling diff --git a/wren-ui/src/pages/api/graphql.ts b/wren-ui/src/pages/api/graphql.ts index 3090226c0..830f5a8d6 100644 --- a/wren-ui/src/pages/api/graphql.ts +++ b/wren-ui/src/pages/api/graphql.ts @@ -108,7 +108,6 @@ const bootstrapServer = async () => { }); const deployService = new DeployService({ wrenAIAdaptor, - wrenEngineAdaptor, deployLogRepository, telemetry, });