[MBQL lib] Put `:fields :all` on joins on deleting last aggregation/breakout (#44321)

Part of #44298.

src/metabase/lib/convert.cljc

@@ -211,7 +211,7 @@
                  (update :stages ->pMBQL))]
     (cond-> join
       (:fields join) (update :fields (fn [fields]
-                                       (if (seqable? fields)
+                                       (if (sequential? fields)
                                          (mapv ->pMBQL fields)
                                          (keyword fields))))
       (not (:alias join)) (assoc :alias legacy-default-join-alias))))

src/metabase/lib/remove_replace.cljc

@@ -197,7 +197,7 @@
       (if location
         (-> query
             (remove-replace-location stage-number query location target-clause remove-replace-fn)
-            normalize-fields-clauses)
+            (normalize-fields-clauses location))
         query))))
 
 (mu/defn remove-clause :- :metabase.lib.schema/query
@@ -563,25 +563,36 @@
          (lib.equality/matching-column-sets? query stage-number fields
                                              (lib.metadata.calculation/default-columns-for-stage query stage-number)))))
 
-(defn- normalize-fields-for-join [query stage-number join]
-  (if (#{:none :all} (:fields join))
+(defn- normalize-fields-for-join [query stage-number removed-location join]
+  (cond
     ;; Nothing to do if it's already a keyword.
-    join
-    (cond-> join
-      (lib.equality/matching-column-sets?
-        query stage-number (:fields join)
-        (lib.metadata.calculation/returned-columns query stage-number (assoc join :fields :all)))
-      (assoc :fields :all))))
-
-(defn- normalize-fields-for-stage [query stage-number]
+    (#{:none :all} (:fields join)) join
+
+    ;; If it's missing, treat it as `:all` unless we just removed a field.
+    ;; TODO: This really should be a different function called by `remove-field`; it also needs to filter on the stage.
+    (and (or (= removed-location [:aggregation])
+             (= removed-location [:breakout]))
+         (not (contains? join :fields)))
+    (assoc join :fields :all)
+
+    (lib.equality/matching-column-sets?
+      query stage-number (:fields join)
+      (lib.metadata.calculation/returned-columns query stage-number (assoc join :fields :all)))
+    (assoc join :fields :all)
+
+    :else join))
+
+(defn- normalize-fields-for-stage [query stage-number removed-location]
   (let [stage (lib.util/query-stage query stage-number)]
     (cond-> query
       (specifies-default-fields? query stage-number)
       (lib.util/update-query-stage stage-number dissoc :fields)
 
-      (:joins stage)
+      (and (empty? (:aggregation stage))
+           (empty? (:breakout stage))
+           (:joins stage))
       (lib.util/update-query-stage stage-number update :joins
-                                   (partial mapv #(normalize-fields-for-join query stage-number %))))))
+                                   (partial mapv #(normalize-fields-for-join query stage-number removed-location %))))))
 
 (mu/defn normalize-fields-clauses :- :metabase.lib.schema/query
   "Check all the `:fields` clauses in the query - on the stages and any joins - and drops them if they are equal to the
@@ -589,5 +600,10 @@
   - For stages, if the `:fields` list is identical to the default fields for this stage.
   - For joins, replace it with `:all` if it's all the fields that are in the join by default.
   - For joins, remove it if the list is empty (the default for joins is no fields)."
-  [query :- :metabase.lib.schema/query]
-  (reduce normalize-fields-for-stage query (range (count (:stages query)))))
+  ([query :- :metabase.lib.schema/query]
+   (normalize-fields-clauses query nil))
+  ([query            :- :metabase.lib.schema/query
+    removed-location :- [:maybe [:sequential :any]]]
+   (reduce #(normalize-fields-for-stage %1 %2 removed-location)
+           query
+           (range (count (:stages query))))))

test/metabase/lib/remove_replace_test.cljc

@@ -1042,6 +1042,53 @@
       (is (empty? (lib/aggregations result)))
       (is (= (lib/breakouts query) (lib/breakouts result))))))
 
+(deftest ^:parallel removing-last-aggregation-brings-back-all-fields-on-joins
+  (testing "Removing the last aggregation puts :fields :all on join clauses"
+    (let [base   (-> lib.tu/venues-query
+                     (lib/join (lib/join-clause (meta/table-metadata :products)
+                                                [(lib/= (meta/field-metadata :orders :product-id)
+                                                        (meta/field-metadata :products :id))])))
+          query  (lib/aggregate base (lib/count))
+          result (lib/remove-clause query (first (lib/aggregations query)))]
+      (is (= :all (-> base :stages first :joins first :fields)))
+      (is (= :all (-> result :stages first :joins first :fields)))
+      (is (=? (map :name (lib/returned-columns base))
+              (map :name (lib/returned-columns result)))))))
+
+(deftest ^:parallel removing-last-breakout-brings-back-all-fields-on-joins
+  (testing "Removing the last breakout puts :fields :all on join clauses"
+    (let [base   (-> lib.tu/venues-query
+                     (lib/join (lib/join-clause (meta/table-metadata :products)
+                                                [(lib/= (meta/field-metadata :orders :product-id)
+                                                        (meta/field-metadata :products :id))])))
+          query  (-> base
+                     (lib/aggregate (lib/count))
+                     (lib/breakout (meta/field-metadata :products :category)))
+          agg    (first (lib/aggregations query))
+          brk    (first (lib/breakouts query))]
+      (is (= :all (-> base :stages first :joins first :fields)))
+
+      (testing "no change to join"
+        (testing "when removing just the aggregation"
+          (is (= (m/dissoc-in query [:stages 0 :aggregation])
+                 (lib/remove-clause query agg))))
+        (testing "when removing just the breakout"
+          (is (= (m/dissoc-in query [:stages 0 :breakout])
+                 (lib/remove-clause query brk)))))
+      (testing "join gets :fields :all"
+        (testing "removing aggregation and then breakout"
+          (is (= :all
+                 (-> query
+                     (lib/remove-clause agg)
+                     (lib/remove-clause brk)
+                     :stages first :joins first :fields))))
+        (testing "removing breakout and then aggregation"
+          (is (= :all
+                 (-> query
+                     (lib/remove-clause brk)
+                     (lib/remove-clause agg)
+                     :stages first :joins first :fields))))))))
+
 (deftest ^:parallel simple-tweak-expression-test
   (let [table (lib/query meta/metadata-provider (meta/table-metadata :orders))
         base (lib/expression table "Tax Rate" (lib// (meta/field-metadata :orders :tax)