[MLv2] Append and drop final stage (#29625)

* [MLv2] Append and drop final stage

* Address review

* order-bys can return nil

* Check exceptions differently in cljs

src/metabase/lib/core.cljc

@@ -145,5 +145,8 @@
    native-query
    query
    saved-question-query]
+  [lib.stage
+   append-stage
+   drop-stage]
   [lib.temporal-bucket
    temporal-bucket])

src/metabase/lib/order_by.cljc

@@ -87,7 +87,7 @@
      (lib.util/update-query-stage query stage-number update :order-by (fn [order-bys]
                                                                         (conj (vec order-bys) new-order-by))))))
 
-(mu/defn order-bys :- [:sequential ::lib.schema.order-by/order-by]
+(mu/defn order-bys :- [:maybe [:sequential ::lib.schema.order-by/order-by]]
   "Get the order-by clauses in a query."
   ([query :- ::lib.schema/query]
    (order-bys query -1))

src/metabase/lib/stage.cljc

@@ -8,6 +8,7 @@
    [metabase.lib.expression :as lib.expression]
    [metabase.lib.metadata :as lib.metadata]
    [metabase.lib.metadata.calculation :as lib.metadata.calculation]
+   [metabase.lib.options :as lib.options]
    [metabase.lib.schema :as lib.schema]
    [metabase.lib.schema.common :as lib.schema.common]
    [metabase.lib.schema.id :as lib.schema.id]
@@ -248,3 +249,15 @@
      (lib.expression/expressions query stage-number)
      columns
      (implicitly-joinable-columns query columns))))
+
+(mu/defn append-stage :- ::lib.schema/query
+  "Adds a new blank stage to the end of the pipeline"
+  [query]
+  (update query :stages conj (lib.options/ensure-uuid {:lib/type :mbql.stage/mbql})))
+
+(mu/defn drop-stage :- ::lib.schema/query
+  "Drops the final stage in the pipeline"
+  [query]
+  (when (= 1 (count (:stages query)))
+    (throw (ex-info (i18n/tru "Cannot drop the only stage") {:stages (:stages query)})))
+  (update query :stages (comp vec butlast)))

test/metabase/lib/stage_test.cljc

@@ -64,3 +64,17 @@
                                :source-table "card__1"}]}]
     (is (= "My Card"
            (lib.metadata.calculation/display-name query -1 query)))))
+
+(deftest ^:parallel adding-and-removing-stages
+  (let [query (lib/query-for-table-name meta/metadata-provider "VENUES")
+        query-with-new-stage (-> query
+                                 lib/append-stage
+                                 (lib/order-by 1 (lib/field "VENUES" "NAME") :asc))]
+    (is (= 0 (count (lib/order-bys query-with-new-stage 0))))
+    (is (= 1 (count (lib/order-bys query-with-new-stage 1))))
+    (is (= query
+           (-> query-with-new-stage
+               (lib/filter (lib/= 1 (lib/field "VENUES" "NAME")))
+               (lib/drop-stage))))
+    (testing "Dropping with 1 stage should error"
+      (is (thrown-with-msg? #?(:cljs :default :clj Exception) #"Cannot drop the only stage" (-> query (lib/drop-stage)))))))