[MLv2] remove replace join conditions and fields (#30373)

* [MLv2] Remove replace join conditions and fields

* Add tests for join fields

* Fix match for other parts

* Address PR feedback

* Fix linter

src/metabase/lib/core.cljc

@@ -145,6 +145,8 @@
   [lib.join
    join
    join-clause
+   join-conditions
+   join-fields
    joins
    with-join-alias
    with-join-fields]

src/metabase/lib/join.cljc

@@ -284,3 +284,13 @@
   [table-name           :- ::lib.schema.common/non-blank-string
    source-field-id-name :- ::lib.schema.common/non-blank-string]
   (lib.util/format "%s__via__%s" table-name source-field-id-name))
+
+(mu/defn join-conditions :- ::lib.schema.join/conditions
+  "Get all join conditions for the given join"
+  [j :- ::lib.schema.join/join]
+  (:conditions j))
+
+(mu/defn join-fields :- [:maybe ::lib.schema/fields]
+  "Get all join conditions for the given join"
+  [j :- ::lib.schema.join/join]
+  (:fields j))

src/metabase/lib/remove_replace.cljc

 (ns metabase.lib.remove-replace
   (:require
    [metabase.lib.common :as lib.common]
+   [metabase.lib.join :as lib.join]
    [metabase.lib.metadata.calculation :as lib.metadata.calculation]
    [metabase.lib.ref :as lib.ref]
    [metabase.lib.util :as lib.util]
    [metabase.mbql.util.match :as mbql.match]
    [metabase.util.malli :as mu]))
 
