| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you under the Apache License, Version 2.0 (the |
| * "License"); you may not use this file except in compliance |
| * with the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, software |
| * distributed under the License is distributed on an "AS IS" BASIS, |
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| * See the License for the specific language governing permissions and |
| * limitations under the License. |
| */ |
| |
| import type { Column } from 'druid-query-toolkit'; |
| import { C, F, SqlColumn, SqlExpression, SqlQuery, SqlStar } from 'druid-query-toolkit'; |
| |
| import { filterMap, filterOrReturn, mapRecordOrReturn } from '../../../utils'; |
| |
| import { ExpressionMeta } from './expression-meta'; |
| import { Measure } from './measure'; |
| import type { OptionValue, ParameterDefinition, Parameters, ParameterValues } from './parameter'; |
| import { evaluateFunctor } from './parameter'; |
| |
| function expressionWithinColumns(ex: SqlExpression, columns: readonly Column[]): boolean { |
| const usedColumns = ex.getUsedColumnNames(); |
| return usedColumns.every(columnName => columns.some(c => c.name === columnName)); |
| } |
| |
| interface QuerySourceValue { |
| query: SqlQuery; |
| baseColumns: readonly Column[]; |
| columns: readonly Column[]; |
| measures: Measure[]; |
| } |
| |
| export class QuerySource { |
| static isSimpleSelect(query: SqlQuery): boolean { |
| return Boolean( |
| !query.hasGroupBy() && !query.unionQuery && query.getFromExpressions().length === 1, |
| ); |
| } |
| |
| static isSingleStarQuery(query: SqlQuery): boolean { |
| const selectExpressions = query.getSelectExpressionsArray(); |
| return selectExpressions.length === 1 && selectExpressions[0] instanceof SqlStar; |
| } |
| |
| static makeLimitZeroIntrospectionQuery(query: SqlQuery): SqlQuery { |
| return SqlQuery.selectStarFrom(query).changeLimitValue(0); |
| } |
| |
| static stripToBaseSource(query: SqlQuery): SqlQuery { |
| if (query.hasGroupBy()) { |
| if (!query.fromClause) throw new Error('must have FROM clause'); |
| return SqlQuery.selectStarFrom(query.fromClause); |
| } else { |
| return query |
| .changeSelectExpressions([SqlStar.PLAIN]) |
| .changeLimitValue(undefined) |
| .changeContextStatements([]); |
| } |
| } |
| |
| static fromIntrospectResult( |
| query: SqlQuery, |
| baseColumns: readonly Column[], |
| columns: readonly Column[], |
| ): QuerySource { |
| let effectiveColumns = columns; |
| if (query.getSelectExpressionsArray().some(ex => ex instanceof SqlStar)) { |
| // The query has a star so carefully pick the columns that make sense |
| effectiveColumns = columns.filter( |
| c => c.sqlType !== 'OTHER' || c.nativeType === 'COMPLEX<json>', |
| ); |
| } |
| |
| let measures = Measure.extractQueryMeasures(query); |
| if (!measures.length) { |
| const countColumn = columns.find(c => c.name === 'count' || c.name === '__count'); |
| measures = [ |
| countColumn |
| ? new Measure({ |
| expression: F.sum(C(countColumn.name)), |
| as: 'Count', |
| }) |
| : Measure.COUNT, |
| ]; |
| |
| measures = [ |
| ...measures, |
| ...filterMap(columns, column => |
| column.nativeType?.startsWith('COMPLEX<') |
| ? Measure.getPossibleMeasuresForColumn(column)[0] |
| : undefined, |
| ), |
| ]; |
| } |
| |
| return new QuerySource({ |
| query, |
| baseColumns, |
| columns: effectiveColumns, |
| measures, |
| }); |
| } |
| |
| public readonly query: SqlQuery; |
| public readonly baseColumns: readonly Column[]; |
| public readonly columns: readonly Column[]; |
| public readonly measures: Measure[]; |
| |
| constructor(value: QuerySourceValue) { |
| this.query = value.query; |
| this.baseColumns = value.baseColumns; |
| this.columns = value.columns; |
| this.measures = value.measures; |
| } |
| |
| public valueOf(): QuerySourceValue { |
| return { |
| query: this.query, |
| baseColumns: this.baseColumns, |
| columns: this.columns, |
| measures: this.measures, |
| }; |
| } |
| |
| public getInitQuery(where?: SqlExpression): SqlQuery { |
| return SqlQuery.from(this.query.as('t')).changeWhereExpression(where); |
| } |
| |
| public getInitBaseQuery(): SqlQuery { |
| return SqlQuery.from(QuerySource.stripToBaseSource(this.query).as('t')); |
| } |
| |
| private materializeStarIfNeeded(): SqlQuery { |
| const { query, columns, measures } = this; |
| let columnsToExpand = columns.map(c => c.name); |
| const selectExpressions = query.getSelectExpressionsArray(); |
| let starCount = 0; |
| for (const selectExpression of selectExpressions) { |
| if (selectExpression instanceof SqlStar) { |
| starCount++; |
| continue; |
| } |
| const outputName = selectExpression.getOutputName(); |
| if (!outputName) continue; |
| columnsToExpand = columnsToExpand.filter(c => c !== outputName); |
| } |
| if (starCount === 0) return query; |
| if (starCount > 1) throw new Error('can not handle multiple stars'); |
| |
| return Measure.addMeasuresToQuery( |
| query |
| .changeSelectExpressions( |
| selectExpressions.flatMap(selectExpression => |
| selectExpression instanceof SqlStar ? columnsToExpand.map(c => C(c)) : selectExpression, |
| ), |
| ) |
| .prettify(), |
| measures, |
| ); |
| } |
| |
| public getFirstAggregateMeasure(): Measure | undefined { |
| return this.measures[0]?.toAggregateBasedMeasure(); |
| } |
| |
| public getFirstAggregateMeasureArray(): Measure[] { |
| return this.measures.length ? [this.measures[0].toAggregateBasedMeasure()] : []; |
| } |
| |
| public getAvailableName(prefix: string): string { |
| let columnName = prefix; |
| let counter = 1; |
| while (this.nameInUse(columnName)) { |
| counter++; |
| columnName = `${prefix}_${counter}`; |
| } |
| return columnName; |
| } |
| |
| public nameInUse(nameToCheck: string): boolean { |
| return this.hasColumnByName(nameToCheck) || this.hasMeasureByName(nameToCheck); |
| } |
| |
| public hasColumnByName(name: string): boolean { |
| return this.columns.some(c => c.name === name); |
| } |
| |
| public hasMeasureByName(name: string): boolean { |
| return this.measures.some(m => m.name === name); |
| } |
| |
| public hasBaseTimeColumn(): boolean { |
| return this.baseColumns.some(column => column.isTimeColumn()); |
| } |
| |
| public getSourceExpressionForColumn(outputName: string): SqlExpression { |
| const selectExpressionsArray = this.query.getSelectExpressionsArray(); |
| |
| const sourceExpression = selectExpressionsArray.find(ex => ex.getOutputName() === outputName); |
| if (sourceExpression) return sourceExpression; |
| |
| const m = /^EXPR\$(\d+)$/.exec(outputName); |
| if (m) { |
| const index = parseInt(m[1], 10); |
| if (selectExpressionsArray[index]) return selectExpressionsArray[index]; |
| } |
| |
| return C(outputName); |
| } |
| |
| private getSourceToBaseSubstitutions(): Map<string, SqlExpression> { |
| return new Map<string, SqlExpression>( |
| filterMap(this.query.getSelectExpressionsArray(), ex => { |
| const outputName = ex.getOutputName(); |
| const underlyingExpression = ex.getUnderlyingExpression(); |
| if (!outputName || underlyingExpression.getOutputName() === outputName) return; |
| return [outputName, underlyingExpression]; |
| }), |
| ); |
| } |
| |
| public transformExpressionToBaseColumns(expression: SqlExpression): SqlExpression { |
| const sourceToBaseSubstitutions = this.getSourceToBaseSubstitutions(); |
| return expression.walk(ex => { |
| if (ex instanceof SqlColumn) { |
| return sourceToBaseSubstitutions.get(ex.getName()) || ex; |
| } |
| return ex; |
| }) as SqlExpression; |
| } |
| |
| public addWhereClause(clause: SqlExpression): SqlQuery { |
| return this.query.addWhere(clause); |
| } |
| |
| public addColumn(newExpression: SqlExpression): SqlQuery { |
| const noStarQuery = this.materializeStarIfNeeded(); |
| return noStarQuery.addSelect(newExpression); |
| } |
| |
| public addColumnAfter(neighborName: string, ...newExpressions: SqlExpression[]): SqlQuery { |
| const noStarQuery = this.materializeStarIfNeeded(); |
| return noStarQuery.changeSelectExpressions( |
| noStarQuery |
| .getSelectExpressionsArray() |
| .flatMap(ex => (ex.getOutputName() === neighborName ? [ex, ...newExpressions] : ex)), |
| ); |
| } |
| |
| public changeColumn(oldName: string, newExpression: SqlExpression): SqlQuery { |
| const noStarQuery = this.materializeStarIfNeeded(); |
| return noStarQuery.changeSelectExpressions( |
| noStarQuery |
| .getSelectExpressionsArray() |
| .map(ex => (ex.getOutputName() === oldName ? newExpression : ex)), |
| ); |
| } |
| |
| public deleteColumn(outputName: string): SqlQuery { |
| const noStarQuery = this.materializeStarIfNeeded(); |
| return noStarQuery.changeSelectExpressions( |
| noStarQuery.getSelectExpressionsArray().filter(ex => ex.getOutputName() !== outputName), |
| ); |
| } |
| |
| public getColumnNameMap(nameTransform: (columnName: string) => string): Map<string, string> { |
| return new Map(this.columns.map(column => [column.name, nameTransform(column.name)])); |
| } |
| |
| public applyColumnNameMap(columnNameMap: Map<string, string>): SqlQuery { |
| const noStarQuery = this.materializeStarIfNeeded(); |
| return noStarQuery.changeSelectExpressions( |
| noStarQuery.getSelectExpressionsArray().map(ex => { |
| const outputName = ex.getOutputName(); |
| if (!outputName) return ex; |
| const newOutputName = columnNameMap.get(outputName); |
| if (!newOutputName || newOutputName === outputName) return ex; |
| return ex.as(newOutputName); |
| }), |
| ); |
| } |
| |
| // ------------------------------------ |
| |
| public addMeasure(measure: Measure): SqlQuery { |
| const noStarQuery = this.materializeStarIfNeeded(); |
| return Measure.addMeasuresToQuery(noStarQuery, this.measures.concat(measure)); |
| } |
| |
| public addMeasureAfter(neighborName: string, newMeasure: Measure): SqlQuery { |
| const noStarQuery = this.materializeStarIfNeeded(); |
| return Measure.addMeasuresToQuery( |
| noStarQuery, |
| this.measures.flatMap(m => (m.name === neighborName ? [m, newMeasure] : m)), |
| ); |
| } |
| |
| public changeMeasure(oldName: string, newMeasure: Measure): SqlQuery { |
| const noStarQuery = this.materializeStarIfNeeded(); |
| return Measure.addMeasuresToQuery( |
| noStarQuery, |
| this.measures.map(m => (m.name === oldName ? newMeasure : m)), |
| ); |
| } |
| |
| public deleteMeasure(measureName: string): SqlQuery { |
| const noStarQuery = this.materializeStarIfNeeded(); |
| return Measure.addMeasuresToQuery( |
| noStarQuery, |
| this.measures.filter(m => m.name !== measureName), |
| ); |
| } |
| |
| // -------------------------------- |
| |
| public restrictWhere(where: SqlExpression): SqlExpression { |
| const { columns } = this; |
| const parts = where.decomposeViaAnd(); |
| const filterParts = parts.filter(ex => expressionWithinColumns(ex, columns)); |
| if (parts.length === filterParts.length) return where; |
| return SqlExpression.and(...filterParts); |
| } |
| |
| public restrictParameterValues( |
| parameterValues: ParameterValues, |
| parameters: Parameters, |
| where: SqlExpression, |
| ): ParameterValues { |
| return mapRecordOrReturn(parameterValues, (parameterValue, k) => { |
| const parameter = parameters[k]; |
| if (!parameter) return; |
| return this.restrictParameterValue(parameterValue, parameter, where, parameterValues); |
| }); |
| } |
| |
| private restrictParameterValue( |
| parameterValue: any, |
| parameter: ParameterDefinition, |
| where: SqlExpression, |
| parameterValues: ParameterValues, |
| ): any { |
| if (typeof parameterValue !== 'undefined') { |
| switch (parameter.type) { |
| case 'number': { |
| if (typeof parameter.min === 'number') { |
| parameterValue = Math.max(parameterValue, parameter.min); |
| } |
| if (typeof parameter.max === 'number') { |
| parameterValue = Math.min(parameterValue, parameter.max); |
| } |
| return parameterValue; |
| } |
| |
| case 'option': { |
| const options = evaluateFunctor(parameter.options, parameterValues, this, where); |
| if (!options || !options.includes(parameterValue as OptionValue)) return; |
| return parameterValue as OptionValue; |
| } |
| |
| case 'options': { |
| const options = evaluateFunctor(parameter.options, {}, this, where) || []; |
| return filterOrReturn<OptionValue>(parameterValue, v => options.includes(v)); |
| } |
| |
| case 'expression': |
| if (!this.validateExpressionMeta(parameterValue)) return; |
| break; |
| |
| case 'measure': |
| if (!this.validateMeasure(parameterValue)) return; |
| break; |
| |
| case 'expressions': |
| return filterOrReturn<ExpressionMeta>(parameterValue, v => |
| this.validateExpressionMeta(v), |
| ); |
| |
| case 'measures': |
| return filterOrReturn<Measure>(parameterValue, v => this.validateMeasure(v)); |
| |
| default: |
| break; |
| } |
| } |
| return parameterValue; |
| } |
| |
| public validateExpressionMeta(e: ExpressionMeta | undefined): e is ExpressionMeta { |
| if (!(e instanceof ExpressionMeta)) return false; |
| return expressionWithinColumns(e.expression, this.columns); |
| } |
| |
| public validateMeasure(m: Measure | undefined): m is Measure { |
| if (!(m instanceof Measure)) return false; |
| |
| const usedAggregates = m.getUsedAggregates(); |
| if (usedAggregates.some(usedAggregate => !this.hasMeasureByName(usedAggregate))) return false; |
| |
| return expressionWithinColumns(m.expression, this.columns); |
| } |
| } |