Show the columns for the correct stage when using combine/extract in the...

Show the columns for the correct stage when using combine/extract in the presence of an aggregation (#43226)

* Do not use asReturned to determine columns for extractions and combine columns

* Add reproduction for #43226 for combine columns

* Add reproduction for #43226 for extract columns

* Remove commented out code

* Only render extraction shortcut when there are extractions

* Only render combination shortcut when there are combinations

* Use hasCombinations from CombineColumns

* Add hasExtractions from ExtractColumn

* Use appendStageIfAggregated over Lib.asReturned

* Only show combine column shortcut when there are two or more columns to be combined

* Lift appendStageIfAggregated to top-level drill

* Add test for extraction on table with just breakout

* Add test for combinations on table with just breakouts

* Disable the 2-column requirement for combinations

* Remove Lib.asReturned in appendStageIfAggregated

* Remove + 1 to stageIndex in appendStageIfAggregated

* Check for empty breakouts too in as-returned

* Switch back to Lib.asReturned

* Fix test for combine column on breakouts

* Reference the correct query and stageIndex

* add CLJS unit tests for as-returned with only aggs, breakouts

---------

Co-authored-by: Braden Shepherdson <braden@metabase.com>

e2e/test/scenarios/custom-column/cc-shortcuts-combine.cy.spec.js

@@ -111,6 +111,26 @@ describe("scenarios > question > custom column > expression shortcuts > combine"
       cy.findByLabelText("Separator").should("have.value", "");
     });
   });
+
+  it("should be possible to edit a previous stages' columns when there is an aggregation (metabase#43226)", () => {
+    openOrdersTable({ mode: "notebook" });
+
+    cy.button("Summarize").click();
+
+    popover().findByText("Count of rows").click();
+
+    addCustomColumn();
+    selectCombineColumns();
+
+    selectColumn("User", "Email");
+
+    expressionEditorWidget().within(() => {
+      cy.findByText("Separated by (empty)").should("exist");
+      cy.findByText(/Separated by/).click();
+
+      cy.findByLabelText("Separator").should("have.value", "");
+    });
+  });
 });
 
 describeWithSnowplow(

e2e/test/scenarios/custom-column/cc-shortcuts-extract.cy.spec.js

@@ -153,6 +153,23 @@ describe("scenarios > question > custom column > expression shortcuts > extract"
       .findByTestId("expression-name")
       .should("have.value", "Hour of day (1)");
   });
+
+  it("should be possible to edit a previous stages' columns when an aggregation is present (metabase#43226)", () => {
+    openOrdersTable({ mode: "notebook", limit: 5 });
+
+    cy.button("Summarize").click();
+    popover().findByText("Count of rows").click();
+
+    addCustomColumn();
+    selectExtractColumn();
+
+    cy.findAllByTestId("dimension-list-item").contains("Created At").click();
+    popover().findAllByRole("button").contains("Hour of day").click();
+
+    expressionEditorWidget()
+      .findByTestId("expression-name")
+      .should("have.value", "Hour of day");
+  });
 });
 
 function selectExtractColumn() {

e2e/test/scenarios/visualizations-tabular/column-shortcuts/column-shortcuts.cy.spec.ts

@@ -290,6 +290,32 @@ describeWithSnowplow("extract shortcut", () => {
       "be.visible",
     );
   });
