From cb6a6849e447900cd867a9b307872bb2de6d2ffb Mon Sep 17 00:00:00 2001
From: Case Nelson <case@metabase.com>
Date: Tue, 28 Mar 2023 21:03:13 -0600
Subject: [PATCH] [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 | 3 +++
src/metabase/lib/order_by.cljc | 2 +-
src/metabase/lib/stage.cljc | 13 +++++++++++++
test/metabase/lib/stage_test.cljc | 14 ++++++++++++++
4 files changed, 31 insertions(+), 1 deletion(-)
diff --git a/src/metabase/lib/core.cljc b/src/metabase/lib/core.cljc
index 849e36c9a50..b1f7100d344 100644
--- a/src/metabase/lib/core.cljc
+++ b/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])
diff --git a/src/metabase/lib/order_by.cljc b/src/metabase/lib/order_by.cljc
index 6d11f6f8425..a8860c94eaf 100644
--- a/src/metabase/lib/order_by.cljc
+++ b/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))
diff --git a/src/metabase/lib/stage.cljc b/src/metabase/lib/stage.cljc
index 74eafaebfbd..f8c6a59099e 100644
--- a/src/metabase/lib/stage.cljc
+++ b/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)))
diff --git a/test/metabase/lib/stage_test.cljc b/test/metabase/lib/stage_test.cljc
index ce32a2cb546..9a7885ccfc3 100644
--- a/test/metabase/lib/stage_test.cljc
+++ b/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)))))))
--
GitLab