[MLv2] Adds stage-count (#35722)

Fixes #35719

src/metabase/lib/core.cljc

@@ -256,6 +256,7 @@
  [lib.query
   can-run
   query
+  stage-count
   with-different-table]
  [lib.ref
   ref]

src/metabase/lib/js.cljs

@@ -996,3 +996,8 @@
     offset-n
     (some-> offset-unit keyword)
     (js->clj options :keywordize-keys true)))
+
+(defn ^:export stage-count
+  "Returns the count of stages in query"
+  [a-query]
+  (lib.core/stage-count a-query))

src/metabase/lib/query.cljc

@@ -11,6 +11,7 @@
    [metabase.lib.metadata.calculation :as lib.metadata.calculation]
    [metabase.lib.normalize :as lib.normalize]
    [metabase.lib.schema :as lib.schema]
+   [metabase.lib.schema.common :as lib.schema.common]
    [metabase.lib.schema.id :as lib.schema.id]
    [metabase.lib.util :as lib.util]
    [metabase.mbql.util :as mbql.u]
@@ -39,6 +40,11 @@
   [query stage-number x style]
   (lib.metadata.calculation/display-name query stage-number (lib.util/query-stage x stage-number) style))
 
+(mu/defn stage-count :- ::lib.schema.common/int-greater-than-or-equal-to-zero
+  "Returns the count of stages in query"
+  [query :- ::lib.schema/query]
+  (count (:stages query)))
+
 (defmulti can-run-method
   "Returns whether the query is runnable based on first stage :lib/type"
   (fn [query]

test/metabase/lib/query_test.cljc

@@ -90,3 +90,8 @@
                                                  "CC"]]]}]}]}
 
               (lib/query meta/metadata-provider converted-query))))))
+
+(deftest ^:parallel stage-count-test
+  (is (= 1 (lib/stage-count lib.tu/venues-query)))
+  (is (= 2 (lib/stage-count (lib/append-stage lib.tu/venues-query))))
+  (is (= 3 (lib/stage-count (lib/append-stage (lib/append-stage lib.tu/venues-query))))))