[MLv2] Properly match aggregations across versions of a query (#37867)

Because aggregation refs are based on `:source-uuid` they are not
comparable between queries, which breaks some logic to find columns
added and removed between two versions of a query.

This adds a fallback based on the newly added `:source-name` option on
aggregation refs.

Fixes #37851.

src/metabase/lib/aggregation.cljc

@@ -23,8 +23,9 @@
 (mu/defn column-metadata->aggregation-ref :- :mbql.clause/aggregation
   "Given `:metadata/column` column metadata for an aggregation, construct an `:aggregation` reference."
   [metadata :- lib.metadata/ColumnMetadata]
-  (let [options {:lib/uuid       (str (random-uuid))
-                 :effective-type ((some-fn :effective-type :base-type) metadata)}
+  (let [options {:lib/uuid        (str (random-uuid))
+                 :effective-type  ((some-fn :effective-type :base-type) metadata)
+                 :lib/source-name (:name metadata)}
         ag-uuid (:lib/source-uuid metadata)]
     (assert ag-uuid "Metadata for an aggregation reference should include :lib/source-uuid")
     [:aggregation options ag-uuid]))

src/metabase/lib/equality.cljc

@@ -286,27 +286,35 @@
     a-ref-or-column :- [:or ::lib.schema.metadata/column ::lib.schema.ref/ref]
     columns         :- [:sequential ::lib.schema.metadata/column]
     opts            :- FindMatchingColumnOptions]
-   (let [[_ref-kind _opts ref-id :as a-ref] (if (lib.util/clause? a-ref-or-column)
-                                              a-ref-or-column
-                                              (lib.ref/ref a-ref-or-column))]
+   (let [[ref-kind ref-opts ref-id :as a-ref] (if (lib.util/clause? a-ref-or-column)
+                                                a-ref-or-column
+                                                (lib.ref/ref a-ref-or-column))]
      (or (find-matching-column a-ref columns opts)
+         ;; Aggregations are matched by `:source-uuid` but if we're comparing old columns to new refs or vice versa
+         ;; the random UUIDs won't match up. This falls back to the `:lib/source-name` option on aggregation refs, if
+         ;; present.
+         (when (and (= ref-kind :aggregation)
+                    (:lib/source-name ref-opts))
+           (m/find-first #(and (= (:lib/source %) :source/aggregations)
+                               (= (:name %) (:lib/source-name ref-opts)))
+                         columns))
          ;; We failed to match by ID, so try again with the column's name. Any columns with `:id` set are dropped.
          ;; Why? Suppose there are two CREATED_AT columns in play - if one has an :id and it failed to match above, then
          ;; it certainly shouldn't match by name just because of the coincidence of column names!
-       (when (and query (number? ref-id))
-         (when-let [no-id-columns (not-empty (remove :id columns))]
-           (when-let [resolved (if (lib.util/clause? a-ref-or-column)
-                                 (resolve-field-id query stage-number ref-id)
-                                 a-ref-or-column)]
-             (find-matching-column (-> (assoc a-ref 2 (or (:lib/desired-column-alias resolved)
-                                                          (:name resolved)))
-                                       ;; make sure the :field ref has a `:base-type`, it's against the rules for a
-                                       ;; nominal :field ref not to have a base-type -- this can fail schema
-                                       ;; validation if it's missing in the Field ID ref we generate the nominal ref
-                                       ;; from.
-                                       (lib.options/update-options (partial merge {:base-type :type/*})))
-                                   no-id-columns
-                                   opts))))))))
+         (when (and query (number? ref-id))
+           (when-let [no-id-columns (not-empty (remove :id columns))]
+             (when-let [resolved (if (lib.util/clause? a-ref-or-column)
+                                   (resolve-field-id query stage-number ref-id)
+                                   a-ref-or-column)]
+               (find-matching-column (-> (assoc a-ref 2 (or (:lib/desired-column-alias resolved)
+                                                            (:name resolved)))
+                                         ;; make sure the :field ref has a `:base-type`, it's against the rules for a
+                                         ;; nominal :field ref not to have a base-type -- this can fail schema
+                                         ;; validation if it's missing in the Field ID ref we generate the nominal ref
+                                         ;; from.
+                                         (lib.options/update-options (partial merge {:base-type :type/*})))
+                                     no-id-columns
+                                     opts))))))))
 
 (defn- ref-id-or-name [[_ref-kind _opts id-or-name]]
   id-or-name)

src/metabase/lib/schema/ref.cljc

@@ -79,7 +79,8 @@
    ::common/options
    [:map
     [:name {:optional true} ::common/non-blank-string]
-    [:display-name {:optional true} ::common/non-blank-string]]])
+    [:display-name {:optional true} ::common/non-blank-string]
+    [:lib/source-name {:optional true} ::common/non-blank-string]]])
 
 (mbql-clause/define-mbql-clause :aggregation
   [:tuple

test/metabase/lib/equality_test.cljc

@@ -459,10 +459,33 @@
   (let [query (-> lib.tu/venues-query
                   (lib/aggregate (lib/count)))
         [ag]  (lib/aggregations query)]
-    (is (=? {:display-name "Count", :lib/source :source/aggregations}
-            (lib.equality/find-matching-column
-             [:aggregation {:lib/uuid (str (random-uuid))} (lib.options/uuid ag)]
-             (lib/returned-columns query))))))
+    (testing "without passing query"
+      (testing "matches with UUID"
+        (is (=? {:display-name "Count", :lib/source :source/aggregations}
+                (lib.equality/find-matching-column
+                  [:aggregation {:lib/uuid (str (random-uuid))} (lib.options/uuid ag)]
+                  (lib/returned-columns query)))))
+      (testing "fails with bad UUID but good source-name"
+        (is (nil? (lib.equality/find-matching-column
+                    [:aggregation {:lib/uuid        (str (random-uuid))
+                                   :lib/source-name "count"}
+                     "this is a bad UUID"]
+                    (lib/returned-columns query))))))
+    (testing "when passing query"
+      (testing "matches with UUID"
+        (is (=? {:display-name "Count", :lib/source :source/aggregations}
+                (lib.equality/find-matching-column
+                  query -1
+                  [:aggregation {:lib/uuid (str (random-uuid))} (lib.options/uuid ag)]
+                  (lib/returned-columns query)))))
+      (testing "matches with bad UUID but good source-name"
+        (is (=? {:display-name "Count", :lib/source :source/aggregations}
+                (lib.equality/find-matching-column
+                  query -1
+                  [:aggregation {:lib/uuid        (str (random-uuid))
+                                 :lib/source-name "count"}
+                   "this is a bad UUID"]
+                  (lib/returned-columns query))))))))
 
 (deftest ^:parallel find-matching-column-expression-test
   (is (=? {:name "expr", :lib/source :source/expressions}