Web console: Improve explore max time cancelation (#18830)

* add a timeout

* hook up cancel for table

* remove console.log

Author: vogievetsky

web-console/src/utils/general.spec.ts

@@ -34,6 +34,7 @@ import {
   OVERLAY_OPEN_SELECTOR,
   parseCsvLine,
   swapElements,
+  wait,
 } from './general';
 
 describe('general', () => {
@@ -231,4 +232,118 @@ describe('general', () => {
       expect(OVERLAY_OPEN_SELECTOR).toEqual('.bp5-portal .bp5-overlay-open');
     });
   });
+
+  describe('wait', () => {
+    beforeEach(() => {
+      jest.useFakeTimers();
+    });
+
+    afterEach(() => {
+      jest.useRealTimers();
+    });
+
+    it('resolves after the specified time', async () => {
+      const promise = wait(100);
+      expect(promise).toBeInstanceOf(Promise);
+
+      jest.advanceTimersByTime(99);
+      await Promise.resolve(); // Let microtasks run
+      expect(promise).not.toBe(await Promise.race([promise, Promise.resolve('pending')]));
+
+      jest.advanceTimersByTime(1);
+      await expect(promise).resolves.toBeUndefined();
+    });
+
+    it('works without a signal (backward compatibility)', async () => {
+      const promise = wait(50);
+      jest.advanceTimersByTime(50);
+      await expect(promise).resolves.toBeUndefined();
+    });
+
+    it('resolves normally when signal does not abort', async () => {
+      const controller = new AbortController();
+      const promise = wait(100, controller.signal);
+
+      jest.advanceTimersByTime(100);
+      await expect(promise).resolves.toBeUndefined();
+    });
+
+    it('rejects when signal aborts before timeout', async () => {
+      const controller = new AbortController();
+      const promise = wait(100, controller.signal);
+
+      jest.advanceTimersByTime(50);
+      controller.abort();
+
+      await expect(promise).rejects.toThrow('Aborted');
+    });
+
+    it('rejects immediately if signal is already aborted', async () => {
+      const controller = new AbortController();
+      controller.abort();
+
+      const promise = wait(100, controller.signal);
+      await expect(promise).rejects.toThrow('Aborted');
+
+      // Timer should not have been created
+      expect(jest.getTimerCount()).toBe(0);
+    });
+
+    it('cleans up timeout when aborted', async () => {
+      const controller = new AbortController();
+      const promise = wait(100, controller.signal);
+
+      expect(jest.getTimerCount()).toBe(1);
+
+      controller.abort();
+
+      try {
+        await promise;
+      } catch {
+        // Expected
+      }
+
+      // Timer should be cleaned up
+      expect(jest.getTimerCount()).toBe(0);
+    });
+
+    it('cleans up event listener when timeout completes', async () => {
+      const controller = new AbortController();
+      const removeEventListenerSpy = jest.spyOn(controller.signal, 'removeEventListener');
+
+      const promise = wait(100, controller.signal);
+      jest.advanceTimersByTime(100);
+      await promise;
+
+      expect(removeEventListenerSpy).toHaveBeenCalledWith('abort', expect.any(Function));
+    });
+
+    it('cleans up event listener when aborted', async () => {
+      const controller = new AbortController();
+      const removeEventListenerSpy = jest.spyOn(controller.signal, 'removeEventListener');
+
+      const promise = wait(100, controller.signal);
+      controller.abort();
+
+      try {
+        await promise;
+      } catch {
+        // Expected
+      }
+
+      expect(removeEventListenerSpy).toHaveBeenCalledWith('abort', expect.any(Function));
+    });
+
+    it('handles multiple waits with same signal', async () => {
+      const controller = new AbortController();
+      const promise1 = wait(100, controller.signal);
+      const promise2 = wait(200, controller.signal);
+
+      jest.advanceTimersByTime(50);
+      controller.abort();
+
+      await expect(promise1).rejects.toThrow('Aborted');
+      await expect(promise2).rejects.toThrow('Aborted');
+    });
+  });
 });

web-console/src/utils/general.tsx

@@ -63,9 +63,29 @@ export function arraysEqualByElement<T>(xs: T[], ys: T[]): boolean {
   return xs.length === ys.length && xs.every((x, i) => x === ys[i]);
 }
 
