[MLv2] remove & replace breakout clause (#29777)

* Add current-breakout function

* Add handling for breakouts to remove and replace clause

* Fix tests and cljs compile

* Fix lints

* Address PR comments

* Adjust naming based on PR comments

* Move remove and replace clause to avoid circular def

* Fix js requires

* Fix JS requires

* Update for filter to filters change

* Address PR comment

src/metabase/lib/breakout.cljc

@@ -31,6 +31,14 @@
      (lib.util/update-query-stage query stage-number update :breakout (fn [breakouts]
                                                                         (conj (vec breakouts) expr))))))
 
+(mu/defn current-breakouts :- [:maybe [:sequential ::lib.schema.expression/expression]]
+  "Return the current breakouts"
+  ([query]
+   (current-breakouts query -1))
+  ([query :- ::lib.schema/query
+    stage-number :- :int]
+   (not-empty (:breakout (lib.util/query-stage query stage-number)))))
+
 (mu/defn breakouts :- [:maybe [:sequential lib.metadata/ColumnMetadata]]
   "Get metadata about the breakouts in a given stage of a `query`."
   [query        :- ::lib.schema/query

src/metabase/lib/core.cljc

@@ -20,6 +20,7 @@
    [metabase.lib.order-by :as lib.order-by]
    [metabase.lib.query :as lib.query]
    [metabase.lib.ref :as lib.ref]
+   [metabase.lib.remove-replace :as lib.remove-replace]
    [metabase.lib.segment :as lib.segment]
    [metabase.lib.stage :as lib.stage]
    [metabase.lib.table :as lib.table]
@@ -64,7 +65,8 @@
    sum
    sum-where]
   [lib.breakout
-   breakout]
+   breakout
+   current-breakouts]
   [lib.dev
    field
    query-for-table-id
@@ -164,11 +166,12 @@
   [lib.query
    native-query
    query
-   remove-clause
-   replace-clause
    saved-question-query]
   [lib.ref
    ref]
+  [lib.remove-replace
+   remove-clause
+   replace-clause]
   [lib.stage
    append-stage
    drop-stage]

src/metabase/lib/js.cljs

@@ -10,6 +10,7 @@
    [metabase.lib.normalize :as lib.normalize]
    [metabase.lib.order-by :as lib.order-by]
    [metabase.lib.query :as lib.query]
+   [metabase.lib.remove-replace :as lib.remove-replace]
    [metabase.mbql.js :as mbql.js]
    [metabase.mbql.normalize :as mbql.normalize]
    [metabase.util :as u]
@@ -145,7 +146,7 @@
   ([a-query clause]
    (remove-clause a-query -1 clause))
   ([a-query stage-number clause]
-   (lib.query/remove-clause
+   (lib.remove-replace/remove-clause
      a-query stage-number
      (lib.normalize/normalize (js->clj clause :keywordize-keys true)))))
 
@@ -154,7 +155,7 @@
   ([a-query target-clause new-clause]
    (replace-clause a-query -1 target-clause new-clause))
   ([a-query stage-number target-clause new-clause]
-   (lib.query/replace-clause
+   (lib.remove-replace/replace-clause
      a-query stage-number
      (lib.normalize/normalize (js->clj target-clause :keywordize-keys true))
      (lib.normalize/normalize (js->clj new-clause :keywordize-keys true)))))

src/metabase/lib/query.cljc

 (ns metabase.lib.query
   (:refer-clojure :exclude [remove])
   (:require
-   [metabase.lib.common :as lib.common]
    [metabase.lib.convert :as lib.convert]
    [metabase.lib.dispatch :as lib.dispatch]
    [metabase.lib.hierarchy :as lib.hierarchy]
@@ -14,45 +13,6 @@
    [metabase.lib.util :as lib.util]
    [metabase.util.malli :as mu]))
 
-(mu/defn replace-clause :- :metabase.lib.schema/query
-  "Replaces the `target-clause` with `new-clause` in the `query` stage."
-  ([query :- :metabase.lib.schema/query
-    target-clause
-    new-clause]
-   (replace-clause query -1 target-clause new-clause))
-  ([query :- :metabase.lib.schema/query
-    stage-number :- :int
-    target-clause
-    new-clause]
-   (let [replacement (lib.common/->op-arg query stage-number new-clause)]
-     ;; Right now, this works for clauses that cannot have dependents.
-     ;; This will change and have different logic depending on `location`
-     ;; `location` should probably be "found" first before iterating.
-     (reduce
-       (fn [query location]
-         (lib.util/update-query-stage query stage-number
-                                      lib.util/replace-clause location target-clause replacement))
-       query
-       [:order-by]))))
-
-(mu/defn remove-clause :- :metabase.lib.schema/query
-  "Removes the `target-clause` in the filter of the `query`."
-  ([query :- :metabase.lib.schema/query
-    target-clause]
-   (remove-clause query -1 target-clause))
-  ([query :- :metabase.lib.schema/query
-    stage-number :- :int
-    target-clause]
-   ;; Right now, this works for clauses that cannot have dependents.
-   ;; This will change and have different logic depending on `location`
-   ;; `location` should probably be "found" first before iterating.
-   (reduce
-     (fn [query location]
-       (lib.util/update-query-stage query stage-number
-                                    lib.util/remove-clause location target-clause))
-     query
-     [:order-by])))
-
 (defmethod lib.normalize/normalize :mbql/query
   [query]
   (lib.normalize/normalize-map

src/metabase/lib/remove_replace.cljc

+(ns metabase.lib.remove-replace
+  (:require
+   [metabase.lib.common :as lib.common]
+   [metabase.lib.ref :as lib.ref]
+   [metabase.lib.stage :as lib.stage]
+   [metabase.lib.util :as lib.util]
+   [metabase.mbql.util.match :as mbql.match]
+   [metabase.util.malli :as mu]))
+
+(defn- find-clause
+  [query stage-number target-clause]
+  (let [stage (lib.util/query-stage query stage-number)
+        [target-type _opts target-id] target-clause]
+    (->> [:order-by :aggregation :breakout :filters :expressions :joins :fields]
+         (keep (fn [top-level-clause]
+                 (mbql.match/match-one (get stage top-level-clause)
+                   [target-type _ target-id] top-level-clause)))
+         (set)
+         (not-empty))))
+
+(defn- target-ref-for-stage
+  "Gets the ref for the target-id exposed by the previous stage"
+  [query stage-number target-id]
+  (->> (lib.stage/visible-columns query stage-number)
+       (some (fn [{:keys [lib/source id] :as column}]
+               (when (and (= :source/previous-stage source) (= target-id id))
+                 (lib.ref/ref column))))))
+
+(defn- check-subsequent-stages-for-invalid-target!
+  "Throws if target-clause is used in a subsequent stage"
+  [previous-query query stage-number target-clause]
+  (let [[_ _ target-id] target-clause]
+    (loop [stage-number stage-number]
+      (when-let [next-stage-number (lib.util/next-stage-number query stage-number)]
+        ;; The target could still be exposed (i.e. removing the last breakout could expose itself through default fields)
+        (when-not (target-ref-for-stage query next-stage-number target-id)
+          ;; Get the ref to look for from the previous-query
+          (let [target-ref (target-ref-for-stage previous-query next-stage-number target-id)]
+            (if-let [found (find-clause query next-stage-number target-ref)]
+              (throw (ex-info "Clause cannot be removed as it has dependents" {:target-clause target-clause
+                                                                               :stage-number next-stage-number
+                                                                               :found found}))
+              (recur next-stage-number))))))))
+
+(mu/defn remove-clause :- :metabase.lib.schema/query
+  "Removes the `target-clause` in the filter of the `query`."
+  ([query :- :metabase.lib.schema/query
+    target-clause]
+   (remove-clause query -1 target-clause))
+  ([query :- :metabase.lib.schema/query
+    stage-number :- :int
+    target-clause]
+   (reduce
+     (fn [query location]
+       (let [result (lib.util/update-query-stage query stage-number
+                                                 lib.util/remove-clause location target-clause)]
+         (when (not= query result)
+           (case location
+             :breakout (check-subsequent-stages-for-invalid-target! query result stage-number (lib.ref/ref target-clause))
+             nil)
+           (reduced result))
+         result))
+     query
+     ;; TODO only these top level clauses are supported at this moment
+     [:order-by :breakout])))
+
+(mu/defn replace-clause :- :metabase.lib.schema/query
+  "Replaces the `target-clause` with `new-clause` in the `query` stage."
+  ([query :- :metabase.lib.schema/query
+    target-clause
+    new-clause]
+   (replace-clause query -1 target-clause new-clause))
+  ([query :- :metabase.lib.schema/query
+    stage-number :- :int
+    target-clause
+    new-clause]
+   (let [replacement (lib.common/->op-arg query stage-number new-clause)]
+     (reduce
+       (fn [query location]
+         (let [result (lib.util/update-query-stage query stage-number
+                                                   lib.util/replace-clause location target-clause replacement)]
+           (when (not= query result)
+             (case location
+               :breakout (check-subsequent-stages-for-invalid-target! query result stage-number (lib.ref/ref target-clause))
+               nil)
+             (reduced result))
+           result))
+       query
+       ;; TODO only these top level clauses are supported at this moment
+       [:order-by :breakout]))))

src/metabase/lib/util.cljc

@@ -3,6 +3,7 @@
   (:require
    [clojure.set :as set]
    [clojure.string :as str]
+   [medley.core :as m]
    [metabase.lib.options :as lib.options]
    [metabase.lib.schema :as lib.schema]
    [metabase.lib.schema.common :as lib.schema.common]
@@ -47,7 +48,7 @@
    If `location` contains no clause with `target-clause` no replacement happens."
   [stage location target-clause new-clause]
   {:pre [(clause? target-clause)]}
-  (update
+  (m/update-existing
     stage
     location
     #(->> (for [clause %]
@@ -57,21 +58,19 @@
           vec)))
 
 (defn remove-clause
-  "Replace the `target-clause` in `stage` `location`.
+  "Remove the `target-clause` in `stage` `location`.
    If a clause has :lib/uuid equal to the `target-clause` it is removed.
    If `location` contains no clause with `target-clause` no removal happens.
    If the the location is empty, dissoc it from stage."
   [stage location target-clause]
   {:pre [(clause? target-clause)]}
-  (let [target-uuid (clause-uuid target-clause)
-        target (get stage location)
-        result (->> target
-                    (remove (comp #{target-uuid} clause-uuid))
-                    vec
-                    not-empty)]
-    (if result
-      (assoc stage location result)
-      (dissoc stage location))))
+  (if-let [target (get stage location)]
+    (let [target-uuid (clause-uuid target-clause)
+          result (into [] (remove (comp #{target-uuid} clause-uuid)) target)]
+      (if (seq result)
+        (assoc stage location result)
+        (dissoc stage location)))
+    stage))
 
 ;;; TODO -- all of this `->pipeline` stuff should probably be merged into [[metabase.lib.convert]] at some point in
 ;;; the near future.
@@ -187,6 +186,15 @@
     (when (pos? stage-number)
       (dec stage-number))))
 
+(defn next-stage-number
+  "The index of the next stage, if there is one. `nil` if there is no next stage."
+  [{:keys [stages], :as _query} stage-number]
+  (let [stage-number (if (neg? stage-number)
+                       (+ (count stages) stage-number)
+                       stage-number)]
+    (when (< (inc stage-number) (count stages))
+      (inc stage-number))))
+
 (mu/defn query-stage :- ::lib.schema/stage
   "Fetch a specific `stage` of a query. This handles negative indices as well, e.g. `-1` will return the last stage of
   the query."

test/metabase/lib/breakout_test.cljc

@@ -25,3 +25,9 @@
            (lib/display-name query query)
            (lib/describe-query query)
            (lib/suggested-name query)))))
+
+(deftest ^:parallel current-breakouts-test
+  (let [query (-> (lib/query-for-table-name meta/metadata-provider "CHECKINS")
+                  (lib/breakout (lib/field (meta/id :checkins :date))))]
+    (is (=? [[:field {} (meta/id :checkins :date)]]
+            (lib/current-breakouts query)))))

test/metabase/lib/query_test.cljc

 (ns metabase.lib.query-test
   (:require
    #?@(:cljs ([metabase.test-runner.assert-exprs.approximately-equal]))
-   [clojure.test :refer [deftest is testing]]
+   [clojure.test :refer [deftest is]]
    [metabase.lib.core :as lib]
    [metabase.lib.metadata.calculation :as lib.metadata.calculation]
    [metabase.lib.test-metadata :as meta]))
@@ -71,36 +71,3 @@
           (lib/query meta/metadata-provider {:database (meta/id)
                                              :type     :query
                                              :query    {:source-query {:source-query {:source-table (meta/id :venues)}}}}))))
