[MLv2] Clean query to remove any parts that fail schema validation (#30096)

* [MLv2] Clean query to remove any parts that fail schema validation

* Fix metric query generation

* Make test parallel

* Fix joins schema

src/metabase/lib/convert.cljc

 (ns metabase.lib.convert
   (:require
    [clojure.set :as set]
+   [malli.core :as mc]
    [medley.core :as m]
    [metabase.lib.dispatch :as lib.dispatch]
    [metabase.lib.options :as lib.options]
-   [metabase.lib.util :as lib.util]))
+   [metabase.lib.schema :as lib.schema]
+   [metabase.lib.util :as lib.util]
+   [metabase.util :as u]))
+
+(defn- clean-location [almost-stage error-type error-location]
+  (let [operate-on-parent? #{:malli.core/missing-key :malli.core/end-of-input}
+        location (if (operate-on-parent? error-type)
+                   (drop-last 2 error-location)
+                   (drop-last 1 error-location))
+        [location-key] (if (operate-on-parent? error-type)
+                         (take-last 2 error-location)
+                         (take-last 1 error-location))]
+    (if (seq location)
+      (update-in almost-stage
+                 location
+                 (fn [error-loc]
+                   (let [result (assoc error-loc location-key nil)]
+                     (cond
+                       (vector? error-loc) (into [] (remove nil?) result)
+                       (map? error-loc) (u/remove-nils result)
+                       :else result))))
+      (dissoc almost-stage location-key))))
+
+(def ^:private stage-keys-to-clean
+  #{:expressions :joins :filters :order-by :aggregation :fields :breakout})
+
+(defn- clean-stage [almost-stage]
+  (loop [almost-stage almost-stage
+         removals []]
+    (if-let [[error-type error-location] (->> (mc/explain ::lib.schema/stage.mbql almost-stage)
+                                              :errors
+                                              (filter (comp stage-keys-to-clean first :in))
+                                              (map (juxt :type :in))
+                                              first)]
+      (let [new-stage (clean-location almost-stage error-type error-location)]
+        (if (= new-stage almost-stage)
+          almost-stage
+          (recur new-stage (conj removals [error-type error-location]))))
+      almost-stage)))
+
+(defn- clean [almost-query]
+  (loop [almost-query almost-query
+         stage-index 0]
+    (let [current-stage (nth (:stages almost-query) stage-index)
+          new-stage (clean-stage current-stage)]
+      (if (= current-stage new-stage)
+        (if (= stage-index (dec (count (:stages almost-query))))
+          almost-query
+          (recur almost-query (inc stage-index)))
+        (recur (update almost-query :stages assoc stage-index new-stage) stage-index)))))
 
 (defmulti ->pMBQL
   "Coerce something to pMBQL (the version of MBQL manipulated by Metabase Lib v2) if it's not already pMBQL."
@@ -59,7 +109,8 @@
   (if (:type m)
     (-> (lib.util/pipeline m)
         (update :stages (fn [stages]
-                          (mapv ->pMBQL stages))))
+                          (mapv ->pMBQL stages)))
+        clean)
     (update-vals m ->pMBQL)))
 
 (defmethod ->pMBQL :field

src/metabase/lib/join.cljc

@@ -299,13 +299,13 @@
      (lib.util/update-query-stage query stage-number update :joins (fn [joins]
                                                                      (conj (vec joins) new-join))))))
 
