Keep breakout and order-by columns as selectable (#30573)

* Keep breakout and order-by columns as selectable

Fixes #30568.

The already selected columns are kept in the lists returned by
breakoutable-columns and orderable-column, but they are marked
with their position in the breakouts and order-bys list, respectively.

* Add test for orderable-columns display-info

* Add test for orderByPosition after a roundtrip

* Cherry pick order-by cleanup fix

* Fix breakoutable and orderable columns after round trip

* Declare breakoutPosition in ColumnDisplayInfo

* Address review feedback

---------

Co-authored-by: Anton Kulyk <kuliks.anton@gmail.com>

frontend/src/metabase-lib/breakout.unit.spec.ts

@@ -21,6 +21,32 @@ describe("breakout", () => {
       expect(breakouts).toHaveLength(1);
       expect(ML.displayName(nextQuery, breakouts[0])).toBe("Title");
     });
+
+    it("should preserve breakout positions between v1-v2 roundtrip", () => {
+      const query = createQuery();
+      const taxColumn = findBreakoutableColumn("ORDERS", "TAX");
+      const nextQuery = ML.breakout(query, taxColumn);
+      const nextQueryColumns = ML.breakoutableColumns(nextQuery);
+      const nextTaxColumn = columnFinder(nextQuery, nextQueryColumns)(
+        "ORDERS",
+        "TAX",
+      );
+
+      expect(ML.displayInfo(nextQuery, nextTaxColumn).breakoutPosition).toBe(0);
+
+      const roundtripQuery = createQuery({
+        query: ML.toLegacyQuery(nextQuery),
+      });
+      const roundtripQueryColumns = ML.breakoutableColumns(roundtripQuery);
+      const roundtripTaxColumn = columnFinder(
+        roundtripQuery,
+        roundtripQueryColumns,
+      )("ORDERS", "TAX");
+
+      expect(
+        ML.displayInfo(roundtripQuery, roundtripTaxColumn).breakoutPosition,
+      ).toBe(0);
+    });
   });
 
   describe("replace breakout", () => {

frontend/src/metabase-lib/order_by.unit.spec.ts

@@ -97,6 +97,32 @@ describe("order by", () => {
         }),
       );
     });
+
+    it("should preserve order-by positions between v1-v2 roundtrip", () => {
+      const query = createQuery();
+      const taxColumn = findOrderableColumn("ORDERS", "TAX");
+      const nextQuery = ML.orderBy(query, taxColumn);
+      const nextQueryColumns = ML.orderableColumns(nextQuery);
+      const nextTaxColumn = columnFinder(nextQuery, nextQueryColumns)(
+        "ORDERS",
+        "TAX",
+      );
+
+      expect(ML.displayInfo(nextQuery, nextTaxColumn).orderByPosition).toBe(0);
+
+      const roundtripQuery = createQuery({
+        query: ML.toLegacyQuery(nextQuery),
+      });
+      const roundtripQueryColumns = ML.orderableColumns(roundtripQuery);
+      const roundtripTaxColumn = columnFinder(
+        roundtripQuery,
+        roundtripQueryColumns,
+      )("ORDERS", "TAX");
+
+      expect(
+        ML.displayInfo(roundtripQuery, roundtripTaxColumn).orderByPosition,
+      ).toBe(0);
+    });
   });
 
   describe("add order by", () => {

frontend/src/metabase-lib/types.ts

@@ -44,6 +44,9 @@ export type ColumnDisplayInfo = {
   isFromJoin: boolean;
   isImplicitlyJoinable: boolean;
   table?: TableInlineDisplayInfo;
+
+  breakoutPosition?: number;
+  orderByPosition?: number;
 };
 
 export type OrderByClauseDisplayInfo = Pick<

src/metabase/lib/breakout.cljc

@@ -69,12 +69,21 @@
 
   ([query        :- ::lib.schema/query
     stage-number :- :int]
-   (let [existing-breakouts (breakouts query stage-number)
-         existing-breakout? (fn [x]
-                              (some (fn [existing-breakout]
-                                      (lib.equality/= (lib.ref/ref x) existing-breakout))
-                                    existing-breakouts))
-         columns            (let [stage (lib.util/query-stage query stage-number)]
-                              (lib.metadata.calculation/visible-columns query stage-number stage))]
+   (let [indexed-breakouts
+         (map-indexed vector (breakouts query stage-number))
+         breakout-pos
+         (fn [x]
+           (some (fn [[pos existing-breakout]]
+                   (let [a-ref (lib.ref/ref x)]
+                     (when (or (lib.equality/= a-ref existing-breakout)
+                               (lib.equality/= a-ref (lib.util/with-default-effective-type existing-breakout)))
+                       pos)))
+                 indexed-breakouts))
+         columns
+         (let [stage (lib.util/query-stage query stage-number)]
+           (lib.metadata.calculation/visible-columns query stage-number stage))]
      (some->> (not-empty columns)
-              (into [] (remove existing-breakout?))))))
+              (into [] (map (fn [col]
+                              (let [pos (breakout-pos col)]
+                                (cond-> col
+                                  pos (assoc :breakout-position pos))))))))))

