Remove dependent joins from subsequent stages too (#34416)

src/metabase/lib/remove_replace.cljc

@@ -338,9 +338,17 @@
   ([query        :- :metabase.lib.schema/query
     stage-number :- :int
     join-spec    :- [:or :metabase.lib.schema.join/join :string :int]]
-   (update-joins query stage-number join-spec (fn [joins join-alias]
-                                                (not-empty (filterv #(not (dependent-join? % join-alias))
-                                                                    joins))))))
+   (try
+     (update-joins query stage-number join-spec (fn [joins join-alias]
+                                                  (not-empty (filterv #(not (dependent-join? % join-alias))
+                                                                      joins))))
+     (catch #?(:clj Exception :cljs :default) e
+       (let [{:keys [error stage join]} (ex-data e)]
+         (if (= error ::lib.util/cannot-remove-final-join-condition)
+           (-> query
+               (remove-join (u/index-of #{stage} (:stages query)) join)
+               (remove-join stage-number join-spec))
+           (throw e)))))))
 
 (mu/defn replace-join :- :metabase.lib.schema/query
   "Replace the join specified by `join-spec` in `query` at `stage-number` with `new-join`.

src/metabase/lib/util.cljc

@@ -112,7 +112,10 @@
 
         (= [:joins :conditions] [first-loc last-loc])
         (throw (ex-info (i18n/tru "Cannot remove the final join condition")
-                        {:conditions (get-in stage location)}))
+                        {:error ::cannot-remove-final-join-condition
+                         :conditions (get-in stage location)
+                         :join (get-in stage (pop location))
+                         :stage stage}))
 
         (= [:joins :fields] [first-loc last-loc])
         (update-in stage (pop location) dissoc last-loc)

test/metabase/lib/remove_replace_test.cljc

@@ -768,7 +768,7 @@
             (is (=? result
                     (lib/remove-clause query' 0 (second (lib/joins query' 0)))))))))))
 
-(deftest ^:parallel remove-dependent-join
+(deftest ^:parallel remove-dependent-join-test
   (let [first-join-clause (-> (lib/join-clause
                                (meta/table-metadata :checkins)
                                [(lib/= (meta/field-metadata :venues :id)
@@ -788,6 +788,32 @@
                                 (lib/with-join-alias "Users"))))]
     (is (nil? (lib/joins (lib/remove-clause query 0 first-join-clause))))))
 
+(deftest ^:parallel remove-dependent-join-from-subsequent-stage-test
+  (testing "Join from previous stage used in join can be removed (#34404)"
+    (let [join-query (-> (lib/query meta/metadata-provider (meta/table-metadata :venues))
+                         (lib/join (-> (lib/join-clause
+                                        (meta/table-metadata :checkins)
+                                        [(lib/= (meta/field-metadata :venues :id)
+                                                (-> (meta/field-metadata :checkins :venue-id)
+                                                    (lib/with-join-alias "Checkins")))])
+                                       (lib/with-join-fields :all)
+                                       (lib/with-join-alias "Checkins"))))
+          breakout-col (m/find-first (comp #{(meta/id :checkins :venue-id)} :id)
+                                     (lib/breakoutable-columns join-query))
+          breakout-query (-> join-query
+                             (lib/breakout breakout-col)
+                             (lib/aggregate (lib/count))
+                             (lib/append-stage))
+          query (-> breakout-query
+                    (lib/join (-> (lib/join-clause
+                                   (meta/table-metadata :checkins)
+                                   [(lib/= (first (lib/returned-columns breakout-query))
+                                           (-> (meta/field-metadata :checkins :venue-id)
+                                               (lib/with-join-alias "Checkins2")))])
+                                  (lib/with-join-fields :all)
+                                  (lib/with-join-alias "Checkins2"))))]
+      (is (nil? (lib/joins (lib/remove-clause query 0 (first (lib/joins query 0)))))))))
+
 (deftest ^:parallel replace-join-test
   (let [query             lib.tu/query-with-join
         expected-original {:stages [{:joins [{:lib/type :mbql/join, :alias "Cat", :fields :all}]}]}