-
-(deftest ^:parallel remove-clause-test
-  (let [query (-> (lib/query-for-table-name meta/metadata-provider "VENUES")
-                  (lib/filter (lib/= "myvenue" (lib/field (meta/id :venues :name))))
-                  (lib/order-by (lib/field (meta/id :venues :name)))
-                  (lib/order-by (lib/field (meta/id :venues :name))))
-        order-bys (lib/order-bys query)]
-    (testing "order-bys"
-      (is (= 2 (count order-bys)))
-      (is (= 1 (-> query
-                   (lib/remove-clause (first order-bys))
-                   (lib/order-bys)
-                   count)))
-      (is (= 0 (-> query
-                   (lib/remove-clause (first order-bys))
-                   (lib/remove-clause (second order-bys))
-                   (lib/order-bys)
-                   count))))))
-
-(deftest ^:parallel replace-clause-test
-  (let [query (-> (lib/query-for-table-name meta/metadata-provider "VENUES")
-                  (lib/filter (lib/= "myvenue" (lib/field (meta/id :venues :name))))
-                  (lib/order-by (lib/field (meta/id :venues :name)))
-                  (lib/order-by (lib/field (meta/id :venues :name))))
-        order-bys (lib/order-bys query)]
-    (testing "order-bys"
-      (is (= 2 (count order-bys)))
-      (let [replaced (-> query
-                         (lib/replace-clause (first order-bys) (lib/order-by-clause (lib/field (meta/id :venues :id)))))
-            replaced-order-bys (lib/order-bys replaced)]
-        (is (not= order-bys replaced-order-bys))
-        (is (= 2 (count replaced-order-bys)))
-        (is (= (second order-bys) (second replaced-order-bys)))))))

