apache · kfaraz · Oct 3, 2024 · Oct 3, 2024
diff --git a/web-console/src/views/explore-view/components/module-pane/module-pane.scss b/web-console/src/views/explore-view/components/module-pane/module-pane.scss
@@ -35,9 +35,17 @@
 
     & > .issue {
       position: absolute;
+      top: 0;
+      left: 0;
       width: 100%;
       height: 100%;
     }
+
+    .loader {
+      position: absolute;
+      top: 0;
+      left: 0;
+    }
   }
 
   .tile-content {

diff --git a/web-console/src/views/explore-view/explore-view.tsx b/web-console/src/views/explore-view/explore-view.tsx
@@ -22,6 +22,7 @@ import { Button, Intent, Menu, MenuDivider, MenuItem } from '@blueprintjs/core';
 import { IconNames } from '@blueprintjs/icons';
 import type { Column, QueryResult, SqlExpression } from '@druid-toolkit/query';
 import { QueryRunner, SqlQuery } from '@druid-toolkit/query';
+import type { CancelToken } from 'axios';
 import classNames from 'classnames';
 import copy from 'copy-to-clipboard';
 import React, { useEffect, useMemo, useRef, useState } from 'react';
@@ -79,27 +80,33 @@ const queryRunner = new QueryRunner({
   },
 });
 
