fix: Drill by with GLOBAL_ASYNC_QUERIES (#27066)

superset-frontend/src/components/Chart/DrillBy/DrillByModal.tsx

@@ -54,11 +54,12 @@ import {
   LOG_ACTIONS_DRILL_BY_MODAL_OPENED,
   LOG_ACTIONS_FURTHER_DRILL_BY,
 } from 'src/logger/LogUtils';
+import { getQuerySettings } from 'src/explore/exploreUtils';
 import { Dataset, DrillByType } from '../types';
 import DrillByChart from './DrillByChart';
 import { ContextMenuItem } from '../ChartContextMenu/ChartContextMenu';
 import { useContextMenu } from '../ChartContextMenu/useContextMenu';
-import { getChartDataRequest } from '../chartAction';
+import { getChartDataRequest, handleChartDataResponse } from '../chartAction';
 import { useDisplayModeToggle } from './useDisplayModeToggle';
 import {
   DrillByBreadcrumb,
@@ -390,13 +391,17 @@ export default function DrillByModal({
 
   useEffect(() => {
     if (drilledFormData) {
+      const [useLegacyApi] = getQuerySettings(drilledFormData);
       setIsChartDataLoading(true);
       setChartDataResult(undefined);
       getChartDataRequest({
         formData: drilledFormData,
       })
-        .then(({ json }) => {
-          setChartDataResult(json.result);
+        .then(({ response, json }) =>
+          handleChartDataResponse(response, json, useLegacyApi),
+        )
+        .then(queriesResponse => {
+          setChartDataResult(queriesResponse);
         })
         .catch(() => {
           addDangerToast(t('Failed to load chart data.'));

superset-frontend/src/components/Chart/chartAction.js

@@ -376,6 +376,29 @@ export function addChart(chart, key) {
   return { type: ADD_CHART, chart, key };
 }
 
+export function handleChartDataResponse(response, json, useLegacyApi) {
+  if (isFeatureEnabled(FeatureFlag.GlobalAsyncQueries)) {
+    // deal with getChartDataRequest transforming the response data
+    const result = 'result' in json ? json.result : json;
+    switch (response.status) {
+      case 200:
+        // Query results returned synchronously, meaning query was already cached.
+        return Promise.resolve(result);
+      case 202:
+        // Query is running asynchronously and we must await the results
+        if (useLegacyApi) {
+          return waitForAsyncData(result[0]);
+        }
+        return waitForAsyncData(result);
+      default:
+        throw new Error(
+          `Received unexpected response status (${response.status}) while fetching chart data`,
+        );
+    }
+  }
+  return json.result;
+}
+
 export function exploreJSON(
   formData,
   force = false,
@@ -411,31 +434,11 @@ export function exploreJSON(
 
     dispatch(chartUpdateStarted(controller, formData, key));
 
+    const [useLegacyApi] = getQuerySettings(formData);
     const chartDataRequestCaught = chartDataRequest
-      .then(({ response, json }) => {
-        if (isFeatureEnabled(FeatureFlag.GLOBAL_ASYNC_QUERIES)) {
-          // deal with getChartDataRequest transforming the response data
-          const result = 'result' in json ? json.result : json;
-          const [useLegacyApi] = getQuerySettings(formData);
-          switch (response.status) {
-            case 200:
-              // Query results returned synchronously, meaning query was already cached.
-              return Promise.resolve(result);
-            case 202:
-              // Query is running asynchronously and we must await the results
-              if (useLegacyApi) {
-                return waitForAsyncData(result[0]);
-              }
-              return waitForAsyncData(result);
-            default:
-              throw new Error(
-                `Received unexpected response status (${response.status}) while fetching chart data`,
-              );
-          }
-        }
-
-        return json.result;
-      })
+      .then(({ response, json }) =>
+        handleChartDataResponse(response, json, useLegacyApi),
+      )
       .then(queriesResponse => {
         queriesResponse.forEach(resultItem =>
           dispatch(
@@ -494,7 +497,6 @@ export function exploreJSON(
       });
 
     // only retrieve annotations when calling the legacy API
-    const [useLegacyApi] = getQuerySettings(formData);
     const annotationLayers = useLegacyApi
       ? formData.annotation_layers || []
       : [];

superset-frontend/src/components/Chart/chartActions.test.js

@@ -21,10 +21,12 @@ import fetchMock from 'fetch-mock';
 import sinon from 'sinon';
 
 import * as chartlib from '@superset-ui/core';
-import { SupersetClient } from '@superset-ui/core';
+import { FeatureFlag, SupersetClient } from '@superset-ui/core';
 import { LOG_EVENT } from 'src/logger/actions';
 import * as exploreUtils from 'src/explore/exploreUtils';
 import * as actions from 'src/components/Chart/chartAction';
+import * as asyncEvent from 'src/middleware/asyncEvent';
+import { handleChartDataResponse } from 'src/components/Chart/chartAction';
 
 describe('chart actions', () => {
   const MOCK_URL = '/mockURL';
@@ -33,6 +35,7 @@ describe('chart actions', () => {
   let getChartDataUriStub;
   let metadataRegistryStub;
   let buildQueryRegistryStub;
+  let waitForAsyncDataStub;
   let fakeMetadata;
 
   const setupDefaultFetchMock = () => {
@@ -66,6 +69,9 @@ describe('chart actions', () => {
           result_format: 'json',
         }),
       }));
+    waitForAsyncDataStub = sinon
+      .stub(asyncEvent, 'waitForAsyncData')
+      .callsFake(data => Promise.resolve(data));
   });
 
   afterEach(() => {
@@ -74,6 +80,11 @@ describe('chart actions', () => {
     fetchMock.resetHistory();
     metadataRegistryStub.restore();
     buildQueryRegistryStub.restore();
+    waitForAsyncDataStub.restore();
+
+    global.featureFlags = {
+      [FeatureFlag.GlobalAsyncQueries]: false,
+    };
   });
 
   describe('v1 API', () => {
@@ -114,6 +125,36 @@ describe('chart actions', () => {
       expect(fetchMock.calls(mockBigIntUrl)).toHaveLength(1);
       expect(json.value.toString()).toEqual(expectedBigNumber);
     });
+
+    it('handleChartDataResponse should return result if GlobalAsyncQueries flag is disabled', async () => {
+      const result = await handleChartDataResponse(
+        { status: 200 },
+        { result: [1, 2, 3] },
+      );
+      expect(result).toEqual([1, 2, 3]);
+    });
+
+    it('handleChartDataResponse should handle responses when GlobalAsyncQueries flag is enabled and results are returned synchronously', async () => {
+      global.featureFlags = {
+        [FeatureFlag.GlobalAsyncQueries]: true,
+      };
+      const result = await handleChartDataResponse(
+        { status: 200 },
+        { result: [1, 2, 3] },
+      );
+      expect(result).toEqual([1, 2, 3]);
+    });
+
+    it('handleChartDataResponse should handle responses when GlobalAsyncQueries flag is enabled and query is running asynchronously', async () => {
+      global.featureFlags = {
+        [FeatureFlag.GlobalAsyncQueries]: true,
+      };
+      const result = await handleChartDataResponse(
+        { status: 202 },
+        { result: [1, 2, 3] },
+      );
+      expect(result).toEqual([1, 2, 3]);
+    });
   });
 
   describe('legacy API', () => {