[MLv2] Remove invalid query parts when summarizing (#30300)

* [MLv2] Remove invalid query parts when summarizing

* Address PR feedback

src/metabase/lib/aggregation.cljc

@@ -208,12 +208,7 @@
   ([query an-aggregate-clause]
    (aggregate query -1 an-aggregate-clause))
   ([query stage-number an-aggregate-clause]
-   (let [stage-number (or stage-number -1)]
-     (lib.util/update-query-stage
-       query stage-number
-       update :aggregation
-       (fn [aggregations]
-         (conj (vec aggregations) (lib.common/->op-arg query stage-number an-aggregate-clause)))))))
+   (lib.util/add-summary-clause query stage-number :aggregation an-aggregate-clause)))
 
 (mu/defn aggregations :- [:maybe [:sequential lib.metadata/ColumnMetadata]]
   "Get metadata about the aggregations in a given stage of a query."

src/metabase/lib/breakout.cljc

@@ -25,11 +25,7 @@
   ([query        :- ::lib.schema/query
     stage-number :- :int
     expr         :- [:or ::lib.schema.expression/expression fn?]]
-   (let [expr (if (fn? expr)
-                (expr query stage-number)
-                expr)]
-     (lib.util/update-query-stage query stage-number update :breakout (fn [breakouts]
-                                                                        (conj (vec breakouts) expr))))))
+   (lib.util/add-summary-clause query stage-number :breakout expr)))
 
 (mu/defn breakouts :- [:maybe [:sequential ::lib.schema.expression/expression]]
   "Return the current breakouts"

src/metabase/lib/join.cljc

 (ns metabase.lib.join
   (:require
    [medley.core :as m]
+   [metabase.lib.common :as lib.common]
    [metabase.lib.dispatch :as lib.dispatch]
    [metabase.lib.hierarchy :as lib.hierarchy]
    [metabase.lib.metadata :as lib.metadata]
@@ -8,7 +9,6 @@
    [metabase.lib.options :as lib.options]
    [metabase.lib.schema :as lib.schema]
    [metabase.lib.schema.common :as lib.schema.common]
-   [metabase.lib.schema.expression :as lib.schema.expression]
    [metabase.lib.schema.join :as lib.schema.join]
    [metabase.lib.util :as lib.util]
    [metabase.shared.util.i18n :as i18n]
@@ -204,41 +204,6 @@
                                            :stage-number stage-number
                                            :f            f})))))
 