src/metabase/lib/metadata/calculation.cljc

@@ -271,7 +271,11 @@
    ;; already joined, but could implicitly join against?
    [:is-implicitly-joinable {:optional true} [:maybe :boolean]]
    ;; For the `:table` field of a Column, is this the source table, or a joined table?
-   [:is-source-table {:optional true} [:maybe :boolean]]])
+   [:is-source-table {:optional true} [:maybe :boolean]]
+   ;; does this column occur in the breakout clause?
+   [:is-breakout-column {:optional true} [:maybe :boolean]]
+   ;; does this column occur in the order-by clause?
+   [:is-order-by-column {:optional true} [:maybe :boolean]]])
 
 (mu/defn display-info :- ::display-info
   "Given some sort of Cljs object, return a map with the info you'd need to implement UI for it. This is mostly meant to
@@ -311,7 +315,8 @@
        {:is-from-previous-stage (= source :source/previous-stage)
         :is-from-join           (= source :source/joins)
         :is-calculated          (= source :source/expressions)
-        :is-implicitly-joinable (= source :source/implicitly-joinable)}))))
+        :is-implicitly-joinable (= source :source/implicitly-joinable)})
+     (select-keys x-metadata [:breakout-position :order-by-position]))))
 
 (defmethod display-info-method :default
   [query stage-number x]

src/metabase/lib/order_by.cljc

@@ -148,13 +148,18 @@
 
   ([query        :- ::lib.schema/query
     stage-number :- :int]
-   (let [existing-order-bys (mapv (fn [[_tag _opts expr]]
-                                    expr)
-                                  (order-bys query stage-number))
-         existing-order-by? (fn [x]
-                              (some (fn [existing-order-by]
-                                      (lib.equality/= (lib.ref/ref x) existing-order-by))
-                                    existing-order-bys))
+   (let [indexed-order-bys (map-indexed (fn [pos [_tag _opts expr]]
+                                          [pos expr])
+                                        (order-bys query stage-number))
+         order-by-pos
+         (fn [x]
+           (some (fn [[pos existing-order-by]]
+                   (let [a-ref (lib.ref/ref x)]
+                     (when (or (lib.equality/= a-ref existing-order-by)
+                               (lib.equality/= a-ref (lib.util/with-default-effective-type existing-order-by)))
+                       pos)))
+                 indexed-order-bys))
+
          breakouts          (not-empty (lib.breakout/breakouts-metadata query stage-number))
          aggregations       (not-empty (lib.aggregation/aggregations query stage-number))
          columns            (if (or breakouts aggregations)
@@ -163,7 +168,10 @@
                                 (lib.metadata.calculation/visible-columns query stage-number stage)))]
      (some->> (not-empty columns)
               (into [] (comp (filter orderable-column?)
-                             (remove existing-order-by?)))))))
+                             (map (fn [col]
+                                    (let [pos (order-by-pos col)]
+                                      (cond-> col
+                                        pos (assoc :order-by-position pos)))))))))))
 
 (def ^:private opposite-direction
   {:asc :desc

src/metabase/lib/util.cljc

@@ -465,3 +465,20 @@
           ;; subvec holds onto references, so create a new vector
           (update :stages (comp #(into [] %) subvec) 0 (inc (non-negative-stage-index stages stage-number))))
       new-query)))
+
+(defn with-default-effective-type
+  "Adds a default :effective-type property if it does not exist and
+  :base-type is known.
+
+  This is needed only because we have to convert queries to the Legacy
+  form.
+  The round trip conversion pMBQL -> legacy MBQL -> pMBQL loses the
+  :effective-type property, but it should be present for the frontend
+  to work. It defaults to the :base-type property."
+  [clause]
+  (let [options (lib.options/options clause)
+        default-effective-type (when-not (:effective-type options)
+                                 (:base-type options))]
+    (cond-> clause
+      default-effective-type
+      (lib.options/update-options assoc :effective-type default-effective-type))))

test/metabase/lib/breakout_test.cljc

@@ -247,17 +247,6 @@
                   (for [col columns]
                     (lib/display-info query col)))))))))
 