-export function wait(ms: number): Promise<void> {
-  return new Promise(resolve => {
-    setTimeout(resolve, ms);
+export function wait(ms: number, signal?: AbortSignal): Promise<void> {
+  return new Promise((resolve, reject) => {
+    if (signal?.aborted) {
+      reject(new Error('Aborted'));
+      return;
+    }
+
+    const timeoutId = setTimeout(() => {
+      cleanup();
+      resolve();
+    }, ms);
+
+    const onAbort = () => {
+      cleanup();
+      reject(new Error('Aborted'));
+    };
+
+    const cleanup = () => {
+      clearTimeout(timeoutId);
+      signal?.removeEventListener('abort', onAbort);
+    };
+
+    signal?.addEventListener('abort', onAbort);
   });
 }
 

web-console/src/views/explore-view/explore-view.tsx

@@ -229,6 +229,7 @@ export const ExploreView = React.memo(function ExploreView({ capabilities }: Exp
 
       const { query: rewrittenQuery, maxTime } = await rewriteMaxDataTime(
         rewriteAggregate(parsedQuery, querySource.measures),
+        signal,
       );
       const results = await runSqlQuery(rewrittenQuery, queryTimezone, signal);
 

web-console/src/views/explore-view/modules/grouping-table-module/grouping-table-module.tsx

@@ -19,7 +19,7 @@
 import { Button } from '@blueprintjs/core';
 import { IconNames } from '@blueprintjs/icons';
 import type { Timezone } from 'chronoshift';
-import type { SqlExpression, SqlOrderByDirection, SqlQuery } from 'druid-query-toolkit';
+import type { SqlExpression, SqlOrderByDirection } from 'druid-query-toolkit';
 import { C, F } from 'druid-query-toolkit';
 import { useMemo } from 'react';
 
@@ -233,10 +233,10 @@ ModuleRepository.registerModule<GroupingTableParameterValues>({
         .changeLimitValue(maxPivotValues);
     }, [querySource, where, moduleWhere, parameterValues]);
 
-    const [pivotValueState, queryManager] = useQueryManager({
+    const [pivotValueState, pivotValueQueryManager] = useQueryManager({
       query: pivotValueQuery,
-      processQuery: async (pivotValueQuery: SqlQuery) => {
-        return (await runSqlQuery(pivotValueQuery)).getColumnByName('v') as string[];
+      processQuery: async (pivotValueQuery, signal) => {
+        return (await runSqlQuery(pivotValueQuery, signal)).getColumnByName('v') as string[];
       },
     });
 
@@ -272,11 +272,12 @@ ModuleRepository.registerModule<GroupingTableParameterValues>({
       };
     }, [querySource.query, timezone, where, parameterValues, pivotValueState.data]);
 
-    const [resultState] = useQueryManager({
+    const [resultState, resultQueryManager] = useQueryManager({
       query: queryAndMore,
       processQuery: async (queryAndMore, signal) => {
         const { timezone, globalWhere, queryAndHints } = queryAndMore;
         const { query, columnHints } = queryAndHints;
+
         let result = await runSqlQuery({ query, timezone }, signal);
         if (result.sqlQuery) {
           result = result.attachQuery(
@@ -328,7 +329,13 @@ ModuleRepository.registerModule<GroupingTableParameterValues>({
           />
         ) : undefined}
         {resultState.loading && (
-          <Loader cancelText="Cancel query" onCancel={() => queryManager.cancelCurrent()} />
+          <Loader
+            cancelText="Cancel query"
+            onCancel={() => {
+              pivotValueQueryManager.cancelCurrent();
+              resultQueryManager.cancelCurrent();
+            }}
+          />
         )}
       </div>
     );

web-console/src/views/explore-view/query-macros/max-data-time.ts

@@ -28,6 +28,7 @@ const tablesForWhichWeCouldNotDetermineMaxTime = new Set<string>();
 
 export async function rewriteMaxDataTime(
   query: SqlQuery,
+  signal?: AbortSignal,
 ): Promise<{ query: SqlQuery; maxTime?: Date }> {
   if (!query.containsFunction('MAX_DATA_TIME')) return { query };
 
@@ -36,7 +37,7 @@ export async function rewriteMaxDataTime(
 
   let maxTime: Date;
   try {
-    maxTime = await getMaxTimeForTable(tableName);
+    maxTime = await getMaxTimeForTable(tableName, signal);
   } catch (error) {
     if (!tablesForWhichWeCouldNotDetermineMaxTime.has(tableName)) {
       tablesForWhichWeCouldNotDetermineMaxTime.add(tableName);

web-console/src/views/explore-view/utils/max-time-for-table.ts

@@ -27,7 +27,7 @@ let lastMaxTimeTable: string | undefined;
 let lastMaxTimeValue: Date | undefined;
 let lastMaxTimeTimestamp = 0;
 
-export async function getMaxTimeForTable(tableName: string): Promise<Date> {
+export async function getMaxTimeForTable(tableName: string, signal?: AbortSignal): Promise<Date> {
   // micro-cache get
   if (
     lastMaxTimeTable === tableName &&
@@ -37,9 +37,15 @@ export async function getMaxTimeForTable(tableName: string): Promise<Date> {
     return lastMaxTimeValue;
   }
 
-  const d = await queryDruidSql({
-    query: sql`SELECT MAX(__time) AS "maxTime" FROM ${T(tableName)}`,
-  });
+  const d = await queryDruidSql(
+    {
+      query: sql`SELECT MAX(__time) AS "maxTime" FROM ${T(tableName)}`,
+      context: {
+        timeout: 2000, // We expect this query to be superfast
+      },
+    },
+    signal,
+  );
 
   const maxTimeRaw = deepGet(d, '0.maxTime');
   const maxTime = new Date(maxTimeRaw);