-(mu/defn joins :- ::lib.schema.join/joins
+(mu/defn joins :- [:maybe ::lib.schema.join/joins]
   "Get all joins in a specific `stage` of a `query`. If `stage` is unspecified, returns joins in the final stage of the
   query."
   ([query]
    (joins query -1))
   ([query        :- ::lib.schema/query
-    stage-number :- ::lib.schema.common/int-greater-than-or-equal-to-zero]
+    stage-number :- :int]
    (not-empty (get (lib.util/query-stage query stage-number) :joins))))
 
 (mu/defn implicit-join-name :- ::lib.schema.common/non-blank-string

src/metabase/lib/metric.cljc

@@ -22,10 +22,11 @@
       definition
       ;; legacy; needs conversion
       (->
-       (lib.convert/legacy-query-from-inner-query nil definition)
-       mbql.normalize/normalize
-       lib.convert/->pMBQL
-       (lib.util/query-stage -1)))))
+        ;; database-id cannot be nil, but gets thrown out
+        (lib.convert/legacy-query-from-inner-query #?(:clj Integer/MAX_VALUE :cljs js/Number.MAX_SAFE_INTEGER) definition)
+        mbql.normalize/normalize
+        lib.convert/->pMBQL
+        (lib.util/query-stage -1)))))
 
 (defmethod lib.metadata.calculation/type-of-method :metadata/metric
   [query stage-number metric-metadata]

test/metabase/lib/convert_test.cljc

@@ -2,6 +2,7 @@
   (:require
    [clojure.test :refer [are deftest is testing]]
    [metabase.lib.convert :as lib.convert]
+   [metabase.lib.core :as lib]
    [metabase.lib.test-metadata :as meta]
    #?@(:cljs ([metabase.test-runner.assert-exprs.approximately-equal]))))
 
@@ -77,7 +78,8 @@
           (lib.convert/->pMBQL
            {:database (meta/id)
             :type     :query
-            :query    {:fields [[:field (meta/id :categories :name) {:join-alias "CATEGORIES__via__CATEGORY_ID"}]]
+            :query    {:source-table (meta/id :categories)
+                       :fields [[:field (meta/id :categories :name) {:join-alias "CATEGORIES__via__CATEGORY_ID"}]]
                        :joins  [{:alias        "CATEGORIES__via__CATEGORY_ID"
                                  :source-table (meta/id :venues)
                                  :condition    [:=
@@ -121,6 +123,7 @@
                                        {:lib/uuid string?, :display-name "Revenue"}
                                        [:field {:lib/uuid string?} 1]]]}]}
           (lib.convert/->pMBQL {:type  :query
+                                :database 5
                                 :query {:source-table 1
                                         :aggregation  [[:aggregation-options
                                                         [:sum [:field 1 nil]]
@@ -188,7 +191,8 @@
 
     {:database 23001
      :type     :query
-     :query    {:fields [[:field 23101 {:join-alias "CATEGORIES__via__CATEGORY_ID"}]]
+     :query    {:source-table 224
+                :fields [[:field 23101 {:join-alias "CATEGORIES__via__CATEGORY_ID"}]]
                 :joins  [{:alias        "CATEGORIES__via__CATEGORY_ID"
                           :source-table 23040
                           :condition    [:=
@@ -199,7 +203,8 @@
 
     {:database 1
      :type     :query
-     :query    {:order-by [[:asc [:field 1 nil]]]}}
+     :query    {:source-table 224
+                :order-by [[:asc [:field 1 nil]]]}}
 
     {:database 5
      :type     :query
@@ -209,3 +214,61 @@
                                 :fields       [[:field 1 {:join-alias "Cat"}]]}]
                 :limit        1
                 :source-table 4}}))
+
+(deftest ^:parallel clean-test
+  (testing "irrecoverable queries"
+    ;; Eventually we should get to a place where ->pMBQL throws an exception here,
+    ;; but legacy e2e tests make this impossible right now
+    (is (= {:type :query
+            :query {}}
+           (lib.convert/->legacy-MBQL
+             (lib.convert/->pMBQL
+               {:type :query}))))
+    (is (= {:type :query
+            :database 1
+            :query {}}
+            (lib.convert/->legacy-MBQL
+              (lib.convert/->pMBQL
+                {:type :query
+                 :database 1}))))
+    (is (= {:type :query
+            :database 1
+            :query {}}
+           (lib.convert/->legacy-MBQL
+             (lib.convert/->pMBQL
+               {:type :query
+                :database 1})))))
+  (testing "recoverable queries"
+    (is (nil? (->
+                {:database 1
+                 :type :query
+                 :query {:source-table 224
+                         :order-by [[:asc [:xfield 1 nil]]]}}
+                lib.convert/->pMBQL
+                lib/order-bys)))
+    (is (nil? (->
+                {:database 1
+                 :type :query
+                 :query {:source-table 224
+                         :filter [:and [:= [:xfield 1 nil]]]}}
+                lib.convert/->pMBQL
+                lib/filters)))
+    (is (nil? (->
+                {:database 5
+                 :type :query
+                 :query {:joins [{:source-table 3
+                                  ;; Invalid condition makes the join invalid
+                                  :condition [:= [:field 2 nil] [:xfield 2 nil]]}]
+                         :source-table 4}}
+                 lib.convert/->pMBQL
+                 lib/joins)))
+    (is (nil? (->
+                {:database 5
+                 :type :query
+                 :query {:joins [{:source-table 3
+                                  :condition [:= [:field 2 nil] [:field 2 nil]]
+                                  ;; Invalid field, the join is still valid
+                                  :fields [[:xfield 2 nil]]}]
+                         :source-table 4}}
+                 lib.convert/->pMBQL
+                 (get-in [:stages 0 :joins 0 :fields]))))))

test/metabase/lib/metric_test.cljc

@@ -16,7 +16,8 @@
     :fields   [(meta/field-metadata :venues :price)]
     :metrics  [{:id         100
                 :name       "My Metric"
-                :definition {:aggregation [[:sum [:field (meta/id :venues :price) nil]]]
+                :definition {:source-table (meta/id :venues)
+                             :aggregation [[:sum [:field (meta/id :venues :price) nil]]]
                              :filter      [:= [:field (meta/id :venues :price) nil] 4]}}]}))
 
 (deftest ^:parallel metric-display-name-test