-(defn- assert-breakout-column-excluded [query query' column]
-  (let [ignore-uuid (map #(dissoc % :metabase.lib.aggregation/aggregation-uuid))
-        columns (disj (into #{}
-                            ignore-uuid
-                            (lib/breakoutable-columns query))
-                      column)
-        columns' (into #{}
-                       ignore-uuid
-                       (lib/breakoutable-columns query'))]
-    (is (= columns columns'))))
-
 (deftest ^:parallel breakoutable-columns-e2e-test
   (testing "Use the metadata returned by `breakoutable-columns` to add a new breakout to a query."
     (let [query (lib/query-for-table-name meta/metadata-provider "VENUES")]
@@ -277,8 +266,7 @@
                                :breakout     [[:field {:lib/uuid string? :base-type :type/Text} (meta/id :venues :name)]]}]}
                   query'))
           (is (=? [[:field {:lib/uuid string? :base-type :type/Text} (meta/id :venues :name)]]
-                  (lib/breakouts query')))
-          (assert-breakout-column-excluded query query' col))))))
+                  (lib/breakouts query'))))))))
 
 (deftest ^:parallel breakoutable-columns-own-and-implicitly-joinable-columns-e2e-test
   (testing "An implicitly joinable column can be broken out by."
@@ -289,7 +277,8 @@
                                       (lib/breakoutable-columns query))
           query' (-> query
                      (lib/breakout cat-name-col)
-                     (lib/breakout ven-price-col))]
+                     (lib/breakout ven-price-col))
+          breakoutables' (lib/breakoutable-columns query')]
       (is (=? {:stages [{:breakout [[:field
                                      {:source-field (meta/id :venues :category-id)}
                                      (meta/id :categories :name)]
@@ -304,8 +293,86 @@
                {:display-name "Category ID", :lib/source :source/table-defaults}
                {:display-name "Latitude",    :lib/source :source/table-defaults}
                {:display-name "Longitude",   :lib/source :source/table-defaults}
-               {:display-name "ID",          :lib/source :source/implicitly-joinable}]
-              (lib/breakoutable-columns query'))))))
+               {:display-name "Price"        :lib/source :source/table-defaults, :breakout-position 1}
+               {:display-name "ID",          :lib/source :source/implicitly-joinable}
+               {:display-name "Name",        :lib/source :source/implicitly-joinable, :breakout-position 0}]
+              breakoutables'))
+      (is (= 2 (count (filter :breakout-position breakoutables'))))
+      (is (=? [{:table {:name "VENUES", :display-name "Venues", :is-source-table true}
+                :semantic-type :type/PK
+                :name "ID"
+                :effective-type :type/BigInteger
+                :is-from-join false
+                :display-name "ID"
+                :is-from-previous-stage false
+                :is-calculated false
+                :is-implicitly-joinable false}
+               {:table {:name "VENUES", :display-name "Venues", :is-source-table true}
+                :semantic-type :type/Name
+                :name "NAME"
+                :effective-type :type/Text
+                :is-from-join false
+                :display-name "Name"
+                :is-from-previous-stage false
+                :is-calculated false
+                :is-implicitly-joinable false}
+               {:table {:name "VENUES", :display-name "Venues", :is-source-table true}
+                :semantic-type :type/FK
+                :name "CATEGORY_ID"
+                :effective-type :type/Integer
+                :is-from-join false
+                :display-name "Category ID"
+                :is-from-previous-stage false
+                :is-calculated false
+                :is-implicitly-joinable false}
+               {:table {:name "VENUES", :display-name "Venues", :is-source-table true}
+                :semantic-type :type/Latitude
+                :name "LATITUDE"
+                :effective-type :type/Float
+                :is-from-join false
+                :display-name "Latitude"
+                :is-from-previous-stage false
+                :is-calculated false
+                :is-implicitly-joinable false}
+               {:table {:name "VENUES", :display-name "Venues", :is-source-table true}
+                :semantic-type :type/Longitude
+                :name "LONGITUDE"
+                :effective-type :type/Float
+                :is-from-join false
+                :display-name "Longitude"
+                :is-from-previous-stage false
+                :is-calculated false
+                :is-implicitly-joinable false}
+               {:table {:name "VENUES", :display-name "Venues", :is-source-table true}
+                :semantic-type :type/Category
+                :name "PRICE"
+                :effective-type :type/Integer
+                :is-from-join false
+                :display-name "Price"
+                :is-from-previous-stage false
+                :is-calculated false
+                :is-implicitly-joinable false
+                :breakout-position 1}
+               {:table {:name "CATEGORIES", :display-name "Categories", :is-source-table false}
+                :semantic-type :type/PK
+                :name "ID"
+                :effective-type :type/BigInteger
+                :is-from-join false
+                :display-name "ID"
+                :is-from-previous-stage false
+                :is-calculated false
+                :is-implicitly-joinable true}
+               {:table {:name "CATEGORIES", :display-name "Categories", :is-source-table false}
+                :semantic-type :type/Name
+                :name "NAME"
+                :effective-type :type/Text
+                :is-from-join false
+                :display-name "Name"
+                :is-from-previous-stage false
+                :is-calculated false
+                :is-implicitly-joinable true
+                :breakout-position 0}]
+              (map #(lib/display-info query' %) breakoutables'))))))
 
 (deftest ^:parallel breakoutable-columns-with-source-card-e2e-test
   (testing "A column that comes from a source Card (Saved Question/Model/etc) can be broken out by."
@@ -325,8 +392,7 @@
                    (lib/describe-query query')))
             (is (= ["User ID"]
                    (for [breakout (lib/breakouts query')]
-                     (lib/display-name query' breakout))))
-            (assert-breakout-column-excluded query query' name-col)))))))
+                     (lib/display-name query' breakout))))))))))
 
 (deftest ^:parallel breakoutable-columns-expression-e2e-test
   (let [query (-> (lib/query-for-table-name meta/metadata-provider "VENUES")
@@ -352,8 +418,7 @@
                 query'))
         (testing "description"
           (is (= "Venues, Grouped by expr"
-                 (lib/describe-query query'))))
-        (assert-breakout-column-excluded query query' expr)))))
+                 (lib/describe-query query'))))))))
 
 (deftest ^:parallel breakoutable-columns-new-stage-e2e-test
   (let [query (-> (lib/query-for-table-name meta/metadata-provider "VENUES")
@@ -375,5 +440,4 @@
                 query'))
         (testing "description"
           (is (= "Grouped by Expr"
-                 (lib/describe-query query'))))
-        (assert-breakout-column-excluded query query' expr)))))
+                 (lib/describe-query query'))))))))

test/metabase/lib/order_by_test.cljc

@@ -498,7 +498,7 @@
                 (lib/orderable-columns))))))
 
 (deftest ^:parallel orderable-columns-exclude-already-sorted-columns-test
-  (testing "orderable-columns should not return normal Fields already included in :order-by (#29807)"
+  (testing "orderable-columns should return position for normal Fields already included in :order-by (#30568)"
     (let [query (lib/query-for-table-name meta/metadata-provider "VENUES")]
       (is (=? [{:display-name "ID",          :lib/source :source/table-defaults}
                {:display-name "Name",        :lib/source :source/table-defaults}
@@ -509,16 +509,22 @@
                {:display-name "ID",          :lib/source :source/implicitly-joinable}
                {:display-name "Name",        :lib/source :source/implicitly-joinable}]
               (lib/orderable-columns query)))
-      (let [query' (lib/order-by query (second (lib/orderable-columns query)))]
-        (is (=? {:stages [{:order-by [[:asc {} [:field {} (meta/id :venues :name)]]]}]}
+      (let [orderable-columns (lib/orderable-columns query)
+            query' (-> query
+                       (lib/order-by (orderable-columns 5))
+                       (lib/order-by (orderable-columns 1)))]
+        (is (=? {:stages [{:order-by [[:asc {} [:field {} (meta/id :venues :price)]]
+                                      [:asc {} [:field {} (meta/id :venues :name)]]]}]}
                 query'))
-        (is (=? [[:asc {} [:field {} (meta/id :venues :name)]]]
+        (is (=? [[:asc {} [:field {} (meta/id :venues :price)]]
+                 [:asc {} [:field {} (meta/id :venues :name)]]]
                 (lib/order-bys query')))
         (is (=? [{:display-name "ID",          :lib/source :source/table-defaults}
+                 {:display-name "Name",        :lib/source :source/table-defaults, :order-by-position 1}
                  {:display-name "Category ID", :lib/source :source/table-defaults}
                  {:display-name "Latitude",    :lib/source :source/table-defaults}
                  {:display-name "Longitude",   :lib/source :source/table-defaults}
-                 {:display-name "Price",       :lib/source :source/table-defaults}
+                 {:display-name "Price",       :lib/source :source/table-defaults, :order-by-position 0}
                  {:display-name "ID",          :lib/source :source/implicitly-joinable}
                  {:display-name "Name",        :lib/source :source/implicitly-joinable}]
                 (lib/orderable-columns query')))
@@ -535,19 +541,24 @@
                     (lib/orderable-columns query'')))))))))
 
 (deftest ^:parallel orderable-columns-exclude-already-sorted-aggregation-test
-  (testing "orderable-columns should not return aggregation refs that are already in :order-by (#29807)"
+  (testing "orderable-columns should return position for aggregation refs that are already in :order-by (#30568)"
     (let [query (-> (lib/query-for-table-name meta/metadata-provider "VENUES")
                     (lib/aggregate (lib/sum (lib/field (meta/id :venues :price))))
                     (lib/aggregate (lib/sum (lib/field (meta/id :venues :id)))))]
-      (is (=? [{:display-name "Sum of Price", :lib/source :source/aggregations}
-               {:display-name "Sum of ID",    :lib/source :source/aggregations}]
-              (lib/orderable-columns query)))
-      (let [query' (lib/order-by query (first (lib/orderable-columns query)))]
-        (is (=? [{:display-name "Sum of ID", :lib/source :source/aggregations}]
-                (lib/orderable-columns query')))))))
+      (let [orderable-columns (lib/orderable-columns query)]
+        (is (=? [{:display-name "Sum of Price", :lib/source :source/aggregations}
+                 {:display-name "Sum of ID",    :lib/source :source/aggregations}]
+                orderable-columns))
+        (is (empty? (filter :order-by-position orderable-columns))))
+      (let [query' (lib/order-by query (first (lib/orderable-columns query)))
+            orderable-columns (lib/orderable-columns query')]
+        (is (=? [{:display-name "Sum of Price", :lib/source :source/aggregations, :order-by-position 0}
+                 {:display-name "Sum of ID", :lib/source :source/aggregations}]
+                orderable-columns))
+        (is (= 1 (count (filter :order-by-position orderable-columns))))))))
 
 (deftest ^:parallel orderable-columns-exclude-already-sorted-joined-columns-test
-  (testing "orderable-columns should not return joined columns that are already in :order-by (#29807)"
+  (testing "orderable-columns should return position for joined columns that are already in :order-by (#30568)"
     (let [query (-> (lib/query-for-table-name meta/metadata-provider "VENUES")
                     (lib/join (-> (lib/table (meta/id :categories))
                                   (lib/with-join-alias "Cat")
@@ -583,11 +594,12 @@
                  {:display-name "Latitude",    :lib/source :source/table-defaults}
                  {:display-name "Longitude",   :lib/source :source/table-defaults}
                  {:display-name "Price",       :lib/source :source/table-defaults}
-                 {:display-name "ID",          :lib/source :source/joins}]
+                 {:display-name "ID",          :lib/source :source/joins}
+                 {:display-name "Name", :lib/source :source/joins, :order-by-position 0}]
                 (lib/orderable-columns query')))))))
 
 (deftest ^:parallel orderable-columns-exclude-already-sorted-implicitly-joinable-columns-test
-  (testing "orderable-columns should not return implicitly joinable columns that are already in :order-by (#29807)"
+  (testing "orderable-columns should return position implicitly joinable columns that are already in :order-by (#30568)"
     (let [query (lib/query-for-table-name meta/metadata-provider "VENUES")
           query (-> query
                     (lib/order-by (m/find-first #(= (:id %) (meta/id :categories :name))
@@ -604,11 +616,12 @@
                {:display-name "Latitude",    :lib/source :source/table-defaults}
                {:display-name "Longitude",   :lib/source :source/table-defaults}
                {:display-name "Price",       :lib/source :source/table-defaults}
-               {:display-name "ID",          :lib/source :source/implicitly-joinable}]
+               {:display-name "ID",          :lib/source :source/implicitly-joinable}
+               {:display-name "Name",        :lib/source :source/implicitly-joinable, :order-by-position 0}]
               (lib/orderable-columns query))))))
 
 (deftest ^:parallel orderable-columns-exclude-already-sorted-expression-test
-  (testing "orderable-columns should not return expressions that are already in :order-by (#29807)"
+  (testing "orderable-columns should return position for expressions that are already in :order-by (#30568)"
     (let [query (-> (lib/query-for-table-name meta/metadata-provider "VENUES")
                     (lib/expression "My Expression" (lib/+ 2 3)))]
       (is (=? [{:display-name "ID",            :lib/source :source/table-defaults}
@@ -629,6 +642,7 @@
                  {:display-name "Latitude",      :lib/source :source/table-defaults}
                  {:display-name "Longitude",     :lib/source :source/table-defaults}
                  {:display-name "Price",         :lib/source :source/table-defaults}
+                 {:display-name "My Expression", :lib/source :source/expressions, :order-by-position 0}
                  {:display-name "ID",            :lib/source :source/implicitly-joinable}
                  {:display-name "Name",          :lib/source :source/implicitly-joinable}]
                 (lib/orderable-columns query')))))))
@@ -737,7 +751,17 @@
               :table          {:name "VENUES", :display-name "Venues"}
               :direction      :asc}]
             (for [order-by (lib/order-bys query')]
-              (lib/display-info query' order-by))))))
+              (lib/display-info query' order-by))))
+    (is (=? [{:display-name "ID"}
+             {:display-name "Name", :order-by-position 0}
+             {:display-name "Category ID"}
+             {:display-name "Latitude"}
+             {:display-name "Longitude"}
+             {:display-name "Price"}
+             {:display-name "ID"}
+             {:display-name "Name"}]
+            (for [orderable-column (lib/orderable-columns query')]
+              (lib/display-info query' orderable-column))))))
 
 (deftest ^:parallel change-direction-test
   (doseq [[dir opposite] {:asc :desc, :desc :asc}]