Add alias info util namespace (#19610)

shared/src/metabase/mbql/schema.cljc

@@ -258,12 +258,15 @@
 
 ;; Expression *references* refer to a something in the `:expressions` clause, e.g. something like
 ;;
-;;    [:+ [:field 1 nil] [:field 2 nil]]`
+;;    [:+ [:field 1 nil] [:field 2 nil]]
+;;
+;; As of 0.42.0 `:expression` references can have an optional options map
 (defclause ^{:requires-features #{:expressions}} expression
-  expression-name helpers/NonBlankString)
+  expression-name helpers/NonBlankString
+  options         (optional (s/pred map? "map")))
 
 (def BinningStrategyName
-  "Schema for a valid value for the `strategy-name` param of a `binning-strategy` clause."
+  "Schema for a valid value for the `strategy-name` param of a [[field]] clause with `:binning` information."
   (s/enum :num-bins :bin-width :default))
 
 (defn- validate-bin-width [schema]
@@ -389,7 +392,7 @@
 (def ^:private Field*
   (one-of expression field))
 
-;; TODO -- consider renaming this FieldOrExpression,
+;; TODO -- consider renaming this FieldOrExpression
 (def Field
   "Schema for either a `:field` clause (reference to a Field) or an `:expression` clause (reference to an expression)."
   (s/recursive #'Field*))
@@ -408,7 +411,11 @@
 ;;
 ;; TODO - it would be nice if we could check that there's actually an aggregation with the corresponding index,
 ;; wouldn't it
-(defclause aggregation, aggregation-clause-index s/Int)
+;;
+;; As of 0.42.0 `:aggregation` references can have an optional options map.
+(defclause aggregation
+  aggregation-clause-index s/Int
+  options                  (optional (s/pred map? "map")))
 
 (def FieldOrAggregationReference
   "Schema for any type of valid Field clause, or for an indexed reference to an aggregation clause."

shared/src/metabase/mbql/util.cljc

@@ -440,7 +440,14 @@
        (unique-name \"A\")])
     ;; -> [\"A\" \"B\" \"A_2\"]
 
-  If idempotence is desired, the function returned by the generator also has a 2 airity version where the first argument is the object for which we are generating the name.
+  By default, unique aliases are generated for each unique `[id original-name]` key pair. By default, a unique `id` is
+  generated for every call, meaning repeated calls to [[unique-name-generator]] with the same `original-name` will
+  return different unique aliases. If idempotence is desired, the function returned by the generator also has a 2
+  airity version with the signature
+
+    (unique-name-fn id original-name)
+
+  for example:
 
     (let [unique-name (unique-name-generator)]
       [(unique-name :x \"A\")
@@ -448,21 +455,76 @@
        (unique-name :x \"A\")
        (unique-name :y \"A\")])
     ;; -> [\"A\" \"B\" \"A\" \"A_2\"]
-  "
-  []
-  (let [identity-objects->aliases (atom {})
-        aliases                   (atom {})]
+
+  Finally, [[unique-name-generator]] accepts the following options to further customize behavior:
+
+  ### `:name-key-fn`
+
+  Generated aliases are unique by the value of `[id (name-key-fn original-name)]`; the default is `identity`, so by
+  default aliases are unique by `[id name-key-fn]`. Specify something custom here if you want to make the unique
+  aliases unique by some other value, for example to make them unique without regards to case:
+
+    (let [f (unique-name-generator :name-key-fn str/lower-case)]
+      [(f \"x\")
+       (f \"X\")
+       (f \"X\")])
+    ;; -> [\"x\" \"X_2\" \"X_3\"]
+
+  This is useful for databases that treat column aliases as case-insensitive (see #19618 for some examples of this).
+
+  ### `:unique-alias-fn`
+
+  The function used to generate a potentially-unique alias given an original alias and unique suffix with the signature
+
+    (unique-alias-fn original suffix)
+
+  By default, combines them like `original_suffix`, but you can supply a custom function if you need to change this
+  behavior:
+
+    (let [f (unique-name-generator :unique-alias-fn (fn [x y] (format \"%s~~%s\" y x)))]
+      [(f \"x\")
+       (f \"x\")])
+  ;; -> [\"x\" \"2~~x\"]
+
+  This is useful if you need to constrain the generated suffix in some way, for example by limiting its length or
+  escaping characters disallowed in a column alias.
+
+  Values generated by this function are recursively checked for uniqueness, and will keep trying values a unique value
+  is generated; for this reason the function *must* return a unique value for every unique input. Use caution when
+  limiting the length of the identifier generated (consider appending a hash in cases like these)."
+  [& {:keys [name-key-fn unique-alias-fn]
+      :or   {name-key-fn     identity
+             unique-alias-fn (fn [original suffix]
+                               (str original \_ suffix))}}]
+  (let [id+original->unique (atom {})   ; map of [id original-alias] -> unique-alias
+        original->count     (atom {})]  ; map of original-alias -> count
     (fn generate-name
-      ([alias] (generate-name (gensym) alias))
-      ([identity-object alias]
-       (or (@identity-objects->aliases [identity-object alias])
-           (loop [maybe-unique alias]
-             (let [total-count (get (swap! aliases update maybe-unique (fnil inc 0)) maybe-unique)]
-               (if (= total-count 1)
-                 (do
-                   (swap! identity-objects->aliases assoc [identity-object alias] maybe-unique)
-                   maybe-unique)
-                 (recur (str maybe-unique \_ total-count))))))))))
+      ([alias]
+       (generate-name (gensym) alias))
+
+      ([id original]
+       (let [name-key (name-key-fn original)]
+         (or
+          ;; if we already have generated an alias for this key (e.g. `[id original]`), return it as-is.
+          (@id+original->unique [id name-key])
+          ;; otherwise generate a new unique alias.
+          ;; see if we're the first to try to use this candidate alias. Update the usage count in `original->count`
+          (let [total-count (get (swap! original->count update name-key (fnil inc 0)) name-key)]
+            (if (= total-count 1)
+              ;; if we are the first to do it, record it in `id+original->unique` and return it.
+              (do
+                (swap! id+original->unique assoc [id name-key] original)
+                original)
+              ;; otherwise prefix the alias by the current total count (e.g. `id` becomes `id_2`) and recur. If `id_2`
+              ;; is unused, it will get returned. Otherwise we'll recursively try `id_2_2`, and so forth.
+              (let [candidate (unique-alias-fn original (str total-count))]
+                ;; double-check that `unique-alias-fn` isn't doing something silly like truncating the generated alias
+                ;; to aggressively or forgetting to include the `suffix` -- otherwise we could end up with an infinite
+                ;; loop
+                (assert (not= candidate original)
+                        (str "unique-alias-fn must return a different string than its input. Input: "
+                             (pr-str candidate)))
+                (recur id candidate))))))))))
 
 (s/defn uniquify-names :- (s/constrained [s/Str] distinct? "sequence of unique strings")
   "Make the names in a sequence of string names unique by adding suffixes such as `_2`.
@@ -591,20 +653,30 @@
     :else
     x))
 
-(s/defn update-field-options :- mbql.s/field
-  "Like `clojure.core/update`, but for the options in a `:field` clause."
-  [[_ id-or-name opts] :- mbql.s/field f & args]
-  [:field id-or-name (remove-empty (apply f opts args))])
+(s/defn update-field-options :- mbql.s/FieldOrAggregationReference
+  "Like [[clojure.core/update]], but for the options in a `:field`, `:expression`, or `:aggregation` clause."
+  {:arglists '([field-or-ag-ref-or-expression-ref f & args])}
+  [[clause-type id-or-name opts] :- mbql.s/FieldOrAggregationReference f & args]
+  (let [opts (not-empty (remove-empty (apply f opts args)))]
+    ;; `:field` clauses should have a `nil` options map if there are no options. `:aggregation` and `:expression`
+    ;; should get the arg removed if it's `nil` or empty. (For now. In the future we may change this if we make the
+    ;; 3-arg versions the "official" normalized versions.)
+    (cond
+      opts                   [clause-type id-or-name opts]
+      (= clause-type :field) [clause-type id-or-name nil]
+      :else                  [clause-type id-or-name])))
 
 (defn assoc-field-options
-  "Like `clojure.core/assoc`, but for the options in a `:field` clause."
-  [field-clause & kvs]
-  (apply update-field-options field-clause assoc kvs))
+  "Like [[clojure.core/assoc]], but for the options in a `:field`, `:expression`, or `:aggregation` clause."
+  [clause & kvs]
+  (apply update-field-options clause assoc kvs))
 
 (defn with-temporal-unit
   "Set the `:temporal-unit` of a `:field` clause to `unit`."
-  [field-clause unit]
-  (assoc-field-options field-clause :temporal-unit unit))
+  [clause unit]
+  ;; it doesn't make sense to call this on an `:expression` or `:aggregation`.
+  (assert (is-clause? :field clause))
+  (assoc-field-options clause :temporal-unit unit))
 
 #?(:clj
    (p/import-vars

shared/test/metabase/mbql/util_test.cljc

 (ns metabase.mbql.util-test
-  (:require [clojure.test :as t]
+  (:require [clojure.string :as str]
+            [clojure.test :as t]
             [metabase.mbql.util :as mbql.u]
             metabase.types))
 
@@ -597,7 +598,26 @@
              (map (mbql.u/unique-name-generator) [:x :y :x :z] ["count" "sum" "count" "count_2"]))))
   (t/testing "Can the same object have multiple aliases"
     (t/is (= ["count" "sum" "count" "count_2"]
-             (map (mbql.u/unique-name-generator) [:x :y :x :x] ["count" "sum" "count" "count_2"])))))
+             (map (mbql.u/unique-name-generator) [:x :y :x :x] ["count" "sum" "count" "count_2"]))))
+
+  (t/testing "idempotence (2-arity calls to generated function)"
+    (let [unique-name (mbql.u/unique-name-generator)]
+      (t/is (= ["A" "B" "A" "A_2"]
+               [(unique-name :x "A")
+                (unique-name :x "B")
+                (unique-name :x "A")
+                (unique-name :y "A")]))))
+
+  (t/testing "options"
+    (t/testing :name-key-fn
+      (let [f (mbql.u/unique-name-generator :name-key-fn str/lower-case)]
+        (t/is (= ["x" "X_2" "X_3"]
+                 (map f ["x" "X" "X"])))))
+
+    (t/testing :unique-alias-fn
+      (let [f (mbql.u/unique-name-generator :unique-alias-fn (fn [x y] (str y "~~" x)))]
+        (t/is (= ["x" "2~~x"]
+                 (map f ["x" "x"])))))))
 
 
 ;;; --------------------------------------------- query->max-rows-limit ----------------------------------------------
@@ -701,4 +721,22 @@
 
   (t/testing "Should normalize the clause"
     (t/is (= [:field 1 nil]
-             (mbql.u/update-field-options [:field 1 {:a {:b 1}}] assoc-in [:a :b] nil)))))
+             (mbql.u/update-field-options [:field 1 {:a {:b 1}}] assoc-in [:a :b] nil))))
+
+  (t/testing "Should work with `:expression` and `:aggregation` references as well"
+    (t/is (= [:expression "wow" {:a 1}]
+             (mbql.u/update-field-options [:expression "wow"] assoc :a 1)))
+    (t/is (= [:expression "wow" {:a 1, :b 2}]
+             (mbql.u/update-field-options [:expression "wow" {:b 2}] assoc :a 1)))
+    (t/is (= [:aggregation 0 {:a 1}]
+             (mbql.u/update-field-options [:aggregation 0] assoc :a 1)))
+    (t/is (= [:aggregation 0 {:a 1, :b 2}]
+             (mbql.u/update-field-options [:aggregation 0 {:b 2}] assoc :a 1)))
+
+    ;; in the future when we make the 3-arg version the normalized/"official" version we will probably want to stop
+    ;; doing this.
+    (t/testing "Remove empty options entirely from `:expression` and `:aggregation` (for now)"
+      (t/is (= [:expression "wow"]
+               (mbql.u/update-field-options [:expression "wow" {:b 2}] dissoc :b)))
+      (t/is (= [:aggregation 0]
+               (mbql.u/update-field-options [:aggregation 0 {:b 2}] dissoc :b))))))

src/metabase/query_processor/util/add_alias_info.clj

+(ns metabase.query-processor.util.add-alias-info
+  (:require [clojure.string :as str]
+            [clojure.walk :as walk]
+            [metabase.driver :as driver]
+            [metabase.mbql.util :as mbql.u]
+            [metabase.query-processor.error-type :as qp.error-type]
+            [metabase.query-processor.store :as qp.store]
+            [metabase.util.i18n :refer [tru]]))
+
+;; these methods were moved from [[metabase.driver.sql.query-processor]] in 0.42.0
+
+(defmulti prefix-field-alias
+  "Create a Field alias by combining a `prefix` string with `field-alias` string. The default implementation just joins
+  the two strings with `__` -- override this if you need to do something different."
+  {:arglists '([driver prefix field-alias]), :added "0.38.1"}
+  driver/dispatch-on-initialized-driver
+  :hierarchy #'driver/hierarchy)
+
+(defmethod prefix-field-alias :default
+  [_driver prefix field-alias]
+  (str prefix "__" field-alias))
+
+(defmulti ^String escape-alias
+  "Return the String that should be emitted in the query for the generated `alias-name`, which will follow the
+  equivalent of a SQL `AS` clause. This is to allow for escaping names that particular databases may not allow as
+  aliases for custom expressions or fields (even when quoted).
+
+  Defaults to identity (i.e. returns `alias-name` unchanged)."
+  {:added "0.41.0" :arglists '([driver alias-name])}
+  driver/dispatch-on-initialized-driver
+  :hierarchy #'driver/hierarchy)
+
+(defn- make-unique-alias-fn
+  "Creates a function with the signature
+
+    (unique-alias position original-alias)
+
+  To return a uniquified version of `original-alias`. Memoized by `position`, so duplicate calls will result in the
+  same unique alias."
+  []
+  (let [unique-name-fn (mbql.u/unique-name-generator
+                        :name-key-fn     str/lower-case
+                        ;; TODO -- we should probably limit the length somehow like we do in
+                        ;; [[metabase.query-processor.middleware.add-implicit-joins/join-alias]], and also update this
+                        ;; function and that one to append a short suffix if we are limited by length. See also
+                        ;; [[escape-alias]] above
+                        :unique-alias-fn (fn [original suffix]
+                                           (escape-alias driver/*driver* (str original \_ suffix))))]
+    (fn unique-alias-fn [position original-alias]
+      (unique-name-fn position (escape-alias driver/*driver* original-alias)))))
+
+;; TODO -- this should probably limit the resulting alias, and suffix a short hash as well if it gets too long. See also
+;; [[unique-alias-fn]] below.
+(defmethod escape-alias :default
+  [_driver alias-name]
+  alias-name)
+
+(defn- remove-namespaced-options [options]
+  (when options
+    (not-empty (into {}
+                     (remove (fn [[k _]]
+                               (when (keyword? k)
+                                 (namespace k))))
+                     options))))
+
+(defn normalize-clause
+  "Normalize a `:field`/`:expression`/`:aggregation` clause by removing extra info so it can serve as a key for
+  `:qp/refs`."
+  [clause]
+  (mbql.u/match-one clause
+    ;; optimization: don't need to rewrite a `:field` clause without any options
+    [:field _ nil]
+    &match
+
+    [:field id-or-name opts]
+    ;; this doesn't use [[mbql.u/update-field-options]] because this gets called a lot and the overhead actually adds up
+    ;; a bit
+    [:field id-or-name (remove-namespaced-options (dissoc opts :source-fields))]
+
+    ;; for `:expression` and `:aggregation` references, remove the options map if they are empty.
+    [:expression expression-name opts]
+    (if-let [opts (remove-namespaced-options opts)]
+      [:expression expression-name opts]
+      [:expression expression-name])
+
+    [:aggregation index opts]
+    (if-let [opts (remove-namespaced-options opts)]
+      [:aggregation index opts]
+      [:aggregation index])
+
+    _
+    &match))
+
+(defn- selected-clauses
+  "Get all the clauses that are returned by this level of the query as a map of normalized-clause -> index of that
+  column in the results."
+  [{:keys [fields breakout aggregation], :as query}]
+  ;; this is cached for the duration of the QP run because it's a little expensive to calculate and caching this speeds
+  ;; up this namespace A LOT
+  (qp.store/cached (select-keys query [:fields :breakout :aggregation])
+    (into
+     {}
+     (comp cat
+           (map-indexed
+            (fn [i clause]
+              [(normalize-clause clause) i])))
+     [breakout
+      (map-indexed
+       (fn [i ag]
+         (mbql.u/replace ag
+           [:aggregation-options wrapped opts]
+           [:aggregation i]
+
+           ;; aggregation clause should be preprocessed into an `:aggregation-options` clause by now.
+           _
+           (throw (ex-info (tru "Expected :aggregation-options clause, got {0}" (pr-str ag))
+                           {:type qp.error-type/qp, :clause ag}))))
+       aggregation)
+      fields])))
+
+(defn- clause->position
+  "Get the position (i.e., column index) `clause` is returned as, if it is returned (i.e. if it is in `:breakout`,
+  `:aggregation`, or `:fields`). Not all clauses are returned."
+  [inner-query clause]
+  ((selected-clauses inner-query) (normalize-clause clause)))
+
+(defn- this-level-join-aliases [{:keys [joins]}]
+  (into #{} (map :alias) joins))
+
+(defn- field-is-from-join-in-this-level? [inner-query [_ _ {:keys [join-alias]}]]
+  (when join-alias
+    ((this-level-join-aliases inner-query) join-alias)))
+
+(defn- field-instance
+  {:arglists '([field-clause])}
+  [[_ id-or-name]]
+  (when (integer? id-or-name)
+    (qp.store/field id-or-name)))
+
+(defn- field-table-id [field-clause]
+  (:table_id (field-instance field-clause)))
+
+(defn- field-source-table-alias
+  "Determine the appropriate `::source-table` alias for a `field-clause`."
+  {:arglists '([inner-query field-clause])}
+  [{:keys [source-table], :as inner-query} [_ _id-or-name {:keys [join-alias]}, :as field-clause]]
+  (let [table-id            (field-table-id field-clause)
+        join-is-this-level? (field-is-from-join-in-this-level? inner-query field-clause)]
+    (cond
+      join-is-this-level?                      join-alias
+      (and table-id (= table-id source-table)) table-id
+      :else                                    ::source)))
+
+(defn- exports [query]
+  (into #{} (mbql.u/match (dissoc query :source-query :source-metadata :joins)
+              [(_ :guard #{:field :expression :aggregation}) _ (_ :guard (every-pred map? ::position))])))
+
+(defn- join-with-alias [{:keys [joins]} join-alias]
+  (some (fn [join]
+          (when (= (:alias join) join-alias)
+            join))
+        joins))
+
+(defn- matching-field-in-join-at-this-level
+  "If `field-clause` is the result of a join *at this level* with a `:source-query`, return the 'source' `:field` clause
+  from that source query."
+  [inner-query [_ _ {:keys [join-alias]} :as field-clause]]
+  (when join-alias
+    (when-let [matching-join-source-query (:source-query (join-with-alias inner-query join-alias))]
+      (let [normalized (mbql.u/update-field-options (normalize-clause field-clause) dissoc :join-alias)]
+        (some (fn [a-clause]
+                (when (and (mbql.u/is-clause? :field a-clause)
+                           (= (mbql.u/update-field-options (normalize-clause a-clause) dissoc :join-alias)
+                              normalized))
+                  a-clause))
+              (exports matching-join-source-query))))))
+
+(defn- field-alias-in-join-at-this-level
+  "If `field-clause` is the result of a join at this level, return the `::desired-alias` from that join (where the Field is
+  introduced). This is the appropriate `::source-alias` for such a Field."
+  [inner-query field-clause]
+  (when-let [[_ _ {::keys [desired-alias]}] (matching-field-in-join-at-this-level inner-query field-clause)]
+    desired-alias))
+
+(defn- matching-field-in-source-query
+  [{:keys [source-query], :as inner-query} [_ _ {:keys [join-alias]}, :as field-clause]]
+  (when (= (field-source-table-alias inner-query field-clause) ::source)
+    (let [normalized (normalize-clause field-clause)]
+      (some (fn [a-clause]
+              (when (and (mbql.u/is-clause? :field a-clause)
+                         (= (normalize-clause a-clause)
+                            normalized))
+                a-clause))
+            (exports source-query)))))
+
+(defn- field-alias-in-source-query
+  [inner-query field-clause]
+  (when-let [[_ _ {::keys [desired-alias]}] (matching-field-in-source-query inner-query field-clause)]
+    desired-alias))
+
+(defn- field-name
+  "*Actual* name of a `:field` from the database or source query (for Field literals)."
+  [_inner-query [_ id-or-name :as field-clause]]
+  (or (:name (field-instance field-clause))
+      (when (string? id-or-name)
+        id-or-name)))
+
+(defn- expensive-field-info
+  "Calculate extra stuff about `field-clause` that's a little expensive to calculate. This is done once so we can pass
+  it around instead of recalculating it a bunch of times."
+  [inner-query field-clause]
+  {:field-name              (field-name inner-query field-clause)
+   :join-is-this-level?     (field-is-from-join-in-this-level? inner-query field-clause)
+   :alias-from-join         (field-alias-in-join-at-this-level inner-query field-clause)
+   :alias-from-source-query (field-alias-in-source-query inner-query field-clause)})
+
+(defn- field-source-alias
+  "Determine the appropriate `::source-alias` for a `field-clause`."
+  {:arglists '([inner-query field-clause expensive-field-info])}
+  [{:keys [source-table], :as inner-query}
+   [_ _id-or-name {:keys [join-alias]}, :as field-clause]
+   {:keys [field-name join-is-this-level? alias-from-join alias-from-source-query]}]
+  (cond
+    (and join-alias (not join-is-this-level?)) (prefix-field-alias driver/*driver* join-alias field-name)
+    (and join-is-this-level? alias-from-join)  alias-from-join
+    alias-from-source-query                    alias-from-source-query
+    :else                                      field-name))
+
+(defn- field-desired-alias
+  "Determine the appropriate `::desired-alias` for a `field-clause`."
+  {:arglists '([inner-query field-clause expensive-field-info])}
+  [inner-query
+   [_ _id-or-name {:keys [join-alias]} :as field-clause]
+   {:keys [field-name alias-from-join alias-from-source-query]}]
+  (cond
+    join-alias              (prefix-field-alias driver/*driver* join-alias (or alias-from-join field-name))
+    alias-from-source-query alias-from-source-query
+    :else                   field-name))
+
+(defmulti ^:private clause-alias-info
+  {:arglists '([inner-query unique-alias-fn clause])}
+  (fn [_ _ [clause-type]]
+    clause-type))
+
+(defmethod clause-alias-info :field
+  [inner-query unique-alias-fn field-clause]
+  (let [expensive-info (expensive-field-info inner-query field-clause)]
+    (merge {::source-table (field-source-table-alias inner-query field-clause)
+            ::source-alias (field-source-alias inner-query field-clause expensive-info)}
+           (when-let [position (clause->position inner-query field-clause)]
+             {::desired-alias (unique-alias-fn position (field-desired-alias inner-query field-clause expensive-info))
+              ::position      position}))))
+
+(defmethod clause-alias-info :aggregation
+  [{aggregations :aggregation, :as inner-query} unique-alias-fn [_ index opts :as ag-ref-clause]]
+  (let [position (clause->position inner-query ag-ref-clause)]
+    ;; an aggregation is ALWAYS returned, so it HAS to have a `position`. If it does not, the aggregation reference
+    ;; is busted.
+    (when-not position
+      (throw (ex-info (tru "Aggregation does not exist at index {0}" index)
+                      {:type   qp.error-type/invalid-query
+                       :clause ag-ref-clause
+                       :query  inner-query})))
+    (let [[_ ag-name _ :as matching-ag] (nth aggregations index)]
+      ;; make sure we have an `:aggregation-options` clause like we expect. This is mostly a precondition check
+      ;; since we should never be running this code on not-preprocessed queries, so it's not i18n'ed
+      (when-not (mbql.u/is-clause? :aggregation-options matching-ag)
+        (throw (ex-info (format "Expected :aggregation-options, got %s. (Query must be fully preprocessed.)"
+                                (pr-str matching-ag))
+                        {:clause ag-ref-clause, :query inner-query})))
+      {::desired-alias (unique-alias-fn position ag-name)
+       ::position      position})))
+
+(defmethod clause-alias-info :expression
+  [inner-query unique-alias-fn [_ expression-name :as expression-ref-clause]]
+  (when-let [position (clause->position inner-query expression-ref-clause)]
+    {::desired-alias (unique-alias-fn position expression-name)
+     ::position      position}))
+
+(defn- add-alias-info* [inner-query]
+  (assert (not (:strategy inner-query)) "add-alias-info* should not be called on a join") ; not user-facing
+  (let [unique-alias-fn (make-unique-alias-fn)]
+    (mbql.u/replace inner-query
+      ;; don't rewrite anything inside any source queries or source metadata.
+      (_ :guard (constantly (some (partial contains? (set &parents))
+                                  [:source-query :source-metadata])))
+      &match
+
+      #{:field :aggregation :expression}
+      (mbql.u/update-field-options &match merge (clause-alias-info inner-query unique-alias-fn &match)))))
+
+(defn add-alias-info
+  "Add extra info to `:field` clauses, `:expression` references, and `:aggregation` references in `query`. `query` must
+  be fully preprocessed.
+
+  Adds some or all of the following keys:
+
+  ### `::source-table`
+
+  String name, integer Table ID, or the keyword `::source`. Use this alias to qualify the clause during compilation.
+  String names are aliases for joins. `::source` means this clause comes from the `:source-query`; the alias to use is
+  theoretically driver-specific but in practice is
+  `source` (see [[metabase.driver.sql.query-processor/source-query-alias]]). An integer Table ID means this comes from
+  the `:source-table` (either directly or indirectly via one or more `:source-query`s; use the Table's schema and name
+  to qualify the clause.
+
+  ### `::source-alias`
+
+  String name to use to refer to this clause during compilation.
+
+  ### `::desired-alias`
+
+  If this clause is 'selected' (i.e., appears in `:fields`, `:aggregation`, or `:breakout`), select the clause `AS`
+  this alias. This alias is guaranteed to be unique.
+
+  ### `::position`
+
+  If this clause is 'selected', this is the position the clause will appear in the results (i.e. the corresponding
+  column index)."
+  [query-or-inner-query]
+  (walk/postwalk
+   (fn [form]
+     (if (and (map? form)
+              ((some-fn :source-query :source-table) form)
+              (not (:strategy form)))
+       (vary-meta (add-alias-info* form) assoc ::transformed true)
+       form))
+   query-or-inner-query))

test/metabase/test_runner/effects.clj

@@ -10,7 +10,6 @@
   "
   (:require [clojure.data :as data]
             [clojure.test :as t]
-            dev.debug-qp
             [metabase.util.date-2 :as date-2]
             [metabase.util.i18n.impl :as i18n.impl]
             [schema.core :as s]))
@@ -46,20 +45,26 @@
        :diffs    (when-not pass?#
                    [[actual# [(s/check schema# actual#) nil]]])})))
 
+(defn query=-report
+  "Impl for [[t/assert-expr]] `query=`."
+  [message expected actual]
+  (let [pass? (= expected actual)]
+    (merge
+     {:type     (if pass? :pass :fail)
+      :message  message
+      :expected expected
+      :actual   actual}
+     ;; don't bother adding names unless the test actually failed
+     (when-not pass?
+       (let [add-names (requiring-resolve 'dev.debug-qp/add-names)]
+         {:expected (add-names expected)
+          :actual   (add-names actual)
+          :diffs    (let [[only-in-actual only-in-expected] (data/diff actual expected)]
+                      [[(add-names actual) [(add-names only-in-expected) (add-names only-in-actual)]]])})))))
+
 ;; basically the same as normal `=` but will add comment forms to MBQL queries for Field clauses and source tables
 ;; telling you the name of the referenced Fields/Tables
 (defmethod t/assert-expr 'query=
-  [message [_ expected actual :as args]]
-  `(let [expected# ~expected
-         actual#   ~actual
-         pass?#    (= expected# actual#)
-         expected# (dev.debug-qp/add-names expected#)
-         actual#   (dev.debug-qp/add-names actual#)]
-     (t/do-report
-      {:type     (if pass?# :pass :fail)
-       :message  ~message
-       :expected expected#
-       :actual   actual#
-       :diffs    (when-not pass?#
-                   (let [[only-in-actual# only-in-expected#] (data/diff actual# expected#)]
-                     [[actual# [only-in-expected# only-in-actual#]]]))})))
+  [message [_ expected actual]]
+  `(t/do-report
+    (query=-report ~message ~expected ~actual)))