diff --git a/docker/.env.example b/docker/.env.example index 03be15ad3..a49f17ed8 100644 --- a/docker/.env.example +++ b/docker/.env.example @@ -5,13 +5,19 @@ PLATFORM=linux/amd64 WREN_ENGINE_PORT=8080 WREN_ENGINE_SQL_PORT=7432 WREN_AI_SERVICE_PORT=5555 +WREN_UI_PORT=3000 +IBIS_SERVER_PORT=8000 + +# service endpoint (for docker-compose-dev.yaml file) +WREN_UI_ENDPOINT=http://docker.for.mac.localhost:3000 # version # CHANGE THIS TO THE LATEST VERSION -WREN_PRODUCT_VERSION=0.3.6 -WREN_ENGINE_VERSION=0.4.6 +WREN_PRODUCT_VERSION=0.4.0-rc.1 +WREN_ENGINE_VERSION=0.4.7 WREN_AI_SERVICE_VERSION=0.4.0 -WREN_UI_VERSION=0.5.7 +WREN_UI_VERSION=epic-ibis-1 +IBIS_SERVER_VERSION=0.4.7 WREN_BOOTSTRAP_VERSION=0.1.4 # keys diff --git a/docker/docker-compose-dev.yaml b/docker/docker-compose-dev.yaml index ad5ffe6a7..5be3568a4 100644 --- a/docker/docker-compose-dev.yaml +++ b/docker/docker-compose-dev.yaml @@ -1,4 +1,4 @@ -version: '3' +version: "3" volumes: data: @@ -46,7 +46,7 @@ services: OPENAI_API_KEY: ${OPENAI_API_KEY} OPENAI_GENERATION_MODEL: ${OPENAI_GENERATION_MODEL} QDRANT_HOST: qdrant - WREN_UI_ENDPOINT: http://wren-ui:${WREN_UI_PORT} + WREN_UI_ENDPOINT: ${WREN_UI_ENDPOINT} REDIS_HOST: ${AI_SERVICE_REDIS_HOST} REDIS_PORT: ${AI_SERVICE_REDIS_PORT} ENABLE_TIMER: ${AI_SERVICE_ENABLE_TIMER} @@ -62,6 +62,20 @@ services: - wren-engine - redis + ibis-server: + image: ghcr.io/canner/wren-engine-ibis:${IBIS_SERVER_VERSION} + pull_policy: always + platform: ${PLATFORM} + expose: + - 8000 + ports: + - ${IBIS_SERVER_PORT}:8000 + environment: + WREN_ENGINE_ENDPOINT: http://wren-engine:${WREN_ENGINE_PORT} + LOG_LEVEL: DEBUG + networks: + - wren + qdrant: image: qdrant/qdrant:v1.7.4 pull_policy: always diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index 37748ce18..f95d6b65a 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -1,4 +1,4 @@ -version: '3' +version: "3" volumes: data: @@ -34,6 +34,17 @@ services: depends_on: - bootstrap + ibis-server: + image: ghcr.io/canner/wren-engine-ibis:${IBIS_SERVER_VERSION} + restart: on-failure + platform: ${PLATFORM} + expose: + - ${IBIS_SERVER_PORT} + environment: + WREN_ENGINE_ENDPOINT: http://wren-engine:${WREN_ENGINE_PORT} + networks: + - wren + wren-ai-service: image: ghcr.io/canner/wren-ai-service:${WREN_AI_SERVICE_VERSION} restart: on-failure @@ -92,6 +103,7 @@ services: SQLITE_FILE: /app/data/db.sqlite3 WREN_ENGINE_ENDPOINT: http://wren-engine:${WREN_ENGINE_PORT} WREN_AI_ENDPOINT: http://wren-ai-service:${WREN_AI_SERVICE_PORT} + IBIS_SERVER_ENDPOINT: http://ibis-server:${IBIS_SERVER_PORT} OPENAI_GENERATION_MODEL: ${OPENAI_GENERATION_MODEL} PG_USERNAME: ${PG_USERNAME} PG_PASSWORD: ${PG_PASSWORD} diff --git a/wren-launcher/utils/docker.go b/wren-launcher/utils/docker.go index cd9ca51ae..d11367973 100644 --- a/wren-launcher/utils/docker.go +++ b/wren-launcher/utils/docker.go @@ -24,7 +24,7 @@ import ( const ( // please change the version when the version is updated - WREN_PRODUCT_VERSION string = "0.3.6" + WREN_PRODUCT_VERSION string = "0.4.0-rc.1" DOCKER_COMPOSE_YAML_URL string = "https://raw.githubusercontent.com/Canner/WrenAI/" + WREN_PRODUCT_VERSION + "/docker/docker-compose.yaml" DOCKER_COMPOSE_ENV_URL string = "https://raw.githubusercontent.com/Canner/WrenAI/" + WREN_PRODUCT_VERSION + "/docker/.env.example" diff --git a/wren-ui/README.md b/wren-ui/README.md index e522c46f9..d9059efa6 100644 --- a/wren-ui/README.md +++ b/wren-ui/README.md @@ -36,6 +36,14 @@ export PG_URL=postgres://user:password@localhost:5432/dbname Step 4. Run the development server: ```bash +# Execute this if you start wren-engine and ibis-server via docker +# Linux or MacOS +export OTHER_SERVICE_USING_DOCKER=true +# Windows +SET OTHER_SERVICE_USING_DOCKER=true + + +# Run the development server yarn dev # or npm run dev diff --git a/wren-ui/migrations/20240327030000_create_ask_table.js b/wren-ui/migrations/20240327030000_create_ask_table.js index b2317dcd6..f48374d87 100644 --- a/wren-ui/migrations/20240327030000_create_ask_table.js +++ b/wren-ui/migrations/20240327030000_create_ask_table.js @@ -37,5 +37,5 @@ exports.up = function (knex) { * @returns { Promise } */ exports.down = function (knex) { - return knex.schema.dropTable('thread').dropTable('thread_response'); + return knex.schema.dropTable('thread_response').dropTable('thread'); }; diff --git a/wren-ui/migrations/20240418000000_update_project_table_pg.js b/wren-ui/migrations/20240418000000_update_project_table_pg.js index 5f7b64704..4f0ed14a3 100644 --- a/wren-ui/migrations/20240418000000_update_project_table_pg.js +++ b/wren-ui/migrations/20240418000000_update_project_table_pg.js @@ -31,6 +31,6 @@ exports.up = function (knex) { */ exports.down = function (knex) { return knex.schema.alterTable('project', (table) => { - table.dropColumns('host', 'port', 'database', 'username', 'password'); + table.dropColumns('host', 'port', 'database', 'user'); }); }; diff --git a/wren-ui/migrations/20240530062133_update_project_table.js b/wren-ui/migrations/20240530062133_update_project_table.js new file mode 100644 index 000000000..8c5168114 --- /dev/null +++ b/wren-ui/migrations/20240530062133_update_project_table.js @@ -0,0 +1,24 @@ +/** + * @param { import("knex").Knex } knex + * @returns { Promise } + */ + +// create connectionInfo column in project table +exports.up = function (knex) { + return knex.schema.table('project', (table) => { + table + .jsonb('connection_info') + .nullable() + .comment('Connection information for the project'); + }); +}; + +/** + * @param { import("knex").Knex } knex + * @returns { Promise } + */ +exports.down = function (knex) { + return knex.schema.table('project', (table) => { + table.dropColumn('connection_info'); + }); +}; diff --git a/wren-ui/migrations/20240530062809_transfer_project_table_data.js b/wren-ui/migrations/20240530062809_transfer_project_table_data.js new file mode 100644 index 000000000..30180e8eb --- /dev/null +++ b/wren-ui/migrations/20240530062809_transfer_project_table_data.js @@ -0,0 +1,82 @@ +/** + * @param { import("knex").Knex } knex + * @returns { Promise } + */ +exports.up = async function (knex) { + const projects = await knex('project').select('*'); + + // bigquery data + const bigqueryConnectionInfo = projects + .filter((project) => project.type === 'BIG_QUERY') + .map((project) => { + return { + id: project.id, + connectionInfo: { + projectId: project.project_id, + datasetId: project.dataset_id, + credentials: project.credentials, + }, + }; + }); + + // duckdb data + const duckdbConnectionInfo = projects + .filter((project) => project.type === 'DUCKDB') + .map((project) => { + return { + id: project.id, + connectionInfo: { + initSql: project.init_sql || '', + configurations: project.configurations || {}, + extensions: project.extensions || [], + }, + }; + }); + + // postgres data + const postgresConnectionInfo = projects + .filter((project) => project.type === 'POSTGRES') + .map((project) => { + const ssl = + project.configurations && project.configurations.ssl ? true : false; + return { + id: project.id, + connectionInfo: { + host: project.host, + port: project.port, + database: project.database, + user: project.user, + password: project.credentials, + ssl, + }, + }; + }); + + // update project table + for (const project of [ + ...bigqueryConnectionInfo, + ...duckdbConnectionInfo, + ...postgresConnectionInfo, + ]) { + const { id, connectionInfo } = project; + if (process.env.DB_TYPE === 'pg') { + // postgres + await knex('project') + .where({ id }) + .update({ connection_info: connectionInfo }); + } else { + // sqlite + await knex('project') + .where({ id }) + .update({ connection_info: JSON.stringify(connectionInfo) }); + } + } +}; + +/** + * @param { import("knex").Knex } knex + * @returns { Promise } + */ +exports.down = async function (knex) { + await knex('project').update({ connection_info: null }); +}; diff --git a/wren-ui/migrations/20240530105955_drop_project_table_columns.js b/wren-ui/migrations/20240530105955_drop_project_table_columns.js new file mode 100644 index 000000000..b7be89adb --- /dev/null +++ b/wren-ui/migrations/20240530105955_drop_project_table_columns.js @@ -0,0 +1,65 @@ +/** + * @param { import("knex").Knex } knex + * @returns { Promise } + */ +exports.up = function (knex) { + return knex.schema.table('project', (table) => { + table.dropColumn('configurations'); + table.dropColumn('credentials'); + table.dropColumn('project_id'); + table.dropColumn('dataset_id'); + table.dropColumn('init_sql'); + table.dropColumn('extensions'); + table.dropColumn('host'); + table.dropColumn('port'); + table.dropColumn('database'); + table.dropColumn('user'); + }); +}; + +/** + * @param { import("knex").Knex } knex + * @returns { Promise } + */ +exports.down = function (knex) { + return knex.schema.table('project', (table) => { + table + .jsonb('configurations') + .nullable() + .comment( + 'duckdb configurations that can be set in session, eg: { "key1": "value1", "key2": "value2" }', + ); + table + .text('credentials') + .nullable() + .comment('database connection credentials'); + table + .string('project_id') + .nullable() + .comment('gcp project id, big query specific'); + table.string('dataset_id').nullable().comment('big query datasetId'); + table.text('init_sql'); + table + .jsonb('extensions') + .nullable() + .comment( + 'duckdb extensions, will be a array-like string like, eg: ["extension1", "extension2"]', + ); + table + .string('host') + .nullable() + .comment('postgresql host, postgresql specific'); + table + .integer('port') + .nullable() + .comment('postgresql port, postgresql specific'); + table + .string('database') + .nullable() + .comment('postgresql database, postgresql specific'); + table + .string('user') + .nullable() + .comment('postgresql user, postgresql specific'); + }); +}; diff --git a/wren-ui/migrations/20240531085916_transfer_model_properties.js b/wren-ui/migrations/20240531085916_transfer_model_properties.js new file mode 100644 index 000000000..817c1b82d --- /dev/null +++ b/wren-ui/migrations/20240531085916_transfer_model_properties.js @@ -0,0 +1,65 @@ +/** + * @param { import("knex").Knex } knex + * @returns { Promise } + */ +exports.up = async function (knex) { + const projects = await knex('project').select('*'); + const models = await knex('model').select('*'); + console.log(`model len:${models.length}`); + for (const model of models) { + const project = projects.find((p) => p.id === model.project_id); + const dataSourceType = project.type; + // get schema & catalog if its available + let schema = null; + let catalog = null; + let table = null; + switch (dataSourceType) { + case 'BIG_QUERY': { + const connectionInfo = + typeof project.connection_info === 'string' + ? JSON.parse(project.connection_info) + : project.connection_info; + const datasetId = connectionInfo.datasetId; + if (!datasetId) continue; + const splitDataSetId = datasetId.split('.'); + schema = splitDataSetId[1]; + catalog = splitDataSetId[0]; + table = model.source_table_name; + break; + } + case 'POSTGRES': { + const connectionInfo = + typeof project.connection_info === 'string' + ? JSON.parse(project.connection_info) + : project.connection_info; + catalog = connectionInfo.database; + schema = model.source_table_name.split('.')[0]; + table = model.source_table_name.split('.')[1]; + break; + } + case 'DUCKDB': { + // already have schema & catalog in properties + table = model.source_table_name; + break; + } + } + const oldProperties = model.properties ? JSON.parse(model.properties) : {}; + const newProperties = { + schema, + catalog, + table, + ...oldProperties, + }; + await knex('model') + .where({ id: model.id }) + .update({ properties: JSON.stringify(newProperties) }); + } +}; + +/** + * @param { import("knex").Knex } knex + * @returns { Promise } + */ +exports.down = function () { + return Promise.resolve(); +}; diff --git a/wren-ui/src/apollo/client/graphql/__types__.ts b/wren-ui/src/apollo/client/graphql/__types__.ts index 59b8a4119..3f663ff04 100644 --- a/wren-ui/src/apollo/client/graphql/__types__.ts +++ b/wren-ui/src/apollo/client/graphql/__types__.ts @@ -326,6 +326,16 @@ export type FieldInfo = { type?: Maybe; }; +export type GetMdlInput = { + hash: Scalars['String']; +}; + +export type GetMdlResult = { + __typename?: 'GetMDLResult'; + hash: Scalars['String']; + mdl?: Maybe; +}; + export type MdlModelSubmitInput = { columns: Array; name: Scalars['String']; @@ -372,8 +382,10 @@ export type Mutation = { deleteThread: Scalars['Boolean']; deleteView: Scalars['Boolean']; deploy: Scalars['JSON']; + getMDL: GetMdlResult; previewData: Scalars['JSON']; previewModelData: Scalars['JSON']; + previewSql: Scalars['JSON']; previewViewData: Scalars['JSON']; resetCurrentProject: Scalars['Boolean']; saveDataSource: DataSource; @@ -458,6 +470,11 @@ export type MutationDeleteViewArgs = { }; +export type MutationGetMdlArgs = { + data?: InputMaybe; +}; + + export type MutationPreviewDataArgs = { where: PreviewDataInput; }; @@ -468,8 +485,13 @@ export type MutationPreviewModelDataArgs = { }; +export type MutationPreviewSqlArgs = { + data?: InputMaybe; +}; + + export type MutationPreviewViewDataArgs = { - where: ViewWhereUniqueInput; + where: PreviewViewDataInput; }; @@ -565,10 +587,22 @@ export type OnboardingStatusResponse = { }; export type PreviewDataInput = { + limit?: InputMaybe; responseId: Scalars['Int']; stepIndex?: InputMaybe; }; +export type PreviewSqlDataInput = { + limit?: InputMaybe; + projectId: Scalars['Int']; + sql: Scalars['String']; +}; + +export type PreviewViewDataInput = { + id: Scalars['Int']; + limit?: InputMaybe; +}; + export type Query = { __typename?: 'Query'; askingTask: AskingTask; diff --git a/wren-ui/src/apollo/client/graphql/view.generated.ts b/wren-ui/src/apollo/client/graphql/view.generated.ts index d41da120c..2642c7cc3 100644 --- a/wren-ui/src/apollo/client/graphql/view.generated.ts +++ b/wren-ui/src/apollo/client/graphql/view.generated.ts @@ -27,10 +27,10 @@ export type GetViewQuery = { __typename?: 'Query', view: { __typename?: 'ViewInf export type ListViewsQueryVariables = Types.Exact<{ [key: string]: never; }>; -export type ListViewsQuery = { __typename?: 'Query', listViews: Array<{ __typename?: 'ViewInfo', id: number, name: string, statement: string }> }; +export type ListViewsQuery = { __typename?: 'Query', listViews: Array<{ __typename?: 'ViewInfo', id: number, name: string, displayName: string, statement: string }> }; export type PreviewViewDataMutationVariables = Types.Exact<{ - where: Types.ViewWhereUniqueInput; + where: Types.PreviewViewDataInput; }>; @@ -152,6 +152,7 @@ export const ListViewsDocument = gql` listViews { id name + displayName statement } } @@ -184,7 +185,7 @@ export type ListViewsQueryHookResult = ReturnType; export type ListViewsLazyQueryHookResult = ReturnType; export type ListViewsQueryResult = Apollo.QueryResult; export const PreviewViewDataDocument = gql` - mutation PreviewViewData($where: ViewWhereUniqueInput!) { + mutation PreviewViewData($where: PreviewViewDataInput!) { previewViewData(where: $where) } `; diff --git a/wren-ui/src/apollo/client/graphql/view.ts b/wren-ui/src/apollo/client/graphql/view.ts index 62ef5eb43..6021ef27a 100644 --- a/wren-ui/src/apollo/client/graphql/view.ts +++ b/wren-ui/src/apollo/client/graphql/view.ts @@ -31,13 +31,14 @@ export const LIST_VIEWS = gql` listViews { id name + displayName statement } } `; export const PREVIEW_VIEW_DATA = gql` - mutation PreviewViewData($where: ViewWhereUniqueInput!) { + mutation PreviewViewData($where: PreviewViewDataInput!) { previewViewData(where: $where) } `; diff --git a/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts b/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts new file mode 100644 index 000000000..aefc26ba1 --- /dev/null +++ b/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts @@ -0,0 +1,249 @@ +import axios, { AxiosResponse } from 'axios'; + +import { getLogger } from '@server/utils/logger'; +import { DataSourceName } from '@server/types'; +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 { snakeCase } from 'lodash'; + +const logger = getLogger('IbisAdaptor'); +logger.level = 'debug'; + +const config = getConfig(); + +export interface IbisPostgresConnectionInfo { + host: string; + port: number; + database: string; + user: string; + password: string; + ssl: boolean; +} + +export interface IbisBigQueryConnectionInfo { + project_id: string; + dataset_id: string; + credentials: string; // base64 encoded +} + +export interface TableResponse { + tables: CompactTable[]; +} + +export enum ValidationRules { + COLUMN_IS_VALID = 'COLUMN_IS_VALID', +} + +export interface ValidationResponse { + valid: boolean; + message: string | null; +} + +export interface QueryOptions { + dataSource: DataSourceName; + connectionInfo: IbisBigQueryConnectionInfo | IbisPostgresConnectionInfo; + mdl: Manifest; +} + +export interface IIbisAdaptor { + query: (query: string, options: QueryOptions) => Promise; + dryRun: (query: string, options: QueryOptions) => Promise; + getTables: ( + dataSource: DataSourceName, + connectionInfo: IbisBigQueryConnectionInfo | IbisPostgresConnectionInfo, + ) => Promise; + getConstraints: ( + dataSource: DataSourceName, + connectionInfo: IbisBigQueryConnectionInfo | IbisPostgresConnectionInfo, + ) => Promise; + + validate: ( + dataSource: DataSourceName, + rule: ValidationRules, + connectionInfo: IbisBigQueryConnectionInfo | IbisPostgresConnectionInfo, + mdl: Manifest, + parameters: Record, + ) => Promise; +} + +export enum SupportedDataSource { + POSTGRES = 'POSTGRES', + BIG_QUERY = 'BIG_QUERY', + SNOWFLAKE = 'SNOWFLAKE', +} + +export interface IbisQueryResponse { + columns: string[]; + data: any[]; + dtypes: Record; +} + +const dataSourceUrlMap: Record = { + [SupportedDataSource.POSTGRES]: 'postgres', + [SupportedDataSource.BIG_QUERY]: 'bigquery', + [SupportedDataSource.SNOWFLAKE]: 'snowflake', +}; + +export class IbisAdaptor implements IIbisAdaptor { + private ibisServerEndpoint: string; + + constructor({ ibisServerEndpoint }: { ibisServerEndpoint: string }) { + this.ibisServerEndpoint = ibisServerEndpoint; + } + + public async query( + query: string, + options: QueryOptions, + ): Promise { + const { dataSource, mdl } = options; + const connectionInfo = this.updateConnectionInfo(options.connectionInfo); + const body = { + sql: query, + connectionInfo, + manifestStr: Buffer.from(JSON.stringify(mdl)).toString('base64'), + }; + logger.debug(`Querying ibis with body: ${JSON.stringify(body, null, 2)}`); + try { + const res = await axios.post( + `${this.ibisServerEndpoint}/v2/ibis/${dataSourceUrlMap[dataSource]}/query`, + body, + ); + const response = res.data; + return response; + } catch (e) { + logger.debug(`Got error when querying ibis: ${e.response.data}`); + + throw Errors.create(Errors.GeneralErrorCodes.IBIS_SERVER_ERROR, { + customMessage: e.response.data || 'Error querying ibis server', + originalError: e, + }); + } + } + + public async dryRun(query: string, options: QueryOptions): Promise { + const { dataSource, mdl } = options; + const connectionInfo = this.updateConnectionInfo(options.connectionInfo); + const body = { + sql: query, + connectionInfo, + manifestStr: Buffer.from(JSON.stringify(mdl)).toString('base64'), + }; + logger.debug(`Dry run ibis with body: ${JSON.stringify(body, null, 2)}`); + try { + await axios.post( + `${this.ibisServerEndpoint}/v2/ibis/${dataSourceUrlMap[dataSource]}/query?dryRun=true`, + body, + ); + return true; + } catch (e) { + logger.debug(`Got error when dry run ibis: ${e.response.data}`); + return false; + } + } + + public async getTables( + dataSource: DataSourceName, + connectionInfo: Record, + ): Promise { + if (config.otherServiceUsingDocker) { + connectionInfo.host = toDockerHost(connectionInfo.host); + logger.debug(`Rewritten host: ${connectionInfo.host}`); + } + const body = { + connectionInfo, + }; + logger.debug(`Getting table with body: ${JSON.stringify(body, null, 2)}`); + try { + const res: AxiosResponse = await axios.post( + `${this.ibisServerEndpoint}/v2/ibis/${dataSourceUrlMap[dataSource]}/metadata/tables`, + body, + ); + return res.data; + } catch (e) { + logger.debug(`Got error when getting table: ${e.response.data}`); + + throw Errors.create(Errors.GeneralErrorCodes.IBIS_SERVER_ERROR, { + customMessage: + e.response.data || 'Error getting table from ibis server', + originalError: e, + }); + } + } + + public async getConstraints( + dataSource: DataSourceName, + connectionInfo: Record, + ): Promise { + if (config.otherServiceUsingDocker) { + connectionInfo.host = toDockerHost(connectionInfo.host); + logger.debug(`Rewritten host: ${connectionInfo.host}`); + } + const body = { + connectionInfo, + }; + logger.debug( + `Getting constraint with body: ${JSON.stringify(body, null, 2)}`, + ); + try { + const res: AxiosResponse = await axios.post( + `${this.ibisServerEndpoint}/v2/ibis/${dataSourceUrlMap[dataSource]}/metadata/constraints`, + body, + ); + return res.data; + } catch (e) { + logger.debug(`Got error when getting constraint: ${e.response.data}`); + + throw Errors.create(Errors.GeneralErrorCodes.IBIS_SERVER_ERROR, { + customMessage: + e.response.data || 'Error getting constraint from ibis server', + originalError: e, + }); + } + } + + public async validate( + dataSource: DataSourceName, + validationRule: ValidationRules, + connectionInfo: Record, + mdl: Manifest, + parameters: Record, + ): Promise { + if (config.otherServiceUsingDocker) { + connectionInfo.host = toDockerHost(connectionInfo.host); + logger.debug(`Rewritten host: ${connectionInfo.host}`); + } + const body = { + connectionInfo, + manifestStr: Buffer.from(JSON.stringify(mdl)).toString('base64'), + parameters, + }; + logger.debug( + `Validating connection with body: ${JSON.stringify(body, null, 2)}`, + ); + try { + await axios.post( + `${this.ibisServerEndpoint}/v2/ibis/${dataSourceUrlMap[dataSource]}/validate/${snakeCase(validationRule)}`, + body, + ); + return { valid: true, message: null }; + } catch (e) { + logger.debug(`Got error when validating connection: ${e.response.data}`); + + return { valid: false, message: e.response.data }; + } + } + + private updateConnectionInfo(connectionInfo: any) { + if ( + config.otherServiceUsingDocker && + Object.hasOwnProperty.call(connectionInfo, 'host') + ) { + connectionInfo.host = toDockerHost(connectionInfo.host); + logger.debug(`Rewritten host: ${connectionInfo.host}`); + } + return connectionInfo; + } +} diff --git a/wren-ui/src/apollo/server/adaptors/tests/ibisAdaptor.test.ts b/wren-ui/src/apollo/server/adaptors/tests/ibisAdaptor.test.ts new file mode 100644 index 000000000..707adad62 --- /dev/null +++ b/wren-ui/src/apollo/server/adaptors/tests/ibisAdaptor.test.ts @@ -0,0 +1,175 @@ +import axios from 'axios'; +import { + IbisAdaptor, + IbisBigQueryConnectionInfo, + IbisPostgresConnectionInfo, + ValidationRules, +} from '../ibisAdaptor'; +import { DataSourceName } from '../../types'; +import { Manifest } from '../../mdl/type'; + +jest.mock('axios'); +const mockedAxios = axios as jest.Mocked; + +describe('IbisAdaptor', () => { + let ibisAdaptor: IbisAdaptor; + const ibisServerEndpoint = 'http://localhost:8080'; + + const mockPostgresConnectionInfo: IbisPostgresConnectionInfo = { + host: 'localhost', + port: 5432, + database: 'my-database', + user: 'my-user', + password: 'my-password', + ssl: false, + }; + + const mockBigQueryConnectionInfo: IbisBigQueryConnectionInfo = { + project_id: 'my-bq-project-id', + dataset_id: 'my-bq-dataset-id', + credentials: 'my-bq-credentials', + }; + + const mockManifest: Manifest = { + catalog: 'wrenai', // eg: "test-catalog" + schema: 'wrenai', // eg: "test-schema" + models: [ + { + name: 'test_table', + tableReference: { + catalog: 'wrenai', + schema: 'wrenai', + table: 'test_table', + }, + properties: { + description: 'test table', + }, + columns: [ + { + name: 'id', + type: 'integer', + properties: {}, + isCalculated: false, + }, + { + name: 'sumId', + type: 'float', + properties: {}, + isCalculated: true, + expression: 'SUM(id)', + }, + ], + cached: false, + }, + ], + relationships: [], + views: [], + }; + + beforeEach(() => { + ibisAdaptor = new IbisAdaptor({ + ibisServerEndpoint: ibisServerEndpoint, + }); + }); + + afterEach(() => { + jest.resetAllMocks(); + }); + + it('should get postgres constraints', async () => { + const mockResponse = { data: [] }; + mockedAxios.post.mockResolvedValue(mockResponse); + + const result = await ibisAdaptor.getConstraints( + DataSourceName.POSTGRES, + mockPostgresConnectionInfo, + ); + + expect(result).toEqual([]); + expect(mockedAxios.post).toHaveBeenCalledWith( + `${ibisServerEndpoint}/v2/ibis/postgres/metadata/constraints`, + { connectionInfo: mockPostgresConnectionInfo }, + ); + }); + + it('should get bigquery constraints', async () => { + const mockResponse = { data: [] }; + mockedAxios.post.mockResolvedValue(mockResponse); + + const result = await ibisAdaptor.getConstraints( + DataSourceName.BIG_QUERY, + mockBigQueryConnectionInfo, + ); + + expect(result).toEqual([]); + expect(mockedAxios.post).toHaveBeenCalledWith( + `${ibisServerEndpoint}/v2/ibis/bigquery/metadata/constraints`, + { connectionInfo: mockBigQueryConnectionInfo }, + ); + }); + + // it('should handle error when getting constraints', async () => { + // const mockError = new Error('Error'); + // mockedAxios.post.mockRejectedValue(mockError); + + // await expect(ibisAdaptor.getConstraints('dataSource', {})).rejects.toThrow( + // 'Error', + // ); + // }); + + it('should validate with rule COLUMN_IS_VALID', async () => { + mockedAxios.post.mockResolvedValue(true); + const parameters = { + modelName: 'test_table', + columnName: 'sumId', + }; + const result = await ibisAdaptor.validate( + DataSourceName.POSTGRES, + ValidationRules.COLUMN_IS_VALID, + mockPostgresConnectionInfo, + mockManifest, + parameters, + ); + + expect(result).toEqual({ valid: true, message: null }); + expect(mockedAxios.post).toHaveBeenCalledWith( + `${ibisServerEndpoint}/v2/ibis/postgres/validate/column_is_valid`, + { + connectionInfo: mockPostgresConnectionInfo, + manifestStr: Buffer.from(JSON.stringify(mockManifest)).toString( + 'base64', + ), + parameters, + }, + ); + }); + + it('should handle error when validating', async () => { + const mockError = { response: { data: 'Error' } }; + const parameters = { + modelName: 'test_table', + columnName: 'sumId', + }; + mockedAxios.post.mockRejectedValue(mockError); + + const result = await ibisAdaptor.validate( + DataSourceName.POSTGRES, + ValidationRules.COLUMN_IS_VALID, + mockPostgresConnectionInfo, + mockManifest, + parameters, + ); + + expect(result).toEqual({ valid: false, message: 'Error' }); + expect(mockedAxios.post).toHaveBeenCalledWith( + `${ibisServerEndpoint}/v2/ibis/postgres/validate/column_is_valid`, + { + connectionInfo: mockPostgresConnectionInfo, + manifestStr: Buffer.from(JSON.stringify(mockManifest)).toString( + 'base64', + ), + parameters, + }, + ); + }); +}); diff --git a/wren-ui/src/apollo/server/adaptors/wrenEngineAdaptor.ts b/wren-ui/src/apollo/server/adaptors/wrenEngineAdaptor.ts index c3bc5b8e0..8e63b4791 100644 --- a/wren-ui/src/apollo/server/adaptors/wrenEngineAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/wrenEngineAdaptor.ts @@ -2,6 +2,7 @@ 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'; const logger = getLogger('WrenEngineAdaptor'); logger.level = 'debug'; @@ -38,7 +39,7 @@ export interface ColumnMetadata { type: string; } -export interface QueryResponse { +export interface EngineQueryResponse { columns: ColumnMetadata[]; data: any[][]; } @@ -66,19 +67,30 @@ export interface ValidationResponse { message?: string; } +export interface DryPlanOption { + modelingOnly?: boolean; + manifest?: Manifest; +} + +export interface DuckDBPrepareOptions { + initSql: string; + sessionProps: Record; +} + export interface IWrenEngineAdaptor { deploy(deployData: deployData): Promise; - initDatabase(sql: string): Promise; + prepareDuckDB(options: DuckDBPrepareOptions): Promise; + listTables(): Promise; putSessionProps(props: Record): Promise; - queryDuckdb(sql: string): Promise; + queryDuckdb(sql: string): Promise; patchConfig(config: Record): Promise; previewData( sql: string, limit?: number, mdl?: Manifest, - ): Promise; + ): Promise; describeStatement(sql: string): Promise; - getNativeSQL(sql: string): Promise; + getNativeSQL(sql: string, options?: DryPlanOption): Promise; validateColumnIsValid( manifest: Manifest, modelName: string, @@ -139,6 +151,21 @@ export class WrenEngineAdaptor implements IWrenEngineAdaptor { } } + public async prepareDuckDB(options: DuckDBPrepareOptions): Promise { + const { initSql, sessionProps } = options; + await this.initDatabase(initSql); + await this.putSessionProps(sessionProps); + } + + public async listTables() { + const sql = + 'SELECT \ + table_catalog, table_schema, table_name, column_name, ordinal_position, is_nullable, data_type\ + FROM INFORMATION_SCHEMA.COLUMNS;'; + const response = await this.queryDuckdb(sql); + return this.formatToCompactTable(response); + } + public async deploy(deployData: deployData): Promise { const { manifest, hash } = deployData; const deployPayload = { manifest, version: hash } as DeployPayload; @@ -166,24 +193,6 @@ export class WrenEngineAdaptor implements IWrenEngineAdaptor { } } - public async initDatabase(sql) { - try { - const url = new URL(this.initSqlUrlPath, this.wrenEngineBaseEndpoint); - logger.debug(`Endpoint: ${url.href}`); - const headers = { - 'Content-Type': 'text/plain; charset=utf-8', - }; - await axios.put(url.href, sql, { headers }); - } catch (err: any) { - logger.debug(`Got error when init database: ${err}`); - throw Errors.create(Errors.GeneralErrorCodes.INIT_SQL_ERROR, { - customMessage: - Errors.errorMessages[Errors.GeneralErrorCodes.INIT_SQL_ERROR], - originalError: err, - }); - } - } - public async putSessionProps(props: Record) { const setSessionStatements = Object.entries(props) .map(([key, value]) => { @@ -210,14 +219,14 @@ export class WrenEngineAdaptor implements IWrenEngineAdaptor { } } - public async queryDuckdb(sql: string): Promise { + public async queryDuckdb(sql: string): Promise { try { const url = new URL(this.queryDuckdbUrlPath, this.wrenEngineBaseEndpoint); const headers = { 'Content-Type': 'text/plain; charset=utf-8', }; const res = await axios.post(url.href, sql, { headers }); - return res.data as QueryResponse; + return res.data as EngineQueryResponse; } catch (err: any) { logger.debug(`Got error when querying duckdb: ${err.message}`); throw err; @@ -247,14 +256,14 @@ export class WrenEngineAdaptor implements IWrenEngineAdaptor { sql: string, limit: number = DEFAULT_PREVIEW_LIMIT, manifest?: Manifest, - ): Promise { + ): Promise { try { const url = new URL(this.previewUrlPath, this.wrenEngineBaseEndpoint); const headers = { 'Content-Type': 'application/json', }; - const res: AxiosResponse = await axios({ + const res: AxiosResponse = await axios({ method: 'get', url: url.href, headers, @@ -285,8 +294,16 @@ export class WrenEngineAdaptor implements IWrenEngineAdaptor { } } - public async getNativeSQL(sql: string): Promise { + public async getNativeSQL( + sql: string, + options: DryPlanOption, + ): Promise { try { + const props = { + modelingOnly: options?.modelingOnly ? true : false, + manifest: options?.manifest, + }; + const url = new URL(this.dryPlanUrlPath, this.wrenEngineBaseEndpoint); const headers = { 'Content-Type': 'application/json' }; @@ -296,7 +313,7 @@ export class WrenEngineAdaptor implements IWrenEngineAdaptor { headers, data: { sql, - modelingOnly: false, + ...props, }, }); @@ -320,4 +337,61 @@ export class WrenEngineAdaptor implements IWrenEngineAdaptor { throw err; } } + + private async initDatabase(sql) { + try { + const url = new URL(this.initSqlUrlPath, this.wrenEngineBaseEndpoint); + logger.debug(`Endpoint: ${url.href}`); + const headers = { + 'Content-Type': 'text/plain; charset=utf-8', + }; + await axios.put(url.href, sql, { headers }); + } catch (err: any) { + logger.debug(`Got error when init database: ${err}`); + throw Errors.create(Errors.GeneralErrorCodes.INIT_SQL_ERROR, { + customMessage: + Errors.errorMessages[Errors.GeneralErrorCodes.INIT_SQL_ERROR], + originalError: err, + }); + } + } + + private formatToCompactTable(columns: EngineQueryResponse): CompactTable[] { + return columns.data.reduce((acc: CompactTable[], row: any) => { + const [ + table_catalog, + table_schema, + table_name, + column_name, + _ordinal_position, + is_nullable, + data_type, + ] = row; + let table = acc.find( + (t) => t.name === table_name && t.properties.schema === table_schema, + ); + if (!table) { + table = { + name: table_name, + description: '', + columns: [], + properties: { + schema: table_schema, + catalog: table_catalog, + table: table_name, + }, + primaryKey: null, + }; + acc.push(table); + } + table.columns.push({ + name: column_name, + type: data_type, + notNull: is_nullable.toLocaleLowerCase() !== 'yes', + description: '', + properties: {}, + }); + return acc; + }, []); + } } diff --git a/wren-ui/src/apollo/server/config.ts b/wren-ui/src/apollo/server/config.ts index 81ef9b07d..b1ddcde1a 100644 --- a/wren-ui/src/apollo/server/config.ts +++ b/wren-ui/src/apollo/server/config.ts @@ -1,6 +1,9 @@ import { pickBy } from 'lodash'; export interface IConfig { + // wren ui + otherServiceUsingDocker: boolean; + // database dbType: string; // pg @@ -18,6 +21,9 @@ export interface IConfig { wrenAIEndpoint: string; openaiGenerationModel?: string; + // ibis server + ibisServerEndpoint: string; + // encryption encryptionPassword: string; encryptionSalt: string; @@ -40,6 +46,9 @@ export interface IConfig { } const defaultConfig = { + // wren ui + otherServiceUsingDocker: false, + // database dbType: 'sqlite', @@ -58,12 +67,18 @@ const defaultConfig = { // wren AI wrenAIEndpoint: 'http://localhost:5555', + // ibis server + ibisServerEndpoint: 'http://localhost:8000', + // encryption encryptionPassword: 'sementic', encryptionSalt: 'layer', }; const config = { + // node + otherServiceUsingDocker: process.env.OTHER_SERVICE_USING_DOCKER === 'true', + // database dbType: process.env.DB_TYPE, // pg @@ -89,6 +104,9 @@ const config = { wrenAIEndpoint: process.env.WREN_AI_ENDPOINT, openaiGenerationModel: process.env.OPENAI_GENERATION_MODEL, + // ibis server + ibisServerEndpoint: process.env.IBIS_SERVER_ENDPOINT, + // encryption encryptionPassword: process.env.ENCRYPTION_PASSWORD, encryptionSalt: process.env.ENCRYPTION_SALT, diff --git a/wren-ui/src/apollo/server/connectors/bqConnector.ts b/wren-ui/src/apollo/server/connectors/bqConnector.ts deleted file mode 100644 index 9a88036f5..000000000 --- a/wren-ui/src/apollo/server/connectors/bqConnector.ts +++ /dev/null @@ -1,166 +0,0 @@ -import { CompactTable } from './connector'; -import { IConnector } from './connector'; -import { BigQuery, BigQueryOptions } from '@google-cloud/bigquery'; -import { getLogger } from '@server/utils'; -import * as Errors from '@server/utils/error'; - -const logger = getLogger('DuckDBConnector'); -logger.level = 'debug'; - -// column type ref: https://cloud.google.com/bigquery/docs/information-schema-columns#schema -export interface BQColumnResponse { - table_catalog: string; - table_schema: string; - table_name: string; - column_name: string; - ordinal_position: number; - is_nullable: string; - data_type: string; - is_generated: string; - generation_expression: any; - is_stored: string; - is_hidden: string; - is_updatable: string; - is_system_defined: string; - is_partitioning_column: string; - clustering_ordinal_position: number; - collation_name: string; - column_default: string; - rounding_mode: string; - column_description: string; - table_description: string; -} - -export interface BQConstraintResponse { - constraintName: string; - constraintType: string; - constraintTable: string; - constraintColumn: string; - constraintedTable: string; - constraintedColumn: string; -} - -export interface BQListTableFilter { - tableName: string; -} -export interface BQListTableOptions { - datasetId: string; - format?: boolean; - filter?: BQListTableFilter; -} -export class BQConnector - implements IConnector -{ - private bq: BigQuery; - - // Not storing the bq client instance because we rarely need to use it - constructor(bqOptions: BigQueryOptions) { - this.bq = new BigQuery(bqOptions); - } - - public async prepare() { - return; - } - - public async connect() { - try { - await this.bq.query('SELECT 1;'); - return true; - } catch (err) { - logger.error(`Error connecting to BigQuery: ${err}`); - throw Errors.create(Errors.GeneralErrorCodes.CONNECTION_ERROR, { - originalError: err, - }); - } - } - - public async listTables(listTableOptions: BQListTableOptions) { - const { datasetId, format, filter } = listTableOptions; - // AND cf.column_name = cf.field_path => filter out the subfield in record - let sql = `SELECT - c.*, - cf.description AS column_description, - table_options.option_value AS table_description - FROM ${datasetId}.INFORMATION_SCHEMA.COLUMNS c - JOIN ${datasetId}.INFORMATION_SCHEMA.COLUMN_FIELD_PATHS cf - ON cf.table_name = c.table_name - AND cf.column_name = c.column_name - LEFT JOIN ${datasetId}.INFORMATION_SCHEMA.TABLE_OPTIONS table_options - ON c.table_name = table_options.table_name - WHERE - NOT REGEXP_CONTAINS(cf.data_type, r'^(STRUCT|ARRAY { - this.bq.query(sql, (err, rows) => { - if (err) { - reject(err); - } else { - resolve(rows); - } - }); - }); - - if (!format) { - return columns as BQColumnResponse[]; - } - return this.formatToCompactTable(columns); - } - - public async listConstraints(options) { - const { datasetId } = options; - // ref: information schema link: https://cloud.google.com/bigquery/docs/information-schema-intro - const constraints = await new Promise((resolve, reject) => { - this.bq.query( - ` - SELECT - ccu.table_name as constraintTable, ccu.column_name constraintColumn, - kcu.table_name as constraintedTable, kcu.column_name as constraintedColumn, - tc.constraint_type as constraintType - FROM ${datasetId}.INFORMATION_SCHEMA.CONSTRAINT_COLUMN_USAGE ccu - JOIN ${datasetId}.INFORMATION_SCHEMA.KEY_COLUMN_USAGE kcu - ON ccu.constraint_name = kcu.constraint_name - JOIN ${datasetId}.INFORMATION_SCHEMA.TABLE_CONSTRAINTS tc - ON ccu.constraint_name = tc.constraint_name - WHERE tc.constraint_type = 'FOREIGN KEY' - `, - (err, rows) => { - if (err) { - reject(err); - } else { - resolve(rows); - } - }, - ); - }); - return constraints as BQConstraintResponse[]; - } - - private formatToCompactTable(columns: any): CompactTable[] { - return columns.reduce((acc: CompactTable[], row: any) => { - let table = acc.find((t) => t.name === row.table_name); - if (!table) { - table = { - name: row.table_name, - description: row.table_description, - columns: [], - }; - acc.push(table); - } - table.columns.push({ - name: row.column_name, - type: row.data_type, - notNull: row.is_nullable.toLocaleLowerCase() !== 'yes', - description: row.column_description, - }); - return acc; - }, []); - } -} diff --git a/wren-ui/src/apollo/server/connectors/connector.ts b/wren-ui/src/apollo/server/connectors/connector.ts deleted file mode 100644 index 6f562b1f0..000000000 --- a/wren-ui/src/apollo/server/connectors/connector.ts +++ /dev/null @@ -1,21 +0,0 @@ -export interface CompactColumn { - name: string; - type: string; - notNull: boolean; - description?: string; - properties?: Record; -} - -export interface CompactTable { - name: string; - columns: CompactColumn[]; - description?: string; - properties?: Record; -} - -export interface IConnector { - prepare(prepareOptions: any): Promise; - connect(): Promise; - listTables(listTableOptions: any): Promise; - listConstraints(listConstraintOptions: any): Promise<[] | C[]>; -} diff --git a/wren-ui/src/apollo/server/connectors/duckdbConnector.ts b/wren-ui/src/apollo/server/connectors/duckdbConnector.ts deleted file mode 100644 index 19b61fcde..000000000 --- a/wren-ui/src/apollo/server/connectors/duckdbConnector.ts +++ /dev/null @@ -1,114 +0,0 @@ -import { - IWrenEngineAdaptor, - QueryResponse, -} from '../adaptors/wrenEngineAdaptor'; -import { CompactTable } from './connector'; -import { IConnector } from './connector'; -import { getLogger } from '@server/utils'; -import * as Errors from '@server/utils/error'; - -const logger = getLogger('DuckDBConnector'); -logger.level = 'debug'; - -export interface DuckDBPrepareOptions { - initSql: string; - sessionProps: Record; -} - -export interface DuckDBListTableOptions { - format?: boolean; -} - -export interface DuckDBColumnResponse { - table_catalog: string; - table_schema: string; - table_name: string; - column_name: string; - ordinal_position: string; - is_nullable: string; - data_type: string; -} - -export class DuckDBConnector - implements IConnector -{ - private wrenEngineAdaptor: IWrenEngineAdaptor; - constructor({ - wrenEngineAdaptor, - }: { - wrenEngineAdaptor: IWrenEngineAdaptor; - }) { - this.wrenEngineAdaptor = wrenEngineAdaptor; - } - public async prepare(prepareOptions: DuckDBPrepareOptions): Promise { - const { initSql, sessionProps } = prepareOptions; - await this.wrenEngineAdaptor.initDatabase(initSql); - await this.wrenEngineAdaptor.putSessionProps(sessionProps); - } - - public async connect(): Promise { - const sql = 'SELECT 1;'; - try { - await this.wrenEngineAdaptor.queryDuckdb(sql); - return true; - } catch (err) { - logger.error(`Error connecting to DuckDB: ${err}`); - throw Errors.create(Errors.GeneralErrorCodes.CONNECTION_ERROR, { - originalError: err, - }); - } - } - - public async listTables(listTableOptions: DuckDBListTableOptions) { - const sql = - 'SELECT \ - table_catalog, table_schema, table_name, column_name, ordinal_position, is_nullable, data_type\ - FROM INFORMATION_SCHEMA.COLUMNS;'; - const response = await this.wrenEngineAdaptor.queryDuckdb(sql); - if (listTableOptions.format) { - return this.formatToCompactTable(response); - } - return response.data; - } - - public async listConstraints(): Promise { - return []; - } - - private formatToCompactTable(columns: QueryResponse): CompactTable[] { - return columns.data.reduce((acc: CompactTable[], row: any) => { - const [ - table_catalog, - table_schema, - table_name, - column_name, - _ordinal_position, - is_nullable, - data_type, - ] = row; - let table = acc.find( - (t) => t.name === table_name && t.properties.schema === table_schema, - ); - if (!table) { - table = { - name: table_name, - description: '', - columns: [], - properties: { - schema: table_schema, - catalog: table_catalog, - }, - }; - acc.push(table); - } - table.columns.push({ - name: column_name, - type: data_type, - notNull: is_nullable.toLocaleLowerCase() !== 'yes', - description: '', - properties: {}, - }); - return acc; - }, []); - } -} diff --git a/wren-ui/src/apollo/server/connectors/postgresConnector.ts b/wren-ui/src/apollo/server/connectors/postgresConnector.ts deleted file mode 100644 index de913c63b..000000000 --- a/wren-ui/src/apollo/server/connectors/postgresConnector.ts +++ /dev/null @@ -1,281 +0,0 @@ -import { CompactTable } from './connector'; -import { IConnector } from './connector'; -import { getLogger } from '@server/utils'; -import { WrenEngineColumnType } from './types'; -import * as Errors from '@server/utils/error'; - -import pg from 'pg'; -import { ConnectionOptions } from 'tls'; -const { Client } = pg; - -const logger = getLogger('PostgresConnector'); -logger.level = 'debug'; - -export interface PostgresConnectionConfig { - user: string; - password: string; - host: string; - database: string; - port: number; - ssl?: boolean | ConnectionOptions; -} - -export interface PostgresColumnResponse { - table_catalog: string; - table_schema: string; - table_name: string; - column_name: string; - ordinal_position: string; - is_nullable: string; - data_type: WrenEngineColumnType; -} - -export interface PostgresConstraintResponse { - constraintName: string; - constraintType: string; - constraintTable: string; - constraintColumn: string; - constraintedTable: string; - constraintedColumn: string; -} - -export interface PostgresListTableOptions { - format?: boolean; -} - -export class PostgresConnector - implements IConnector -{ - private config: PostgresConnectionConfig; - private client?: pg.Client; - - constructor(config: PostgresConnectionConfig) { - this.config = config; - } - - public async prepare() { - return; - } - - public async connect(): Promise { - try { - await this.prepareClient(); - // query to check if connection is successful - await this.client.query('SELECT 1;'); - return true; - } catch (err) { - logger.error(`Error connecting to Postgres: ${err}`); - const errCode = - err.code === 'ECONNREFUSED' - ? Errors.GeneralErrorCodes.CONNECTION_REFUSED - : Errors.GeneralErrorCodes.CONNECTION_ERROR; - throw Errors.create(errCode, { - originalError: err, - }); - } - } - - public async listTables(options: PostgresListTableOptions) { - const sql = ` - SELECT - t.table_catalog, - t.table_schema, - t.table_name, - c.column_name, - c.data_type, - c.is_nullable, - c.ordinal_position - FROM - information_schema.tables t - JOIN - information_schema.columns c ON t.table_schema = c.table_schema AND t.table_name = c.table_name - WHERE - t.table_type in ('BASE TABLE', 'VIEW') - and t.table_schema not in ('information_schema', 'pg_catalog') - ORDER BY - t.table_schema, - t.table_name, - c.ordinal_position; - `; - await this.prepareClient(); - const res = await this.client.query(sql); - const columns = res.rows.map((row) => { - return { - table_catalog: row.table_catalog, - table_schema: row.table_schema, - table_name: row.table_name, - column_name: row.column_name, - ordinal_position: row.ordinal_position, - is_nullable: row.is_nullable, - data_type: this.transformColumnType(row.data_type), - }; - }) as PostgresColumnResponse[]; - - return options.format ? this.formatToCompactTable(columns) : columns; - } - - public async listConstraints() { - const sql = ` - SELECT - tc.table_schema, - tc.constraint_name, - tc.table_name, - kcu.column_name, - ccu.table_schema AS foreign_table_schema, - ccu.table_name AS foreign_table_name, - ccu.column_name AS foreign_column_name - FROM information_schema.table_constraints AS tc - JOIN information_schema.key_column_usage AS kcu - ON tc.constraint_name = kcu.constraint_name - AND tc.table_schema = kcu.table_schema - JOIN information_schema.constraint_column_usage AS ccu - ON ccu.constraint_name = tc.constraint_name - WHERE tc.constraint_type = 'FOREIGN KEY' - `; - await this.prepareClient(); - const res = await this.client.query(sql); - const constraints = res.rows.map((row) => { - return { - constraintName: row.constraint_name, - constraintType: 'FOREIGN KEY', - constraintTable: this.formatCompactTableName( - row.table_name, - row.table_schema, - ), - constraintColumn: row.column_name, - constraintedTable: this.formatCompactTableName( - row.foreign_table_name, - row.foreign_table_schema, - ), - constraintedColumn: row.foreign_column_name, - }; - }) as PostgresConstraintResponse[]; - return constraints; - } - - private transformColumnType(dataType: string) { - // lower case the dataType - dataType = dataType.toLowerCase(); - - // all possible types listed here: https://www.postgresql.org/docs/current/datatype.html#DATATYPE-TABLE - - switch (dataType) { - case 'text': - return WrenEngineColumnType.TEXT; - case 'char': - case 'character': - case 'bpchar': - case 'name': - return WrenEngineColumnType.CHAR; - case 'character varying': - return WrenEngineColumnType.VARCHAR; - case 'bigint': - return WrenEngineColumnType.BIGINT; - case 'int': - case 'integer': - return WrenEngineColumnType.INTEGER; - case 'smallint': - return WrenEngineColumnType.SMALLINT; - case 'real': - return WrenEngineColumnType.REAL; - case 'double precision': - return WrenEngineColumnType.DOUBLE; - case 'numeric': - case 'decimal': - return WrenEngineColumnType.DECIMAL; - case 'boolean': - return WrenEngineColumnType.BOOLEAN; - case 'timestamp': - case 'timestamp without time zone': - return WrenEngineColumnType.TIMESTAMP; - case 'timestamp with time zone': - return WrenEngineColumnType.TIMESTAMPTZ; - case 'date': - return WrenEngineColumnType.DATE; - case 'interval': - return WrenEngineColumnType.INTERVAL; - case 'json': - return WrenEngineColumnType.JSON; - case 'bytea': - return WrenEngineColumnType.BYTEA; - case 'uuid': - return WrenEngineColumnType.UUID; - case 'inet': - return WrenEngineColumnType.INET; - case 'oid': - return WrenEngineColumnType.OID; - default: - return WrenEngineColumnType.UNKNOWN; - } - } - - private formatToCompactTable( - columns: PostgresColumnResponse[], - ): CompactTable[] { - return columns.reduce( - (acc: CompactTable[], row: PostgresColumnResponse) => { - const { - table_catalog, - table_schema, - table_name, - column_name, - is_nullable, - data_type, - } = row; - const tableName = this.formatCompactTableName(table_name, table_schema); - let table = acc.find((t) => t.name === tableName); - if (!table) { - table = { - name: tableName, - description: '', - columns: [], - properties: { - schema: table_schema, - catalog: table_catalog, - }, - }; - acc.push(table); - } - table.columns.push({ - name: column_name, - type: data_type, - notNull: is_nullable.toLocaleLowerCase() !== 'yes', - description: '', - properties: {}, - }); - return acc; - }, - [], - ); - } - - public async close() { - if (this.client) { - await this.client.end(); - this.client = null; - } - } - - public formatCompactTableName(tableName: string, schema: string) { - return `${schema}.${tableName}`; - } - - public parseCompactTableName(compactTableName: string) { - const [schema, tableName] = compactTableName.split('.'); - return { schema, tableName }; - } - - private async prepareClient() { - if (this.client) { - return; - } - - // bypass server certificate validation - if (this.config.ssl) { - this.config.ssl = { rejectUnauthorized: false }; - } - - this.client = new Client(this.config); - await this.client.connect(); - } -} diff --git a/wren-ui/src/apollo/server/connectors/types.ts b/wren-ui/src/apollo/server/connectors/types.ts deleted file mode 100644 index 22ce2105b..000000000 --- a/wren-ui/src/apollo/server/connectors/types.ts +++ /dev/null @@ -1,59 +0,0 @@ -export enum WrenEngineColumnType { - // Boolean Types - BOOLEAN = 'BOOLEAN', - - // Numeric Types - TINYINT = 'TINYINT', - - INT2 = 'INT2', - SMALLINT = 'SMALLINT', // alias for INT2 - - INT4 = 'INT4', - INTEGER = 'INTEGER', // alias for INT4 - - INT8 = 'INT8', - BIGINT = 'BIGINT', // alias for INT8 - - NUMERIC = 'NUMERIC', - DECIMAL = 'DECIMAL', - - // Floating-Point Types - FLOAT4 = 'FLOAT4', - REAL = 'REAL', // alias for FLOAT4 - - FLOAT8 = 'FLOAT8', - DOUBLE = 'DOUBLE', // alias for FLOAT8 - - // Character Types - VARCHAR = 'VARCHAR', - CHAR = 'CHAR', - BPCHAR = 'BPCHAR', // BPCHAR is fixed-length, blank padded string - TEXT = 'TEXT', // alias for VARCHAR - STRING = 'STRING', // alias for VARCHAR - NAME = 'NAME', // alias for VARCHAR - - // Date/Time Types - TIMESTAMP = 'TIMESTAMP', - TIMESTAMPTZ = 'TIMESTAMP WITH TIME ZONE', - DATE = 'DATE', - INTERVAL = 'INTERVAL', - - // JSON Types - JSON = 'JSON', - - // Object identifiers (OIDs) are used internally by PostgreSQL as primary keys for various system tables. - // https://www.postgresql.org/docs/current/datatype-oid.html - OID = 'OID', - - // Binary Data Types - BYTEA = 'BYTEA', - - // UUID Type - UUID = 'UUID', - - // Network Address Types - INET = 'INET', - - // Unknown Type - UNKNOWN = 'UNKNOWN', -} diff --git a/wren-ui/src/apollo/server/factories/bqStrategy.ts b/wren-ui/src/apollo/server/factories/bqStrategy.ts deleted file mode 100644 index f75d8c099..000000000 --- a/wren-ui/src/apollo/server/factories/bqStrategy.ts +++ /dev/null @@ -1,437 +0,0 @@ -import { BigQueryOptions } from '@google-cloud/bigquery'; -import { capitalize } from 'lodash'; -import { IConnector } from '../connectors/connector'; -import { Model, ModelColumn, Project } from '../repositories'; -import { - AnalysisRelationInfo, - BigQueryDataSourceProperties, - DataSourceName, - IContext, - RelationType, -} from '../types'; -import { - BQColumnResponse, - BQConnector, - BQListTableOptions, -} from '../connectors/bqConnector'; -import { Encryptor, toBase64, getLogger } from '../utils'; -import { IDataSourceStrategy } from './dataSourceStrategy'; -import { - findColumnsToUpdate, - updateModelPrimaryKey, - transformInvalidColumnName, -} from './util'; - -const logger = getLogger('BigQueryStrategy'); -logger.level = 'debug'; - -export class BigQueryStrategy implements IDataSourceStrategy { - connector: IConnector; - project: Project; - ctx: IContext; - - constructor({ ctx, project }: { ctx: IContext; project?: Project }) { - if (project) { - this.project = project; - } - this.ctx = ctx; - } - - public async createDataSource(properties: BigQueryDataSourceProperties) { - const { displayName, projectId, datasetId, credentials } = properties; - const { config } = this.ctx; - - await this.testConnection({ projectId, datasetId, credentials }); - - await this.patchConfigToWrenEngine({ projectId, credentials }); - - // save DataSource to database - const encryptor = new Encryptor(config); - const encryptedCredentials = encryptor.encrypt(credentials); - const project = await this.ctx.projectRepository.createOne({ - displayName, - schema: 'public', - catalog: 'wrenai', - type: DataSourceName.BIG_QUERY, - projectId, - datasetId, - credentials: encryptedCredentials, - }); - return project; - } - - public async updateDataSource( - properties: BigQueryDataSourceProperties, - ): Promise { - const { displayName, credentials: newCredentials } = properties; - const { config } = this.ctx; - const { - projectId, - datasetId, - credentials: oldEncryptedCredentials, - } = this.project; - - const encryptor = new Encryptor(config); - const oldCredentials = JSON.parse( - encryptor.decrypt(oldEncryptedCredentials), - ); - - const credentials = newCredentials || oldCredentials; - - await this.testConnection({ projectId, datasetId, credentials }); - - await this.patchConfigToWrenEngine({ projectId, credentials }); - - // update DataSource to database - const encryptedCredentials = encryptor.encrypt(credentials); - const project = await this.ctx.projectRepository.updateOne( - this.project.id, - { - displayName, - credentials: encryptedCredentials, - }, - ); - return project; - } - - public async listTable({ formatToCompactTable }) { - const connector = await this.getBQConnector(); - const listTableOptions = { - datasetId: this.project.datasetId, - format: formatToCompactTable, - } as BQListTableOptions; - const tables = await connector.listTables(listTableOptions); - return tables; - } - - public async saveModels(tables: string[]) { - const connector = await this.getBQConnector(); - const listTableOptions = { - datasetId: this.project.datasetId, - format: false, - } as BQListTableOptions; - const dataSourceColumns = (await connector.listTables( - listTableOptions, - )) as BQColumnResponse[]; - - const models = await this.createModels( - this.project, - tables, - dataSourceColumns, - ); - // create columns - const columns = await this.createAllColumns( - tables, - models, - dataSourceColumns as BQColumnResponse[], - ); - return { models, columns }; - } - - public async saveModel( - table: string, - columns: string[], - primaryKey?: string, - ) { - const connector = await this.getBQConnector(); - const listTableOptions = { - datasetId: this.project.datasetId, - format: false, - } as BQListTableOptions; - const dataSourceColumns = (await connector.listTables( - listTableOptions, - )) as BQColumnResponse[]; - - const models = await this.createModels( - this.project, - [table], - dataSourceColumns, - ); - const model = models[0]; - const modelColumns = await this.createColumns( - columns, - model, - dataSourceColumns, - primaryKey, - ); - return { model, columns: modelColumns }; - } - - public async updateModel( - model: Model, - columns: string[], - primaryKey: string, - ) { - // get current column in the database - const existingColumns = await this.ctx.modelColumnRepository.findAllBy({ - modelId: model.id, - }); - const connector = await this.getBQConnector(); - const listTableOptions = { - datasetId: this.project.datasetId, - format: false, - } as BQListTableOptions; - const dataSourceColumns = (await connector.listTables( - listTableOptions, - )) as BQColumnResponse[]; - - const { toDeleteColumnIds, toCreateColumns } = findColumnsToUpdate( - columns, - existingColumns, - ); - await updateModelPrimaryKey( - this.ctx.modelColumnRepository, - model.id, - primaryKey, - ); - if (toCreateColumns.length) { - await this.createColumns( - toCreateColumns, - model, - dataSourceColumns, - primaryKey, - ); - } - if (toDeleteColumnIds.length) { - await this.ctx.modelColumnRepository.deleteMany(toDeleteColumnIds); - } - } - - public async analysisRelation(models, columns) { - const connector = await this.getBQConnector(); - const listConstraintOptions = { - datasetId: this.project.datasetId, - }; - const constraints = await connector.listConstraints(listConstraintOptions); - const relations = []; - for (const constraint of constraints) { - const { - constraintTable, - constraintColumn, - constraintedTable, - constraintedColumn, - } = constraint; - // validate tables and columns exists in our models and model columns - const fromModel = models.find( - (m) => m.sourceTableName === constraintTable, - ); - const toModel = models.find( - (m) => m.sourceTableName === constraintedTable, - ); - if (!fromModel || !toModel) { - continue; - } - const fromColumn = columns.find( - (c) => - c.modelId === fromModel.id && c.sourceColumnName === constraintColumn, - ); - const toColumn = columns.find( - (c) => - c.modelId === toModel.id && c.sourceColumnName === constraintedColumn, - ); - if (!fromColumn || !toColumn) { - continue; - } - // create relation - const relation: AnalysisRelationInfo = { - // upper case the first letter of the sourceTableName - name: - capitalize(fromModel.sourceTableName) + - capitalize(toModel.sourceTableName), - fromModelId: fromModel.id, - fromModelReferenceName: fromModel.referenceName, - fromColumnId: fromColumn.id, - fromColumnReferenceName: fromColumn.referenceName, - toModelId: toModel.id, - toModelReferenceName: toModel.referenceName, - toColumnId: toColumn.id, - toColumnReferenceName: toColumn.referenceName, - // TODO: add join type - type: RelationType.ONE_TO_MANY, - }; - relations.push(relation); - } - return relations; - } - - private async testConnection(args: { - projectId: string; - datasetId: string; - credentials: JSON; - }) { - const { projectId, datasetId, credentials } = args; - const { config } = this.ctx; - - // check DataSource is valid and can connect to it - const filePath = this.ctx.projectService.writeCredentialFile( - credentials, - config.persistCredentialDir, - ); - - const connectionOption: BigQueryOptions = { - projectId, - keyFilename: filePath, - }; - const connector = new BQConnector(connectionOption); - await connector.prepare(); - - // check can connect to bigquery - const connected = await connector.connect(); - if (!connected) { - throw new Error('Can not connect to data source'); - } - // check this credential have permission can list dataset table - try { - await connector.listTables({ datasetId }); - } catch (_e) { - throw new Error('Can not list tables in dataset'); - } - - return true; - } - - private async patchConfigToWrenEngine(args: { - projectId: string; - credentials: JSON; - }) { - const { projectId, credentials } = args; - - // update wren-engine config - const wrenEngineConfig = { - 'wren.datasource.type': 'bigquery', - 'bigquery.project-id': projectId, - 'bigquery.credentials-key': toBase64(JSON.stringify(credentials)), - }; - await this.ctx.wrenEngineAdaptor.patchConfig(wrenEngineConfig); - } - - private async getBQConnector() { - const filePath = await this.ctx.projectService.getCredentialFilePath( - this.project, - ); - // fetch tables - const { projectId } = this.project; - const connectionOption: BigQueryOptions = { - projectId, - keyFilename: filePath, - }; - return new BQConnector(connectionOption); - } - - private async createModels( - project: Project, - tables: string[], - dataSourceColumns: BQColumnResponse[], - ) { - const projectId = this.project.id; - const tableDescriptionMap = dataSourceColumns - .filter((col) => col.table_description) - .reduce((acc, column) => { - acc[column.table_name] = column.table_description; - return acc; - }, {}); - const modelValues = tables.map((tableName) => { - const description = tableDescriptionMap[tableName]; - const properties = description ? JSON.stringify({ description }) : null; - const model = { - projectId, - displayName: tableName, //use table name as displayName, referenceName and tableName - referenceName: tableName, - sourceTableName: tableName, - refSql: `select * from "${project.datasetId}".${tableName}`, - cached: false, - refreshTime: null, - properties, - } as Partial; - return model; - }); - - const models = await this.ctx.modelRepository.createMany(modelValues); - return models; - } - - private async createColumns( - columns: string[], - model: Model, - dataSourceColumns: BQColumnResponse[], - primaryKey?: string, - ) { - const columnValues = columns.reduce((acc, columnName) => { - const tableColumn = dataSourceColumns.find( - (col) => col.column_name === columnName, - ); - if (!tableColumn) { - throw new Error(`Column not found: ${columnName}`); - } - const properties = tableColumn.column_description - ? JSON.stringify({ - description: tableColumn.column_description, - }) - : null; - const columnValue = { - modelId: model.id, - isCalculated: false, - displayName: columnName, - sourceColumnName: columnName, - referenceName: transformInvalidColumnName(columnName), - type: tableColumn?.data_type || 'string', - notNull: tableColumn.is_nullable.toLocaleLowerCase() !== 'yes', - isPk: primaryKey === columnName, - properties, - } as Partial; - acc.push(columnValue); - return acc; - }, []); - const modelColumns = - await this.ctx.modelColumnRepository.createMany(columnValues); - return modelColumns; - } - - private async createAllColumns( - tables: string[], - models: Model[], - dataSourceColumns: BQColumnResponse[], - ) { - const columnValues = tables.reduce((acc, tableName) => { - const modelId = models.find((m) => m.sourceTableName === tableName)?.id; - if (!modelId) { - throw new Error('Model not found'); - } - const tableColumns = dataSourceColumns.filter( - (col) => col.table_name === tableName, - ); - for (const tableColumn of tableColumns) { - const columnName = tableColumn.column_name; - const properties = tableColumn.column_description - ? JSON.stringify({ - description: tableColumn.column_description, - }) - : null; - const columnValue = { - modelId, - isCalculated: false, - displayName: columnName, - sourceColumnName: columnName, - referenceName: transformInvalidColumnName(columnName), - type: tableColumn?.data_type || 'string', - notNull: tableColumn.is_nullable.toLocaleLowerCase() !== 'yes', - isPk: false, - properties, - } as Partial; - acc.push(columnValue); - } - return acc; - }, []); - const batch = 100; - const columns = []; - for (let i = 0; i < columnValues.length; i += batch) { - logger.debug(`Creating columns: ${i} - ${i + batch}`); - const res = await this.ctx.modelColumnRepository.createMany( - columnValues.slice(i, i + batch), - ); - columns.push(...res); - } - - return columns; - } -} diff --git a/wren-ui/src/apollo/server/factories/dataSourceStrategy.ts b/wren-ui/src/apollo/server/factories/dataSourceStrategy.ts deleted file mode 100644 index 9ee7e2d6c..000000000 --- a/wren-ui/src/apollo/server/factories/dataSourceStrategy.ts +++ /dev/null @@ -1,38 +0,0 @@ -import { CompactTable } from '../connectors/connector'; -import { Model, ModelColumn } from '../repositories'; -import { AnalysisRelationInfo, DataSourceProperties } from '../types'; - -export interface IDataSourceStrategy { - createDataSource(properties: DataSourceProperties): Promise; - updateDataSource(properties: DataSourceProperties): Promise; - listTable({ - formatToCompactTable, - }: { - formatToCompactTable: boolean; - }): Promise; - - /** - * Save multiple models, all the supported column in the table will be created - * @param models : array of table names in the datasource - */ - saveModels(models: any): Promise; - /** - * Save single model, only the column specified in the columns array will be created - * @param table : source table name - * @param columns : source column name of the table - */ - saveModel( - table: string, - columns: string[], - primaryKey?: string, - ): Promise; - updateModel( - model: Model, - columns: string[], - primaryKey?: string, - ): Promise; - analysisRelation( - models: Model[], - columns: ModelColumn[], - ): Promise; -} diff --git a/wren-ui/src/apollo/server/factories/duckdbStrategy.ts b/wren-ui/src/apollo/server/factories/duckdbStrategy.ts deleted file mode 100644 index 90eb9d20a..000000000 --- a/wren-ui/src/apollo/server/factories/duckdbStrategy.ts +++ /dev/null @@ -1,335 +0,0 @@ -import { IConnector, CompactTable } from '../connectors/connector'; -import { Model, ModelColumn, Project } from '../repositories'; -import { DataSourceName, DuckDBDataSourceProperties, IContext } from '../types'; -import { - DuckDBConnector, - DuckDBListTableOptions, - DuckDBPrepareOptions, -} from '../connectors/duckdbConnector'; -import { trim } from '../utils'; -import { IDataSourceStrategy } from './dataSourceStrategy'; -import { - findColumnsToUpdate, - updateModelPrimaryKey, - transformInvalidColumnName, -} from './util'; -import { getLogger } from '@server/utils'; - -const logger = getLogger('DuckDBStrategy'); -logger.level = 'debug'; - -export class DuckDBStrategy implements IDataSourceStrategy { - connector: IConnector; - project: Project; - ctx: IContext; - - constructor({ ctx, project }: { ctx: IContext; project?: Project }) { - if (project) { - this.project = project; - } - this.ctx = ctx; - } - - public async createDataSource(properties: DuckDBDataSourceProperties) { - const { displayName, initSql, extensions, configurations } = properties; - const initSqlWithExtensions = this.concatInitSql(initSql, extensions); - - await this.testConnection({ - initSql: initSqlWithExtensions, - configurations, - }); - - await this.patchConfigToWrenEngine(); - - // save DataSource to database - const project = await this.ctx.projectRepository.createOne({ - displayName, - schema: 'public', - catalog: 'wrenai', - type: DataSourceName.DUCKDB, - initSql: trim(initSql), - configurations, - extensions, - }); - return project; - } - - public async updateDataSource( - properties: DuckDBDataSourceProperties, - ): Promise { - const { displayName, initSql, extensions, configurations } = properties; - const initSqlWithExtensions = this.concatInitSql(initSql, extensions); - - await this.testConnection({ - initSql: initSqlWithExtensions, - configurations, - }); - - await this.patchConfigToWrenEngine(); - - const project = await this.ctx.projectRepository.updateOne( - this.project.id, - { - displayName, - initSql: trim(initSql), - configurations, - extensions, - }, - ); - return project; - } - - public async listTable({ formatToCompactTable }) { - const connector = new DuckDBConnector({ - wrenEngineAdaptor: this.ctx.wrenEngineAdaptor, - }); - const listTableOptions = { - format: formatToCompactTable, - } as DuckDBListTableOptions; - const tables = (await connector.listTables( - listTableOptions, - )) as CompactTable[]; - return tables; - } - - public async saveModels(tables: string[]) { - const connector = new DuckDBConnector({ - wrenEngineAdaptor: this.ctx.wrenEngineAdaptor, - }); - const listTableOptions = { format: true } as DuckDBListTableOptions; - const dataSourceColumns = (await connector.listTables( - listTableOptions, - )) as CompactTable[]; - const models = await this.createModels( - tables, - dataSourceColumns as CompactTable[], - ); - // create columns - const columns = await this.createAllColumns( - tables, - models, - dataSourceColumns as CompactTable[], - ); - return { models, columns }; - } - - public async saveModel( - table: string, - columns: string[], - primaryKey?: string, - ) { - const connector = new DuckDBConnector({ - wrenEngineAdaptor: this.ctx.wrenEngineAdaptor, - }); - const listTableOptions = { format: true } as DuckDBListTableOptions; - const dataSourceColumns = (await connector.listTables( - listTableOptions, - )) as CompactTable[]; - const model = await this.createModels( - [table], - dataSourceColumns as CompactTable[], - ); - const modelColumns = await this.createColumns( - columns, - model[0], - dataSourceColumns, - primaryKey, - ); - return { model, columns: modelColumns }; - } - - public async updateModel( - model: Model, - columns: string[], - primaryKey?: string, - ) { - const connector = new DuckDBConnector({ - wrenEngineAdaptor: this.ctx.wrenEngineAdaptor, - }); - const listTableOptions = { format: true } as DuckDBListTableOptions; - const dataSourceColumns = (await connector.listTables( - listTableOptions, - )) as CompactTable[]; - const existingColumns = await this.ctx.modelColumnRepository.findAllBy({ - modelId: model.id, - }); - const { toDeleteColumnIds, toCreateColumns } = findColumnsToUpdate( - columns, - existingColumns, - ); - await updateModelPrimaryKey( - this.ctx.modelColumnRepository, - model.id, - primaryKey, - ); - if (toCreateColumns.length) { - await this.createColumns( - toCreateColumns, - model, - dataSourceColumns, - primaryKey, - ); - } - if (toDeleteColumnIds.length) { - await this.ctx.modelColumnRepository.deleteMany(toDeleteColumnIds); - } - } - - public async analysisRelation(_models, _columns) { - return []; - } - - private async testConnection(args: { - configurations: Record; - initSql: string; - }) { - const { initSql, configurations } = args; - - const connector = new DuckDBConnector({ - wrenEngineAdaptor: this.ctx.wrenEngineAdaptor, - }); - - // prepare duckdb environment in wren-engine - const prepareOption = { - sessionProps: configurations, - initSql, - } as DuckDBPrepareOptions; - await connector.prepare(prepareOption); - - // check DataSource is valid and can connect to it - const connected = await connector.connect(); - if (!connected) { - throw new Error('Can not connect to data source'); - } - // check can list dataset table - try { - await connector.listTables({ format: false }); - } catch (_e) { - throw new Error('Can not list tables in dataset'); - } - } - - private async patchConfigToWrenEngine() { - // update wren-engine config - const config = { - 'wren.datasource.type': 'duckdb', - }; - await this.ctx.wrenEngineAdaptor.patchConfig(config); - } - - private concatInitSql(initSql: string, extensions: string[]) { - const installExtensions = extensions - .map((ext) => `INSTALL ${ext};`) - .join('\n'); - return trim(`${installExtensions}\n${initSql}`); - } - - private async createModels(tables: string[], compactTables: CompactTable[]) { - const projectId = this.project.id; - - const modelValues = tables.map((tableName) => { - const compactTable = compactTables.find( - (table) => table.name === tableName, - ); - const properties = compactTable.properties - ? JSON.stringify(compactTable.properties) - : null; - const model = { - projectId, - displayName: tableName, //use table name as displayName, referenceName and tableName - referenceName: tableName, - sourceTableName: tableName, - refSql: `select * from ${compactTable.properties.schema}.${tableName}`, - cached: false, - refreshTime: null, - properties, - } as Partial; - return model; - }); - - const models = await this.ctx.modelRepository.createMany(modelValues); - return models; - } - - private async createColumns( - columns: string[], - model: Model, - compactTables: CompactTable[], - primaryKey?: string, - ) { - const columnValues = columns.reduce((acc, columnName) => { - const compactColumns = compactTables.find( - (table) => table.name === model.sourceTableName, - )?.columns; - if (!compactColumns) { - throw new Error('Table not found'); - } - const compactColumn = compactColumns.find( - (column) => column.name === columnName, - ); - if (!compactColumn) { - throw new Error('Column not found'); - } - const columnValue = { - modelId: model.id, - isCalculated: false, - displayName: columnName, - sourceColumnName: columnName, - referenceName: transformInvalidColumnName(columnName), - type: compactColumn.type || 'string', - notNull: compactColumn.notNull, - isPk: primaryKey === columnName, - properties: JSON.stringify(compactColumn.properties), - } as Partial; - acc.push(columnValue); - return acc; - }, []); - const res = await this.ctx.modelColumnRepository.createMany(columnValues); - return res; - } - - private async createAllColumns( - tables: string[], - models: Model[], - compactTables: CompactTable[], - ) { - const columnValues = tables.reduce((acc, tableName) => { - const modelId = models.find((m) => m.sourceTableName === tableName)?.id; - if (!modelId) { - throw new Error(`Model not found: ${tableName}`); - } - const compactColumns = compactTables.find( - (table) => table.name === tableName, - )?.columns; - if (!compactColumns) { - throw new Error('Table not found'); - } - for (const compactColumn of compactColumns) { - const columnName = compactColumn.name; - const columnValue = { - modelId, - isCalculated: false, - displayName: columnName, - sourceColumnName: columnName, - referenceName: transformInvalidColumnName(columnName), - type: compactColumn.type || 'string', - notNull: compactColumn.notNull, - isPk: false, - properties: JSON.stringify(compactColumn.properties), - } as Partial; - acc.push(columnValue); - } - return acc; - }, []); - let columns = []; - const batch = 100; - for (let i = 0; i < columnValues.length; i += batch) { - logger.debug(`Creating columns: ${i} - ${i + batch}`); - const columnValueChunk = columnValues.slice(i, i + batch); - const columnChunk = - await this.ctx.modelColumnRepository.createMany(columnValueChunk); - columns = columns.concat(columnChunk); - } - return columns; - } -} diff --git a/wren-ui/src/apollo/server/factories/onboardingFactory.ts b/wren-ui/src/apollo/server/factories/onboardingFactory.ts deleted file mode 100644 index 943e0131a..000000000 --- a/wren-ui/src/apollo/server/factories/onboardingFactory.ts +++ /dev/null @@ -1,20 +0,0 @@ -import { DataSourceName } from '../types'; -import { BigQueryStrategy } from './bqStrategy'; -import { IDataSourceStrategy } from './dataSourceStrategy'; -import { DuckDBStrategy } from './duckdbStrategy'; -import { PostgresStrategy } from './postgresStrategy'; - -export class DataSourceStrategyFactory { - static create(dataSourceType: string, options: any): IDataSourceStrategy { - switch (dataSourceType) { - case DataSourceName.BIG_QUERY: - return new BigQueryStrategy(options); - case DataSourceName.DUCKDB: - return new DuckDBStrategy(options); - case DataSourceName.POSTGRES: - return new PostgresStrategy(options); - default: - throw new Error(`Unsupported data source type: ${dataSourceType}`); - } - } -} diff --git a/wren-ui/src/apollo/server/factories/postgresStrategy.ts b/wren-ui/src/apollo/server/factories/postgresStrategy.ts deleted file mode 100644 index 192ab1730..000000000 --- a/wren-ui/src/apollo/server/factories/postgresStrategy.ts +++ /dev/null @@ -1,406 +0,0 @@ -import { capitalize } from 'lodash'; -import { IDataSourceStrategy } from './dataSourceStrategy'; -import { - AnalysisRelationInfo, - DataSourceName, - IContext, - RelationType, - PGDataSourceProperties, -} from '../types'; -import { Model, ModelColumn, Project } from '../repositories'; -import { - PostgresColumnResponse, - PostgresConnector, -} from '../connectors/postgresConnector'; -import { Encryptor, getLogger } from '../utils'; -import { - findColumnsToUpdate, - updateModelPrimaryKey, - transformInvalidColumnName, -} from './util'; - -const logger = getLogger('PostgresStrategy'); -logger.level = 'debug'; - -export class PostgresStrategy implements IDataSourceStrategy { - private project?: Project; - private ctx: IContext; - - constructor({ ctx, project }: { ctx: IContext; project?: Project }) { - this.project = project; - this.ctx = ctx; - } - - public async createDataSource(properties: PGDataSourceProperties) { - const { displayName, host, port, database, user, password, ssl } = - properties; - - await this.testConnection(properties); - - await this.patchConfigToWrenEngine(properties); - - // save DataSource to database - const credentials = { password } as any; - const encryptor = new Encryptor(this.ctx.config); - const encryptedCredentials = encryptor.encrypt(credentials); - - const project = await this.ctx.projectRepository.createOne({ - displayName, - schema: 'public', - catalog: 'wrenai', - type: DataSourceName.POSTGRES, - host, - port, - database, - user, - credentials: encryptedCredentials, - configurations: { ssl }, - }); - return project; - } - - public async updateDataSource( - properties: PGDataSourceProperties, - ): Promise { - const { displayName, user, password: newPassword, ssl } = properties; - const { - host, - port, - database, - credentials: oldEncryptedCredentials, - } = this.project; - - const encryptor = new Encryptor(this.ctx.config); - const { password: oldPassword } = JSON.parse( - encryptor.decrypt(oldEncryptedCredentials), - ); - const password = newPassword || oldPassword; - - const newProperties = { - host, - port, - database, - user, - password, - ssl, - }; - - await this.testConnection(newProperties); - - await this.patchConfigToWrenEngine(newProperties); - - const credentials = { password } as any; - const encryptedCredentials = encryptor.encrypt(credentials); - const project = await this.ctx.projectRepository.updateOne( - this.project.id, - { - displayName, - user, - credentials: encryptedCredentials, - }, - ); - return project; - } - - public async listTable({ - formatToCompactTable, - }: { - formatToCompactTable: boolean; - }) { - const connector = this.getPGConnector(); - - // list tables - const listTableOptions = { - format: formatToCompactTable, - }; - - const tables = await connector.listTables(listTableOptions); - return tables; - } - - public async saveModels(tables: string[]) { - const connector = this.getPGConnector(); - const dataSourceColumns = (await connector.listTables({ - format: false, - })) as PostgresColumnResponse[]; - - const models = await this.createModels(tables, connector); - // create columns - const columns = await this.createAllColumns( - tables, - models, - dataSourceColumns as PostgresColumnResponse[], - connector, - ); - return { models, columns }; - } - - public async saveModel( - table: string, - columns: string[], - primaryKey?: string, - ) { - const connector = this.getPGConnector(); - const dataSourceColumns = (await connector.listTables({ - format: false, - })) as PostgresColumnResponse[]; - - const models = await this.createModels([table], connector); - const model = models[0]; - // create columns - const modelColumns = await this.createColumns( - columns, - model, - dataSourceColumns as PostgresColumnResponse[], - primaryKey, - ); - return { model, columns: modelColumns }; - } - - public async updateModel( - model: Model, - columns: string[], - primaryKey?: string, - ) { - const connector = this.getPGConnector(); - const dataSourceColumns = (await connector.listTables({ - format: false, - })) as PostgresColumnResponse[]; - const existingColumns = await this.ctx.modelColumnRepository.findAllBy({ - modelId: model.id, - }); - const { toDeleteColumnIds, toCreateColumns } = findColumnsToUpdate( - columns, - existingColumns, - ); - await updateModelPrimaryKey( - this.ctx.modelColumnRepository, - model.id, - primaryKey, - ); - if (toCreateColumns.length) { - await this.createColumns( - toCreateColumns, - model, - dataSourceColumns as PostgresColumnResponse[], - primaryKey, - ); - } - if (toDeleteColumnIds.length) { - await this.ctx.modelColumnRepository.deleteMany(toDeleteColumnIds); - } - } - - public async analysisRelation(models: Model[], columns: ModelColumn[]) { - const connector = this.getPGConnector(); - const constraints = await connector.listConstraints(); - const relations = []; - for (const constraint of constraints) { - const { - constraintTable, - constraintColumn, - constraintedTable, - constraintedColumn, - } = constraint; - // validate tables and columns exists in our models and model columns - const fromModel = models.find( - (m) => m.sourceTableName === constraintTable, - ); - const toModel = models.find( - (m) => m.sourceTableName === constraintedTable, - ); - if (!fromModel || !toModel) { - continue; - } - const fromColumn = columns.find( - (c) => - c.modelId === fromModel.id && c.sourceColumnName === constraintColumn, - ); - const toColumn = columns.find( - (c) => - c.modelId === toModel.id && c.sourceColumnName === constraintedColumn, - ); - if (!fromColumn || !toColumn) { - continue; - } - // create relation - const relation: AnalysisRelationInfo = { - // upper case the first letter of the sourceTableName - name: - capitalize(fromModel.sourceTableName) + - capitalize(toModel.sourceTableName), - fromModelId: fromModel.id, - fromModelReferenceName: fromModel.referenceName, - fromColumnId: fromColumn.id, - fromColumnReferenceName: fromColumn.referenceName, - toModelId: toModel.id, - toModelReferenceName: toModel.referenceName, - toColumnId: toColumn.id, - toColumnReferenceName: toColumn.referenceName, - // TODO: add join type - type: RelationType.ONE_TO_MANY, - }; - relations.push(relation); - } - return relations; - } - - private async testConnection(properties: any) { - const connector = new PostgresConnector(properties); - - // check DataSource is valid and can connect to it - await connector.connect(); - - // check can list dataset table - try { - await connector.listTables({ format: false }); - } catch (_e) { - throw new Error('Can not list tables in dataset'); - } - } - - private async patchConfigToWrenEngine(properties: any) { - const { host, port, database, user, password, ssl } = properties; - const sslMode = ssl ? '?sslmode=require' : ''; - // update wren-engine config - const jdbcUrl = `jdbc:postgresql://${host}:${port}/${database}${sslMode}`; - const config = { - 'wren.datasource.type': 'postgres', - 'postgres.jdbc.url': jdbcUrl, - 'postgres.user': user, - 'postgres.password': password, - }; - await this.ctx.wrenEngineAdaptor.patchConfig(config); - } - - private getPGConnector() { - // get credentials decrypted - const { credentials: encryptedCredentials } = this.project; - const encryptor = new Encryptor(this.ctx.config); - const credentials = JSON.parse(encryptor.decrypt(encryptedCredentials)); - - // connect to data source - const connector = new PostgresConnector({ - user: this.project.user, - password: credentials.password, - host: this.project.host, - database: this.project.database, - port: this.project.port, - ssl: this.project.configurations?.ssl, - }); - return connector; - } - - private async createModels(tables: string[], connector: PostgresConnector) { - const projectId = this.project.id; - const modelValues = tables.map((compactTableName) => { - const { schema, tableName } = - connector.parseCompactTableName(compactTableName); - // make referenceName = schema + _ + tableName - const referenceName = `${schema}_${tableName}`; - const model = { - projectId, - displayName: compactTableName, // use table name as displayName, referenceName and tableName - referenceName, - sourceTableName: compactTableName, - refSql: `select * from "${schema}"."${tableName}"`, - cached: false, - refreshTime: null, - } as Partial; - return model; - }); - - const models = await this.ctx.modelRepository.createMany(modelValues); - return models; - } - - private async createColumns( - columns: string[], - model: Model, - dataSourceColumns: PostgresColumnResponse[], - primaryKey?: string, - ) { - const columnValues = columns.reduce((acc, columnName) => { - const tableColumn = dataSourceColumns.find( - (col) => col.column_name === columnName, - ); - if (!tableColumn) { - throw new Error(`Column not found: ${columnName}`); - } - const columnValue = { - modelId: model.id, - isCalculated: false, - displayName: columnName, - sourceColumnName: columnName, - referenceName: transformInvalidColumnName(columnName), - type: tableColumn?.data_type || 'string', - notNull: tableColumn.is_nullable.toLocaleLowerCase() !== 'yes', - isPk: primaryKey === columnName, - } as Partial; - acc.push(columnValue); - return acc; - }, []); - const modelColumns = await Promise.all( - columnValues.map( - async (column) => - await this.ctx.modelColumnRepository.createOne(column), - ), - ); - return modelColumns; - } - - private async createAllColumns( - tables: string[], - models: Model[], - dataSourceColumns: PostgresColumnResponse[], - connector: PostgresConnector, - ) { - const columnValues = tables.reduce((acc, compactTableName) => { - // sourceTableName is the same as compactTableName when we create models - const modelId = models.find( - (m) => m.sourceTableName === compactTableName, - )?.id; - - if (!modelId) { - throw new Error('Model not found'); - } - - // get columns of the table - // format the table_name & table_schema of the column to match compactTableName - const tableColumns = dataSourceColumns.filter( - (col) => - connector.formatCompactTableName(col.table_name, col.table_schema) === - compactTableName, - ); - - // create column for each column in the table - // and add it to accumulated columnValues - // columnValues will be used to create columns in database - for (const tableColumn of tableColumns) { - const columnName = tableColumn.column_name; - const columnValue = { - modelId, - isCalculated: false, - displayName: columnName, - sourceColumnName: columnName, - referenceName: transformInvalidColumnName(columnName), - type: tableColumn?.data_type || 'string', - notNull: tableColumn.is_nullable.toLocaleLowerCase() !== 'yes', - isPk: false, - } as Partial; - acc.push(columnValue); - } - return acc; - }, []); - const batch = 100; - const columns = []; - for (let i = 0; i < columnValues.length; i += batch) { - logger.debug(`Creating columns: ${i} - ${i + batch}`); - const res = await this.ctx.modelColumnRepository.createMany( - columnValues.slice(i, i + batch), - ); - columns.push(...res); - } - return columns; - } -} diff --git a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts index 6b99edbd5..f313234b5 100644 --- a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts +++ b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts @@ -6,7 +6,7 @@ import { RelationInfo, View, } from '../repositories'; -import { Manifest, ModelMDL } from './type'; +import { Manifest, ModelMDL, TableReference } from './type'; import { getLogger } from '@server/utils'; const logger = getLogger('MDLBuilder'); @@ -93,14 +93,20 @@ export class MDLBuilder implements IMDLBuilder { if (model.displayName) { properties.displayName = model.displayName; } + const tableReference = this.buildTableReference(model); return { name: model.referenceName, columns: [], - refSql: model.refSql, + tableReference, + // can only have one of refSql or tableReference + refSql: tableReference ? null : model.refSql, cached: model.cached, refreshTime: model.refreshTime, - properties, + properties: { + displayName: model.displayName, + description: properties.description, + }, primaryKey: '', // will be modified in addColumn } as ModelMDL; }); @@ -386,4 +392,19 @@ export class MDLBuilder implements IMDLBuilder { relation; return `"${fromModelName}".${fromColumnName} = "${toModelName}".${toColumnName}`; } + + private buildTableReference(model: Model): TableReference | null { + const modelProps = + model.properties && typeof model.properties === 'string' + ? JSON.parse(model.properties) + : {}; + if (!modelProps.table) { + return null; + } + return { + catalog: modelProps.catalog || null, + schema: modelProps.schema || null, + table: modelProps.table, + }; + } } diff --git a/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts b/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts index 6b9ff1358..e51243fa5 100644 --- a/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts +++ b/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts @@ -1,9 +1,11 @@ +import { DataSourceName } from '@server/types'; import { Model, Project, ModelColumn, RelationInfo, View, + BIG_QUERY_CONNECTION_INFO, } from '../../repositories'; import { MDLBuilder, MDLBuilderBuildFromOptions } from '../mdlBuilder'; import { ModelMDL, RelationMDL, ViewMDL } from '../type'; @@ -32,11 +34,13 @@ describe('MDLBuilder', () => { // Arrange const project = { id: 1, - type: 'bigquery', + type: DataSourceName.BIG_QUERY, displayName: 'my project', - projectId: 'bq-project-id', - datasetId: 'my-dataset', - credentials: 'my-credential', + connectionInfo: { + projectId: 'bq-project-id', + datasetId: 'bq-project-id.my-dataset', + credentials: 'my-credential', + } as BIG_QUERY_CONNECTION_INFO, catalog: 'wrenai', schema: 'public', sampleDataset: null, @@ -51,7 +55,12 @@ describe('MDLBuilder', () => { refSql: 'SELECT * FROM order', cached: false, refreshTime: null, - properties: JSON.stringify({ description: 'foo table' }), + properties: JSON.stringify({ + description: 'foo table', + schema: 'my-dataset', + catalog: 'bq-project-id', + table: 'order', + }), }, { id: 2, @@ -62,7 +71,11 @@ describe('MDLBuilder', () => { refSql: 'SELECT * FROM customer', cached: false, refreshTime: null, - properties: null, + properties: JSON.stringify({ + schema: null, + catalog: null, + table: 'customer', + }), }, ] as Model[]; const columns = [ @@ -134,7 +147,12 @@ describe('MDLBuilder', () => { const expectedModels = [ { name: 'order', - refSql: 'SELECT * FROM order', + tableReference: { + schema: 'my-dataset', + catalog: 'bq-project-id', + table: 'order', + }, + refSql: null, columns: [ { name: 'orderKey', @@ -156,11 +174,19 @@ describe('MDLBuilder', () => { cached: false, refreshTime: null, primaryKey: 'orderKey', - properties: { description: 'foo table', displayName: 'order' }, + properties: { + description: 'foo table', + displayName: 'order', + }, }, { name: 'customer', - refSql: 'SELECT * FROM customer', + tableReference: { + schema: null, + catalog: null, + table: 'customer', + }, + refSql: null, columns: [ { name: 'orderKey', @@ -182,7 +208,9 @@ describe('MDLBuilder', () => { primaryKey: '', cached: false, refreshTime: null, - properties: { displayName: 'customer' }, + properties: { + displayName: 'customer', + }, }, ] as ModelMDL[]; @@ -207,11 +235,13 @@ describe('MDLBuilder', () => { // Arrange const project = { id: 1, - type: 'bigquery', + type: DataSourceName.BIG_QUERY, displayName: 'my project', - projectId: 'bq-project-id', - datasetId: 'my-dataset', - credentials: 'my-credential', + connectionInfo: { + projectId: 'bq-project-id', + datasetId: 'my-dataset', + credentials: 'my-credential', + } as BIG_QUERY_CONNECTION_INFO, catalog: 'wrenai', schema: 'public', sampleDataset: null, @@ -226,7 +256,12 @@ describe('MDLBuilder', () => { refSql: 'SELECT * FROM order', cached: false, refreshTime: null, - properties: JSON.stringify({ description: 'foo table' }), + properties: JSON.stringify({ + description: 'foo table', + catalog: 'bq-project-id', + schema: 'my-dataset', + table: 'order', + }), }, { id: 2, @@ -237,7 +272,11 @@ describe('MDLBuilder', () => { refSql: 'SELECT * FROM customer', cached: false, refreshTime: null, - properties: null, + properties: JSON.stringify({ + catalog: 'bq-project-id', + schema: 'my-dataset', + table: 'customer', + }), }, ] as Model[]; const columns = [ @@ -324,7 +363,12 @@ describe('MDLBuilder', () => { const expectedModels = [ { name: 'order', - refSql: 'SELECT * FROM order', + refSql: null, + tableReference: { + schema: 'my-dataset', + catalog: 'bq-project-id', + table: 'order', + }, columns: [ { name: 'orderKey', @@ -350,7 +394,12 @@ describe('MDLBuilder', () => { }, { name: 'customer', - refSql: 'SELECT * FROM customer', + refSql: null, + tableReference: { + schema: 'my-dataset', + catalog: 'bq-project-id', + table: 'customer', + }, columns: [ { name: 'orderKey', diff --git a/wren-ui/src/apollo/server/mdl/type.ts b/wren-ui/src/apollo/server/mdl/type.ts index 351c25824..53889d941 100644 --- a/wren-ui/src/apollo/server/mdl/type.ts +++ b/wren-ui/src/apollo/server/mdl/type.ts @@ -14,6 +14,7 @@ export interface ColumnMDL { export interface ModelMDL { name: string; // eg: "OrdersModel", "LineitemModel" refSql?: string; // eg: "select * from orders", "select * from lineitem" + tableReference?: TableReference; columns?: ColumnMDL[]; primaryKey?: string; // eg: "orderkey", "custkey" cached: boolean; // eg true, false @@ -74,3 +75,9 @@ export interface Manifest { enumDefinitions?: EnumDefinition[]; views?: ViewMDL[]; } + +export interface TableReference { + schema?: string; + catalog?: string; + table: string; +} diff --git a/wren-ui/src/apollo/server/models/model.ts b/wren-ui/src/apollo/server/models/model.ts index 82b30654d..f601b07f3 100644 --- a/wren-ui/src/apollo/server/models/model.ts +++ b/wren-ui/src/apollo/server/models/model.ts @@ -81,3 +81,10 @@ export interface CheckCalculatedFieldCanQueryData { expression: ExpressionName; lineage: number[]; } + +export interface PreviewSQLData { + sql: string; + projectId?: number; + limit?: number; + dryRun?: boolean; +} diff --git a/wren-ui/src/apollo/server/repositories/deployLogRepository.ts b/wren-ui/src/apollo/server/repositories/deployLogRepository.ts index 4c498bb72..6e2e0471f 100644 --- a/wren-ui/src/apollo/server/repositories/deployLogRepository.ts +++ b/wren-ui/src/apollo/server/repositories/deployLogRepository.ts @@ -1,5 +1,6 @@ import { Knex } from 'knex'; import { BaseRepository, IBasicRepository } from './baseRepository'; +import { camelCase, isPlainObject, mapKeys, mapValues } from 'lodash'; export interface Deploy { id: number; // ID @@ -55,4 +56,19 @@ export class DeployLogRepository .first(); return (res && this.transformFromDBData(res)) || null; } + + public override transformFromDBData: (data: any) => Deploy = (data: any) => { + if (!isPlainObject(data)) { + throw new Error('Unexpected dbdata'); + } + const camelCaseData = mapKeys(data, (_value, key) => camelCase(key)); + const formattedData = mapValues(camelCaseData, (value, key) => { + if (['manifest'].includes(key)) { + // sqlite return a string for json field, but postgres return an object + return typeof value === 'string' ? JSON.parse(value) : value; + } + return value; + }); + return formattedData as Deploy; + }; } diff --git a/wren-ui/src/apollo/server/repositories/projectRepository.ts b/wren-ui/src/apollo/server/repositories/projectRepository.ts index 97f5ad627..38e3ae0c5 100644 --- a/wren-ui/src/apollo/server/repositories/projectRepository.ts +++ b/wren-ui/src/apollo/server/repositories/projectRepository.ts @@ -8,31 +8,39 @@ import { snakeCase, isEmpty, } from 'lodash'; +import { DataSourceName } from '@/apollo/client/graphql/__types__'; + +export interface BIG_QUERY_CONNECTION_INFO { + projectId: string; + datasetId: string; + credentials: string; +} +export interface POSTGRES_CONNECTION_INFO { + host: string; + port: number; + user: string; + password: string; + database: string; + ssl: boolean; +} + +export interface DUCKDB_CONNECTION_INFO { + initSql: string; + extensions: Array; + configurations: Record; +} export interface Project { id: number; // ID - type: string; // Project datasource type. ex: bigquery, mysql, postgresql, mongodb, etc + type: DataSourceName; // Project datasource type. ex: bigquery, mysql, postgresql, mongodb, etc displayName: string; // Project display name - credentials: string; // Database credentials. General purpose field for storing credentials - configurations: Record; // Project connection configurations - - // bq - projectId: string; // BigQuery project id - datasetId: string; // BigQuery datasetId - - // duckdb - initSql: string; // DuckDB init sql - extensions: string[]; - - // pg - host: string; // Host - port: number; // Port - database: string; // Database - user: string; // User - catalog: string; // Catalog name schema: string; // Schema name sampleDataset: string; // Sample dataset name + connectionInfo: + | BIG_QUERY_CONNECTION_INFO + | POSTGRES_CONNECTION_INFO + | DUCKDB_CONNECTION_INFO; } export interface IProjectRepository extends IBasicRepository { @@ -60,16 +68,16 @@ export class ProjectRepository public override transformFromDBData: (data: any) => Project = (data: any) => { if (!isPlainObject(data)) { - throw new Error('Unexpected dbdata'); + throw new Error('Unexpected db data'); } const camelCaseData = mapKeys(data, (_value, key) => camelCase(key)); const formattedData = mapValues(camelCaseData, (value, key) => { - if (key === 'configurations') { + if (key === 'connectionInfo' && typeof value === 'string') { // should return {} if value is null / {}, use value ? {} : JSON.parse(value) will throw error when value is null return isEmpty(value) ? {} : JSON.parse(value); } - if (key === 'extensions') { - return isEmpty(value) ? [] : JSON.parse(value); + if (key === 'type') { + return DataSourceName[value]; } return value; }); @@ -80,11 +88,11 @@ export class ProjectRepository data: Project, ) => { if (!isPlainObject(data)) { - throw new Error('Unexpected dbdata'); + throw new Error('Unexpected db data'); } const snakeCaseData = mapKeys(data, (_value, key) => snakeCase(key)); const formattedData = mapValues(snakeCaseData, (value, key) => { - if (['configurations', 'extensions'].includes(key)) { + if (key === 'connectionInfo') { return JSON.stringify(value); } return value; diff --git a/wren-ui/src/apollo/server/resolvers.ts b/wren-ui/src/apollo/server/resolvers.ts index 07d90b28d..79ff3a2c9 100644 --- a/wren-ui/src/apollo/server/resolvers.ts +++ b/wren-ui/src/apollo/server/resolvers.ts @@ -44,6 +44,7 @@ const resolvers = { }, Mutation: { deploy: modelResolver.deploy, + getMDL: modelResolver.getMDL, saveDataSource: projectResolver.saveDataSource, startSampleDataset: projectResolver.startSampleDataset, saveTables: projectResolver.saveTables, @@ -86,6 +87,9 @@ const resolvers = { // Settings resetCurrentProject: projectResolver.resetCurrentProject, updateDataSource: projectResolver.updateDataSource, + + // preview + previewSql: modelResolver.previewSql, }, ThreadResponse: askingResolver.getThreadResponseNestedResolver(), DetailStep: askingResolver.getDetailStepNestedResolver(), diff --git a/wren-ui/src/apollo/server/resolvers/askingResolver.ts b/wren-ui/src/apollo/server/resolvers/askingResolver.ts index 2133c8072..d76ee928c 100644 --- a/wren-ui/src/apollo/server/resolvers/askingResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/askingResolver.ts @@ -164,7 +164,6 @@ export class AskingResolver { const askingService = ctx.askingService; const responses = await askingService.getResponsesWithThread(threadId); - // reduce responses to group by thread id const thread = reduce( responses, @@ -278,12 +277,12 @@ export class AskingResolver { public async previewData( _root: any, - args: { where: { responseId: number; stepIndex?: number } }, + args: { where: { responseId: number; stepIndex?: number; limit?: number } }, ctx: IContext, ): Promise { - const { responseId, stepIndex } = args.where; + const { responseId, stepIndex, limit } = args.where; const askingService = ctx.askingService; - const data = await askingService.previewData(responseId, stepIndex); + const data = await askingService.previewData(responseId, stepIndex, limit); return data; } @@ -292,22 +291,25 @@ export class AskingResolver { */ public getThreadResponseNestedResolver = () => ({ detail: async (parent: ThreadResponse, _args: any, ctx: IContext) => { + if (!parent.detail) { + return null; + } // extend view & sql to detail + + // handle sql + const sql = format(constructCteSql(parent.detail.steps)); + + // handle view + let view = null; const viewId = parent?.detail?.viewId; - if (!viewId) return parent.detail; - const view = viewId - ? await ctx.viewRepository.findOneBy({ id: viewId }) - : null; - const displayName = view.properties - ? JSON.parse(view.properties)?.displayName - : view.name; - return parent.detail - ? { - ...parent.detail, - sql: format(constructCteSql(parent.detail.steps)), - view: { ...view, displayName }, - } - : null; + if (viewId) { + view = await ctx.viewRepository.findOneBy({ id: viewId }); + const displayName = view.properties + ? JSON.parse(view.properties)?.displayName + : view.name; + view = { ...view, displayName }; + } + return { ...parent.detail, sql, view }; }, }); diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 86f2ae4fe..7bb2fc764 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -5,17 +5,23 @@ import { CreateCalculatedFieldData, UpdateCalculatedFieldData, UpdateViewMetadataInput, + PreviewSQLData, } from '../models'; import { IContext, RelationData, UpdateRelationData } from '../types'; -import { getLogger } from '@server/utils'; -import { CompactTable } from '../connectors/connector'; +import { getLogger, transformInvalidColumnName } from '@server/utils'; import { DeployResponse } from '../services/deployService'; import { constructCteSql } from '../services/askingService'; import { format } from 'sql-formatter'; import { isEmpty, isNil } from 'lodash'; -import { DataSourceStrategyFactory } from '../factories/onboardingFactory'; import { replaceAllowableSyntax, validateDisplayName } from '../utils/regex'; import * as Errors from '@server/utils/error'; +import { Model, ModelColumn } from '../repositories'; +import { + findColumnsToUpdate, + replaceInvalidReferenceName, + updateModelPrimaryKey, +} from '../utils/model'; +import { CompactTable, PreviewDataResponse } from '@server/services'; const logger = getLogger('ModelResolver'); logger.level = 'debug'; @@ -38,6 +44,7 @@ export class ModelResolver { this.deleteModel = this.deleteModel.bind(this); this.updateModelMetadata = this.updateModelMetadata.bind(this); this.deploy = this.deploy.bind(this); + this.getMDL = this.getMDL.bind(this); this.checkModelSync = this.checkModelSync.bind(this); // view @@ -51,6 +58,7 @@ export class ModelResolver { // preview this.previewModelData = this.previewModelData.bind(this); this.previewViewData = this.previewViewData.bind(this); + this.previewSql = this.previewSql.bind(this); this.getNativeSql = this.getNativeSql.bind(this); // calculated field @@ -135,14 +143,13 @@ export class ModelResolver { } public async checkModelSync(_root: any, _args: any, ctx: IContext) { - const project = await ctx.projectService.getCurrentProject(); + const { id } = await ctx.projectService.getCurrentProject(); const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); - const currentHash = ctx.deployService.createMDLHash(manifest); - const lastDeployHash = await ctx.deployService.getLastDeployment( - project.id, - ); + const currentHash = ctx.deployService.createMDLHash(manifest, id); + const lastDeploy = await ctx.deployService.getLastDeployment(id); + const lastDeployHash = lastDeploy.hash; const inProgressDeployment = - await ctx.deployService.getInProgressDeployment(project.id); + await ctx.deployService.getInProgressDeployment(id); if (inProgressDeployment) { return { status: SyncStatusEnum.IN_PROGRESS }; } @@ -156,14 +163,22 @@ export class ModelResolver { _args: any, ctx: IContext, ): Promise { - const project = await ctx.projectService.getCurrentProject(); + const { id } = await ctx.projectService.getCurrentProject(); const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); - return await ctx.deployService.deploy(manifest, project.id); + return await ctx.deployService.deploy(manifest, id); + } + + public async getMDL(_root: any, args: any, ctx: IContext) { + const { hash } = args.data; + const mdl = await ctx.deployService.getMDLByHash(hash); + return { + hash, + mdl, + }; } public async listModels(_root: any, _args: any, ctx: IContext) { - const project = await ctx.projectService.getCurrentProject(); - const projectId = project.id; + const { id: projectId } = await ctx.projectService.getCurrentProject(); const models = await ctx.modelRepository.findAllBy({ projectId }); const modelIds = models.map((m) => m.id); const modelColumnList = @@ -231,26 +246,51 @@ export class ModelResolver { const { sourceTableName, fields, primaryKey } = args.data; const project = await ctx.projectService.getCurrentProject(); - const dataSourceType = project.type; - const strategyOptions = { - ctx, - project, - }; - const strategy = DataSourceStrategyFactory.create( - dataSourceType, - strategyOptions, - ); - const dataSourceTables = await strategy.listTable({ - formatToCompactTable: true, - }); + const dataSourceTables = + await ctx.projectService.getProjectDataSourceTables(project); this.validateTableExist(sourceTableName, dataSourceTables); this.validateColumnsExist(sourceTableName, fields, dataSourceTables); - const { model, _columns } = await strategy.saveModel( - sourceTableName, - fields, - primaryKey, + // create model + const dataSourceTable = dataSourceTables.find( + (table) => table.name === sourceTableName, ); + if (!dataSourceTable) { + throw new Error('Table not found in the data source'); + } + const properties = dataSourceTable?.properties; + const modelValue = { + projectId: project.id, + displayName: sourceTableName, //use table name as displayName, referenceName and tableName + referenceName: replaceInvalidReferenceName(sourceTableName), + sourceTableName: sourceTableName, + cached: false, + refreshTime: null, + properties: properties ? JSON.stringify(properties) : null, + } as Partial; + const model = await ctx.modelRepository.createOne(modelValue); + + const columnValues = []; + fields.forEach((field) => { + const compactColumn = dataSourceTable.columns.find( + (c) => c.name === field, + ); + const columnValue = { + modelId: model.id, + isCalculated: false, + displayName: compactColumn.name, + referenceName: transformInvalidColumnName(compactColumn.name), + sourceColumnName: compactColumn.name, + type: compactColumn.type || 'string', + notNull: compactColumn.notNull || false, + isPk: primaryKey === field, + properties: compactColumn.properties + ? JSON.stringify(compactColumn.properties) + : null, + } as Partial; + columnValues.push(columnValue); + }); + await ctx.modelColumnRepository.createMany(columnValues); logger.info(`Model created: ${JSON.stringify(model)}`); return model; @@ -264,24 +304,53 @@ export class ModelResolver { const { fields, primaryKey } = args.data; const project = await ctx.projectService.getCurrentProject(); - const dataSourceType = project.type; - const strategyOptions = { - ctx, - project, - }; - const strategy = DataSourceStrategyFactory.create( - dataSourceType, - strategyOptions, - ); - const dataSourceTables = await strategy.listTable({ - formatToCompactTable: true, - }); + const dataSourceTables = + await ctx.projectService.getProjectDataSourceTables(project); const model = await ctx.modelRepository.findOneBy({ id: args.where.id }); + const existingColumns = await ctx.modelColumnRepository.findAllBy({ + modelId: model.id, + }); const { sourceTableName } = model; this.validateTableExist(sourceTableName, dataSourceTables); this.validateColumnsExist(sourceTableName, fields, dataSourceTables); - await strategy.updateModel(model, fields, primaryKey); + const sourceTableColumns = dataSourceTables.find( + (table) => table.name === sourceTableName, + )?.columns; + const { toDeleteColumnIds, toCreateColumns } = findColumnsToUpdate( + fields, + existingColumns, + ); + await updateModelPrimaryKey( + ctx.modelColumnRepository, + model.id, + primaryKey, + ); + if (toCreateColumns.length) { + const columnValues = toCreateColumns.map((columnName) => { + const sourceTableColumn = sourceTableColumns.find( + (col) => col.name === columnName, + ); + if (!sourceTableColumn) { + throw new Error(`Column not found: ${columnName}`); + } + const columnValue = { + modelId: model.id, + isCalculated: false, + displayName: columnName, + sourceColumnName: columnName, + referenceName: transformInvalidColumnName(columnName), + type: sourceTableColumn.type || 'string', + notNull: sourceTableColumn.notNull, + isPk: primaryKey === columnName, + } as Partial; + return columnValue; + }); + await ctx.modelColumnRepository.createMany(columnValues); + } + if (toDeleteColumnIds.length) { + await ctx.modelColumnRepository.deleteMany(toDeleteColumnIds); + } logger.info(`Model created: ${JSON.stringify(model)}`); return model; @@ -434,9 +503,14 @@ export class ModelResolver { // list views public async listViews(_root: any, _args: any, ctx: IContext) { - const project = await ctx.projectService.getCurrentProject(); - const views = await ctx.viewRepository.findAllBy({ projectId: project.id }); - return views; + const { id } = await ctx.projectService.getCurrentProject(); + const views = await ctx.viewRepository.findAllBy({ projectId: id }); + return views.map((view) => ({ + ...view, + displayName: view.properties + ? JSON.parse(view.properties)?.displayName + : view.name, + })); } public async getView(_root: any, args: any, ctx: IContext) { @@ -445,7 +519,10 @@ export class ModelResolver { if (!view) { throw new Error('View not found'); } - return view; + const displayName = view.properties + ? JSON.parse(view.properties)?.displayName + : view.name; + return { ...view, displayName }; } // validate a view name @@ -466,6 +543,8 @@ export class ModelResolver { // create view const project = await ctx.projectService.getCurrentProject(); + const deployment = await ctx.deployService.getLastDeployment(project.id); + const mdl = deployment.manifest; // get sql statement of a response const response = await ctx.askingService.getResponse(responseId); @@ -478,8 +557,13 @@ export class ModelResolver { const statement = format(constructCteSql(steps)); // describe columns - const { columns } = - await ctx.wrenEngineAdaptor.describeStatement(statement); + const { columns } = await ctx.queryService.describeStatement(statement, { + project, + limit: PREVIEW_MAX_OUTPUT_ROW, + modelingOnly: false, + mdl, + }); + if (isEmpty(columns)) { throw new Error('Failed to describe statement'); } @@ -520,7 +604,7 @@ export class ModelResolver { displayName, }); - return view; + return { ...view, displayName }; } // delete view @@ -540,32 +624,59 @@ export class ModelResolver { if (!model) { throw new Error('Model not found'); } - - // pass the current mdl to wren engine to preview data, prevent the model is not deployed + const project = await ctx.projectService.getCurrentProject(); const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); const sql = `select * from ${model.referenceName}`; - const data = await ctx.wrenEngineAdaptor.previewData( - sql, - PREVIEW_MAX_OUTPUT_ROW, - manifest, - ); + + const data = (await ctx.queryService.preview(sql, { + project, + limit: PREVIEW_MAX_OUTPUT_ROW, + modelingOnly: false, + mdl: manifest, + })) as PreviewDataResponse; + return data; } public async previewViewData(_root: any, args: any, ctx: IContext) { - const viewId = args.where.id; + const { id: viewId, limit } = args.where; const view = await ctx.viewRepository.findOneBy({ id: viewId }); if (!view) { throw new Error('View not found'); } + const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); + const project = await ctx.projectService.getCurrentProject(); - const data = await ctx.wrenEngineAdaptor.previewData( - view.statement, - PREVIEW_MAX_OUTPUT_ROW, - ); + const data = (await ctx.queryService.preview(view.statement, { + project, + limit: limit | PREVIEW_MAX_OUTPUT_ROW, + mdl: manifest, + modelingOnly: false, + })) as PreviewDataResponse; return data; } + public async previewSql( + _root: any, + args: { data: PreviewSQLData }, + ctx: IContext, + ) { + const { sql, projectId, limit, dryRun } = args.data; + 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 previewRes = await ctx.queryService.preview(sql, { + project, + limit: limit || PREVIEW_MAX_OUTPUT_ROW, + modelingOnly: false, + mdl, + dryRun, + }); + return dryRun ? { dryRun: previewRes } : previewRes; + } + public async getNativeSql( _root: any, args: { responseId: number }, @@ -574,8 +685,7 @@ export class ModelResolver { const { responseId } = args; // If using a sample dataset, native SQL is not supported - const project = await ctx.projectService.getCurrentProject(); - const sampleDataset = project.sampleDataset; + const { sampleDataset } = await ctx.projectService.getCurrentProject(); if (sampleDataset) { throw new Error(`Doesn't support Native SQL`); } @@ -677,8 +787,8 @@ export class ModelResolver { } const referenceName = replaceAllowableSyntax(viewDisplayName); // check if view name is duplicated - const project = await ctx.projectService.getCurrentProject(); - const views = await ctx.viewRepository.findAllBy({ projectId: project.id }); + const { id } = await ctx.projectService.getCurrentProject(); + const views = await ctx.viewRepository.findAllBy({ projectId: id }); if (views.find((v) => v.name === referenceName && v.id !== selfView)) { return { valid: false, @@ -691,8 +801,11 @@ export class ModelResolver { }; } - private validateTableExist(tableName: string, columns: CompactTable[]) { - if (!columns.find((c) => c.name === tableName)) { + private validateTableExist( + tableName: string, + dataSourceTables: CompactTable[], + ) { + if (!dataSourceTables.find((c) => c.name === tableName)) { throw new Error(`Table ${tableName} not found in the data Source`); } } @@ -700,9 +813,11 @@ export class ModelResolver { private validateColumnsExist( tableName: string, fields: string[], - columns: CompactTable[], + dataSourceTables: CompactTable[], ) { - const tableColumns = columns.find((c) => c.name === tableName)?.columns; + const tableColumns = dataSourceTables.find( + (c) => c.name === tableName, + )?.columns; for (const field of fields) { if (!tableColumns.find((c) => c.name === field)) { throw new Error( diff --git a/wren-ui/src/apollo/server/resolvers/projectResolver.ts b/wren-ui/src/apollo/server/resolvers/projectResolver.ts index b7e5c2ef5..8278a7de4 100644 --- a/wren-ui/src/apollo/server/resolvers/projectResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/projectResolver.ts @@ -1,13 +1,22 @@ import { + AnalysisRelationInfo, DataSource, DataSourceName, DataSourceProperties, IContext, RelationData, + RelationType, SampleDatasetData, } from '../types'; -import { getLogger } from '@server/utils'; -import { Model, ModelColumn, Project } from '../repositories'; +import { getLogger, replaceInvalidReferenceName, trim } from '@server/utils'; +import { + BIG_QUERY_CONNECTION_INFO, + DUCKDB_CONNECTION_INFO, + Model, + ModelColumn, + POSTGRES_CONNECTION_INFO, + Project, +} from '../repositories'; import { SampleDatasetName, SampleDatasetRelationship, @@ -16,7 +25,9 @@ import { sampleDatasets, } from '@server/data'; import { snakeCase } from 'lodash'; -import { DataSourceStrategyFactory } from '../factories/onboardingFactory'; +import { CompactTable, ProjectData } from '../services'; +import { replaceAllowableSyntax } from '../utils/regex'; +import { DuckDBPrepareOptions } from '@server/adaptors/wrenEngineAdaptor'; const logger = getLogger('DataSourceResolver'); logger.level = 'debug'; @@ -57,20 +68,21 @@ export class ProjectResolver { } public async resetCurrentProject(_root: any, _arg: any, ctx: IContext) { - let project: Project; + let id: number; try { - project = await ctx.projectService.getCurrentProject(); + const project = await ctx.projectService.getCurrentProject(); + id = project.id; } catch { // no project found return true; } - await ctx.deployService.deleteAllByProjectId(project.id); - await ctx.askingService.deleteAllByProjectId(project.id); - await ctx.modelService.deleteAllViewsByProjectId(project.id); - await ctx.modelService.deleteAllModelsByProjectId(project.id); + await ctx.deployService.deleteAllByProjectId(id); + await ctx.askingService.deleteAllByProjectId(id); + await ctx.modelService.deleteAllViewsByProjectId(id); + await ctx.modelService.deleteAllModelsByProjectId(id); - await ctx.projectService.deleteProject(project.id); + await ctx.projectService.deleteProject(id); return true; } @@ -175,8 +187,50 @@ export class ProjectResolver { // Currently only can create one project await this.resetCurrentProject(_root, args, ctx); - const strategy = DataSourceStrategyFactory.create(type, { ctx }); - const project = await strategy.createDataSource(properties); + const { displayName, ...connectionInfo } = properties; + const project = await ctx.projectService.createProject({ + displayName, + type, + connectionInfo, + } as ProjectData); + logger.debug(`Created project: ${JSON.stringify(project)}`); + // try to connect to the data source + try { + if (type === DataSourceName.DUCKDB) { + // handle duckdb connection + connectionInfo as DUCKDB_CONNECTION_INFO; + const { initSql, extensions } = connectionInfo; + const initSqlWithExtensions = this.concatInitSql(initSql, extensions); + + // prepare duckdb environment in wren-engine + await ctx.wrenEngineAdaptor.prepareDuckDB({ + sessionProps: connectionInfo.configurations, + initSql: initSqlWithExtensions, + } as DuckDBPrepareOptions); + + // check can list dataset table + await ctx.wrenEngineAdaptor.listTables(); + + // patch wren-engine config + const config = { + 'wren.datasource.type': 'duckdb', + }; + await ctx.wrenEngineAdaptor.patchConfig(config); + } else { + const tables = + await ctx.projectService.getProjectDataSourceTables(project); + logger.debug( + `Can connect to the data source, tables: ${JSON.stringify(tables[0])}...`, + ); + } + } catch (err) { + logger.error( + 'Failed to get project tables', + JSON.stringify(err, null, 2), + ); + await ctx.projectRepository.deleteOne(project.id); + throw err; + } // telemetry ctx.telemetry.send_event('save_data_source', { dataSourceType: type }); @@ -194,13 +248,54 @@ export class ProjectResolver { ctx: IContext, ) { const { properties } = args.data; + const { displayName, ...connectionInfo } = properties; const project = await ctx.projectService.getCurrentProject(); + const dataSourceType = project.type; - const strategy = DataSourceStrategyFactory.create(project.type, { - ctx, - project, + // only new connection info needed to encrypt + const toUpdateConnectionInfo = + ctx.projectService.encryptSensitiveConnectionInfo(connectionInfo as any); + + if (dataSourceType === DataSourceName.DUCKDB) { + // prepare duckdb environment in wren-engine + const { initSql, extensions } = toUpdateConnectionInfo; + const initSqlWithExtensions = this.concatInitSql(initSql, extensions); + await ctx.wrenEngineAdaptor.prepareDuckDB({ + sessionProps: toUpdateConnectionInfo.configurations, + initSql: initSqlWithExtensions, + } as DuckDBPrepareOptions); + + // check can list dataset table + try { + await ctx.wrenEngineAdaptor.listTables(); + } catch (_e) { + throw new Error('Can not list tables in dataset'); + } + + // patch wren-engine config + const config = { + 'wren.datasource.type': 'duckdb', + }; + await ctx.wrenEngineAdaptor.patchConfig(config); + } else { + const updatedProject = { + ...project, + displayName, + connectionInfo: { + ...project.connectionInfo, + ...toUpdateConnectionInfo, + }, + } as Project; + const tables = + await ctx.projectService.getProjectDataSourceTables(updatedProject); + logger.debug( + `Can connect to the data source, tables: ${JSON.stringify(tables[0])}...`, + ); + } + const nextProject = await ctx.projectRepository.updateOne(project.id, { + displayName, + connectionInfo: { ...project.connectionInfo, ...toUpdateConnectionInfo }, }); - const nextProject = await strategy.updateDataSource(properties); return { type: nextProject.type, properties: this.getDataSourceProperties(nextProject), @@ -208,13 +303,7 @@ export class ProjectResolver { } public async listDataSourceTables(_root: any, _arg, ctx: IContext) { - const project = await ctx.projectService.getCurrentProject(); - const dataSourceType = project.type; - const strategy = DataSourceStrategyFactory.create(dataSourceType, { - ctx, - project, - }); - return await strategy.listTable({ formatToCompactTable: true }); + return await ctx.projectService.getProjectDataSourceTables(); } public async saveTables( @@ -258,13 +347,58 @@ export class ProjectResolver { const modelIds = models.map((m) => m.id); const columns = await ctx.modelColumnRepository.findColumnsByModelIds(modelIds); + const constraints = + await ctx.projectService.getProjectSuggestedConstraint(project); // generate relation - const strategy = DataSourceStrategyFactory.create(project.type, { - project, - ctx, - }); - const relations = await strategy.analysisRelation(models, columns); + const relations = []; + for (const constraint of constraints) { + const { + constraintTable, + constraintColumn, + constraintedTable, + constraintedColumn, + } = constraint; + // validate tables and columns exists in our models and model columns + const fromModel = models.find( + (m) => m.sourceTableName === constraintTable, + ); + const toModel = models.find( + (m) => m.sourceTableName === constraintedTable, + ); + if (!fromModel || !toModel) { + continue; + } + const fromColumn = columns.find( + (c) => + c.modelId === fromModel.id && c.sourceColumnName === constraintColumn, + ); + const toColumn = columns.find( + (c) => + c.modelId === toModel.id && c.sourceColumnName === constraintedColumn, + ); + if (!fromColumn || !toColumn) { + continue; + } + // create relation + const relation: AnalysisRelationInfo = { + // upper case the first letter of the sourceTableName + name: constraint.constraintName, + fromModelId: fromModel.id, + fromModelReferenceName: fromModel.referenceName, + fromColumnId: fromColumn.id, + fromColumnReferenceName: fromColumn.referenceName, + toModelId: toModel.id, + toModelReferenceName: toModel.referenceName, + toColumnId: toColumn.id, + toColumnReferenceName: toColumn.referenceName, + // TODO: add join type + type: RelationType.ONE_TO_MANY, + }; + relations.push(relation); + } + logger.debug({ relations }); + // group by model return models.map(({ id, displayName, referenceName }) => ({ id, displayName, @@ -294,9 +428,9 @@ export class ProjectResolver { } private async deploy(ctx: IContext) { - const project = await ctx.projectService.getCurrentProject(); + const { id } = await ctx.projectService.getCurrentProject(); const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); - return await ctx.deployService.deploy(manifest, project.id); + return await ctx.deployService.deploy(manifest, id); } private buildRelationInput( @@ -353,11 +487,57 @@ export class ProjectResolver { await ctx.modelService.deleteAllModelsByProjectId(project.id); // create model and columns - const strategy = DataSourceStrategyFactory.create(project.type, { - ctx, - project, + const compactTables: CompactTable[] = + await ctx.projectService.getProjectDataSourceTables(project); + + const modelValues = tables.map((tableName) => { + const compactTable = compactTables.find( + (table) => table.name === tableName, + ); + if (!compactTable) { + throw new Error(`Table not found in data source: ${tableName}`); + } + const properties = compactTable?.properties; + // compactTable contain schema and catalog, these information are for building tableReference in mdl + const model = { + projectId: project.id, + displayName: tableName, //use table name as displayName, referenceName and tableName + referenceName: replaceInvalidReferenceName(tableName), + sourceTableName: tableName, + cached: false, + refreshTime: null, + properties: properties ? JSON.stringify(properties) : null, + } as Partial; + return model; }); - const { models, columns } = await strategy.saveModels(tables); + const models = await ctx.modelRepository.createMany(modelValues); + + const columnValues = []; + tables.forEach((tableName) => { + const compactTable = compactTables.find( + (table) => table.name === tableName, + ); + const compactColumns = compactTable.columns; + const primaryKey = compactTable.primaryKey; + const model = models.find((m) => m.sourceTableName === compactTable.name); + compactColumns.forEach((column) => { + const columnValue = { + modelId: model.id, + isCalculated: false, + displayName: column.name, + referenceName: this.transformInvalidColumnName(column.name), + sourceColumnName: column.name, + type: column.type || 'string', + notNull: column.notNull || false, + isPk: primaryKey === column.name, + properties: column.properties + ? JSON.stringify(column.properties) + : null, + } as Partial; + columnValues.push(columnValue); + }); + }); + const columns = await ctx.modelColumnRepository.createMany(columnValues); return { models, columns }; } @@ -369,20 +549,43 @@ export class ProjectResolver { } as DataSourceProperties; if (dataSourceType === DataSourceName.BIG_QUERY) { - properties.projectId = project.projectId; - properties.datasetId = project.datasetId; + const { projectId, datasetId } = + project.connectionInfo as BIG_QUERY_CONNECTION_INFO; + properties.projectId = projectId; + properties.datasetId = datasetId; } else if (dataSourceType === DataSourceName.DUCKDB) { - properties.initSql = project.initSql; - properties.extensions = project.extensions; - properties.configurations = project.configurations; + const { initSql, extensions, configurations } = + project.connectionInfo as DUCKDB_CONNECTION_INFO; + properties.initSql = initSql; + properties.extensions = extensions; + properties.configurations = configurations; } else if (dataSourceType === DataSourceName.POSTGRES) { - properties.host = project.host; - properties.port = project.port; - properties.database = project.database; - properties.user = project.user; - properties.ssl = project.configurations?.ssl; + const { host, port, database, user, ssl } = + project.connectionInfo as POSTGRES_CONNECTION_INFO; + properties.host = host; + properties.port = port; + properties.database = database; + properties.user = user; + properties.ssl = ssl; } return properties; } + + private transformInvalidColumnName(columnName: string) { + let referenceName = replaceAllowableSyntax(columnName); + // If the reference name does not start with a letter, add a prefix + const startWithLetterRegex = /^[A-Za-z]/; + if (!startWithLetterRegex.test(referenceName)) { + referenceName = `col_${referenceName}`; + } + return referenceName; + } + + private concatInitSql(initSql: string, extensions: string[]) { + const installExtensions = extensions + .map((ext) => `INSTALL ${ext};`) + .join('\n'); + return trim(`${installExtensions}\n${initSql}`); + } } diff --git a/wren-ui/src/apollo/server/schema.ts b/wren-ui/src/apollo/server/schema.ts index 940656c5b..d480c7a01 100644 --- a/wren-ui/src/apollo/server/schema.ts +++ b/wren-ui/src/apollo/server/schema.ts @@ -306,6 +306,13 @@ export const typeDefs = gql` id: Int! } + input PreviewViewDataInput { + id: Int! + # It will return default 100 rows if not specified limit + # refer: PREVIEW_MAX_OUTPUT_ROW + limit: Int + } + input CreateViewInput { name: String! responseId: Int! @@ -361,7 +368,7 @@ export const typeDefs = gql` displayName: String! referenceName: String! sourceTableName: String! - refSql: String! + refSql: String cached: Boolean! refreshTime: String description: String @@ -515,6 +522,9 @@ export const typeDefs = gql` responseId: Int! # Optional, only used for preview data of a single step stepIndex: Int + # It will return default 500 rows if not specified limit + # refer: DEFAULT_PREVIEW_LIMIT + limit: Int } type DetailStep { @@ -588,6 +598,22 @@ export const typeDefs = gql` dataSource: DataSource! } + input GetMDLInput { + hash: String! + } + + type GetMDLResult { + hash: String! + mdl: String + } + + input PreviewSQLDataInput { + sql: String! + projectId: Int + limit: Int + dryRun: Boolean + } + # Query and Mutation type Query { # On Boarding Steps @@ -627,6 +653,7 @@ export const typeDefs = gql` saveTables(data: SaveTablesInput!): JSON! saveRelations(data: SaveRelationInput!): JSON! deploy: JSON! + getMDL(data: GetMDLInput): GetMDLResult! # Modeling Page createModel(data: CreateModelInput!): JSON! @@ -663,7 +690,7 @@ export const typeDefs = gql` # View createView(data: CreateViewInput!): ViewInfo! deleteView(where: ViewWhereUniqueInput!): Boolean! - previewViewData(where: ViewWhereUniqueInput!): JSON! + previewViewData(where: PreviewViewDataInput!): JSON! validateView(data: ValidateViewInput!): ViewValidationResponse! # Ask @@ -688,5 +715,8 @@ export const typeDefs = gql` # Settings resetCurrentProject: Boolean! updateDataSource(data: UpdateDataSourceInput!): DataSource! + + # preview + previewSql(data: PreviewSQLDataInput): JSON! } `; diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index 4993d10d3..a9f40db5f 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -14,13 +14,10 @@ import { } from '../repositories/threadResponseRepository'; import { getLogger } from '@server/utils'; import { isEmpty, isNil } from 'lodash'; -import { - IWrenEngineAdaptor, - QueryResponse, -} from '../adaptors/wrenEngineAdaptor'; import { format } from 'sql-formatter'; import { Telemetry } from '../telemetry/telemetry'; import { IViewRepository, View } from '../repositories'; +import { IQueryService, PreviewDataResponse } from './queryService'; const logger = getLogger('AskingService'); logger.level = 'debug'; @@ -69,7 +66,11 @@ export interface IAskingService { threadId: number, ): Promise; getResponse(responseId: number): Promise; - previewData(responseId: number, stepIndex?: number): Promise; + previewData( + responseId: number, + stepIndex?: number, + limit?: number, + ): Promise; deleteAllByProjectId(projectId: number): Promise; } @@ -233,42 +234,42 @@ class BackgroundTracker { export class AskingService implements IAskingService { private wrenAIAdaptor: IWrenAIAdaptor; - private wrenEngineAdaptor: IWrenEngineAdaptor; private deployService: IDeployService; private projectService: IProjectService; private viewRepository: IViewRepository; private threadRepository: IThreadRepository; private threadResponseRepository: IThreadResponseRepository; private backgroundTracker: BackgroundTracker; + private queryService: IQueryService; private telemetry: Telemetry; constructor({ telemetry, wrenAIAdaptor, - wrenEngineAdaptor, deployService, projectService, viewRepository, threadRepository, threadResponseRepository, + queryService, }: { telemetry: Telemetry; wrenAIAdaptor: IWrenAIAdaptor; - wrenEngineAdaptor: IWrenEngineAdaptor; deployService: IDeployService; projectService: IProjectService; viewRepository: IViewRepository; threadRepository: IThreadRepository; threadResponseRepository: IThreadResponseRepository; + queryService: IQueryService; }) { this.wrenAIAdaptor = wrenAIAdaptor; - this.wrenEngineAdaptor = wrenEngineAdaptor; this.deployService = deployService; this.projectService = projectService; this.viewRepository = viewRepository; this.threadRepository = threadRepository; this.threadResponseRepository = threadResponseRepository; this.telemetry = telemetry; + this.queryService = queryService; this.backgroundTracker = new BackgroundTracker({ telemetry, wrenAIAdaptor, @@ -346,9 +347,9 @@ export class AskingService implements IAskingService { }); // 2. create a thread and the first thread response - const project = await this.projectService.getCurrentProject(); + const { id } = await this.projectService.getCurrentProject(); const thread = await this.threadRepository.createOne({ - projectId: project.id, + projectId: id, sql: input.sql, summary: input.summary, }); @@ -370,8 +371,8 @@ export class AskingService implements IAskingService { } public async listThreads(): Promise { - const project = await this.projectService.getCurrentProject(); - return await this.threadRepository.listAllTimeDescOrder(project.id); + const { id } = await this.projectService.getCurrentProject(); + return await this.threadRepository.listAllTimeDescOrder(id); } public async updateThread( @@ -459,16 +460,25 @@ export class AskingService implements IAskingService { public async previewData( responseId: number, stepIndex?: number, - ): Promise { + limit?: number, + ): Promise { const response = await this.getResponse(responseId); if (!response) { throw new Error(`Thread response ${responseId} not found`); } - + const project = await this.projectService.getCurrentProject(); + const deployment = await this.deployService.getLastDeployment(project.id); + const mdl = deployment.manifest; const steps = response.detail.steps; const sql = format(constructCteSql(steps, stepIndex)); + const data = (await this.queryService.preview(sql, { + project, + mdl, + limit, + })) as PreviewDataResponse; + this.telemetry.send_event('preview_data', { sql }); - return this.wrenEngineAdaptor.previewData(sql); + return data; } public async deleteAllByProjectId(projectId: number): Promise { @@ -477,9 +487,9 @@ export class AskingService implements IAskingService { } private async getDeployId() { - const project = await this.projectService.getCurrentProject(); - const lastDeploy = await this.deployService.getLastDeployment(project.id); - return lastDeploy; + const { id } = await this.projectService.getCurrentProject(); + const lastDeploy = await this.deployService.getLastDeployment(id); + return lastDeploy.hash; } /** @@ -510,9 +520,9 @@ export class AskingService implements IAskingService { } const properties = JSON.parse(view.properties) || {}; - const project = await this.projectService.getCurrentProject(); + const { id } = await this.projectService.getCurrentProject(); const thread = await this.threadRepository.createOne({ - projectId: project.id, + projectId: id, sql: view.statement, summary: properties.summary, }); diff --git a/wren-ui/src/apollo/server/services/deployService.ts b/wren-ui/src/apollo/server/services/deployService.ts index ec4b7f9b0..8e08f146a 100644 --- a/wren-ui/src/apollo/server/services/deployService.ts +++ b/wren-ui/src/apollo/server/services/deployService.ts @@ -30,9 +30,10 @@ export interface MDLSyncResponse { export interface IDeployService { deploy(manifest: Manifest, projectId: number): Promise; - getLastDeployment(projectId: number): Promise; + getLastDeployment(projectId: number): Promise; getInProgressDeployment(projectId: number): Promise; - createMDLHash(manifest: Manifest): string; + createMDLHash(manifest: Manifest, projectId: number): string; + getMDLByHash(hash: string): Promise; deleteAllByProjectId(projectId: number): Promise; } @@ -65,7 +66,7 @@ export class DeployService implements IDeployService { if (!lastDeploy) { return null; } - return lastDeploy.hash; + return lastDeploy; } public async getInProgressDeployment(projectId) { @@ -76,7 +77,7 @@ export class DeployService implements IDeployService { public async deploy(manifest, projectId) { // generate hash of manifest - const hash = this.createMDLHash(manifest); + const hash = this.createMDLHash(manifest, projectId); logger.debug(`Deploying model, hash: ${hash}`); logger.debug(JSON.stringify(manifest)); @@ -116,12 +117,22 @@ export class DeployService implements IDeployService { return { status, error }; } - public createMDLHash(manifest: Manifest) { - const content = JSON.stringify(manifest); + public createMDLHash(manifest: Manifest, projectId: number) { + const manifestStr = JSON.stringify(manifest); + const content = `${projectId} ${manifestStr}`; const hash = createHash('sha1').update(content).digest('hex'); return hash; } + public async getMDLByHash(hash: string) { + const deploy = await this.deployLogRepository.findOneBy({ hash }); + if (!deploy) { + return null; + } + // return base64 encoded manifest + return Buffer.from(JSON.stringify(deploy.manifest)).toString('base64'); + } + public async deleteAllByProjectId(projectId: number): Promise { // delete all deploy logs await this.deployLogRepository.deleteAllBy({ projectId }); diff --git a/wren-ui/src/apollo/server/services/index.ts b/wren-ui/src/apollo/server/services/index.ts new file mode 100644 index 000000000..658dffaf7 --- /dev/null +++ b/wren-ui/src/apollo/server/services/index.ts @@ -0,0 +1,7 @@ +export * from './askingService'; +export * from './deployService'; +export * from './mdlService'; +export * from './modelService'; +export * from './projectService'; +export * from './queryService'; +export * from './metadataService'; diff --git a/wren-ui/src/apollo/server/services/metadataService.ts b/wren-ui/src/apollo/server/services/metadataService.ts new file mode 100644 index 000000000..c78e79b0d --- /dev/null +++ b/wren-ui/src/apollo/server/services/metadataService.ts @@ -0,0 +1,138 @@ +/** + This class is responsible for handling the retrieval of metadata from the data source. + For DuckDB, we control the access logic and directly query the WrenEngine. + For PostgreSQL and BigQuery, we will use the Ibis server API. + */ + +import { + IbisBigQueryConnectionInfo, + IIbisAdaptor, + IbisPostgresConnectionInfo, +} from '../adaptors/ibisAdaptor'; +import { IWrenEngineAdaptor } from '../adaptors/wrenEngineAdaptor'; +import { getConfig } from '@server/config'; +import { + BIG_QUERY_CONNECTION_INFO, + POSTGRES_CONNECTION_INFO, + Project, +} from '../repositories'; +import { DataSourceName } from '../types'; +import { Encryptor, getLogger } from '@server/utils'; + +const logger = getLogger('MetadataService'); +logger.level = 'debug'; + +const config = getConfig(); + +export interface CompactColumn { + name: string; + type: string; + notNull: boolean; + description?: string; + properties?: Record; +} + +export enum ConstraintType { + PRIMARY_KEY = 'PRIMARY KEY', + FOREIGN_KEY = 'FOREIGN KEY', + UNIQUE = 'UNIQUE', +} + +export interface CompactTable { + name: string; + columns: CompactColumn[]; + description?: string; + properties?: Record; + primaryKey?: string; +} + +export interface RecommendConstraint { + constraintName: string; + constraintType: ConstraintType; + constraintTable: string; + constraintColumn: string; + constraintedTable: string; + constraintedColumn: string; +} + +export interface IDataSourceMetadataService { + listTables(project: Project): Promise; + listConstraints(project: Project): Promise; +} + +export class DataSourceMetadataService implements IDataSourceMetadataService { + private readonly ibisAdaptor: IIbisAdaptor; + private readonly wrenEngineAdaptor: IWrenEngineAdaptor; + + constructor({ + ibisAdaptor, + wrenEngineAdaptor, + }: { + ibisAdaptor: IIbisAdaptor; + wrenEngineAdaptor: IWrenEngineAdaptor; + }) { + this.ibisAdaptor = ibisAdaptor; + this.wrenEngineAdaptor = wrenEngineAdaptor; + } + + public async listTables(project): Promise { + const { type: datasource } = project; + if (datasource === DataSourceName.DUCKDB) { + const tables = await this.wrenEngineAdaptor.listTables(); + return tables; + } else { + const { connectionInfo } = this.transformToIbisConnectionInfo(project); + return await this.ibisAdaptor.getTables(datasource, connectionInfo); + } + } + + public async listConstraints( + project: Project, + ): Promise { + const { type: datasource } = project; + if (datasource === DataSourceName.DUCKDB) { + return []; + } else { + const { connectionInfo } = this.transformToIbisConnectionInfo(project); + return await this.ibisAdaptor.getConstraints(datasource, connectionInfo); + } + } + + // transform connection info to ibis connection info format + private transformToIbisConnectionInfo(project: Project) { + const { type } = project; + switch (type) { + case DataSourceName.POSTGRES: { + const connectionInfo = + project.connectionInfo as POSTGRES_CONNECTION_INFO; + const encryptor = new Encryptor(config); + const decryptedCredentials = encryptor.decrypt(connectionInfo.password); + const { password } = JSON.parse(decryptedCredentials); + return { + connectionInfo: { + ...connectionInfo, + password: password, + } as IbisPostgresConnectionInfo, + }; + } + case DataSourceName.BIG_QUERY: { + const connectionInfo = + project.connectionInfo as BIG_QUERY_CONNECTION_INFO; + const encryptor = new Encryptor(config); + const decryptedCredentials = encryptor.decrypt( + connectionInfo.credentials, + ); + const credential = Buffer.from(decryptedCredentials).toString('base64'); + return { + connectionInfo: { + project_id: connectionInfo.projectId, + dataset_id: connectionInfo.datasetId, + credentials: credential, + } as IbisBigQueryConnectionInfo, + }; + } + default: + throw new Error(`Unsupported project type: ${type}`); + } + } +} diff --git a/wren-ui/src/apollo/server/services/modelService.ts b/wren-ui/src/apollo/server/services/modelService.ts index 491d9a5c1..797c85034 100644 --- a/wren-ui/src/apollo/server/services/modelService.ts +++ b/wren-ui/src/apollo/server/services/modelService.ts @@ -1,4 +1,4 @@ -import { SampleDatasetTable } from '../data'; +import { SampleDatasetTable } from '@server/data'; import { IModelColumnRepository, IModelRepository, @@ -7,21 +7,27 @@ import { Model, ModelColumn, Relation, -} from '../repositories'; +} from '@server/repositories'; import { getLogger } from '@server/utils'; -import { RelationData, UpdateRelationData } from '../types'; +import { RelationData, UpdateRelationData } from '@server/types'; import { IProjectService } from './projectService'; import { CreateCalculatedFieldData, ExpressionName, UpdateCalculatedFieldData, CheckCalculatedFieldCanQueryData, -} from '../models'; +} from '@server/models'; import { IMDLService } from './mdlService'; import { IWrenEngineAdaptor } from '../adaptors/wrenEngineAdaptor'; +import { ValidationRules } from '@server/adaptors/ibisAdaptor'; import { isEmpty, capitalize } from 'lodash'; -import { replaceAllowableSyntax, validateDisplayName } from '../utils/regex'; +import { + replaceAllowableSyntax, + validateDisplayName, +} from '@server/utils/regex'; import * as Errors from '@server/utils/error'; +import { DataSourceName } from '@server/types'; +import { IQueryService } from './queryService'; const logger = getLogger('ModelService'); logger.level = 'debug'; @@ -69,6 +75,7 @@ export class ModelService implements IModelService { private viewRepository: IViewRepository; private mdlService: IMDLService; private wrenEngineAdaptor: IWrenEngineAdaptor; + private queryService: IQueryService; constructor({ projectService, @@ -78,6 +85,7 @@ export class ModelService implements IModelService { viewRepository, mdlService, wrenEngineAdaptor, + queryService, }: { projectService: IProjectService; modelRepository: IModelRepository; @@ -86,6 +94,7 @@ export class ModelService implements IModelService { viewRepository: IViewRepository; mdlService: IMDLService; wrenEngineAdaptor: IWrenEngineAdaptor; + queryService: IQueryService; }) { this.projectService = projectService; this.modelRepository = modelRepository; @@ -94,6 +103,7 @@ export class ModelService implements IModelService { this.viewRepository = viewRepository; this.mdlService = mdlService; this.wrenEngineAdaptor = wrenEngineAdaptor; + this.queryService = queryService; } public async createCalculatedField( @@ -224,9 +234,9 @@ export class ModelService implements IModelService { public async updatePrimaryKeys(tables: SampleDatasetTable[]) { logger.debug('start update primary keys'); - const project = await this.projectService.getCurrentProject(); + const { id } = await this.projectService.getCurrentProject(); const models = await this.modelRepository.findAllBy({ - projectId: project.id, + projectId: id, }); const tableToUpdate = tables.filter((t) => t.primaryKey); for (const table of tableToUpdate) { @@ -243,9 +253,9 @@ export class ModelService implements IModelService { public async batchUpdateModelProperties(tables: SampleDatasetTable[]) { logger.debug('start batch update model description'); - const project = await this.projectService.getCurrentProject(); + const { id } = await this.projectService.getCurrentProject(); const models = await this.modelRepository.findAllBy({ - projectId: project.id, + projectId: id, }); await Promise.all([ @@ -267,9 +277,9 @@ export class ModelService implements IModelService { public async batchUpdateColumnProperties(tables: SampleDatasetTable[]) { logger.debug('start batch update column description'); - const project = await this.projectService.getCurrentProject(); + const { id } = await this.projectService.getCurrentProject(); const models = await this.modelRepository.findAllBy({ - projectId: project.id, + projectId: id, }); const sourceColumns = (await this.modelColumnRepository.findColumnsByModelIds( @@ -329,10 +339,10 @@ export class ModelService implements IModelService { if (isEmpty(relations)) { return []; } - const project = await this.projectService.getCurrentProject(); + const { id } = await this.projectService.getCurrentProject(); const models = await this.modelRepository.findAllBy({ - projectId: project.id, + projectId: id, }); const columnIds = relations @@ -355,7 +365,7 @@ export class ModelService implements IModelService { } const relationName = this.generateRelationName(relation, models, columns); return { - projectId: project.id, + projectId: id, name: relationName, fromColumnId: relation.fromColumnId, toColumnId: relation.toColumnId, @@ -370,7 +380,7 @@ export class ModelService implements IModelService { } public async createRelation(relation: RelationData): Promise { - const project = await this.projectService.getCurrentProject(); + const { id } = await this.projectService.getCurrentProject(); const modelIds = [relation.fromModelId, relation.toModelId]; const models = await this.modelRepository.findAllByIds(modelIds); const columnIds = [relation.fromColumnId, relation.toColumnId]; @@ -387,7 +397,7 @@ export class ModelService implements IModelService { } const relationName = this.generateRelationName(relation, models, columns); const savedRelation = await this.relationRepository.createOne({ - projectId: project.id, + projectId: id, name: relationName, fromColumnId: relation.fromColumnId, toColumnId: relation.toColumnId, @@ -589,6 +599,7 @@ export class ModelService implements IModelService { modelName: string, data: CheckCalculatedFieldCanQueryData, ) { + const project = await this.projectService.getCurrentProject(); const { mdlBuilder } = await this.mdlService.makeCurrentModelMDL(); const { referenceName, expression, lineage } = data; const inputFieldId = lineage[lineage.length - 1]; @@ -621,16 +632,24 @@ export class ModelService implements IModelService { ?.columns.find((c) => c.name === referenceName); logger.debug(`Calculated field MDL: ${JSON.stringify(calculatedField)}`); - const { valid, message } = - await this.wrenEngineAdaptor.validateColumnIsValid( + + // validate calculated field can query + const dataSource = project.type; + if (dataSource === DataSourceName.DUCKDB) { + return await this.wrenEngineAdaptor.validateColumnIsValid( manifest, modelName, referenceName, ); - if (!valid) { - logger.debug(`Calculated field can not query: ${message}`); + } else { + const parameters = { modelName, columnName: referenceName }; + return await this.queryService.validate( + project, + ValidationRules.COLUMN_IS_VALID, + manifest, + parameters, + ); } - return { valid, message }; } private async validateCreateRelation( diff --git a/wren-ui/src/apollo/server/services/projectService.ts b/wren-ui/src/apollo/server/services/projectService.ts index d5d27c028..bc5966718 100644 --- a/wren-ui/src/apollo/server/services/projectService.ts +++ b/wren-ui/src/apollo/server/services/projectService.ts @@ -2,17 +2,56 @@ import crypto from 'crypto'; import * as fs from 'fs'; import path from 'path'; import { Encryptor, getLogger } from '@server/utils'; -import { IProjectRepository } from '../repositories'; +import { + BIG_QUERY_CONNECTION_INFO, + DUCKDB_CONNECTION_INFO, + POSTGRES_CONNECTION_INFO, + IProjectRepository, +} from '../repositories'; import { Project } from '../repositories'; import { getConfig } from '../config'; +import { + CompactTable, + IDataSourceMetadataService, + RecommendConstraint, +} from './metadataService'; +import { DataSourceName } from '../types'; const config = getConfig(); const logger = getLogger('ProjectService'); logger.level = 'debug'; +const SENSITIVE_PROPERTY_NAME = new Set(['credentials', 'password']); +export interface ProjectData { + displayName: string; + type: DataSourceName; + connectionInfo: + | BIG_QUERY_CONNECTION_INFO + | POSTGRES_CONNECTION_INFO + | DUCKDB_CONNECTION_INFO; +} + export interface IProjectService { + createProject: (projectData: ProjectData) => Promise; + getSensitiveConnectionInfo: () => Set; + encryptSensitiveConnectionInfo: ( + connectionInfo: + | BIG_QUERY_CONNECTION_INFO + | POSTGRES_CONNECTION_INFO + | DUCKDB_CONNECTION_INFO, + ) => any; + getProjectDataSourceTables: ( + project?: Project, + projectId?: number, + ) => Promise; + getProjectSuggestedConstraint: ( + project?: Project, + projectId?: number, + ) => Promise; + getCurrentProject: () => Promise; + getProjectById: (projectId: number) => Promise; getCredentialFilePath: (project?: Project) => Promise; writeCredentialFile: ( credentials: JSON, @@ -23,20 +62,73 @@ export interface IProjectService { export class ProjectService implements IProjectService { private projectRepository: IProjectRepository; + private metadataService: IDataSourceMetadataService; - constructor({ projectRepository }: { projectRepository: any }) { + constructor({ + projectRepository, + metadataService, + }: { + projectRepository: IProjectRepository; + metadataService: IDataSourceMetadataService; + }) { this.projectRepository = projectRepository; + this.metadataService = metadataService; } public async getCurrentProject() { return await this.projectRepository.getCurrentProject(); } + public async getProjectById(projectId: number) { + return await this.projectRepository.findOneBy({ id: projectId }); + } + + public async getProjectDataSourceTables( + project?: Project, + projectId?: number, + ) { + const usedProject = project + ? project + : projectId + ? await this.getProjectById(projectId) + : await this.getCurrentProject(); + return await this.metadataService.listTables(usedProject); + } + + public async getProjectSuggestedConstraint( + project?: Project, + projectId?: number, + ) { + const usedProject = project + ? project + : projectId + ? await this.getProjectById(projectId) + : await this.getCurrentProject(); + return await this.metadataService.listConstraints(usedProject); + } + + public async createProject(projectData: ProjectData) { + const projectValue = { + displayName: projectData.displayName, + type: projectData.type, + catalog: 'wrenai', + schema: 'public', + connectionInfo: this.encryptSensitiveConnectionInfo( + projectData.connectionInfo, + ), + }; + logger.debug('Creating project...'); + logger.debug({ projectValue }); + const project = await this.projectRepository.createOne(projectValue); + return project; + } + public async getCredentialFilePath(project?: Project) { if (!project) { project = await this.getCurrentProject(); } - const { credentials: encryptedCredentials } = project; + const connectionInfo = project.connectionInfo as BIG_QUERY_CONNECTION_INFO; + const encryptedCredentials = connectionInfo.credentials; const encryptor = new Encryptor(config); const credentials = encryptor.decrypt(encryptedCredentials); const filePath = this.writeCredentialFile( @@ -73,4 +165,28 @@ export class ProjectService implements IProjectService { public async deleteProject(projectId: number): Promise { await this.projectRepository.deleteOne(projectId); } + + public getSensitiveConnectionInfo() { + return SENSITIVE_PROPERTY_NAME; + } + + public encryptSensitiveConnectionInfo( + connectionInfo: + | BIG_QUERY_CONNECTION_INFO + | POSTGRES_CONNECTION_INFO + | DUCKDB_CONNECTION_INFO, + ) { + const encryptor = new Encryptor(config); + const encryptConnectionInfo = Object.entries(connectionInfo).reduce( + (acc, [key, value]) => { + if (SENSITIVE_PROPERTY_NAME.has(key)) { + const toEncrypt = key === 'password' ? { password: value } : value; + acc[key] = encryptor.encrypt(toEncrypt); + } + return acc; + }, + connectionInfo, + ); + return encryptConnectionInfo; + } } diff --git a/wren-ui/src/apollo/server/services/queryService.ts b/wren-ui/src/apollo/server/services/queryService.ts new file mode 100644 index 000000000..8f2efa10d --- /dev/null +++ b/wren-ui/src/apollo/server/services/queryService.ts @@ -0,0 +1,257 @@ +import { DataSourceName } from '@server/types'; +import { Manifest } from '@server/mdl/type'; +import { IWrenEngineAdaptor } from '../adaptors/wrenEngineAdaptor'; +import { + SupportedDataSource, + IIbisAdaptor, + IbisQueryResponse, + IbisPostgresConnectionInfo, + IbisBigQueryConnectionInfo, + ValidationRules, + QueryOptions, +} from '../adaptors/ibisAdaptor'; +import { Encryptor, getLogger } from '@server/utils'; +import { + BIG_QUERY_CONNECTION_INFO, + POSTGRES_CONNECTION_INFO, + Project, +} from '../repositories'; +import { getConfig } from '../config'; + +const logger = getLogger('QueryService'); +logger.level = 'debug'; + +const config = getConfig(); + +export interface ColumnMetadata { + name: string; + type: string; +} + +export interface PreviewDataResponse { + columns: ColumnMetadata[]; + data: any[][]; +} + +export interface DescribeStatementResponse { + columns: ColumnMetadata[]; +} + +export interface PreviewOptions { + project: Project; + modelingOnly?: boolean; + mdl: Manifest; + limit?: number; + dryRun?: boolean; + // if not given, will use the deployed manifest +} + +export interface SqlValidateOptions { + project: Project; + mdl: Manifest; + modelingOnly?: boolean; +} + +export interface ComposeConnectionInfoResult { + datasource: DataSourceName; + connectionInfo?: IbisPostgresConnectionInfo | IbisBigQueryConnectionInfo; +} + +export interface ValidateResponse { + valid: boolean; + message?: string; +} + +export interface IQueryService { + preview( + sql: string, + options: PreviewOptions, + ): Promise; + + describeStatement( + sql: string, + options: PreviewOptions, + ): Promise; + + validate( + project: Project, + rule: ValidationRules, + manifest: Manifest, + parameters: Record, + ): Promise; +} + +export class QueryService implements IQueryService { + private readonly ibisAdaptor: IIbisAdaptor; + private readonly wrenEngineAdaptor: IWrenEngineAdaptor; + + constructor({ + ibisAdaptor, + wrenEngineAdaptor, + }: { + ibisAdaptor: IIbisAdaptor; + wrenEngineAdaptor: IWrenEngineAdaptor; + }) { + this.ibisAdaptor = ibisAdaptor; + this.wrenEngineAdaptor = wrenEngineAdaptor; + } + + public async preview( + sql: string, + options: PreviewOptions, + ): Promise { + const { project, mdl, limit, dryRun } = options; + + const dataSource = project.type; + if (this.useEngine(dataSource)) { + logger.debug('Using wren engine for preview'); + const data = await this.wrenEngineAdaptor.previewData(sql, limit, mdl); + return data as PreviewDataResponse; + } else { + logger.debug('Use ibis adaptor for 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; + const { connectionInfo } = this.transformToIbisConnectionInfo(project); + this.checkDataSourceIsSupported(dataSource); + const queryOptions = { + dataSource, + connectionInfo, + mdl, + dryRun, + } as QueryOptions; + if (dryRun) { + return await this.tryDryRun(rewrittenSql, queryOptions); + } else { + const data = await this.ibisAdaptor.query(rewrittenSql, queryOptions); + return this.transformDataType(data); + } + } + } + + public async describeStatement( + sql: string, + options: PreviewOptions, + ): Promise { + try { + // preview data with limit 1 to get column metadata + options.limit = 1; + const res = (await this.preview(sql, options)) as PreviewDataResponse; + return { columns: res.columns }; + } catch (err: any) { + logger.debug(`Got error when describing statement: ${err.message}`); + throw err; + } + } + + public async validate( + project, + rule: ValidationRules, + manifest: Manifest, + parameters: Record, + ): Promise { + const { connectionInfo } = this.transformToIbisConnectionInfo(project); + const dataSource = project.type; + const res = await this.ibisAdaptor.validate( + dataSource, + rule, + connectionInfo, + manifest, + parameters, + ); + return res; + } + + // transform connection info to ibis connection info format + private transformToIbisConnectionInfo(project: Project) { + const { type } = project; + switch (type) { + case DataSourceName.POSTGRES: { + const connectionInfo = + project.connectionInfo as POSTGRES_CONNECTION_INFO; + const encryptor = new Encryptor(config); + const decryptedCredentials = encryptor.decrypt(connectionInfo.password); + const { password } = JSON.parse(decryptedCredentials); + return { + connectionInfo: { + ...connectionInfo, + password, + } as IbisPostgresConnectionInfo, + }; + } + case DataSourceName.BIG_QUERY: { + const connectionInfo = + project.connectionInfo as BIG_QUERY_CONNECTION_INFO; + const encryptor = new Encryptor(config); + const decryptedCredentials = encryptor.decrypt( + connectionInfo.credentials, + ); + const credential = Buffer.from(decryptedCredentials).toString('base64'); + return { + connectionInfo: { + project_id: connectionInfo.projectId, + dataset_id: connectionInfo.datasetId, + credentials: credential, + } as IbisBigQueryConnectionInfo, + }; + } + default: + throw new Error(`Unsupported project type: ${type}`); + } + } + + private useEngine(dataSource: DataSourceName): boolean { + if (dataSource === DataSourceName.DUCKDB) { + return true; + } else { + return false; + } + } + + private transformDataType(data: IbisQueryResponse): PreviewDataResponse { + const columns = data.columns; + const dtypes = data.dtypes; + const transformedColumns = columns.map((column) => { + let type = 'unknown'; + if (dtypes && dtypes[column]) { + type = dtypes[column] === 'object' ? 'string' : dtypes[column]; + } + if (type === 'unknown') { + logger.debug(`Did not find type mapping for "${column}"`); + logger.debug( + `dtypes mapping: ${dtypes ? JSON.stringify(dtypes, null, 2) : 'undefined'} `, + ); + } + return { + name: column, + type, + } as ColumnMetadata; + }); + return { + columns: transformedColumns, + data: data.data, + } as PreviewDataResponse; + } + + private checkDataSourceIsSupported(dataSource: DataSourceName) { + if ( + !Object.prototype.hasOwnProperty.call(SupportedDataSource, dataSource) + ) { + throw new Error(`Unsupported datasource for ibis: "${dataSource}"`); + } + } + + private async tryDryRun( + rewrittenSql: string, + queryOptions: QueryOptions, + ): Promise { + try { + await this.ibisAdaptor.dryRun(rewrittenSql, queryOptions); + return true; + } catch (_err) { + return false; + } + } +} 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 0b35fd192..dd783d080 100644 --- a/wren-ui/src/apollo/server/services/tests/deployService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/deployService.test.ts @@ -67,7 +67,7 @@ describe('DeployService', () => { const projectId = 1; mockDeployLogRepository.findLastProjectDeployLog.mockResolvedValue({ - hash: deployService.createMDLHash(manifest), + hash: deployService.createMDLHash(manifest, 1), }); const response = await deployService.deploy(manifest, projectId); diff --git a/wren-ui/src/apollo/server/services/tests/projectService.test.ts b/wren-ui/src/apollo/server/services/tests/projectService.test.ts new file mode 100644 index 000000000..4b14b2f5d --- /dev/null +++ b/wren-ui/src/apollo/server/services/tests/projectService.test.ts @@ -0,0 +1,91 @@ +import { IDataSourceMetadataService, ProjectService } from '@server/services'; +import { + BIG_QUERY_CONNECTION_INFO, + DUCKDB_CONNECTION_INFO, + POSTGRES_CONNECTION_INFO, +} from '@server/repositories'; +import { Encryptor } from '@server/utils/encryptor'; + +jest.mock('@server/utils/encryptor'); + +describe('ProjectService', () => { + let projectService; + let mockProjectRepository; + let mockMetadataService; + let mockEncryptor; + + beforeEach(() => { + mockProjectRepository = { + getCurrentProject: jest.fn(), + }; + mockMetadataService = new (jest.fn())(); + mockEncryptor = { + encrypt: jest.fn().mockReturnValue('encrypted string'), + decrypt: jest.fn().mockReturnValue('decrypted string'), + }; + + projectService = new ProjectService({ + projectRepository: mockProjectRepository, + metadataService: mockMetadataService, + }); + + (Encryptor as jest.Mock).mockImplementation(() => mockEncryptor); + }); + + it('should encrypt sensitive connection info for BigQuery connection info', async () => { + const connectionInfo = { + credentials: 'some-credentials', + datasetId: 'my-bq-dataset-id', + projectId: 'my-bq-project-id', + } as BIG_QUERY_CONNECTION_INFO; + + const encryptedConnectionInfo = + await projectService.encryptSensitiveConnectionInfo(connectionInfo); + + expect(encryptedConnectionInfo).toEqual({ + credentials: 'encrypted string', + datasetId: 'my-bq-dataset-id', + projectId: 'my-bq-project-id', + }); + }); + + it('should encrypt sensitive connection info for Postgres connection info', async () => { + const connectionInfo = { + host: 'localhost', + port: 5432, + database: 'my-database', + user: 'my-user', + password: 'my-password', + ssl: false, + } as POSTGRES_CONNECTION_INFO; + + const encryptedConnectionInfo = + await projectService.encryptSensitiveConnectionInfo(connectionInfo); + + expect(encryptedConnectionInfo).toEqual({ + host: 'localhost', + port: 5432, + database: 'my-database', + user: 'my-user', + password: 'encrypted string', + ssl: false, + }); + }); + + it('should encrypt sensitive connection info for DuckDB connection info', async () => { + const connectionInfo = { + initSql: 'some-sql', + extensions: ['extension1', 'extension2'], + configurations: { key: 'value' }, + } as DUCKDB_CONNECTION_INFO; + + const encryptedConnectionInfo = + await projectService.encryptSensitiveConnectionInfo(connectionInfo); + + expect(encryptedConnectionInfo).toEqual({ + initSql: 'some-sql', + extensions: ['extension1', 'extension2'], + configurations: { key: 'value' }, + }); + }); +}); diff --git a/wren-ui/src/apollo/server/tests/connector.test.ts b/wren-ui/src/apollo/server/tests/connector.test.ts deleted file mode 100644 index 672d4fc57..000000000 --- a/wren-ui/src/apollo/server/tests/connector.test.ts +++ /dev/null @@ -1,281 +0,0 @@ -import { CompactTable } from '@server/connectors/connector'; -import { - PostgresColumnResponse, - PostgresConnector, -} from '@server/connectors/postgresConnector'; -import { TestingEnv } from './env'; -import { TestingPostgres } from './testingDatabase/postgres'; -import { WrenEngineColumnType } from '@server/connectors/types'; - -describe('Connector', () => { - let connector: PostgresConnector; - let testingEnv: TestingEnv; - let testingDatabase: TestingPostgres; - - // expected result - const tpchTables = [ - 'customer', - 'lineitem', - 'nation', - 'orders', - 'part', - 'partsupp', - 'region', - 'supplier', - ]; - - const tpchCustomerColumns = [ - { - name: 'c_custkey', - type: WrenEngineColumnType.INTEGER, - notNull: true, - description: '', - properties: { - ordinalPosition: 1, - }, - }, - { - name: 'c_name', - type: WrenEngineColumnType.VARCHAR, - notNull: true, - description: '', - properties: { - ordinalPosition: 2, - }, - }, - { - name: 'c_address', - type: WrenEngineColumnType.VARCHAR, - notNull: true, - description: '', - properties: { - ordinalPosition: 3, - }, - }, - { - name: 'c_nationkey', - type: WrenEngineColumnType.INTEGER, - notNull: true, - description: '', - properties: { - ordinalPosition: 4, - }, - }, - { - name: 'c_phone', - type: WrenEngineColumnType.CHAR, - notNull: true, - description: '', - properties: { - ordinalPosition: 5, - }, - }, - { - name: 'c_acctbal', - type: WrenEngineColumnType.DECIMAL, - notNull: true, - description: '', - properties: { - ordinalPosition: 6, - }, - }, - { - name: 'c_mktsegment', - type: WrenEngineColumnType.CHAR, - notNull: true, - description: '', - properties: { - ordinalPosition: 7, - }, - }, - { - name: 'c_comment', - type: WrenEngineColumnType.VARCHAR, - notNull: true, - description: '', - properties: { - ordinalPosition: 8, - }, - }, - ]; - const expectedConstraints = [ - { - constraintName: 'supplier_nation_fkey', - constraintType: 'FOREIGN KEY', - constraintTable: 'public.supplier', - constraintColumn: 's_nationkey', - constraintedTable: 'public.nation', - constraintedColumn: 'n_nationkey', - }, - { - constraintName: 'partsupp_part_fkey', - constraintType: 'FOREIGN KEY', - constraintTable: 'public.partsupp', - constraintColumn: 'ps_partkey', - constraintedTable: 'public.part', - constraintedColumn: 'p_partkey', - }, - { - constraintName: 'partsupp_supplier_fkey', - constraintType: 'FOREIGN KEY', - constraintTable: 'public.partsupp', - constraintColumn: 'ps_suppkey', - constraintedTable: 'public.supplier', - constraintedColumn: 's_suppkey', - }, - { - constraintName: 'customer_nation_fkey', - constraintType: 'FOREIGN KEY', - constraintTable: 'public.customer', - constraintColumn: 'c_nationkey', - constraintedTable: 'public.nation', - constraintedColumn: 'n_nationkey', - }, - { - constraintName: 'orders_customer_fkey', - constraintType: 'FOREIGN KEY', - constraintTable: 'public.orders', - constraintColumn: 'o_custkey', - constraintedTable: 'public.customer', - constraintedColumn: 'c_custkey', - }, - { - constraintName: 'lineitem_orders_fkey', - constraintType: 'FOREIGN KEY', - constraintTable: 'public.lineitem', - constraintColumn: 'l_orderkey', - constraintedTable: 'public.orders', - constraintedColumn: 'o_orderkey', - }, - { - constraintName: 'nation_region_fkey', - constraintType: 'FOREIGN KEY', - constraintTable: 'public.nation', - constraintColumn: 'n_regionkey', - constraintedTable: 'public.region', - constraintedColumn: 'r_regionkey', - }, - ]; - - beforeAll(async () => { - testingEnv = new TestingEnv(); - testingDatabase = new TestingPostgres(); - - // initialize - await testingEnv.initialize(); - await testingDatabase.initialize(testingEnv.context); - - // create connector - const container = testingDatabase.getContainer(); - connector = new PostgresConnector({ - user: container.getUsername(), - password: container.getPassword(), - host: container.getHost(), - database: container.getDatabase(), - port: container.getPort(), - }); - }, 60000); - - afterAll(async () => { - // close connector - await connector.close(); - - // finalize testing database - await testingDatabase.finalize(); - }); - - it('should test connect', async () => { - const connected = await connector.connect(); - expect(connected).toBeTruthy(); - }); - - it('should list tables with format: true', async () => { - const tables = (await connector.listTables({ - format: true, - })) as CompactTable[]; - - // check if tables include tpch tables - for (const table of tpchTables) { - const found = tables.find( - (t: CompactTable) => t.name === `public.${table}`, - ); - expect(found).toBeTruthy(); - } - - // check if customer table has correct columns - const customerTable = tables.find( - (t: CompactTable) => t.name === 'public.customer', - ); - expect(customerTable).toBeTruthy(); - expect(customerTable.columns).toBeTruthy(); - expect(customerTable.columns.length).toBe(tpchCustomerColumns.length); - - for (const column of tpchCustomerColumns) { - const found = customerTable.columns.find( - (c) => c.name === column.name && c.type === column.type, - ); - expect(found).toBeTruthy(); - // check type, notNull, and ordinalPosition - expect(found.type).toBe(column.type); - expect(found.notNull).toBe(column.notNull); - - // properties will be an empty object - expect(found.properties).toStrictEqual({}); - } - }); - - it('should list tables with format: false', async () => { - const columns = (await connector.listTables({ - format: false, - })) as PostgresColumnResponse[]; - - // check if columns not null and has length - expect(columns).toBeTruthy(); - expect(columns.length).toBeGreaterThan(0); - - // check the format of the columns - /* - the format of the columns should be like this: - { - table_catalog: 'test', - table_schema: 'public', - table_name: 'supplier', - column_name: 's_comment', - ordinal_position: 7, - is_nullable: 'NO', - data_type: 'character varying' - } - */ - const container = testingDatabase.getContainer(); - const expectedCatalog = container.getDatabase(); - const expectedSchema = 'public'; - const column = columns[0]; - expect(column.table_catalog).toBe(expectedCatalog); - expect(column.table_schema).toBe(expectedSchema); - expect(column.table_name).toBeTruthy(); - expect(column.column_name).toBeTruthy(); - expect(column.ordinal_position).toBeTruthy(); - expect(column.is_nullable).toBeTruthy(); - expect(column.data_type).toBeTruthy(); - }); - - it('should list constraints', async () => { - const constraints = await connector.listConstraints(); - - // compare the constraints with the expected constraints - expect(constraints).toBeTruthy(); - expect(constraints.length).toBe(expectedConstraints.length); - for (const constraint of expectedConstraints) { - const found = constraints.find( - (c) => - c.constraintName === constraint.constraintName && - c.constraintType === constraint.constraintType && - c.constraintTable === constraint.constraintTable && - c.constraintColumn === constraint.constraintColumn && - c.constraintedTable === constraint.constraintedTable && - c.constraintedColumn === constraint.constraintedColumn, - ); - expect(found).toBeTruthy(); - } - }); -}); diff --git a/wren-ui/src/apollo/server/tests/env.ts b/wren-ui/src/apollo/server/tests/env.ts deleted file mode 100644 index eafab1ace..000000000 --- a/wren-ui/src/apollo/server/tests/env.ts +++ /dev/null @@ -1,62 +0,0 @@ -import { TestingContext } from './testingDatabase/database'; -import { Database } from 'duckdb-async'; -import { getLogger } from '@server/utils'; -import * as tmp from 'tmp'; - -const logger = getLogger('TestingEnv'); -logger.level = 'debug'; - -export class TestingEnv { - public context: TestingContext; - - constructor() { - this.context = { - tpch: {}, - }; - } - - public async initialize(): Promise { - await this.prepareTpchData(); - } - - private async prepareTpchData(): Promise { - logger.info('Preparing TPCH data'); - // run duckdb to load tpch data - const db = await Database.create(':memory:'); - await db.run('INSTALL tpch'); - await db.run('LOAD tpch'); - await db.run('CALL dbgen(sf = 0.001)'); - - // output tpch data as csv to tmp dir - const tmpDir = tmp.dirSync(); - const rows = await db.all('SHOW TABLES'); - /* - rows will be like - [ - { name: 'customer' }, - { name: 'lineitem' }, - { name: 'nation' }, - { name: 'orders' }, - { name: 'part' }, - { name: 'partsupp' }, - { name: 'region' }, - { name: 'supplier' } - ] - */ - for (const row of rows) { - const table = row.name; - // run COPY command to output csv - // COPY customer TO 'output.csv' (HEADER, DELIMITER ',') - logger.info(`Exporting ${table} as csv to ${tmpDir.name}/${table}.csv`); - await db.run( - `COPY ${table} TO '${tmpDir.name}/${table}.csv' (HEADER, DELIMITER ',')`, - ); - } - - // close db - await db.close(); - - // set context - this.context.tpch.dataDir = tmpDir.name; - } -} diff --git a/wren-ui/src/apollo/server/tests/testingDatabase/database.ts b/wren-ui/src/apollo/server/tests/testingDatabase/database.ts deleted file mode 100644 index 99b01a7c4..000000000 --- a/wren-ui/src/apollo/server/tests/testingDatabase/database.ts +++ /dev/null @@ -1,11 +0,0 @@ -export interface TestingContext { - tpch: { - // /path/to/tpch/data - dataDir?: string; - }; -} - -export interface TestingDatabase { - initialize(context: TestingContext): Promise; - getContainer(): C; -} diff --git a/wren-ui/src/apollo/server/tests/testingDatabase/postgres.ts b/wren-ui/src/apollo/server/tests/testingDatabase/postgres.ts deleted file mode 100644 index 87e0f9d6b..000000000 --- a/wren-ui/src/apollo/server/tests/testingDatabase/postgres.ts +++ /dev/null @@ -1,196 +0,0 @@ -import { - PostgreSqlContainer, - StartedPostgreSqlContainer, -} from '@testcontainers/postgresql'; -import { TestingContext, TestingDatabase } from './database'; -import { Client } from 'pg'; -import { getLogger } from '@server/utils'; - -const logger = getLogger('TestingPostgres'); -logger.level = 'debug'; - -const psqlInitCommands = (dataDir: string) => ` -CREATE TABLE NATION ( - N_NATIONKEY INT PRIMARY KEY, - N_NAME CHAR(25) NOT NULL, - N_REGIONKEY INT NOT NULL, - N_COMMENT VARCHAR(152) -); - -CREATE TABLE REGION ( - R_REGIONKEY INT PRIMARY KEY, - R_NAME CHAR(25) NOT NULL, - R_COMMENT VARCHAR(152) -); - -CREATE TABLE PART ( - P_PARTKEY INT PRIMARY KEY, - P_NAME VARCHAR(55) NOT NULL, - P_MFGR CHAR(25) NOT NULL, - P_BRAND CHAR(10) NOT NULL, - P_TYPE VARCHAR(25) NOT NULL, - P_SIZE INT NOT NULL, - P_CONTAINER CHAR(10) NOT NULL, - P_RETAILPRICE DECIMAL(15,2) NOT NULL, - P_COMMENT VARCHAR(23) NOT NULL -); - -CREATE TABLE SUPPLIER ( - S_SUPPKEY INT PRIMARY KEY, - S_NAME CHAR(25) NOT NULL, - S_ADDRESS VARCHAR(40) NOT NULL, - S_NATIONKEY INT NOT NULL, - S_PHONE CHAR(15) NOT NULL, - S_ACCTBAL DECIMAL(15,2) NOT NULL, - S_COMMENT VARCHAR(101) NOT NULL -); - -CREATE TABLE PARTSUPP ( - PS_PARTKEY INT NOT NULL, - PS_SUPPKEY INT NOT NULL, - PS_AVAILQTY INT NOT NULL, - PS_SUPPLYCOST DECIMAL(15,2) NOT NULL, - PS_COMMENT VARCHAR(199) NOT NULL - -- weird that csv that duckdb produced has conflicting primary key - -- PRIMARY KEY (PS_PARTKEY, PS_SUPPKEY) -); - -CREATE TABLE CUSTOMER ( - C_CUSTKEY INT PRIMARY KEY, - C_NAME VARCHAR(25) NOT NULL, - C_ADDRESS VARCHAR(40) NOT NULL, - C_NATIONKEY INT NOT NULL, - C_PHONE CHAR(15) NOT NULL, - C_ACCTBAL DECIMAL(15,2) NOT NULL, - C_MKTSEGMENT CHAR(10) NOT NULL, - C_COMMENT VARCHAR(117) NOT NULL -); - -CREATE TABLE ORDERS ( - O_ORDERKEY INT PRIMARY KEY, - O_CUSTKEY INT NOT NULL, - O_ORDERSTATUS CHAR(1) NOT NULL, - O_TOTALPRICE DECIMAL(15,2) NOT NULL, - O_ORDERDATE DATE NOT NULL, - O_ORDERPRIORITY CHAR(15) NOT NULL, - O_CLERK CHAR(15) NOT NULL, - O_SHIPPRIORITY INT NOT NULL, - O_COMMENT VARCHAR(79) NOT NULL -); - -CREATE TABLE LINEITEM ( - L_ORDERKEY INT NOT NULL, - L_PARTKEY INT NOT NULL, - L_SUPPKEY INT NOT NULL, - L_LINENUMBER INT NOT NULL, - L_QUANTITY DECIMAL(15,2) NOT NULL, - L_EXTENDEDPRICE DECIMAL(15,2) NOT NULL, - L_DISCOUNT DECIMAL(15,2) NOT NULL, - L_TAX DECIMAL(15,2) NOT NULL, - L_RETURNFLAG CHAR(1) NOT NULL, - L_LINESTATUS CHAR(1) NOT NULL, - L_SHIPDATE DATE NOT NULL, - L_COMMITDATE DATE NOT NULL, - L_RECEIPTDATE DATE NOT NULL, - L_SHIPINSTRUCT CHAR(25) NOT NULL, - L_SHIPMODE CHAR(10) NOT NULL, - L_COMMENT VARCHAR(44) NOT NULL - -- PRIMARY KEY (L_ORDERKEY, L_LINENUMBER) -); - --- Set Foreign Key -ALTER TABLE SUPPLIER - ADD CONSTRAINT supplier_nation_fkey - FOREIGN KEY (S_NATIONKEY) REFERENCES NATION(N_NATIONKEY); - -ALTER TABLE PARTSUPP - ADD CONSTRAINT partsupp_part_fkey - FOREIGN KEY (PS_PARTKEY) REFERENCES PART(P_PARTKEY); - -ALTER TABLE PARTSUPP - ADD CONSTRAINT partsupp_supplier_fkey - FOREIGN KEY (PS_SUPPKEY) REFERENCES SUPPLIER(S_SUPPKEY); - -ALTER TABLE CUSTOMER - ADD CONSTRAINT customer_nation_fkey - FOREIGN KEY (C_NATIONKEY) REFERENCES NATION(N_NATIONKEY); - -ALTER TABLE ORDERS - ADD CONSTRAINT orders_customer_fkey - FOREIGN KEY (O_CUSTKEY) REFERENCES CUSTOMER(C_CUSTKEY); - -ALTER TABLE LINEITEM - ADD CONSTRAINT lineitem_orders_fkey - FOREIGN KEY (L_ORDERKEY) REFERENCES ORDERS(O_ORDERKEY); - --- weird that csv that duckdb produced has conflicting primary key --- skip this for now. --- ALTER TABLE LINEITEM --- ADD CONSTRAINT lineitem_partsupp_fkey --- FOREIGN KEY (L_PARTKEY, L_SUPPKEY) --- REFERENCES PARTSUPP(PS_PARTKEY, PS_SUPPKEY); - -ALTER TABLE NATION - ADD CONSTRAINT nation_region_fkey - FOREIGN KEY (N_REGIONKEY) REFERENCES REGION(R_REGIONKEY); - --- COPY -COPY REGION FROM '${dataDir}/region.csv' DELIMITER ',' CSV HEADER; -COPY NATION FROM '${dataDir}/nation.csv' DELIMITER ',' CSV HEADER; -COPY SUPPLIER FROM '${dataDir}/supplier.csv' DELIMITER ',' CSV HEADER; -COPY CUSTOMER FROM '${dataDir}/customer.csv' DELIMITER ',' CSV HEADER; -COPY ORDERS FROM '${dataDir}/orders.csv' DELIMITER ',' CSV HEADER; -COPY PART FROM '${dataDir}/part.csv' DELIMITER ',' CSV HEADER; -COPY PARTSUPP FROM '${dataDir}/partsupp.csv' DELIMITER ',' CSV HEADER; -COPY LINEITEM FROM '${dataDir}/lineitem.csv' DELIMITER ',' CSV HEADER; -`; - -export class TestingPostgres - implements TestingDatabase -{ - private container: StartedPostgreSqlContainer; - - public async initialize(context: TestingContext): Promise { - const { dataDir } = context.tpch; - // dataDir that copied into container - const containerDataDir = '/etc/testing_data'; - - logger.info('Initializing TestingPostgres'); - const container = await new PostgreSqlContainer() - .withExposedPorts({ - container: 5432, - host: 8432, - }) - .withCopyDirectoriesToContainer([ - { - source: dataDir, - target: containerDataDir, - }, - ]) - .start(); - - // running init commands - const client = new Client({ - connectionString: container.getConnectionUri(), - }); - - // connect to container and run init commands - logger.info('Running init commands'); - await client.connect(); - await client.query(psqlInitCommands(containerDataDir)); - await client.end(); - - // assign container to instance - logger.info('Container started'); - this.container = container; - } - - public getContainer(): StartedPostgreSqlContainer { - return this.container; - } - - public async finalize(): Promise { - logger.info('Stopping container'); - await this.container.stop(); - } -} diff --git a/wren-ui/src/apollo/server/types/context.ts b/wren-ui/src/apollo/server/types/context.ts index b74e684ea..0345b9842 100644 --- a/wren-ui/src/apollo/server/types/context.ts +++ b/wren-ui/src/apollo/server/types/context.ts @@ -1,3 +1,4 @@ +import { IIbisAdaptor } from '../adaptors/ibisAdaptor'; import { IWrenEngineAdaptor } from '../adaptors/wrenEngineAdaptor'; import { IConfig } from '../config'; import { @@ -8,11 +9,14 @@ import { IViewRepository, } from '../repositories'; import { IDeployLogRepository } from '../repositories/deployLogRepository'; -import { IAskingService } from '../services/askingService'; -import { IDeployService } from '../services/deployService'; -import { IMDLService } from '../services/mdlService'; -import { IModelService } from '../services/modelService'; -import { IProjectService } from '../services/projectService'; +import { + IQueryService, + IAskingService, + IDeployService, + IModelService, + IMDLService, + IProjectService, +} from '../services'; export interface IContext { config: IConfig; @@ -21,6 +25,7 @@ export interface IContext { // adaptor wrenEngineAdaptor: IWrenEngineAdaptor; + ibisServerAdaptor: IIbisAdaptor; // services projectService: IProjectService; @@ -28,6 +33,7 @@ export interface IContext { mdlService: IMDLService; deployService: IDeployService; askingService: IAskingService; + queryService: IQueryService; // repository projectRepository: IProjectRepository; diff --git a/wren-ui/src/apollo/server/types/diagram.ts b/wren-ui/src/apollo/server/types/diagram.ts index 3aaaef836..3df220d90 100644 --- a/wren-ui/src/apollo/server/types/diagram.ts +++ b/wren-ui/src/apollo/server/types/diagram.ts @@ -40,7 +40,7 @@ export interface DiagramModel { displayName: string; referenceName: string; sourceTableName: string; - refSql: string; + refSql?: string; cached: boolean; refreshTime: string; description: string; diff --git a/wren-ui/src/apollo/server/utils/docker.ts b/wren-ui/src/apollo/server/utils/docker.ts new file mode 100644 index 000000000..85f70b6e4 --- /dev/null +++ b/wren-ui/src/apollo/server/utils/docker.ts @@ -0,0 +1,16 @@ +export const toDockerHost = (host: string) => { + // if host is localhost or 127.0.0.1, rewrite it to docker.for.{platform}.localhost + if (host === 'localhost' || host === '127.0.0.1') { + const platform = process.platform; + switch (platform) { + case 'darwin': + return 'docker.for.mac.localhost'; + case 'linux': + return 'docker.for.linux.localhost'; + default: + // windows and others... + return 'host.docker.internal'; + } + } + return host; +}; diff --git a/wren-ui/src/apollo/server/utils/error.ts b/wren-ui/src/apollo/server/utils/error.ts index ebb62c80e..4ee19d884 100644 --- a/wren-ui/src/apollo/server/utils/error.ts +++ b/wren-ui/src/apollo/server/utils/error.ts @@ -11,6 +11,9 @@ export enum GeneralErrorCodes { // Exception error for AI service (e.g., network connection error) AI_SERVICE_UNDEFINED_ERROR = 'OTHERS', + // IBIS Error + IBIS_SERVER_ERROR = 'IBIS_SERVER_ERROR', + // Connector errors CONNECTION_ERROR = 'CONNECTION_ERROR', // duckdb @@ -50,6 +53,10 @@ export const errorMessages = { [GeneralErrorCodes.CONNECTION_REFUSED]: 'Connection refused by the server, Please check your connection settings and try again.', + // ibis service errors + [GeneralErrorCodes.IBIS_SERVER_ERROR]: + 'Error occurred while querying ibis server, please try again later.', + // calculated field validation [GeneralErrorCodes.DUPLICATED_FIELD_NAME]: 'This field name already exists', [GeneralErrorCodes.INVALID_EXPRESSION]: @@ -67,6 +74,7 @@ export const shortMessages = { [GeneralErrorCodes.NO_RELEVANT_DATA]: 'No relevant data', [GeneralErrorCodes.NO_RELEVANT_SQL]: 'No relevant SQL', [GeneralErrorCodes.CONNECTION_ERROR]: 'Failed to connect', + [GeneralErrorCodes.IBIS_SERVER_ERROR]: 'Ibis server error', [GeneralErrorCodes.INIT_SQL_ERROR]: 'Invalid initializing SQL', [GeneralErrorCodes.SESSION_PROPS_ERROR]: 'Invalid session properties', [GeneralErrorCodes.CONNECTION_REFUSED]: 'Connection refused', diff --git a/wren-ui/src/apollo/server/utils/index.ts b/wren-ui/src/apollo/server/utils/index.ts index 7ab43fffa..b6ca2132d 100644 --- a/wren-ui/src/apollo/server/utils/index.ts +++ b/wren-ui/src/apollo/server/utils/index.ts @@ -2,3 +2,5 @@ export * from './logger'; export * from './encryptor'; export * from './encode'; export * from './string'; +export * from './docker'; +export * from './model'; diff --git a/wren-ui/src/apollo/server/factories/util.ts b/wren-ui/src/apollo/server/utils/model.ts similarity index 82% rename from wren-ui/src/apollo/server/factories/util.ts rename to wren-ui/src/apollo/server/utils/model.ts index 964a2a288..768709c76 100644 --- a/wren-ui/src/apollo/server/factories/util.ts +++ b/wren-ui/src/apollo/server/utils/model.ts @@ -1,5 +1,20 @@ -import { IModelColumnRepository, ModelColumn } from '../repositories'; -import { replaceAllowableSyntax } from '@server/utils/regex'; +import { IModelColumnRepository, ModelColumn } from '@server/repositories'; +import { replaceAllowableSyntax } from './regex'; + +export function transformInvalidColumnName(columnName: string) { + let referenceName = replaceAllowableSyntax(columnName); + // If the reference name does not start with a letter, add a prefix + const startWithLetterRegex = /^[A-Za-z]/; + if (!startWithLetterRegex.test(referenceName)) { + referenceName = `col_${referenceName}`; + } + return referenceName; +} + +export function replaceInvalidReferenceName(referenceName: string) { + // replace dot with underscore + return referenceName.replace(/\./g, '_'); +} export function findColumnsToUpdate( columns: string[], @@ -36,13 +51,3 @@ export async function updateModelPrimaryKey( await repository.setModelPrimaryKey(modelId, primaryKey); } } - -export function transformInvalidColumnName(columnName: string) { - let referenceName = replaceAllowableSyntax(columnName); - // If the reference name does not start with a letter, add a prefix - const startWithLetterRegex = /^[A-Za-z]/; - if (!startWithLetterRegex.test(referenceName)) { - referenceName = `col_${referenceName}`; - } - return referenceName; -} diff --git a/wren-ui/src/pages/api/graphql.ts b/wren-ui/src/pages/api/graphql.ts index a6906f5b1..76697709b 100644 --- a/wren-ui/src/pages/api/graphql.ts +++ b/wren-ui/src/pages/api/graphql.ts @@ -27,6 +27,11 @@ import { ThreadRepository } from '@/apollo/server/repositories/threadRepository' import { ThreadResponseRepository } from '@/apollo/server/repositories/threadResponseRepository'; import { defaultApolloErrorHandler } from '@/apollo/server/utils/error'; import { Telemetry } from '@/apollo/server/telemetry/telemetry'; +import { IbisAdaptor } from '@/apollo/server/adaptors/ibisAdaptor'; +import { + DataSourceMetadataService, + QueryService, +} from '@/apollo/server/services'; const serverConfig = getConfig(); const logger = getLogger('APOLLO'); @@ -65,8 +70,19 @@ const bootstrapServer = async () => { const wrenAIAdaptor = new WrenAIAdaptor({ wrenAIBaseEndpoint: serverConfig.wrenAIEndpoint, }); + const ibisAdaptor = new IbisAdaptor({ + ibisServerEndpoint: serverConfig.ibisServerEndpoint, + }); + + const metadataService = new DataSourceMetadataService({ + ibisAdaptor, + wrenEngineAdaptor, + }); - const projectService = new ProjectService({ projectRepository }); + const projectService = new ProjectService({ + projectRepository, + metadataService, + }); const mdlService = new MDLService({ projectRepository, modelRepository, @@ -74,6 +90,10 @@ const bootstrapServer = async () => { relationRepository, viewRepository, }); + const queryService = new QueryService({ + ibisAdaptor, + wrenEngineAdaptor, + }); const modelService = new ModelService({ projectService, modelRepository, @@ -82,6 +102,7 @@ const bootstrapServer = async () => { viewRepository, mdlService, wrenEngineAdaptor, + queryService, }); const deployService = new DeployService({ wrenAIAdaptor, @@ -93,12 +114,12 @@ const bootstrapServer = async () => { const askingService = new AskingService({ telemetry, wrenAIAdaptor, - wrenEngineAdaptor, deployService, projectService, viewRepository, threadRepository, threadResponseRepository, + queryService, }); // initialize services @@ -138,6 +159,7 @@ const bootstrapServer = async () => { telemetry, // adaptor wrenEngineAdaptor, + ibisServerAdaptor: ibisAdaptor, // services projectService, @@ -145,6 +167,7 @@ const bootstrapServer = async () => { mdlService, deployService, askingService, + queryService, // repository projectRepository,