-(defn- find-clause
+(defn- find-reference-clause
   [query stage-number target-clause]
   (let [stage (lib.util/query-stage query stage-number)
         [target-type _opts target-id] target-clause]
@@ -21,11 +22,11 @@
 (defn- target-ref-for-stage
   "Gets the ref for the target-id exposed by the previous stage"
   [query stage-number target-id]
-  (->> (let [stage (lib.util/query-stage query stage-number)]
-         (lib.metadata.calculation/visible-columns query stage-number stage))
-       (some (fn [{:keys [lib/source id] :as column}]
-               (when (and (= :source/previous-stage source) (= target-id id))
-                 (lib.ref/ref column))))))
+  (let [stage (lib.util/query-stage query stage-number)]
+    (->> (lib.metadata.calculation/visible-columns query stage-number stage)
+         (some (fn [{:keys [lib/source id] :as column}]
+                 (when (and (= :source/previous-stage source) (= target-id id))
+                   (lib.ref/ref column)))))))
 
 (defn- check-subsequent-stages-for-invalid-target!
   "Throws if target-clause is used in a subsequent stage"
@@ -37,28 +38,36 @@
         (when-not (target-ref-for-stage query next-stage-number target-id)
           ;; Get the ref to look for from the previous-query
           (let [target-ref (target-ref-for-stage previous-query next-stage-number target-id)]
-            (if-let [found (find-clause query next-stage-number target-ref)]
+            (if-let [found (find-reference-clause query next-stage-number target-ref)]
               (throw (ex-info "Clause cannot be removed as it has dependents" {:target-clause target-clause
                                                                                :stage-number next-stage-number
                                                                                :found found}))
               (recur next-stage-number))))))))
 
 (defn- remove-replace* [query stage-number target-clause remove-replace-fn]
-  (reduce
-     (fn [query location]
-       (let [target-clause (lib.common/->op-arg query stage-number target-clause)
-             result (lib.util/update-query-stage query stage-number
-                                                 remove-replace-fn location target-clause)]
-         (when (not= query result)
-           (case location
-             :breakout (check-subsequent-stages-for-invalid-target! query result stage-number (lib.ref/ref target-clause))
-             :fields (check-subsequent-stages-for-invalid-target! query result stage-number (lib.ref/ref target-clause))
-             nil)
-           (reduced result))
-         result))
-     query
-     ;; TODO only these top level clauses are supported at this moment
-     [:order-by :breakout :filters :fields]))
+  (let [join-indices (range (count (lib.join/joins query stage-number)))
+        join-condition-paths (for [idx join-indices]
+                               [:joins idx :conditions])
+        join-field-paths (for [idx join-indices]
+                           [:joins idx :fields])]
+    (reduce
+      (fn [query location]
+        (let [target-clause (lib.common/->op-arg query stage-number target-clause)
+              result (lib.util/update-query-stage query stage-number
+                                                  remove-replace-fn location target-clause)]
+          (when (not= query result)
+            (mbql.match/match location
+              (:or
+                [:breakout]
+                [:fields]
+                [:joins _ :fields]) (check-subsequent-stages-for-invalid-target! query result stage-number (lib.ref/ref target-clause)))
+            (reduced result))
+          result))
+      query
+      ;; TODO only these top level clauses are supported at this moment
+      (concat [[:order-by] [:breakout] [:filters] [:fields]]
+              join-field-paths
+              join-condition-paths))))
 
 (mu/defn remove-clause :- :metabase.lib.schema/query
   "Removes the `target-clause` in the filter of the `query`."
@@ -81,4 +90,4 @@
     target-clause
     new-clause]
    (let [replacement (lib.common/->op-arg query stage-number new-clause)]
-     (remove-replace* query stage-number target-clause #(lib.util/replace-clause %1 %2 %3  replacement)))))
+     (remove-replace* query stage-number target-clause #(lib.util/replace-clause %1 %2 %3 replacement)))))

src/metabase/lib/schema/join.cljc

@@ -37,12 +37,15 @@
    {:gen/fmap #(str % "-" (random-uuid))}
    ::common/non-blank-string])
 
+(mr/def ::conditions
+  [:sequential {:min 1} [:ref ::expression/boolean]])
+
 (mr/def ::join
   [:map
    [:lib/type    [:= :mbql/join]]
    [:lib/options ::common/options]
    [:stages      [:ref :metabase.lib.schema/stages]]
-   [:conditions  [:sequential {:min 1} [:ref ::expression/boolean]]]
+   [:conditions  ::conditions]
    [:fields {:optional true} ::fields]
    [:alias  {:optional true} ::alias]])
 

src/metabase/lib/util.cljc

@@ -53,7 +53,7 @@
    If `location` contains no clause with `target-clause` no replacement happens."
   [stage location target-clause new-clause]
   {:pre [(clause? target-clause)]}
-  (m/update-existing
+  (m/update-existing-in
     stage
     location
     #(->> (for [clause %]
@@ -69,13 +69,18 @@
    If the the location is empty, dissoc it from stage."
   [stage location target-clause]
   {:pre [(clause? target-clause)]}
-  (if-let [target (get stage location)]
+  (if-let [target (get-in stage location)]
     (let [target-uuid (clause-uuid target-clause)
           result (into [] (remove (comp #{target-uuid} clause-uuid)) target)]
       (if (seq result)
-        (assoc stage location result)
-        (dissoc stage location)))
-    stage))
+        (assoc-in stage location result)
+        (if (= 1 (count location))
+          (dissoc stage (first location))
+          (case [(first location) (last location)]
+            [:joins :conditions] (throw (ex-info (i18n/tru "Cannot remove the final join condition")
+                                                 {:conditions (get-in stage location)}))
+            [:joins :fields] (update-in stage (pop location) dissoc (peek location))))))
+            stage))
 
 ;;; TODO -- all of this `->pipeline` stuff should probably be merged into [[metabase.lib.convert]] at some point in
 ;;; the near future.

test/metabase/lib/remove_replace_test.cljc

@@ -31,11 +31,30 @@
                  (lib/remove-clause (first filters))
                  (lib/filters)
                  count)))
-    (is (= 0 (-> query
-                 (lib/remove-clause (first filters))
-                 (lib/remove-clause (second filters))
-                 (lib/filters)
-                 count)))))
+    (is (nil? (-> query
+                  (lib/remove-clause (first filters))
+                  (lib/remove-clause (second filters))
+                  (lib/filters))))))
+
+(deftest ^:parallel remove-clause-join-conditions-test
+  (let [query (-> (lib/query-for-table-name meta/metadata-provider "VENUES")
+                  (lib/join (lib/query-for-table-name meta/metadata-provider "CATEGORIES")
+                            [(lib/= (lib/field "VENUES" "PRICE") 4)
+                             (lib/= (lib/field "VENUES" "NAME") "x")]))
+        conditions (lib/join-conditions (first (lib/joins query)))]
+    (is (= 2 (count conditions)))
+    (is (= [(second conditions)]
+           (-> query
+               (lib/remove-clause (first conditions))
+               lib/joins
+               first
+               lib/join-conditions)))
+    (is (thrown-with-msg?
+          #?(:clj Exception :cljs js/Error)
+          #"Cannot remove the final join condition"
+          (-> query
+              (lib/remove-clause (first conditions))
+              (lib/remove-clause (second conditions)))))))
 
 (deftest ^:parallel remove-clause-breakout-test
   (let [query (-> (lib/query-for-table-name meta/metadata-provider "VENUES")
@@ -94,7 +113,7 @@
                  (lib/remove-clause (second fields))
                  (lib/fields)
                  count)))
-    (testing "removing breakout with dependent should not be allowed"
+    (testing "removing field with dependent should not be allowed"
       (is (thrown-with-msg?
             #?(:clj Exception :cljs js/Error)
             #"Clause cannot be removed as it has dependents"
@@ -122,6 +141,46 @@
                    (lib/fields 0)
                    count))))))
 
