import { QueryError, QueryResult, QueryService, ConnectionTestResult } from '@tooljet-plugins/common'; import { SourceOptions, QueryOptions } from './types'; import { BigQuery } from '@google-cloud/bigquery'; const JSON5 = require('json5'); const _ = require('lodash'); export default class Bigquery implements QueryService { async run(sourceOptions: SourceOptions, queryOptions: QueryOptions, dataSourceId: string): Promise { const operation = queryOptions.operation; const client = await this.getConnection(sourceOptions); let result = {}; try { switch (operation) { case 'list_datasets': { const [datasets] = await client.getDatasets(); result = this.sanitizeResponse(datasets, ['metadata.datasetReference']); break; } case 'list_tables': { const [tables] = await client.dataset(queryOptions.datasetId).getTables(); result = this.sanitizeResponse(tables, ['metadata.tableReference']); break; } case 'create_table': { const [table] = await client .dataset(queryOptions.datasetId) .createTable(queryOptions.tableId, this.parseJSON(queryOptions.options)); result = { tableId: table.id }; break; } case 'delete_table': { await client.dataset(queryOptions.datasetId).table(queryOptions.tableId).delete(); result = `Table ${queryOptions.tableId} deleted.`; break; } case 'create_view': { const query = `CREATE VIEW ${queryOptions.datasetId}.${queryOptions.view_name} AS SELECT ${queryOptions.viewcolumns} FROM ${queryOptions.datasetId}.${queryOptions.tableId} ${queryOptions.condition ? `WHERE ${queryOptions.condition}` : 'WHERE TRUE'}`; const [job] = await client.createQueryJob({ ...this.parseJSON(queryOptions.queryOptions), query: query, }); const [rows] = await job.getQueryResults(this.parseJSON(queryOptions.queryResultsOptions)); result = rows; break; } case 'query': { const [job] = await client.createQueryJob({ ...this.parseJSON(queryOptions.queryOptions), query: queryOptions.query, }); const [rows] = await job.getQueryResults(this.parseJSON(queryOptions.queryResultsOptions)); result = rows; break; } case 'delete_record': { const query = `DELETE FROM ${queryOptions.datasetId}.${queryOptions.tableId} ${ queryOptions.condition ? `WHERE ${queryOptions.condition}` : 'WHERE TRUE' }`; const [job] = await client.createQueryJob({ ...this.parseJSON(queryOptions.queryOptions), query: query, }); const [rows] = await job.getQueryResults(this.parseJSON(queryOptions.queryResultsOptions)); result = rows; break; } case 'insert_record': { const rows = await client .dataset(queryOptions.datasetId) .table(queryOptions.tableId) .insert(this.parseJSON(queryOptions.rows)); result = { ...rows[0], records: (this.parseJSON(queryOptions.rows) as []).length }; break; } case 'update_record': { let columString = ''; columString = await this.columnBuilder(queryOptions); const query = `UPDATE ${queryOptions.datasetId}.${queryOptions.tableId} SET ${columString} ${ queryOptions.condition ? `WHERE ${queryOptions.condition}` : 'WHERE TRUE' }`; const [job] = await client.createQueryJob({ ...this.parseJSON(queryOptions.queryOptions), query: query, }); const [rows] = await job.getQueryResults(this.parseJSON(queryOptions.queryResultsOptions)); result = rows; break; } } } catch (error) { console.log(error); throw new QueryError('Query could not be completed', error.message, {}); } return { status: 'ok', data: result, }; } async columnBuilder(queryOptions: any): Promise { const columString = []; const columns = queryOptions.columns; for (const [key, value] of Object.entries(columns)) { const primaryKeyValue = typeof value === 'string' ? `'${value}'` : value; columString.push(`${key}=${primaryKeyValue}`); } return columString.join(','); } async getConnection(sourceOptions: any, _options?: object): Promise { const privateKey = this.getPrivateKey(sourceOptions?.private_key); return new BigQuery({ projectId: privateKey?.project_id, credentials: { client_email: privateKey?.client_email, private_key: privateKey?.private_key, }, }); } async testConnection(sourceOptions: SourceOptions): Promise { const privateKey = this.getPrivateKey(sourceOptions?.private_key); const client = new BigQuery({ projectId: privateKey?.project_id, credentials: { client_email: privateKey?.client_email, private_key: privateKey?.private_key, }, }); if (!client) { throw new Error('Invalid credentials'); } await client.getDatasets(); return { status: 'ok', }; } private parseJSON(json?: string): object { if (!json) return {}; return JSON5.parse(json); } private getPrivateKey(configs?: string): { project_id?: string; client_email?: string; private_key?: string; } { return this.parseJSON(configs); } private sanitizeResponse(response: object | [], pickFields: string[]): object | [] { if (!response) return response; if (Array.isArray(response)) { return response.map((item) => this.sanitizeResponse(item, pickFields)); } const pickedKeyValue = pickFields.map((field) => _.result(response, field)); if (pickedKeyValue.length === 1) { return pickedKeyValue[0]; } return pickedKeyValue; } }