-async function runSqlQuery(query: string | SqlQuery): Promise<QueryResult> {
+async function runSqlQuery(
+  query: string | SqlQuery,
+  cancelToken?: CancelToken,
+): Promise<QueryResult> {
   try {
     return await queryRunner.runQuery({
       query,
       defaultQueryContext: {
         sqlStringifyArrays: false,
       },
+      cancelToken,
     });
   } catch (e) {
     throw new DruidError(e);
   }
 }
 
-async function introspectSource(source: string): Promise<QuerySource> {
+async function introspectSource(source: string, cancelToken?: CancelToken): Promise<QuerySource> {
   const query = SqlQuery.parse(source);
   const introspectResult = await runSqlQuery(QuerySource.makeLimitZeroIntrospectionQuery(query));
 
+  cancelToken?.throwIfRequested();
   const baseIntrospectResult = QuerySource.isSingleStarQuery(query)
     ? introspectResult
     : await runSqlQuery(
         QuerySource.makeLimitZeroIntrospectionQuery(QuerySource.stripToBaseSource(query)),
+        cancelToken,
       );
 
   return QuerySource.fromIntrospectResult(
@@ -238,11 +245,15 @@ export const ExploreView = React.memo(function ExploreView() {
   const querySource = querySourceState.getSomeData();
 
   const runSqlPlusQuery = useMemo(() => {
-    return async (query: string | SqlQuery) => {
+    return async (query: string | SqlQuery, cancelToken?: CancelToken) => {
       if (!querySource) throw new Error('no querySource');
-      return await runSqlQuery(
-        await rewriteMaxDataTime(rewriteAggregate(SqlQuery.parse(query), querySource.measures)),
-      );
+      const parsedQuery = SqlQuery.parse(query);
+      return (
+        await runSqlQuery(
+          await rewriteMaxDataTime(rewriteAggregate(parsedQuery, querySource.measures)),
+          cancelToken,
+        )
+      ).attachQuery({ query: '' }, parsedQuery);
     };
   }, [querySource]);
 

diff --git a/web-console/src/views/explore-view/models/query-source.ts b/web-console/src/views/explore-view/models/query-source.ts
@@ -44,31 +44,6 @@ export class QuerySource {
     );
   }
 
-  static materializeStarIfNeeded(query: SqlQuery, columns: readonly Column[]): SqlQuery {
-    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 query
-      .changeSelectExpressions(
-        selectExpressions.flatMap(selectExpression =>
-          selectExpression instanceof SqlStar ? columnsToExpand.map(c => C(c)) : selectExpression,
-        ),
-      )
-      .prettify();
-  }
-
   static isSingleStarQuery(query: SqlQuery): boolean {
     const selectExpressions = query.getSelectExpressionsArray();
     return selectExpressions.length === 1 && selectExpressions[0] instanceof SqlStar;
@@ -151,6 +126,35 @@ export class QuerySource {
     };
   }
 
+  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();
   }
@@ -226,12 +230,12 @@ export class QuerySource {
   }
 
   public addColumn(newExpression: SqlExpression): SqlQuery {
-    const noStarQuery = QuerySource.materializeStarIfNeeded(this.query, this.columns);
+    const noStarQuery = this.materializeStarIfNeeded();
     return noStarQuery.addSelect(newExpression);
   }
 
   public addColumnAfter(neighborName: string, ...newExpressions: SqlExpression[]): SqlQuery {
-    const noStarQuery = QuerySource.materializeStarIfNeeded(this.query, this.columns);
+    const noStarQuery = this.materializeStarIfNeeded();
     return noStarQuery.changeSelectExpressions(
       noStarQuery
         .getSelectExpressionsArray()
@@ -240,7 +244,7 @@ export class QuerySource {
   }
 
   public changeColumn(oldName: string, newExpression: SqlExpression): SqlQuery {
-    const noStarQuery = QuerySource.materializeStarIfNeeded(this.query, this.columns);
+    const noStarQuery = this.materializeStarIfNeeded();
     return noStarQuery.changeSelectExpressions(
       noStarQuery
         .getSelectExpressionsArray()
@@ -249,7 +253,7 @@ export class QuerySource {
   }
 
   public deleteColumn(outputName: string): SqlQuery {
-    const noStarQuery = QuerySource.materializeStarIfNeeded(this.query, this.columns);
+    const noStarQuery = this.materializeStarIfNeeded();
     return noStarQuery.changeSelectExpressions(
       noStarQuery.getSelectExpressionsArray().filter(ex => ex.getOutputName() !== outputName),
     );
@@ -260,7 +264,7 @@ export class QuerySource {
   }
 
   public applyColumnNameMap(columnNameMap: Map<string, string>): SqlQuery {
-    const noStarQuery = QuerySource.materializeStarIfNeeded(this.query, this.columns);
+    const noStarQuery = this.materializeStarIfNeeded();
     return noStarQuery.changeSelectExpressions(
       noStarQuery.getSelectExpressionsArray().map(ex => {
         const outputName = ex.getOutputName();
@@ -275,28 +279,28 @@ export class QuerySource {
   // ------------------------------------
 
   public addMeasure(measure: Measure): SqlQuery {
-    const noStarQuery = QuerySource.materializeStarIfNeeded(this.query, this.columns);
+    const noStarQuery = this.materializeStarIfNeeded();
     return Measure.addMeasuresToQuery(noStarQuery, this.measures.concat(measure));
   }
 
   public addMeasureAfter(neighborName: string, newMeasure: Measure): SqlQuery {
-    const noStarQuery = QuerySource.materializeStarIfNeeded(this.query, this.columns);
+    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 = QuerySource.materializeStarIfNeeded(this.query, this.columns);
+    const noStarQuery = this.materializeStarIfNeeded();
     return Measure.addMeasuresToQuery(
       noStarQuery,
       this.measures.map(m => (m.name === oldName ? newMeasure : m)),
     );
   }
 
   public deleteMeasure(measureName: string): SqlQuery {
-    const noStarQuery = QuerySource.materializeStarIfNeeded(this.query, this.columns);
+    const noStarQuery = this.materializeStarIfNeeded();
     return Measure.addMeasuresToQuery(
       noStarQuery,
       this.measures.filter(m => m.name !== measureName),

diff --git a/web-console/src/views/explore-view/module-repository/module-repository.ts b/web-console/src/views/explore-view/module-repository/module-repository.ts
@@ -17,6 +17,7 @@
  */
 
 import type { QueryResult, SqlExpression, SqlQuery } from '@druid-toolkit/query';
+import type { CancelToken } from 'axios';
 
 import type { ParameterDefinition, QuerySource, Stage } from '../models';
 
@@ -34,7 +35,7 @@ interface ModuleComponentProps<P> {
   setWhere(where: SqlExpression): void;
   parameterValues: P;
   setParameterValues: (parameters: Partial<P>) => void;
-  runSqlQuery(query: string | SqlQuery): Promise<QueryResult>;
+  runSqlQuery(query: string | SqlQuery, cancelToken?: CancelToken): Promise<QueryResult>;
 }
 
 export class ModuleRepository {

diff --git a/web-console/src/views/explore-view/modules/bar-chart-module.tsx b/web-console/src/views/explore-view/modules/bar-chart-module.tsx
@@ -21,6 +21,7 @@ import type { ECharts } from 'echarts';
 import * as echarts from 'echarts';
 import React, { useEffect, useMemo, useRef } from 'react';
 
+import { Loader } from '../../../components';
 import { useQueryManager } from '../../../hooks';
 import { formatEmpty } from '../../../utils';
 import { Issue } from '../components';
@@ -90,10 +91,10 @@ ModuleRepository.registerModule<BarChartParameterValues>({
         .changeLimitValue(limit);
     }, [querySource, where, splitColumn, measure, measureToSort, limit]);
 
-    const [sourceDataState] = useQueryManager({
+    const [sourceDataState, queryManager] = useQueryManager({
       query: dataQuery,
-      processQuery: async (query: SqlQuery) => {
-        return (await runSqlQuery(query)).toObjectArray();
+      processQuery: async (query, cancelToken) => {
+        return (await runSqlQuery(query, cancelToken)).toObjectArray();
       },
     });
 
@@ -203,6 +204,9 @@ ModuleRepository.registerModule<BarChartParameterValues>({
           }}
         />
         {errorMessage && <Issue issue={errorMessage} />}
+        {sourceDataState.loading && (
+          <Loader cancelText="Cancel query" onCancel={() => queryManager.cancelCurrent()} />
+        )}
       </div>
     );
   },

diff --git a/web-console/src/views/explore-view/modules/grouping-table-module.tsx b/web-console/src/views/explore-view/modules/grouping-table-module.tsx
@@ -17,7 +17,7 @@
  */
 
 import { Button } from '@blueprintjs/core';
-import type { SqlOrderByDirection } from '@druid-toolkit/query';
+import type { SqlExpression, SqlOrderByDirection } from '@druid-toolkit/query';
 import { C, F, SqlQuery } from '@druid-toolkit/query';
 import React, { useMemo } from 'react';
 
@@ -43,6 +43,11 @@ import './grouping-table-module.scss';
 // when ordering on non __time is more robust
 const NEEDS_GROUPING_TO_ORDER = true;
 
+interface QueryAndMore {
+  originalWhere: SqlExpression;
+  queryAndHints: QueryAndHints;
+}
+
 interface GroupingTableParameterValues {
   splitColumns: ExpressionMeta[];
   timeBucket: string;
@@ -216,47 +221,58 @@ ModuleRepository.registerModule<GroupingTableParameterValues>({
         .changeLimitValue(maxPivotValues);
     }, [querySource.query, parameterValues]);
 
-    const [pivotValueState] = useQueryManager({
+    const [pivotValueState, queryManager] = useQueryManager({
       query: pivotValueQuery,
       processQuery: async (pivotValueQuery: SqlQuery) => {
         return (await runSqlQuery(pivotValueQuery)).getColumnByName('v') as string[];
       },
     });
 
-    const queryAndHints = useMemo((): QueryAndHints | undefined => {
+    const queryAndMore = useMemo((): QueryAndMore | undefined => {
       const pivotValues = pivotValueState.data;
       if (parameterValues.pivotColumn && !pivotValues) return;
       const { orderByColumn, orderByDirection } = parameterValues;
       const orderBy = orderByColumn
         ? C(orderByColumn).toOrderByExpression(orderByDirection)
         : undefined;
 
-      return makeTableQueryAndHints({
-        source: querySource.query,
-        where,
-        splitColumns: parameterValues.splitColumns,
-        timeBucket: parameterValues.timeBucket,
-        showColumns: parameterValues.showColumns,
-        multipleValueMode: parameterValues.multipleValueMode,
-        pivotColumn: parameterValues.pivotColumn,
-        pivotValues,
-        measures: parameterValues.measures,
-        compares: parameterValues.compares || [],
-        compareStrategy: parameterValues.compareStrategy,
-        compareTypes: parameterValues.compareTypes,
-        restrictTop: parameterValues.restrictTop,
-        maxRows: parameterValues.maxRows,
-        orderBy,
-        useGroupingToOrderSubQueries: NEEDS_GROUPING_TO_ORDER,
-      });
+      return {
+        originalWhere: where,
+        queryAndHints: makeTableQueryAndHints({
+          source: querySource.query,
+          where,
+          splitColumns: parameterValues.splitColumns,
+          timeBucket: parameterValues.timeBucket,
+          showColumns: parameterValues.showColumns,
+          multipleValueMode: parameterValues.multipleValueMode,
+          pivotColumn: parameterValues.pivotColumn,
+          pivotValues,
+          measures: parameterValues.measures,
+          compares: parameterValues.compares || [],
+          compareStrategy: parameterValues.compareStrategy,
+          compareTypes: parameterValues.compareTypes,
+          restrictTop: parameterValues.restrictTop,
+          maxRows: parameterValues.maxRows,
+          orderBy,
+          useGroupingToOrderSubQueries: NEEDS_GROUPING_TO_ORDER,
+        }),
+      };
     }, [querySource.query, where, parameterValues, pivotValueState.data]);
 
     const [resultState] = useQueryManager({
-      query: queryAndHints,
-      processQuery: async (queryAndHints: QueryAndHints) => {
+      query: queryAndMore,
+      processQuery: async (queryAndMore, cancelToken) => {
+        const { originalWhere, queryAndHints } = queryAndMore;
         const { query, columnHints } = queryAndHints;
+        let result = await runSqlQuery(query, cancelToken);
+        if (result.sqlQuery) {
+          result = result.attachQuery(
+            { query: '' },
+            result.sqlQuery.changeWhereExpression(originalWhere),
+          );
+        }
         return {
-          result: await runSqlQuery(query),
+          result,
           columnHints,
         };
       },
@@ -297,7 +313,9 @@ ModuleRepository.registerModule<GroupingTableParameterValues>({
             initPageSize={calculateInitPageSize(stage.height)}
           />
         ) : undefined}
-        {resultState.loading && <Loader />}
+        {resultState.loading && (
+          <Loader cancelText="Cancel query" onCancel={() => queryManager.cancelCurrent()} />
+        )}
       </div>
     );
   },