-;; TODO this is basically the same as lib.common/->op-args,
-;; but requiring lib.common leads to crircular dependencies:
-;; join -> common -> field -> join.
-(defmulti ^:private ->join-condition
-  {:arglists '([query stage-number x])}
-  (fn [_query _stage-number x]
-    (lib.dispatch/dispatch-value x))
-  :hierarchy lib.hierarchy/hierarchy)
-
-(defmethod ->join-condition :default
-  [_query _stage-number x]
-  x)
-
-(defmethod ->join-condition :lib/external-op
-  [query stage-number {:keys [operator options args] :or {options {}}}]
-  (->join-condition query stage-number
-                    (lib.options/ensure-uuid (into [operator options] args))))
-
-(defmethod ->join-condition :dispatch-type/fn
-  [query stage-number f]
-  (->join-condition query stage-number (f query stage-number)))
-
-(mu/defn join-condition :- [:or
-                            fn?
-                            ::lib.schema.expression/boolean]
-  "Create a MBQL condition expression to include in the `:conditions` in a join map.
-
-  - One arity: return a function that will be resolved later once we have `query` and `stage-number.`
-  - Three arity: return the join condition expression immediately."
-  ([x]
-   (fn [query stage-number]
-     (join-condition query stage-number x)))
-  ([query stage-number x]
-   (->join-condition query stage-number x)))
-
 (defn join-clause
   "Create an MBQL join map from something that can conceptually be joined against. A `Table`? An MBQL or native query? A
   Saved Question? You should be able to join anything, and this should return a sensible MBQL join map."
@@ -255,7 +220,7 @@
 
   ([query stage-number x conditions]
    (cond-> (join-clause query stage-number x)
-     conditions (assoc :conditions (mapv #(join-condition query stage-number %) conditions)))))
+     conditions (assoc :conditions (mapv #(lib.common/->op-arg query stage-number %) conditions)))))
 
 (defmulti with-join-fields-method
   "Impl for [[with-join-fields]]."
@@ -267,7 +232,9 @@
 (defmethod with-join-fields-method :dispatch-type/fn
   [f fields]
   (fn [query stage-number]
-    (with-join-fields-method (f query stage-number) fields)))
+    (with-join-fields-method (f query stage-number) (if (keyword? fields)
+                                                      fields
+                                                      (mapv #(lib.common/->op-arg query stage-number %) fields)))))
 
 (defmethod with-join-fields-method :mbql/join
   [join fields]
@@ -276,7 +243,7 @@
 (mu/defn with-join-fields
   "Update a join (or a function that will return a join) to include `:fields`, either `:all`, `:none`, or a sequence of
   references."
-  [x fields :- ::lib.schema.join/fields]
+  [x fields]
   (with-join-fields-method x fields))
 
 (mu/defn join :- ::lib.schema/query

src/metabase/lib/util.cljc

 (ns metabase.lib.util
   (:refer-clojure :exclude [format])
   (:require
+   #?@(:clj
+       ([potemkin :as p]))
+   #?@(:cljs
+       (["crc-32" :as CRC32]
+        [goog.string :as gstring]
+        [goog.string.format :as gstring.format]))
    [clojure.set :as set]
    [clojure.string :as str]
    [medley.core :as m]
+   [metabase.lib.common :as lib.common]
    [metabase.lib.options :as lib.options]
    [metabase.lib.schema :as lib.schema]
    [metabase.lib.schema.common :as lib.schema.common]
@@ -11,13 +18,7 @@
    [metabase.mbql.util :as mbql.u]
    [metabase.shared.util.i18n :as i18n]
    [metabase.util :as u]
-   [metabase.util.malli :as mu]
-   #?@(:clj
-       ([potemkin :as p]))
-   #?@(:cljs
-       (["crc-32" :as CRC32]
-        [goog.string :as gstring]
-        [goog.string.format :as gstring.format]))))
+   [metabase.util.malli :as mu]))
 
 #?(:clj
    (set! *warn-on-reflection* true))
@@ -422,3 +423,31 @@
   (-> display-name
       (str/replace strip-id-regex "")
       str/trim))
+
+(mu/defn add-summary-clause :- ::lib.schema/query
+  "If the given stage has no summary, it will drop :fields, :order-by, and :join :fields from it,
+   as well as any subsequent stages."
+  [query :- ::lib.schema/query
+   stage-number :- :int
+   location :- [:enum :breakout :aggregation]
+   a-summary-clause]
+  (let [{:keys [stages] :as query} (pipeline query)
+        stage-number (or stage-number -1)
+        stage (query-stage query stage-number)
+        new-summary? (not (or (seq (:aggregation stage)) (seq (:breakout stage))))
+        new-query (update-query-stage
+                    query stage-number
+                    update location
+                    (fn [summary-clauses]
+                      (conj (vec summary-clauses) (lib.common/->op-arg query stage-number a-summary-clause))))]
+    (if new-summary?
+      (-> new-query
+          (update-query-stage
+            stage-number
+            (fn [stage]
+              (-> stage
+                  (dissoc :order-by :fields)
+                  (m/update-existing :joins (fn [joins] (mapv #(dissoc % :fields) joins))))))
+          ;; subvec holds onto references, so create a new vector
+          (update :stages (comp #(into [] %) subvec) 0 (inc (non-negative-stage-index stages stage-number))))
+      new-query)))

test/metabase/lib/aggregation_test.cljc

 (ns metabase.lib.aggregation-test
   (:require
+   #?@(:cljs ([metabase.test-runner.assert-exprs.approximately-equal]))
    [clojure.test :refer [are deftest is testing]]
    [metabase.lib.convert :as lib.convert]
    [metabase.lib.core :as lib]
@@ -9,7 +10,7 @@
    [metabase.lib.schema.expression :as lib.schema.expression]
    [metabase.lib.test-metadata :as meta]
    [metabase.lib.test-util :as lib.tu]
-   #?@(:cljs ([metabase.test-runner.assert-exprs.approximately-equal]))))
+   [metabase.lib.util :as lib.util]))
 
 #?(:cljs (comment metabase.test-runner.assert-exprs.approximately-equal/keep-me))
 
@@ -281,3 +282,32 @@
     (is (=? {:display_name   "Average of Price + 1"
              :effective_type :type/Float}
             (lib.metadata.calculation/display-info query ag-ref)))))
+
+(deftest ^:parallel aggregate-should-drop-invalid-parts
+  (let [query (-> (lib/query-for-table-name meta/metadata-provider "VENUES")
+                  (lib/with-fields [(lib/field "VENUES" "PRICE")])
+                  (lib/order-by (lib/field "VENUES" "PRICE"))
+                  (lib/join (-> (lib/join-clause (meta/table-metadata :categories)
+                                                 [(lib/=
+                                                    (lib/field "VENUES" "CATEGORY_ID")
+                                                    (lib/with-join-alias (lib/field "CATEGORIES" "ID") "Cat"))])
+                                (lib/with-join-fields [(lib/field "CATEGORIES" "ID")])))
+                  (lib/append-stage)
+                  (lib/with-fields [(lib/field "VENUES" "PRICE")])
+                  (lib/aggregate 0 (lib/sum (lib/field "VENUES" "CATEGORY_ID"))))
+        first-stage (lib.util/query-stage query 0)
+        first-join (first (lib/joins query 0))]
+    (is (= 1 (count (:stages query))))
+    (is (not (contains? first-stage :fields)))
+    (is (not (contains? first-stage :order-by)))
+    (is (= 1 (count (lib/joins query 0))))
+    (is (not (contains? first-join :fields))))
+  (testing "Already summarized query should be left alone"
+    (let [query (-> (lib/query-for-table-name meta/metadata-provider "VENUES")
+                    (lib/breakout (lib/field "VENUES" "CATEGORY_ID"))
+                    (lib/order-by (lib/field "VENUES" "CATEGORY_ID"))
+                    (lib/append-stage)
+                    (lib/aggregate 0 (lib/sum (lib/field "VENUES" "CATEGORY_ID"))))
+          first-stage (lib.util/query-stage query 0)]
+      (is (= 2 (count (:stages query))))
+      (is (contains? first-stage :order-by)))))

test/metabase/lib/breakout_test.cljc

 (ns metabase.lib.breakout-test
   (:require
-   [clojure.test :refer [deftest is]]
+   #?@(: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]
-   #?@(:cljs ([metabase.test-runner.assert-exprs.approximately-equal]))))
+   [metabase.lib.util :as lib.util]))
 
 #?(:cljs (comment metabase.test-runner.assert-exprs.approximately-equal/keep-me))
 
@@ -30,3 +31,32 @@
                   (lib/breakout (lib/field (meta/id :checkins :date))))]
     (is (=? [[:field {} (meta/id :checkins :date)]]
             (lib/breakouts query)))))
+
+(deftest ^:parallel breakout-should-drop-invalid-parts
+  (let [query (-> (lib/query-for-table-name meta/metadata-provider "VENUES")
+                  (lib/with-fields [(lib/field "VENUES" "PRICE")])
+                  (lib/order-by (lib/field "VENUES" "PRICE"))
+                  (lib/join (-> (lib/join-clause (meta/table-metadata :categories)
+                                                 [(lib/=
+                                                    (lib/field "VENUES" "CATEGORY_ID")
+                                                    (lib/with-join-alias (lib/field "CATEGORIES" "ID") "Cat"))])
+                                (lib/with-join-fields [(lib/field "CATEGORIES" "ID")])))
+                  (lib/append-stage)
+                  (lib/with-fields [(lib/field "VENUES" "PRICE")])
+                  (lib/breakout 0 (lib/field "VENUES" "CATEGORY_ID")))
+        first-stage (lib.util/query-stage query 0)
+        first-join (first (lib/joins query 0))]
+    (is (= 1 (count (:stages query))))
+    (is (not (contains? first-stage :fields)))
+    (is (not (contains? first-stage :order-by)))
+    (is (= 1 (count (lib/joins query 0))))
+    (is (not (contains? first-join :fields))))
+  (testing "Already summarized query should be left alone"
+    (let [query (-> (lib/query-for-table-name meta/metadata-provider "VENUES")
+                    (lib/breakout (lib/field "VENUES" "CATEGORY_ID"))
+                    (lib/order-by (lib/field "VENUES" "CATEGORY_ID"))
+                    (lib/append-stage)
+                    (lib/breakout 0 (lib/field "VENUES" "PRICE")))
+          first-stage (lib.util/query-stage query 0)]
+      (is (= 2 (count (:stages query))))
+      (is (contains? first-stage :order-by)))))