+(deftest ^:parallel remove-clause-join-fields-test
+  (let [query (-> (lib/query-for-table-name meta/metadata-provider "CATEGORIES")
+                  (lib/join (-> (lib/join-clause (lib/query-for-table-name meta/metadata-provider "VENUES")
+                                                 [(lib/= (lib/field "VENUES" "PRICE") 4)])
+                                (lib/with-join-fields [(lib/field "VENUES" "PRICE")
+                                                       (lib/field "VENUES" "ID")]))))
+        fields (lib/join-fields (first (lib/joins query)))]
+    (is (= 2 (count fields)))
+    (is (= [(second fields)]
+           (-> query
+               (lib/remove-clause (first fields))
+               lib/joins
+               first
+               lib/join-fields)))
+    (is (nil? (-> query
+                  (lib/remove-clause (first fields))
+                  (lib/remove-clause (second fields))
+                  lib/joins
+                  first
+                  lib/join-fields)))
+    (testing "removing field with dependent should not be allowed"
+      (is (thrown-with-msg?
+            #?(:clj Exception :cljs js/Error)
+            #"Clause cannot be removed as it has dependents"
+            (-> query
+                (lib/append-stage)
+                ;; TODO Should be able to create a ref with lib/field [#29763]
+                (lib/filter (lib/= [:field {:lib/uuid (str (random-uuid)) :base-type :type/Integer} "PRICE"] 1))
+                (lib/remove-clause 0 (first fields)))))
+      (is (thrown-with-msg?
+            #?(:clj Exception :cljs js/Error)
+            #"Clause cannot be removed as it has dependents"
+            (-> query
+                (lib/append-stage)
+                (lib/with-fields [[:field {:lib/uuid (str (random-uuid)) :base-type :type/Integer} "PRICE"]])
+                (lib/append-stage)
+                ;; TODO Should be able to create a ref with lib/field [#29763]
+                (lib/filter (lib/= [:field {:lib/uuid (str (random-uuid)) :base-type :type/Integer} "PRICE"] 1))
+                (lib/remove-clause 0 (first fields))))))))
+
 (deftest ^:parallel replace-clause-order-by-test
   (let [query (-> (lib/query-for-table-name meta/metadata-provider "VENUES")
                   (lib/filter (lib/= "myvenue" (lib/field (meta/id :venues :name))))
@@ -153,6 +212,39 @@
       (is (= 2 (count replaced-filters)))
       (is (= (second filters) (second replaced-filters))))))
 
+(deftest ^:parallel replace-clause-join-conditions-test
+  (let [query (-> (lib/query-for-table-name meta/metadata-provider "VENUES")
+                  (lib/join (lib/query-for-table-name meta/metadata-provider "CATEGORIES")
+                            [(lib/= (lib/field "VENUES" "PRICE") 4)]))
+        conditions (lib/join-conditions (first (lib/joins query)))]
+    (is (= 1 (count conditions)))
+    (let [replaced (-> query
+                       (lib/replace-clause (first conditions) (lib/= (lib/field (meta/id :venues :id)) 1)))
+          replaced-conditions (lib/join-conditions (first (lib/joins replaced)))]
+      (is (not= conditions replaced-conditions))
+      (is (=? [:= {} [:field {} (meta/id :venues :id)] 1]
+              (first replaced-conditions)))
+      (is (= 1 (count replaced-conditions)))
+      (is (= (second conditions) (second replaced-conditions))))))
+
+(deftest ^:parallel replace-clause-join-fields-test
+  (let [query (-> (lib/query-for-table-name meta/metadata-provider "VENUES")
+                  (lib/join
+                    (-> (lib/join-clause (lib/query-for-table-name meta/metadata-provider "CATEGORIES")
+                                         [(lib/= (lib/field "VENUES" "PRICE") 4)])
+                        (lib/with-join-fields
+                          [(lib/field "CATEGORIES" "ID")]))))
+        fields (lib/join-fields (first (lib/joins query)))]
+    (is (= 1 (count fields)))
+    (let [replaced (-> query
+                       (lib/replace-clause (first fields) (lib/field "CATEGORIES" "NAME")))
+          replaced-fields (lib/join-fields (first (lib/joins replaced)))]
+      (is (not= fields replaced-fields))
+      (is (=? [:field {} (meta/id :categories :name)]
+              (first replaced-fields)))
+      (is (= 1 (count fields)))
+      (is (= 1 (count replaced-fields))))))
+
 (deftest ^:parallel replace-clause-breakout-by-test
   (let [query (-> (lib/query-for-table-name meta/metadata-provider "VENUES")
                   (lib/breakout (lib/field (meta/id :venues :id)))