perf(sqllab): reduce bootstrap data delay by queries (apache#27488)

qleroy · Apr 28, 2024 · 24e18ac · 24e18ac
1 parent f1492a2
commit 24e18ac
Show file tree

Hide file tree

Showing 11 changed files with 577 additions and 76 deletions.
diff --git a/superset-frontend/src/SqlLab/components/QueryHistory/QueryHistory.test.tsx b/superset-frontend/src/SqlLab/components/QueryHistory/QueryHistory.test.tsx
@@ -17,7 +17,10 @@
  * under the License.
  */
 import React from 'react';
-import { render, screen } from 'spec/helpers/testing-library';
+import fetchMock from 'fetch-mock';
+import * as uiCore from '@superset-ui/core';
+import { FeatureFlag, QueryState } from '@superset-ui/core';
+import { render, screen, waitFor } from 'spec/helpers/testing-library';
 import QueryHistory from 'src/SqlLab/components/QueryHistory';
 import { initialState } from 'src/SqlLab/fixtures';
 
@@ -27,18 +30,72 @@ const mockedProps = {
   latestQueryId: 'yhMUZCGb',
 };
 
+const fakeApiResult = {
+  count: 4,
+  ids: [692],
+  result: [
+    {
+      changed_on: '2024-03-12T20:01:02.497775',
+      client_id: 'b0ZDzRYzn',
+      database: {
+        database_name: 'examples',
+        id: 1,
+      },
+      end_time: '1710273662496.047852',
+      error_message: null,
+      executed_sql: 'SELECT * from "FCC 2018 Survey"\nLIMIT 1001',
+      id: 692,
+      limit: 1000,
+      limiting_factor: 'DROPDOWN',
+      progress: 100,
+      results_key: null,
+      rows: 443,
+      schema: 'main',
+      select_as_cta: false,
+      sql: 'SELECT * from "FCC 2018 Survey" ',
+      sql_editor_id: '22',
+      start_time: '1710273662445.992920',
+      status: QueryState.Success,
+      tab_name: 'Untitled Query 16',
+      tmp_table_name: null,
+      tracking_url: null,
+      user: {
+        first_name: 'admin',
+        id: 1,
+        last_name: 'user',
+      },
+    },
+  ],
+};
+
 const setup = (overrides = {}) => (
   <QueryHistory {...mockedProps} {...overrides} />
 );
 
-describe('QueryHistory', () => {
-  it('Renders an empty state for query history', () => {
-    render(setup(), { useRedux: true, initialState });
+test('Renders an empty state for query history', () => {
+  render(setup(), { useRedux: true, initialState });
+
+  const emptyStateText = screen.getByText(
+    /run a query to display query history/i,
+  );
+
+  expect(emptyStateText).toBeVisible();
+});
 
-    const emptyStateText = screen.getByText(
-      /run a query to display query history/i,
+test('fetches the query history when the persistence mode is enabled', async () => {
+  const isFeatureEnabledMock = jest
+    .spyOn(uiCore, 'isFeatureEnabled')
+    .mockImplementation(
+      featureFlag => featureFlag === FeatureFlag.SqllabBackendPersistence,
     );
 
-    expect(emptyStateText).toBeVisible();
-  });
+  const editorQueryApiRoute = `glob:*/api/v1/query/?q=*`;
+  fetchMock.get(editorQueryApiRoute, fakeApiResult);
+  render(setup(), { useRedux: true, initialState });
+  await waitFor(() =>
+    expect(fetchMock.calls(editorQueryApiRoute).length).toBe(1),
+  );
+  const queryResultText = screen.getByText(fakeApiResult.result[0].rows);
+  expect(queryResultText).toBeInTheDocument();
+  isFeatureEnabledMock.mockClear();
 });
diff --git a/superset-frontend/src/SqlLab/components/QueryHistory/index.tsx b/superset-frontend/src/SqlLab/components/QueryHistory/index.tsx
@@ -16,12 +16,23 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import React, { useMemo } from 'react';
+import React, { useEffect, useMemo, useState } from 'react';
 import { shallowEqual, useSelector } from 'react-redux';
+import { useInView } from 'react-intersection-observer';
+import { omit } from 'lodash';
 import { EmptyStateMedium } from 'src/components/EmptyState';
-import { t, styled } from '@superset-ui/core';
+import {
+  t,
+  styled,
+  css,
+  FeatureFlag,
+  isFeatureEnabled,
+} from '@superset-ui/core';
 import QueryTable from 'src/SqlLab/components/QueryTable';
 import { SqlLabRootState } from 'src/SqlLab/types';
+import { useEditorQueriesQuery } from 'src/hooks/apiResources/queries';
+import { Skeleton } from 'src/components';
+import useEffectEvent from 'src/hooks/useEffectEvent';
 
 interface QueryHistoryProps {
   queryEditorId: string | number;
@@ -40,39 +51,92 @@ const StyledEmptyStateWrapper = styled.div`
   }
 `;
 
+const getEditorQueries = (
+  queries: SqlLabRootState['sqlLab']['queries'],
+  queryEditorId: string | number,
+) =>
+  Object.values(queries).filter(
+    ({ sqlEditorId }) => String(sqlEditorId) === String(queryEditorId),
+  );
+
 const QueryHistory = ({
   queryEditorId,
   displayLimit,
   latestQueryId,
 }: QueryHistoryProps) => {
+  const [ref, hasReachedBottom] = useInView({ threshold: 0 });
+  const [pageIndex, setPageIndex] = useState(0);
   const queries = useSelector(
     ({ sqlLab: { queries } }: SqlLabRootState) => queries,
     shallowEqual,
   );
+  const { data, isLoading, isFetching } = useEditorQueriesQuery(
+    { editorId: `${queryEditorId}`, pageIndex },
+    {
+      skip: !isFeatureEnabled(FeatureFlag.SqllabBackendPersistence),
+    },
+  );
   const editorQueries = useMemo(
     () =>
-      Object.values(queries).filter(
-        ({ sqlEditorId }) => String(sqlEditorId) === String(queryEditorId),
-      ),
-    [queries, queryEditorId],
+      data
+        ? getEditorQueries(
+            omit(
+              queries,
+              data.result.map(({ id }) => id),
+            ),
+            queryEditorId,
+          )
+            .concat(data.result)
+            .reverse()
+        : getEditorQueries(queries, queryEditorId),
+    [queries, data, queryEditorId],
   );
 
+  const loadNext = useEffectEvent(() => {
+    setPageIndex(pageIndex + 1);
+  });
+
+  const loadedDataCount = data?.result.length || 0;
+  const totalCount = data?.count || 0;
+
+  useEffect(() => {
+    if (hasReachedBottom && loadedDataCount < totalCount) {
+      loadNext();
+    }
+  }, [hasReachedBottom, loadNext, loadedDataCount, totalCount]);
+
+  if (!editorQueries.length && isLoading) {
+    return <Skeleton active />;
+  }
+
   return editorQueries.length > 0 ? (
-    <QueryTable
-      columns={[
-        'state',
-        'started',
-        'duration',
-        'progress',
-        'rows',
-        'sql',
-        'results',
-        'actions',
-      ]}
-      queries={editorQueries}
-      displayLimit={displayLimit}
-      latestQueryId={latestQueryId}
-    />
+    <>
+      <QueryTable
+        columns={[
+          'state',
+          'started',
+          'duration',
+          'progress',
+          'rows',
+          'sql',
+          'results',
+          'actions',
+        ]}
+        queries={editorQueries}
+        displayLimit={displayLimit}
+        latestQueryId={latestQueryId}
+      />
+      {data && loadedDataCount < totalCount && (
+        <div
+          ref={ref}
+          css={css`
+            position: relative;
+            top: -150px;
+          `}
+        />
+      )}
+      {isFetching && <Skeleton active />}
+    </>
   ) : (
     <StyledEmptyStateWrapper>
       <EmptyStateMedium

diff --git a/superset-frontend/src/SqlLab/reducers/getInitialState.test.ts b/superset-frontend/src/SqlLab/reducers/getInitialState.test.ts
@@ -25,7 +25,6 @@ const apiData = {
   common: DEFAULT_COMMON_BOOTSTRAP_DATA,
   tab_state_ids: [],
   databases: [],
-  queries: {},
   user: {
     userId: 1,
     username: 'some name',
@@ -220,18 +219,20 @@ describe('getInitialState', () => {
         }),
       );
 
+      const latestQuery = {
+        ...runningQuery,
+        id: 'latestPersisted',
+        startDttm: Number(startDttmInStr),
+        endDttm: Number(endDttmInStr),
+      };
       const initializedQueries = getInitialState({
-        ...apiData,
-        queries: {
-          backendPersisted: {
-            ...runningQuery,
-            id: 'backendPersisted',
-            startDttm: startDttmInStr,
-            endDttm: endDttmInStr,
-          },
+        ...apiDataWithTabState,
+        active_tab: {
+          ...apiDataWithTabState.active_tab,
+          latest_query: latestQuery,
         },
       }).sqlLab.queries;
-      expect(initializedQueries.backendPersisted).toEqual(
+      expect(initializedQueries.latestPersisted).toEqual(
         expect.objectContaining({
           startDttm: Number(startDttmInStr),
           endDttm: Number(endDttmInStr),

diff --git a/superset-frontend/src/SqlLab/reducers/getInitialState.ts b/superset-frontend/src/SqlLab/reducers/getInitialState.ts
@@ -136,7 +136,12 @@ export default function getInitialState({
       });
   }
 
-  const queries = { ...queries_ };
+  const queries = {
+    ...queries_,
+    ...(activeTab?.latest_query && {
+      [activeTab.latest_query.id]: activeTab.latest_query,
+    }),
+  };
 
   /**
    * If the `SQLLAB_BACKEND_PERSISTENCE` feature flag is off, or if the user