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

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<IbisQueryResponse>;
-  dryRun: (query: string, options: IbisQueryOptions) => Promise<boolean>;
+  dryRun: (query: string, options: IbisBaseOptions) => Promise<boolean>;
   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;

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<DeployResponse>;
+  // duckdb data source related
   prepareDuckDB(options: DuckDBPrepareOptions): Promise<void>;
-  listTables(): Promise<CompactTable[]>;
-  putSessionProps(props: Record<string, any>): Promise<void>;
   queryDuckdb(sql: string): Promise<EngineQueryResponse>;
+  putSessionProps(props: Record<string, any>): Promise<void>;
+
+  // metadata related, used to fetch metadata of duckdb
+  listTables(): Promise<CompactTable[]>;
+
+  // config wren engine
   patchConfig(config: Record<string, any>): Promise<void>;
+
+  // query
   previewData(
     sql: string,
+    mdl: Manifest,
     limit?: number,
-    mdl?: Manifest,
   ): Promise<EngineQueryResponse>;
-  describeStatement(sql: string): Promise<DescribeStatementResponse>;
   getNativeSQL(sql: string, options?: DryPlanOption): Promise<string>;
   validateColumnIsValid(
     manifest: Manifest,
@@ -182,33 +165,6 @@ export class WrenEngineAdaptor implements IWrenEngineAdaptor {
     return this.formatToCompactTable(response);
   }
 
-  public async deploy(deployData: deployData): Promise<DeployResponse> {
-    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<string, any>) {
     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<EngineQueryResponse> {
     try {
       const url = new URL(this.previewUrlPath, this.wrenEngineBaseEndpoint);
@@ -297,19 +253,6 @@ export class WrenEngineAdaptor implements IWrenEngineAdaptor {
     }
   }
 
-  public async describeStatement(
-    sql: string,
-  ): Promise<DescribeStatementResponse> {
-    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,

wren-ui/src/apollo/server/data/sample.ts

@@ -657,7 +657,7 @@ export const sampleDatasets: Record<string, SampleDataset> = {
           },
           {
             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<string, SampleDataset> = {
           },
           {
             name: 'BlkAway',
-            properties: { description: 'number of bloacks by the away team.' },
+            properties: { description: 'number of blocks by the away team.' },
           },
           {
             name: 'TovAway',

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(

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
   }
 

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;
 

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) {

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<PreviewDataResponse | boolean> {
-    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);
       }
     }

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

wren-ui/src/pages/api/graphql.ts

@@ -108,7 +108,6 @@ const bootstrapServer = async () => {
   });
   const deployService = new DeployService({
     wrenAIAdaptor,
-    wrenEngineAdaptor,
     deployLogRepository,
     telemetry,
   });