+
+  it("should be possible to extract columns from table with breakouts", () => {
+    createQuestion(
+      {
+        query: {
+          "source-table": ORDERS_ID,
+          limit: 5,
+          breakout: [
+            ["field", ORDERS.CREATED_AT, { "temporal-unit": "month" }],
+          ],
+        },
+      },
+      {
+        visitQuestion: true,
+      },
+    );
+
+    extractColumnAndCheck({
+      column: "Created At: Month",
+      option: "Month of year",
+    });
+
+    cy.findAllByRole("columnheader", { name: "Month of year" }).should(
+      "be.visible",
+    );
+  });
 });
 
 function extractColumnAndCheck({

e2e/test/scenarios/visualizations-tabular/column-shortcuts/combine-column-shortcut.cy.spec.ts

@@ -12,7 +12,7 @@ import {
   expectGoodSnowplowEvent,
 } from "e2e/support/helpers";
 
-const { PEOPLE, PEOPLE_ID, ORDERS_ID, ORDERS } = SAMPLE_DATABASE;
+const { PEOPLE, PEOPLE_ID, ORDERS_ID, ORDERS, PRODUCTS } = SAMPLE_DATABASE;
 
 describeWithSnowplow("scenarios > visualizations > combine shortcut", () => {
   beforeEach(() => {
@@ -81,6 +81,36 @@ describeWithSnowplow("scenarios > visualizations > combine shortcut", () => {
       newValue: "0 766",
     });
   });
+
+  it("should allow combining columns on a table with just breakouts", () => {
+    createQuestion(
+      {
+        query: {
+          "source-table": ORDERS_ID,
+          limit: 1,
+          breakout: [
+            ["field", ORDERS.CREATED_AT, { "temporal-unit": "hour-of-day" }],
+            [
+              "field",
+              PRODUCTS.CATEGORY,
+              { "base-type": "type/text", "source-field": ORDERS.PRODUCT_ID },
+            ],
+          ],
+        },
+      },
+      {
+        visitQuestion: true,
+      },
+    );
+
+    cy.findByTestId("TableInteractive-root").should("exist");
+    combineColumns({
+      columns: ["Created At: Hour of day", "Category"],
+      newColumn: "Combined Created At: Hour of day, Category",
+      example: "2042-01-01 12:34:56.789 text",
+      newValue: "0 Doohickey",
+    });
+  });
 });
 
 function combineColumns({

frontend/src/metabase/query_builder/components/expressions/CombineColumns/CombineColumns.tsx

@@ -32,12 +32,7 @@ type State = {
 
 const initialDefaultSeparator = " ";
 
-export function CombineColumns({
-  query: originalQuery,
-  stageIndex: originalStageIndex,
-  onSubmit,
-  width,
-}: Props) {
+export function CombineColumns({ query, stageIndex, onSubmit, width }: Props) {
   const [state, setState] = useState<State>({
     columnsAndSeparators: [
       {
@@ -55,11 +50,6 @@ export function CombineColumns({
 
   const { columnsAndSeparators, isUsingDefaultSeparator } = state;
 
-  const { query, stageIndex } = Lib.asReturned(
-    originalQuery,
-    originalStageIndex,
-  );
-
   const expressionableColumns = Lib.expressionableColumns(query, stageIndex);
 
   const handleRowChange = (

frontend/src/metabase/query_builder/components/expressions/CombineColumns/index.tsx

 export { CombineColumns } from "./CombineColumns";
+export { hasCombinations } from "./util";

frontend/src/metabase/query_builder/components/expressions/CombineColumns/util.ts

@@ -155,3 +155,7 @@ const getColumnExample = (
 
   return "text";
 };
+
+export function hasCombinations(query: Lib.Query, stageIndex: number) {
+  return Lib.expressionableColumns(query, stageIndex).length > 0;
+}

frontend/src/metabase/query_builder/components/expressions/ExpressionWidget.tsx

@@ -14,7 +14,7 @@ import {
   trackColumnExtractViaShortcut,
 } from "../../analytics";
 
-import { CombineColumns } from "./CombineColumns";
+import { CombineColumns, hasCombinations } from "./CombineColumns";
 import { ExpressionEditorTextfield } from "./ExpressionEditorTextfield";
 import {
   ActionButtonsWrapper,
@@ -27,7 +27,7 @@ import {
 } from "./ExpressionWidget.styled";
 import { ExpressionWidgetHeader } from "./ExpressionWidgetHeader";
 import { ExpressionWidgetInfo } from "./ExpressionWidgetInfo";
-import { ExtractColumn } from "./ExtractColumn";
+import { ExtractColumn, hasExtractions } from "./ExtractColumn";
 
 export type ExpressionWidgetProps<Clause = Lib.ExpressionClause> = {
   query: Lib.Query;
@@ -211,20 +211,22 @@ export const ExpressionWidget = <Clause extends object = Lib.ExpressionClause>(
           onCommit={handleCommit}
           onError={(errorMessage: string) => setError(errorMessage)}
           shortcuts={[
-            !startRule && {
-              shortcut: true,
-              name: t`Combine columns`,
-              action: () => setIsCombiningColumns(true),
-              group: "shortcuts",
-              icon: "combine",
-            },
-            !startRule && {
-              shortcut: true,
-              name: t`Extract columns`,
-              icon: "arrow_split",
-              group: "shortcuts",
-              action: () => setIsExtractingColumn(true),
-            },
+            !startRule &&
+              hasCombinations(query, stageIndex) && {
+                shortcut: true,
+                name: t`Combine columns`,
+                action: () => setIsCombiningColumns(true),
+                group: "shortcuts",
+                icon: "combine",
+              },
+            !startRule &&
+              hasExtractions(query, stageIndex) && {
+                shortcut: true,
+                name: t`Extract columns`,
+                icon: "arrow_split",
+                group: "shortcuts",
+                action: () => setIsExtractingColumn(true),
+              },
           ].filter(Boolean)}
         />
       </ExpressionFieldWrapper>

frontend/src/metabase/query_builder/components/expressions/ExtractColumn/ExtractColumn.tsx

@@ -22,18 +22,13 @@ type Props = {
 };
 
 export function ExtractColumn({
-  query: originalQuery,
-  stageIndex: originalStageIndex,
+  query,
+  stageIndex,
   onCancel,
   onSubmit,
 }: Props) {
   const [column, setColumn] = useState<Lib.ColumnMetadata | null>(null);
 
-  const { query, stageIndex } = Lib.asReturned(
-    originalQuery,
-    originalStageIndex,
-  );
-
   function handleSelect(column: Lib.ColumnMetadata) {
     setColumn(column);
   }

frontend/src/metabase/query_builder/components/expressions/ExtractColumn/index.tsx

 export { ExtractColumn } from "./ExtractColumn";
+export { hasExtractions } from "./util";

frontend/src/metabase/query_builder/components/expressions/ExtractColumn/util.ts

@@ -48,3 +48,13 @@ export function getName(
 
   return getNextName(columnNames, info.displayName, 0);
 }
+
+export function hasExtractions(query: Lib.Query, stageIndex: number) {
+  for (const column of Lib.expressionableColumns(query, stageIndex)) {
+    if (Lib.columnExtractions(query, column).length > 0) {
+      return true;
+    }
+  }
+
+  return false;
+}

frontend/src/metabase/querying/utils/drills/combine-columns-drill/combine-columns-drill.tsx

@@ -2,6 +2,7 @@ import { t } from "ttag";
 
 import { useDispatch } from "metabase/lib/redux";
 import { setUIControls } from "metabase/query_builder/actions";
+import { hasCombinations } from "metabase/query_builder/components/expressions/CombineColumns";
 import { trackColumnCombineViaColumnHeader } from "metabase/querying/analytics";
 import type {
   ClickActionPopoverProps,
@@ -13,11 +14,16 @@ import { CombineColumnsDrill } from "./components";
 
 export const combineColumnsDrill: Drill<Lib.CombineColumnsDrillThruInfo> = ({
   question,
-  query,
-  stageIndex,
+  query: originalQuery,
+  stageIndex: originalStageIndex,
   clicked,
 }) => {
-  if (!clicked.column) {
+  const { query, stageIndex } = Lib.asReturned(
+    originalQuery,
+    originalStageIndex,
+  );
+
+  if (!clicked.column || !hasCombinations(query, stageIndex)) {
     return [];
   }
 

frontend/src/metabase/querying/utils/drills/combine-columns-drill/components/CombineColumnsDrill/CombineColumnsDrill.tsx

@@ -40,15 +40,11 @@ interface Props {
 
 export const CombineColumnsDrill = ({
   column,
-  query: originalQuery,
-  stageIndex: originalStageIndex,
+  query,
+  stageIndex,
   onSubmit,
 }: Props) => {
-  const columnInfo = Lib.displayInfo(originalQuery, originalStageIndex, column);
-  const { query, stageIndex } = Lib.asReturned(
-    originalQuery,
-    originalStageIndex,
-  );
+  const columnInfo = Lib.displayInfo(query, stageIndex, column);
   const expressionableColumns = Lib.expressionableColumns(query, stageIndex);
   const defaultSeparator = getDefaultSeparator(column);
   const [isUsingDefaultSeparator, setIsUsingDefaultSeparator] = useState(true);

frontend/src/metabase/visualizations/click-actions/actions/CombineColumnsAction/CombineColumnsAction.tsx

@@ -3,24 +3,24 @@ import { t } from "ttag";
 import { useDispatch } from "metabase/lib/redux";
 import { setUIControls } from "metabase/query_builder/actions";
 import { trackColumnCombineViaPlusModal } from "metabase/query_builder/analytics";
-import { CombineColumns } from "metabase/query_builder/components/expressions/CombineColumns";
+import {
+  CombineColumns,
+  hasCombinations,
+} from "metabase/query_builder/components/expressions/CombineColumns";
 import type { LegacyDrill } from "metabase/visualizations/types";
 import type { ClickActionPopoverProps } from "metabase/visualizations/types/click-actions";
 import * as Lib from "metabase-lib";
 
 export const CombineColumnsAction: LegacyDrill = ({ question, clicked }) => {
   const { query, stageIndex } = Lib.asReturned(question.query(), -1);
-
   const { isEditable } = Lib.queryDisplayInfo(query);
-  const isExpressionable =
-    Lib.expressionableColumns(query, stageIndex).length > 0;
 
   if (
     !clicked ||
     clicked.value !== undefined ||
     !clicked.columnShortcuts ||
     !isEditable ||
-    !isExpressionable
+    !hasCombinations(query, stageIndex)
   ) {
     return [];
   }

frontend/src/metabase/visualizations/click-actions/actions/ExtractColumn/ExtractColumn.tsx

@@ -3,7 +3,10 @@ import { t } from "ttag";
 import { useDispatch } from "metabase/lib/redux";
 import { setUIControls } from "metabase/query_builder/actions";
 import { trackColumnExtractViaPlusModal } from "metabase/query_builder/analytics";
-import { ExtractColumn } from "metabase/query_builder/components/expressions/ExtractColumn";
+import {
+  ExtractColumn,
+  hasExtractions,
+} from "metabase/query_builder/components/expressions/ExtractColumn";
 import { rem, Box } from "metabase/ui";
 import type { LegacyDrill } from "metabase/visualizations/types";
 import type { ClickActionPopoverProps } from "metabase/visualizations/types/click-actions";
@@ -13,18 +16,13 @@ export const ExtractColumnAction: LegacyDrill = ({ question, clicked }) => {
   const { query, stageIndex } = Lib.asReturned(question.query(), -1);
 
   const { isEditable } = Lib.queryDisplayInfo(query);
-  const expressionableColumns = Lib.expressionableColumns(query, stageIndex);
-  const isExtractable =
-    expressionableColumns.reduce(function (sum, column) {
-      return sum + Lib.columnExtractions(query, column).length;
-    }, 0) > 0;
 
   if (
     !clicked ||
     clicked.value !== undefined ||
     !clicked.columnShortcuts ||
     !isEditable ||
-    !isExtractable
+    !hasExtractions(query, stageIndex)
   ) {
     return [];
   }

src/metabase/lib/js.cljs

@@ -248,7 +248,9 @@
 
   > **Code health:** Healthy"
   [a-query stage-number]
-  (if (empty? (lib.core/aggregations a-query stage-number))
+  (if (and
+        (empty? (lib.core/aggregations a-query stage-number))
+        (empty? (lib.core/breakouts a-query stage-number)))
     ;; No extra stage needed with no aggregations.
     #js {:query      a-query
          :stageIndex stage-number}

test/metabase/lib/js_test.cljs

@@ -611,7 +611,7 @@
                             (lib/filter (lib/> (m/find-first #(= (:name "count") %) base-cols)
                                                100)))
           two-stage-agg (lib/aggregate two-stage (lib/count))]
-      (testing "does not change a query with no aggregations"
+      (testing "does not change a query with no aggregations or breakouts"
         (doseq [stage [0 -1]]
           (let [obj (lib.js/as-returned simple-query stage)]
             (is (=? simple-query (.-query obj)))
@@ -635,4 +635,36 @@
         (let [obj (lib.js/as-returned two-stage-agg 1)]
           (is (=? (lib/append-stage two-stage-agg)
                   (.-query obj)))
-          (is (=? -1 (.-stageIndex obj))))))))
+          (is (=? -1 (.-stageIndex obj)))))
+
+      (testing "only breakouts"
+        (let [brk-only  (-> (lib/query meta/metadata-provider (meta/table-metadata :orders))
+                            (lib/breakout (lib/with-temporal-bucket (meta/field-metadata :orders :created-at) :month)))
+              two-stage (-> brk-only
+                            lib/append-stage
+                            (lib/filter (lib/> (first (lib/returned-columns brk-only)) 100)))]
+          (testing "uses an existing later stage if it exists"
+            (let [obj (lib.js/as-returned two-stage 0)]
+              (is (=? two-stage (.-query obj)))
+              (is (=? 1         (.-stageIndex obj)))))
+          (testing "appends a new stage if necessary"
+            (let [obj (lib.js/as-returned brk-only 0)]
+              (is (=? (lib/append-stage brk-only)
+                      (.-query obj)))
+              (is (=? -1 (.-stageIndex obj)))))))
+
+      (testing "only aggregations"
+        (let [agg-only  (-> (lib/query meta/metadata-provider (meta/table-metadata :orders))
+                            (lib/aggregate (lib/count)))
+              two-stage (-> agg-only
+                            lib/append-stage
+                            (lib/filter (lib/> (first (lib/returned-columns agg-only)) 100)))]
+          (testing "uses an existing later stage if it exists"
+            (let [obj (lib.js/as-returned two-stage 0)]
+              (is (=? two-stage (.-query obj)))
+              (is (=? 1         (.-stageIndex obj)))))
+          (testing "appends a new stage if necessary"
+            (let [obj (lib.js/as-returned agg-only 0)]
+              (is (=? (lib/append-stage agg-only)
+                      (.-query obj)))
+              (is (=? -1 (.-stageIndex obj))))))))))