test/metabase/lib/remove_replace_test.cljc

+(ns metabase.lib.remove-replace-test
+  (:require
+   #?@(:cljs ([metabase.test-runner.assert-exprs.approximately-equal]))
+   [clojure.test :refer [deftest is testing]]
+   [metabase.lib.core :as lib]
+   [metabase.lib.test-metadata :as meta]))
+
+(deftest ^:parallel remove-clause-order-bys-test
+  (let [query (-> (lib/query-for-table-name meta/metadata-provider "VENUES")
+                  (lib/order-by (lib/field (meta/id :venues :name)))
+                  (lib/order-by (lib/field (meta/id :venues :name))))
+        order-bys (lib/order-bys query)]
+    (is (= 2 (count order-bys)))
+    (is (= 1 (-> query
+                 (lib/remove-clause (first order-bys))
+                 (lib/order-bys)
+                 count)))
+    (is (= 0 (-> query
+                 (lib/remove-clause (first order-bys))
+                 (lib/remove-clause (second order-bys))
+                 (lib/order-bys)
+                 count)))))
+
+(deftest ^:parallel remove-clause-breakout-test
+  (let [query (-> (lib/query-for-table-name meta/metadata-provider "VENUES")
+                  (lib/breakout (lib/field "VENUES" "ID"))
+                  (lib/breakout (lib/field "VENUES" "NAME")))
+        breakouts (lib/current-breakouts query)]
+    (is (= 2 (count breakouts)))
+    (is (= 1 (-> query
+                 (lib/remove-clause (first breakouts))
+                 (lib/current-breakouts)
+                 count)))
+    (is (= 0 (-> query
+                 (lib/remove-clause (first breakouts))
+                 (lib/remove-clause (second breakouts))
+                 (lib/current-breakouts)
+                 count)))
+    (testing "removing breakout with dependent should not be allowed"
+      (is (thrown-with-msg?
+            #?(:clj Exception :cljs js/Error)
+            #"Clause cannot be removed as it has dependents"
+            (-> query
+                (lib/append-stage)
+                ;; TODO Should be able to create a ref with lib/field [#29763]
+                (lib/filter (lib/= [:field {:lib/uuid (str (random-uuid)) :base-type :type/Integer} "ID"] 1))
+                (lib/remove-clause 0 (first breakouts)))))
+      (is (thrown-with-msg?
+            #?(:clj Exception :cljs js/Error)
+            #"Clause cannot be removed as it has dependents"
+            (-> query
+                (lib/append-stage)
+                (lib/fields [[:field {:lib/uuid (str (random-uuid)) :base-type :type/Integer} "ID"]])
+                (lib/append-stage)
+                ;; TODO Should be able to create a ref with lib/field [#29763]
+                (lib/filter (lib/= [:field {:lib/uuid (str (random-uuid)) :base-type :type/Integer} "ID"] 1))
+                (lib/remove-clause 0 (first breakouts)))))
+      (is (= 0 (-> query
+                   (lib/remove-clause 0 (second breakouts))
+                   (lib/append-stage)
+                   ;; TODO Should be able to create a ref with lib/field [#29763]
+                   (lib/filter (lib/= [:field {:lib/uuid (str (random-uuid)) :base-type :type/Integer} "ID"] 1))
+                   (lib/remove-clause 0 (first breakouts))
+                   (lib/current-breakouts 0)
+                   count))))))
+
+(deftest ^:parallel replace-clause-order-by-test
+  (let [query (-> (lib/query-for-table-name meta/metadata-provider "VENUES")
+                  (lib/filter (lib/= "myvenue" (lib/field (meta/id :venues :name))))
+                  (lib/order-by (lib/field (meta/id :venues :name)))
+                  (lib/order-by (lib/field (meta/id :venues :name))))
+        order-bys (lib/order-bys query)]
+    (is (= 2 (count order-bys)))
+    (let [replaced (-> query
+                       (lib/replace-clause (first order-bys) (lib/order-by-clause (lib/field (meta/id :venues :id)))))
+          replaced-order-bys (lib/order-bys replaced)]
+      (is (not= order-bys replaced-order-bys))
+      (is (=? [:asc {} [:field {} (meta/id :venues :id)]]
+              (first replaced-order-bys)))
+      (is (= 2 (count replaced-order-bys)))
+      (is (= (second order-bys) (second replaced-order-bys))))))
+
+(deftest ^:parallel replace-clause-breakout-by-test
+  (let [query (-> (lib/query-for-table-name meta/metadata-provider "VENUES")
+                  (lib/breakout (lib/field (meta/id :venues :id)))
+                  (lib/breakout (lib/field (meta/id :venues :name))))
+        breakouts (lib/current-breakouts query)
+        replaced (-> query
+                     (lib/replace-clause (first breakouts) (lib/field (meta/id :venues :price))))
+        replaced-breakouts (lib/current-breakouts replaced)]
+    (is (= 2 (count breakouts)))
+    (is (=? [:field {} (meta/id :venues :price)]
+            (first replaced-breakouts)))
+    (is (not= breakouts replaced-breakouts))
+    (is (= 2 (count replaced-breakouts)))
+    (is (= (second breakouts) (second replaced-breakouts)))
+    (testing "replacing breakout with dependent should not be allowed"
+      (is (thrown-with-msg?
+            #?(:clj Exception :cljs js/Error)
+            #"Clause cannot be removed as it has dependents"
+            (-> query
+                (lib/append-stage)
+                ;; TODO Should be able to create a ref with lib/field [#29763]
+                (lib/filter (lib/= [:field {:lib/uuid (str (random-uuid)) :base-type :type/Integer} "ID"] 1))
+                (lib/replace-clause 0 (first breakouts) (lib/field "VENUES" "PRICE")))))
+      (is (not= breakouts (-> query
+                              (lib/append-stage)
+                              ;; TODO Should be able to create a ref with lib/field [#29763]
+                              (lib/filter (lib/= [:field {:lib/uuid (str (random-uuid)) :base-type :type/Integer} "ID"] 1))
+                              (lib/replace-clause 0 (second breakouts) (lib/field "VENUES" "PRICE"))
+                              (lib/current-breakouts 0)))))))