QB: evaluate notebook preview lazy (#39194)

* QB: evaluate notebook preview lazily

* fixup! QB: evaluate notebook preview lazily

* fixup! fixup! QB: evaluate notebook preview lazily

* fixup! fixup! fixup! QB: evaluate notebook preview lazily

frontend/src/metabase/query_builder/components/notebook/NotebookStep/NotebookStep.tsx

@@ -104,7 +104,7 @@ function NotebookStep({
   } = STEP_UI[step.type] || {};
 
   const color = getColor();
-  const canPreview = Boolean(step.previewQuery);
+  const canPreview = Boolean(step.getPreviewQuery);
   const hasPreviewButton = !isPreviewOpen && canPreview;
   const canRevert = typeof step.revert === "function" && !readOnly;
 

frontend/src/metabase/query_builder/components/notebook/NotebookStepPreview.jsx

@@ -37,7 +37,8 @@ class NotebookStepPreview extends Component {
   };
 
   getPreviewQuestion(step) {
-    const { previewQuery: query, stageIndex } = step;
+    const { getPreviewQuery, stageIndex } = step;
+    const query = getPreviewQuery();
     const limit = Lib.currentLimit(query, stageIndex);
     const hasSuitableLimit = limit !== null && limit <= PREVIEW_ROWS_LIMIT;
     const queryWithLimit = hasSuitableLimit

frontend/src/metabase/query_builder/components/notebook/lib/steps.ts

@@ -235,10 +235,10 @@ function getStageSteps(
     const id = getId(STEP, itemIndex);
     const active = STEP.active(query, stageIndex, itemIndex ?? undefined);
     const step: NotebookStep = {
-      id: id,
+      id,
       type: STEP.type,
-      stageIndex: stageIndex,
-      itemIndex: itemIndex,
+      stageIndex,
+      itemIndex,
       question,
       query,
       valid: STEP.valid(query, stageIndex, metadata),
@@ -253,9 +253,8 @@ function getStageSteps(
             return revert(query, stageIndex, itemIndex ?? undefined);
           }
         : null,
-      // `actions`, `previewQuery`, `next` and `previous` will be set later
+      // `actions`, `next` and `previous` will be set later
       actions: [],
-      previewQuery: null,
       next: null,
       previous: null,
     };
@@ -263,22 +262,17 @@ function getStageSteps(
   }
 
   // get the currently visible steps, flattening "items"
-  const steps = _.flatten(
-    STEPS.map(STEP => {
-      if (STEP.subSteps) {
-        // add 1 for the initial or next action button
-        const itemIndexes = _.range(0, STEP.subSteps(query, stageIndex) + 1);
-        return itemIndexes.map(itemIndex => getStep(STEP, itemIndex));
-      } else {
-        return [getStep(STEP)];
-      }
-    }),
-  );
+  const steps = STEPS.flatMap(STEP => {
+    if (STEP.subSteps) {
+      // add 1 for the initial or next action button
+      const itemIndexes = _.range(0, STEP.subSteps(query, stageIndex) + 1);
+      return itemIndexes.map(itemIndex => getStep(STEP, itemIndex));
+    }
 
-  let previewQuery = _.range(Lib.stageCount(query) - 1, stageIndex, -1).reduce(
-    Lib.dropStage,
-    query,
-  );
+    return [getStep(STEP)];
+  });
+  // keep original steps to be able to get a step by index to preview query
+  const originalSteps = [...steps];
 
   let actions = [];
   // iterate over steps in reverse so we can revert query for previewing and accumulate valid actions
@@ -286,7 +280,11 @@ function getStageSteps(
     const step = steps[i];
     if (step.visible) {
       // only include previewQuery if the section would be visible (i.e. excluding "openSteps")
-      step.previewQuery = step.active ? previewQuery : null;
+      if (step.active) {
+        step.getPreviewQuery = () =>
+          getPreviewQuery(query, stageIndex, originalSteps, i);
+      }
+
       // add any accumulated actions and reset
       step.actions = actions;
       actions = [];
@@ -304,15 +302,31 @@ function getStageSteps(
       }
       steps.splice(i, 1);
     }
-    // revert the previewQuery for this step
-    if (step.revert) {
-      previewQuery = step.revert(
-        previewQuery,
-        step.stageIndex,
-        step.itemIndex ?? undefined,
-      );
-    }
   }
 
   return { steps, actions };
 }
+function getPreviewQuery(
+  query: Lib.Query,
+  stageIndex: number,
+  steps: NotebookStep[],
+  activeStepIndex: number,
+): Lib.Query {
+  const queryWithoutNextStages = _.range(
+    Lib.stageCount(query) - 1,
+    stageIndex,
+    -1,
+  ).reduce(Lib.dropStage, query);
+
+  const queryWithoutNextStagesAndSteps = steps
+    .slice(activeStepIndex + 1, steps.length)
+    .reduceRight((query: Lib.Query, step: NotebookStep) => {
+      if (step.revert) {
+        return step.revert(query, step.stageIndex, step.itemIndex ?? undefined);
+      }
+
+      return query;
+    }, queryWithoutNextStages);
+
+  return queryWithoutNextStagesAndSteps;
+}

frontend/src/metabase/query_builder/components/notebook/lib/steps.unit.spec.ts

@@ -93,9 +93,9 @@ describe("filtered and summarized query", () => {
     });
   });
 
-  describe("previewQuery", () => {
+  describe("getPreviewQuery", () => {
     it("shouldn't include filter, summarize for data step", () => {
-      const previewQuery = checkNotNull(dataStep.previewQuery);
+      const previewQuery = checkNotNull(dataStep.getPreviewQuery)();
 
       expect(Lib.aggregations(previewQuery, 0)).toHaveLength(0);
       expect(Lib.breakouts(previewQuery, 0)).toHaveLength(0);
@@ -103,7 +103,7 @@ describe("filtered and summarized query", () => {
     });
 
     it("shouldn't include summarize for filter step", () => {
-      const previewQuery = checkNotNull(filterStep.previewQuery);
+      const previewQuery = checkNotNull(filterStep.getPreviewQuery)();
 
       expect(Lib.aggregations(previewQuery, 0)).toHaveLength(0);
       expect(Lib.breakouts(previewQuery, 0)).toHaveLength(0);
@@ -159,9 +159,9 @@ describe("filtered and summarized query with post-aggregation filter", () => {
     });
   });
 
-  describe("previewQuery", () => {
+  describe("getPreviewQuery", () => {
     it("shouldn't include filter, summarize, or post-aggregation filter for data step", () => {
-      const previewQuery = checkNotNull(dataStep.previewQuery);
+      const previewQuery = checkNotNull(dataStep.getPreviewQuery)();
 
       expect(Lib.stageCount(previewQuery)).toBe(1);
       expect(Lib.aggregations(previewQuery, 0)).toHaveLength(0);
@@ -170,7 +170,7 @@ describe("filtered and summarized query with post-aggregation filter", () => {
     });
 
     it("shouldn't include summarize or post-aggregation filter for filter step", () => {
-      const previewQuery = checkNotNull(filterStep.previewQuery);
+      const previewQuery = checkNotNull(filterStep.getPreviewQuery)();
 
       expect(Lib.stageCount(previewQuery)).toBe(1);
       expect(Lib.aggregations(previewQuery, 0)).toHaveLength(0);
@@ -179,7 +179,7 @@ describe("filtered and summarized query with post-aggregation filter", () => {
     });
 
     it("shouldn't include filters from the next stages for summarizeStep", () => {
-      const previewQuery = checkNotNull(summarizeStep.previewQuery);
+      const previewQuery = checkNotNull(summarizeStep.getPreviewQuery)();
 
       expect(Lib.stageCount(previewQuery)).toBe(1);
       expect(Lib.aggregations(previewQuery, 0)).toHaveLength(1);
@@ -188,7 +188,9 @@ describe("filtered and summarized query with post-aggregation filter", () => {
     });
 
     it("shouldn't include aggregations for post-aggregation filter step", () => {
-      const previewQuery = checkNotNull(postAggregationFilterStep.previewQuery);
+      const previewQuery = checkNotNull(
+        postAggregationFilterStep.getPreviewQuery,
+      )();
 
       expect(Lib.stageCount(previewQuery)).toBe(2);
       expect(Lib.aggregations(previewQuery, 0)).toHaveLength(1);

frontend/src/metabase/query_builder/components/notebook/test-utils.ts

@@ -36,7 +36,6 @@ export function createMockNotebookStep({
     active: true,
     visible: true,
     actions: [],
-    previewQuery: null,
     next: null,
     previous: null,
     revert: jest.fn(),

frontend/src/metabase/query_builder/components/notebook/types.ts

@@ -27,9 +27,9 @@ export interface NotebookStep {
   testID: string;
   revert: RevertFn | null;
   actions: NotebookStepAction[];
-  previewQuery: Query | null;
   next: NotebookStep | null;
   previous: NotebookStep | null;
+  getPreviewQuery?: (() => Query) | undefined;
 }
 
 export interface NotebookStepAction {