[MLv2] Full `:column`s for `:dimensions` in `available-drill-thrus` (#34810)

Previously only the `:column-name` was provided, and there is a
self-proclaimed "icky hack" to guess at the full columns. The FE
provides the (JSON) columns, so this PR converts and uses them.

The icky hack can be removed.

`:column-ref` is included because sometimes it's important that the refs
use the same UUID.

frontend/src/metabase-lib/drills.unit.spec.ts

@@ -181,6 +181,7 @@ const AGGREGATED_ORDERS_COLUMNS = {
         "temporal-unit": "month",
       },
     ],
+    unit: "month",
   }),
 
   count: createMockColumn({

src/metabase/lib/convert.cljc

@@ -533,3 +533,30 @@
                             :legacy-ref               legacy-ref
                             :legacy-index->pMBQL-uuid *legacy-index->pMBQL-uuid*}
                            e))))))))
+
+(defn- from-json [query-fragment]
+  #?(:cljs (if (object? query-fragment)
+             (js->clj query-fragment :keywordize-keys true)
+             query-fragment)
+     :clj  query-fragment))
+
+(defn js-legacy-query->pMBQL
+  "Given a JSON-formatted legacy MBQL query, transform it to pMBQL.
+
+  If you have only the inner query map (`{:source-table 2 ...}` or similar), call [[js-legacy-inner-query->pMBQL]]
+  instead."
+  [query-map]
+  (-> query-map
+      from-json
+      (u/assoc-default :type :query)
+      mbql.normalize/normalize
+      ->pMBQL))
+
+(defn js-legacy-inner-query->pMBQL
+  "Given a JSON-formatted *inner* query, transform it to pMBQL.
+
+  If you have a complete legacy query (`{:type :query, :query {...}}` or similar), call [[js-legacy-query->pMBQL]]
+  instead."
+  [inner-query]
+  (js-legacy-query->pMBQL {:type  :query
+                           :query (from-json inner-query)}))

src/metabase/lib/drill_thru.cljc

 (ns metabase.lib.drill-thru
   (:require
-   [metabase.lib.aggregation :as lib.aggregation]
    [metabase.lib.drill-thru.column-filter :as lib.drill-thru.column-filter]
    [metabase.lib.drill-thru.common :as lib.drill-thru.common]
    [metabase.lib.drill-thru.distribution :as lib.drill-thru.distribution]
@@ -16,12 +15,9 @@
    [metabase.lib.drill-thru.underlying-records :as lib.drill-thru.underlying-records]
    [metabase.lib.drill-thru.zoom :as lib.drill-thru.zoom]
    [metabase.lib.drill-thru.zoom-in-timeseries :as lib.drill-thru.zoom-in-timeseries]
-   [metabase.lib.field :as lib.field]
    [metabase.lib.metadata.calculation :as lib.metadata.calculation]
-   [metabase.lib.options :as lib.options]
    [metabase.lib.schema :as lib.schema]
    [metabase.lib.schema.drill-thru :as lib.schema.drill-thru]
-   [metabase.lib.util :as lib.util]
    [metabase.util.malli :as mu]))
 
 (comment
@@ -39,28 +35,6 @@
 
 ;; TODO: ActionMode, PublicMode, MetabotMode need to be captured in the FE before calling `available-drill-thrus`.
 
-(defn- icky-hack-add-source-uuid-to-aggregation-column-metadata
-  "This is an evil HACK -- the FE calls [[available-drill-thrus]] with query results metadata as produced
-  by [[metabase.query-processor.middleware.annotate]], which does not include the `:lib/source-uuid` for aggregation
-  columns (since it's still using legacy MBQL at this point), which
-  means [[metabase.lib.aggregation/column-metadata->aggregation-ref]] can't generate references (since it requires
-  `:lib/source-uuid`)... so we have to add it back in manually. I've added `:aggregation-index` to the annotate
-  results (see [[metabase.query-processor.middleware.annotate/cols-for-ags-and-breakouts]]), and we can use that to
-  determine the correct `:lib/source-uuid`.
-
-  There's probably a more general place we can be doing this, but it's escaping me, so I guess this will have to do
-  for now. It doesn't seem like the FE generally uses result metadata in most other places so this is not an issue
-  elsewhere.
-
-  Once we're using MLv2 queries everywhere and stop converting back to legacy we should be able to remove this ICKY
-  HACK."
-  [query stage-number {{:keys [aggregation-index], :as column} :column, :as context}]
-  (or (when (and aggregation-index
-                 (not (:lib/source-uuid column)))
-        (when-let [ag (lib.aggregation/aggregation-at-index query stage-number aggregation-index)]
-          (assoc-in context [:column :lib/source-uuid] (lib.options/uuid ag))))
-      context))
-
 ;;; TODO: Missing drills: automatic insights, format.
 (def ^:private available-drill-thru-fns
   "Some drill thru functions are expected to return drills for just the specified `:column`; others are expected to
@@ -70,30 +44,22 @@
    {:f #'lib.drill-thru.column-filter/column-filter-drill,                       :return-drills-for-dimensions? true}
    {:f #'lib.drill-thru.foreign-key/foreign-key-drill,                           :return-drills-for-dimensions? false}
    {:f #'lib.drill-thru.object-details/object-detail-drill,                      :return-drills-for-dimensions? false}
-   {:f #'lib.drill-thru.pivot/pivot-drill,                                       :return-drills-for-dimensions? true}
+   {:f #'lib.drill-thru.pivot/pivot-drill,                                       :return-drills-for-dimensions? false}
    {:f #'lib.drill-thru.quick-filter/quick-filter-drill,                         :return-drills-for-dimensions? false}
    {:f #'lib.drill-thru.sort/sort-drill,                                         :return-drills-for-dimensions? true}
    {:f #'lib.drill-thru.summarize-column/summarize-column-drill,                 :return-drills-for-dimensions? true}
    {:f #'lib.drill-thru.summarize-column-by-time/summarize-column-by-time-drill, :return-drills-for-dimensions? true}
    {:f #'lib.drill-thru.underlying-records/underlying-records-drill,             :return-drills-for-dimensions? false}
-   {:f #'lib.drill-thru.zoom-in-timeseries/zoom-in-timeseries-drill,             :return-drills-for-dimensions? true}])
+   {:f #'lib.drill-thru.zoom-in-timeseries/zoom-in-timeseries-drill,             :return-drills-for-dimensions? false}])
 
 (mu/defn ^:private dimension-contexts :- [:maybe [:sequential {:min 1} ::lib.schema.drill-thru/context]]
   "Create new context maps (with updated `:column` and `:value` keys) for each of the `:dimensions` passed in. Some
   drill thru functions are expected to return drills for each of these columns, while others are expected to ignore
   them. Why? Who knows."
-  [query                             :- ::lib.schema/query
-   stage-number                      :- :int
-   {:keys [dimensions], :as context} :- ::lib.schema.drill-thru/context]
-  (when (seq dimensions)
-    (let [stage            (lib.util/query-stage query stage-number)
-          returned-columns (lib.metadata.calculation/returned-columns query stage-number stage)]
-      (for [{:keys [column-name value]} dimensions
-            :let                        [col (lib.field/resolve-column-name-in-metadata column-name returned-columns)]
-            :when                       col]
-        (assoc context
-               :column col
-               :value value)))))
+  [{:keys [dimensions], :as context} :- ::lib.schema.drill-thru/context]
+  (not-empty
+    (for [dimension dimensions]
+      (merge context dimension))))
 
 (mu/defn available-drill-thrus :- [:sequential [:ref ::lib.schema.drill-thru/drill-thru]]
   "Get a list (possibly empty) of available drill-thrus for a column, or a column + value pair.
@@ -109,13 +75,12 @@
     stage-number :- :int
     context      :- ::lib.schema.drill-thru/context]
    (try
-     (let [primary-context (icky-hack-add-source-uuid-to-aggregation-column-metadata query stage-number context)
-           dim-contexts    (dimension-contexts query stage-number primary-context)]
+     (let [dim-contexts (dimension-contexts context)]
        (into []
              (for [{:keys [f return-drills-for-dimensions?]} available-drill-thru-fns
-                   context                                   (if (and return-drills-for-dimensions? (seq dim-contexts))
+                   context                                   (if (and return-drills-for-dimensions? dim-contexts)
                                                                dim-contexts
-                                                               [primary-context])
+                                                               [context])
                    :let                                      [drill (f query stage-number context)]
                    :when                                     drill]
                drill)))

src/metabase/lib/drill_thru/object_details.cljc

 (ns metabase.lib.drill-thru.object-details
   (:require
+   [medley.core :as m]
    [metabase.lib.aggregation :as lib.aggregation]
    [metabase.lib.drill-thru.common :as lib.drill-thru.common]
    [metabase.lib.metadata.calculation :as lib.metadata.calculation]
@@ -24,8 +25,7 @@
            (empty? (lib.aggregation/aggregations query stage-number)))
       (let [[pk-column] (lib.metadata.calculation/primary-keys query) ; Already know there's only one.
             pk-value    (->> row
-                             (filter #(= (:column-name %) (:name pk-column)))
-                             first
+                             (m/find-first #(-> % :column :name (= (:name pk-column))))
                              :value)]
         (when (and pk-value
                    ;; Only recurse if this is a different column - otherwise it's an infinite loop.

src/metabase/lib/drill_thru/quick_filter.cljc

@@ -52,7 +52,7 @@
   (when (and (lib.drill-thru.common/mbql-stage? query stage-number)
              ;; (editable? query stage-number)
              column
-             (some? value)
+             (some? value) ; Deliberately allows value :null, only a missing value should fail this test.
              (not (lib.types.isa/primary-key? column))
              (not (lib.types.isa/foreign-key? column)))
     ;; for aggregate columns, we want to introduce a new stage when applying the drill-thru, `:new-stage?` is used to

src/metabase/lib/drill_thru/zoom_in_timeseries.cljc

 (ns metabase.lib.drill-thru.zoom-in-timeseries
   (:require
-   [medley.core :as m]
    [metabase.lib.breakout :as lib.breakout]
    [metabase.lib.drill-thru.common :as lib.drill-thru.common]
+   [metabase.lib.equality :as lib.equality]
    [metabase.lib.filter :as lib.filter]
-   [metabase.lib.join.util :as lib.join.util]
    [metabase.lib.metadata :as lib.metadata]
    [metabase.lib.remove-replace :as lib.remove-replace]
    [metabase.lib.schema :as lib.schema]
@@ -25,26 +24,18 @@
   (zipmap (drop-last valid-current-units)
           (drop 1 valid-current-units)))
 
-(defn- is-ref-for-source-column? [a-ref column]
-  (and (lib.util/clause-of-type? a-ref :field)
-       (let [[_field _opts id-or-name] a-ref]
-         (if (integer? id-or-name)
-           (= id-or-name (:id column))
-           (and (if-let [join-alias (lib.join.util/current-join-alias a-ref)]
-                  (= join-alias (lib.join.util/current-join-alias column))
-                  true)
-                (= id-or-name (:lib/source-column-alias column)))))))
-
-(mu/defn ^:private matching-breakout-ref :- [:maybe :mbql.clause/field]
+(mu/defn ^:private matching-breakout-dimension :- [:maybe ::lib.schema.drill-thru/context.row.value]
   [query        :- ::lib.schema/query
    stage-number :- :int
-   column       :- lib.metadata/ColumnMetadata]
-  (let [breakouts (lib.breakout/breakouts query stage-number)]
-    (m/find-first (fn [breakout]
-                    (and (is-ref-for-source-column? breakout column)
-                         (= (lib.temporal-bucket/temporal-bucket breakout)
-                            (lib.temporal-bucket/temporal-bucket column))))
-                  breakouts)))
+   dimensions   :- [:sequential ::lib.schema.drill-thru/context.row.value]]
+  (first (for [breakout (lib.breakout/breakouts query stage-number)
+               :when (and (lib.util/clause-of-type? breakout :field)
+                          (lib.temporal-bucket/temporal-bucket breakout))
+               {:keys [column] :as dimension} dimensions
+               :when (and (lib.equality/find-matching-column breakout [column])
+                          (= (lib.temporal-bucket/temporal-bucket breakout)
+                             (lib.temporal-bucket/temporal-bucket column)))]
+           (assoc dimension :column-ref breakout))))
 
 (mu/defn ^:private next-breakout-unit :- [:maybe ::lib.schema.temporal-bucketing/unit.date-time.truncate]
   [column :- lib.metadata/ColumnMetadata]
@@ -69,27 +60,28 @@
 
   This is different from the `:drill-thru/zoom` type, which is for showing the details of a single object."
   ;; TODO: This naming is confusing. Fix it?
-  [query                  :- ::lib.schema/query
-   stage-number           :- :int
-   {:keys [column value]} :- ::lib.schema.drill-thru/context]
+  [query                             :- ::lib.schema/query
+   stage-number                      :- :int
+   {:keys [column dimensions value]} :- ::lib.schema.drill-thru/context]
   (when (and (lib.drill-thru.common/mbql-stage? query stage-number)
              column
-             (some? value)
-             (matching-breakout-ref query stage-number column))
-    (when-let [next-unit (next-breakout-unit column)]
-      {:lib/type     :metabase.lib.drill-thru/drill-thru
-       :display-name (describe-next-unit next-unit)
-       :type         :drill-thru/zoom-in.timeseries
-       :column       column
-       :value        value
-       :next-unit    next-unit})))
+             (not-empty dimensions)
+             (some? value))
+    (when-let [dimension (matching-breakout-dimension query stage-number dimensions)]
+      (when-let [next-unit (next-breakout-unit (:column dimension))]
+        {:lib/type     :metabase.lib.drill-thru/drill-thru
+         :display-name (describe-next-unit next-unit)
+         :type         :drill-thru/zoom-in.timeseries
+         :dimension    dimension
+         :next-unit    next-unit}))))
 
 (mu/defmethod lib.drill-thru.common/drill-thru-method :drill-thru/zoom-in.timeseries
-  [query                            :- ::lib.schema/query
-   stage-number                     :- :int
-   {:keys [column value next-unit]} :- ::lib.schema.drill-thru/drill-thru.zoom-in.timeseries]
-  (let [breakout     (matching-breakout-ref query stage-number column)
-        new-breakout (lib.temporal-bucket/with-temporal-bucket breakout next-unit)]
+  [query                         :- ::lib.schema/query
+   stage-number                  :- :int
+   {:keys [dimension next-unit]} :- ::lib.schema.drill-thru/drill-thru.zoom-in.timeseries]
+  (let [{:keys [column value]} dimension
+        old-breakout           (:column-ref dimension)
+        new-breakout           (lib.temporal-bucket/with-temporal-bucket old-breakout next-unit)]
     (-> query
-      (lib.filter/filter stage-number (lib.filter/= column value))
-      (lib.remove-replace/replace-clause stage-number breakout new-breakout))))
+        (lib.filter/filter stage-number (lib.filter/= column value))
+        (lib.remove-replace/replace-clause stage-number old-breakout new-breakout))))

src/metabase/lib/js.cljs

@@ -23,7 +23,6 @@
    [metabase.lib.stage :as lib.stage]
    [metabase.lib.util :as lib.util]
    [metabase.mbql.js :as mbql.js]
-   [metabase.mbql.normalize :as mbql.normalize]
    [metabase.shared.util.time :as shared.ut]
    [metabase.util :as u]
    [metabase.util.log :as log]))
@@ -69,15 +68,6 @@
   [query]
   (lib.core/suggested-name query))
 
-(defn- pMBQL [query-map]
-  (as-> query-map <>
-    (js->clj <> :keywordize-keys true)
-    (if (:type <>)
-      <>
-      (assoc <> :type :query))
-    (mbql.normalize/normalize <>)
-    (lib.convert/->pMBQL <>)))
-
 (defn ^:export metadataProvider
   "Convert metadata to a metadata provider if it is not one already."
   [database-id metadata]
@@ -88,7 +78,7 @@
 (defn ^:export query
   "Coerce a plain map `query` to an actual query object that you can use with MLv2."
   [database-id metadata query-map]
-  (let [query-map (pMBQL query-map)]
+  (let [query-map (lib.convert/js-legacy-query->pMBQL query-map)]
     (log/debugf "query map: %s" (pr-str query-map))
     (lib.core/query (metadataProvider database-id metadata) query-map)))
 
@@ -853,16 +843,31 @@
   [a-query stage-number join-condition bucketing-option]
   (lib.core/join-condition-update-temporal-bucketing a-query stage-number join-condition bucketing-option))
 
+(defn- fix-column-with-ref [a-ref column]
+  (cond-> column
+    ;; Sometimes the FE has result metadata from the QP, without the required :lib/source-uuid on it.
+    ;; We have the UUID for the aggregation in its ref, so use that here.
+    (some-> a-ref first (= :aggregation)) (assoc :lib/source-uuid (last a-ref))))
+
 (defn- js-cells-by
   "Given a `col-fn`, returns a function that will extract a JS object like
-  `{col: {name: \"ID\", ...}, value: 12}` into a CLJS map like `{:column-name \"ID\", :value 12}`.
+  `{col: {name: \"ID\", ...}, value: 12}` into a CLJS map like
+  ```
+  {:column     {:lib/type :metadata/column ...}
+   :column-ref [:field ...]
+   :value 12}
+  ```
 
   The spelling of the column key differs between multiple JS objects of this same general shape
   (`col` on data rows, `column` on dimensions), etc., hence the abstraction."
   [col-fn]
   (fn [^js cell]
-    {:column-name (.-name (col-fn cell))
-     :value       (.-value cell)}))
+    (let [column     (js.metadata/parse-column (col-fn cell))
+          column-ref (when-let [a-ref (:field-ref column)]
+                       (legacy-ref->pMBQL a-ref))]
+      {:column     (fix-column-with-ref column-ref column)
+       :column-ref column-ref
+       :value      (.-value cell)})))
 
 (def ^:private row-cell       (js-cells-by #(.-col ^js %)))
 (def ^:private dimension-cell (js-cells-by #(.-column ^js %)))
@@ -874,15 +879,19 @@
   - Nullable data row (the array of `{col, value}` pairs from `clicked.data`)
   - Nullable dimensions list (`{column, value}` pairs from `clicked.dimensions`)"
   [a-query stage-number column value row dimensions]
-  (->> (merge {:column (js.metadata/parse-column column)
-               :value  (cond
-                         (undefined? value) nil   ; Missing a value, ie. a column click
-                         (nil? value)       :null ; Provided value is null, ie. database NULL
-                         :else              value)}
-              (when row                    {:row        (mapv row-cell       row)})
-              (when (not-empty dimensions) {:dimensions (mapv dimension-cell dimensions)}))
-       (lib.core/available-drill-thrus a-query stage-number)
-       to-array))
+  (lib.convert/with-aggregation-list (lib.core/aggregations a-query stage-number)
+    (let [column-ref (when-let [a-ref (.-field_ref ^js column)]
+                       (legacy-ref->pMBQL a-ref))]
+      (->> (merge {:column     (fix-column-with-ref column-ref (js.metadata/parse-column column))
+                   :column-ref column-ref
+                   :value      (cond
+                                 (undefined? value) nil   ; Missing a value, ie. a column click
+                                 (nil? value)       :null ; Provided value is null, ie. database NULL
+                                 :else              value)}
+                  (when row                    {:row        (mapv row-cell       row)})
+                  (when (not-empty dimensions) {:dimensions (mapv dimension-cell dimensions)}))
+           (lib.core/available-drill-thrus a-query stage-number)
+           to-array))))
 
 (defn ^:export drill-thru
   "Applies the given `drill-thru` to the specified query and stage. Returns the updated query.

src/metabase/lib/js/metadata.cljs

@@ -211,7 +211,8 @@
 
 (defmethod rename-key-fn :field
   [_object-type]
-  {:source :lib/source})
+  {:source :lib/source
+   :unit   :metabase.lib.field/temporal-unit})
 
 (defn- parse-field-id
   [id]
@@ -224,19 +225,20 @@
   [_object-type]
   (fn [k v]
     (case k
-      :base-type         (keyword v)
-      :coercion-strategy (keyword v)
-      :effective-type    (keyword v)
-      :fingerprint       (if (map? v)
-                           (walk/keywordize-keys v)
-                           (js->clj v :keywordize-keys true))
-      :has-field-values  (keyword v)
-      :lib/source        (if (= v "aggregation")
-                           :source/aggregations
-                           (keyword "source" v))
-      :semantic-type     (keyword v)
-      :visibility-type   (keyword v)
-      :id                (parse-field-id v)
+      :base-type                        (keyword v)
+      :coercion-strategy                (keyword v)
+      :effective-type                   (keyword v)
+      :fingerprint                      (if (map? v)
+                                          (walk/keywordize-keys v)
+                                          (js->clj v :keywordize-keys true))
+      :has-field-values                 (keyword v)
+      :lib/source                       (if (= v "aggregation")
+                                          :source/aggregations
+                                          (keyword "source" v))
+      :metabase.lib.field/temporal-unit (keyword v)
+      :semantic-type                    (keyword v)
+      :visibility-type                  (keyword v)
+      :id                               (parse-field-id v)
       v)))
 
 (defmethod parse-objects :field

src/metabase/lib/schema/drill_thru.cljc

@@ -10,6 +10,7 @@
    [metabase.lib.schema.id :as lib.schema.id]
    [metabase.lib.schema.metadata :as lib.schema.metadata]
    [metabase.lib.schema.order-by :as lib.schema.order-by]
+   [metabase.lib.schema.ref :as lib.schema.ref]
    [metabase.lib.schema.temporal-bucketing
     :as lib.schema.temporal-bucketing]
    [metabase.util.malli.registry :as mr]))
@@ -157,8 +158,7 @@
    ::drill-thru.common
    [:map
     [:type      [:= :drill-thru/zoom-in.timeseries]]
-    [:column    [:ref ::lib.schema.metadata/column]]
-    [:value     some?]
+    [:dimension [:ref ::context.row.value]]
     [:next-unit [:ref ::drill-thru.zoom-in.timeseries.next-unit]]]])
 
 (mr/def ::drill-thru
@@ -182,34 +182,19 @@
     [:drill-thru/automatic-insights       ::drill-thru.automatic-insights]
     [:drill-thru/zoom-in.timeseries       ::drill-thru.zoom-in.timeseries]]])
 
-;;; Frontend passes in something that looks like this. Why this shape? Who knows.
-(comment
-  {:column     {:lib/type            :metadata/column
-                :remapped-from-index nil
-                :base-type           :type/BigInteger
-                :semantic-type       :type/Quantity
-                :name                "count"
-                :lib/source          :source/aggregations
-                :aggregation-index   0
-                :effective-type      :type/BigInteger
-                :display-name        "Count"
-                :remapping           nil}
-   :value      457
-   :row        [{:column-name "CREATED_AT", :value "2024-01-01T00:00:00Z"}
-                {:column-name "count", :value 457}]
-   :dimensions [{:column-name "CREATED_AT", :value "2024-01-01T00:00:00Z"}]})
-
 (mr/def ::context.row.value
   [:map
-   [:column-name string?]
-   [:value       :any]])
+   [:column     [:ref ::lib.schema.metadata/column]]
+   [:column-ref [:ref ::lib.schema.ref/ref]]
+   [:value      :any]])
 
 (mr/def ::context.row
   [:sequential [:ref ::context.row.value]])
 
 (mr/def ::context
   [:map
-   [:column [:ref ::lib.schema.metadata/column]]
-   [:value  [:maybe :any]]
+   [:column     [:ref ::lib.schema.metadata/column]]
+   [:column-ref [:ref ::lib.schema.ref/ref]]
+   [:value      [:maybe :any]]
    [:row        {:optional true} [:ref ::context.row]]
    [:dimensions {:optional true} [:maybe [:ref ::context.row]]]])

test/metabase/lib/drill_thru/distribution_test.cljc

@@ -18,8 +18,9 @@
           count-col (m/find-first (fn [col]
                                     (= (:display-name col) "Count"))
                                   (lib/returned-columns query))
-          context   {:column count-col
-                     :value  nil}]
+          context   {:column     count-col
+                     :column-ref (lib/ref count-col)
+                     :value      nil}]
       (is (some? count-col))
       (is (nil? (lib.drill-thru.distribution/distribution-drill query -1 context))))))
 

test/metabase/lib/drill_thru/sort_test.cljc

@@ -16,8 +16,9 @@
   (let [query (lib/query meta/metadata-provider (meta/table-metadata :orders))
         drill (lib.drill-thru.sort/sort-drill query
                                               -1
-                                              {:column (meta/field-metadata :orders :id)
-                                               :value  nil})]
+                                              {:column     (meta/field-metadata :orders :id)
+                                               :column-ref (lib/ref (meta/field-metadata :orders :id))
+                                               :value      nil})]
     (is (=? {:type            :drill-thru/sort
              :column          {:id (meta/id :orders :id)}
              :sort-directions [:asc :desc]}
@@ -55,8 +56,9 @@
           count-col (m/find-first (fn [col]
                                     (= (:display-name col) "Count"))
                                   (lib/returned-columns query))
-          context   {:column count-col
-                     :value  nil}]
+          context   {:column     count-col
+                     :column-ref (lib/ref count-col)
+                     :value      nil}]
       (is (some? count-col))
       (let [drill (lib.drill-thru.sort/sort-drill query -1 context)]
         (is (=? {:lib/type        :metabase.lib.drill-thru/drill-thru
@@ -79,8 +81,10 @@
     (let [query   (-> (lib/query meta/metadata-provider (meta/table-metadata :orders))
                       (lib/order-by (meta/field-metadata :orders :user-id))
                       (lib/order-by (meta/field-metadata :orders :id)))
-          context {:column (meta/field-metadata :orders :user-id)
-                   :value  nil}
+          user-id (meta/field-metadata :orders :user-id)
+          context {:column     user-id
+                   :column-ref (lib/ref user-id)
+                   :value      nil}
           drill   (lib.drill-thru.sort/sort-drill query -1 context)]
       (is (=? {:stages
                [{:order-by [[:asc {} [:field {} (meta/id :orders :user-id)]]

test/metabase/lib/drill_thru/summarize_column_by_time_test.cljc

@@ -19,8 +19,9 @@
           count-col (m/find-first (fn [col]
                                     (= (:display-name col) "Count"))
                                   (lib/returned-columns query))
-          context   {:column count-col
-                     :value  nil}]
+          context   {:column     count-col
+                     :column-ref (lib/ref count-col)
+                     :value      nil}]
       (is (some? count-col))
       (is (nil? (lib.drill-thru.summarize-column-by-time/summarize-column-by-time-drill query -1 context))))))
 

test/metabase/lib/drill_thru/test_util.cljc

@@ -85,9 +85,9 @@
    row       :- Row
    {:keys [column-name click-type query-type], :as _test-case} :- TestCase]
   (let [cols       (lib/returned-columns query -1 (lib.util/query-stage query -1))
-        col        (m/find-first (fn [col]
-                                   (= (:name col) column-name))
-                                 cols)
+        by-name    (m/index-by :name cols)
+        col        (get by-name column-name)
+        refs       (update-vals by-name lib/ref)
         _          (assert col (lib.util/format "No column found named %s; found: %s"
                                                 (pr-str column-name)
                                                 (pr-str (map :name cols))))
@@ -98,14 +98,20 @@
                      (for [col   cols
                            :when (and (= (:lib/source col) :source/breakouts)
                                       (not= (:name col) column-name))]
-                       {:column-name (:name col), :value (get row (:name col))}))]
+                       {:column     col
+                        :column-ref (get refs (:name col))
+                        :value      (get row (:name col))}))]
     (merge
-     {:column col
-      :value  nil}
+     {:column     col
+      :column-ref (get refs column-name)
+      :value      nil}
      (when (= click-type :cell)
        {:value      value
-        :row        (for [[column-name value] row]
-                      {:column-name column-name, :value value})
+        :row        (for [[column-name value] row
+                          :let [column (by-name column-name)]]
+                      {:column     column
+                       :column-ref (get refs column-name)
+                       :value      value})
         :dimensions dimensions}))))
 
 (def ^:private AvailableDrillsTestCase

test/metabase/lib/drill_thru/zoom_in_timeseries_test.cljc

@@ -20,18 +20,26 @@
                        :breakout    [[:field {:temporal-unit :day} (meta/id :orders :created-at)]
                                      [:field {:temporal-unit :year} (meta/id :orders :created-at)]]}]}
             query))
-    (let [created-at (m/find-first #(and (= (:id %) (meta/id :orders :created-at))
+    (let [columns    (lib/returned-columns query)
+          created-at (m/find-first #(and (= (:id %) (meta/id :orders :created-at))
                                          (= (lib.temporal-bucket/raw-temporal-bucket %) :year))
                                    (lib/returned-columns query))
           _          (assert created-at)
-          drill      (lib.drill-thru.zoom-in-timeseries/zoom-in-timeseries-drill query
-                                                                                 -1
-                                                                                 {:column created-at
-                                                                                  :value  2022})]
+          count-col  (m/find-first #(= (:name %) "count") columns)
+          _          (assert count-col)
+          drill      (lib.drill-thru.zoom-in-timeseries/zoom-in-timeseries-drill
+                       query -1
+                       {:column     count-col
+                        :column-ref (lib/ref count-col)
+                        :value 200
+                        :dimensions [{:column     created-at
+                                      :column-ref (lib/ref created-at)
+                                      :value      2022}]})]
       (is (=? {:type         :drill-thru/zoom-in.timeseries
-               :column       {:id                               (meta/id :orders :created-at)
-                              :metabase.lib.field/temporal-unit :year}
-               :value        2022
+               :dimension    {:column     {:id                               (meta/id :orders :created-at)
+                                           :metabase.lib.field/temporal-unit :year}
+                              :column-ref [:field {} (meta/id :orders :created-at)]
+                              :value      2022}
                :next-unit    :quarter
                :display-name "See this year by quarter"}
               drill))
@@ -46,18 +54,24 @@
                                           [:field {:temporal-unit :year} (meta/id :orders :created-at)]
                                           2022]]}]}
                 query'))
-        (let [created-at (m/find-first #(and (= (:id %) (meta/id :orders :created-at))
+        (let [columns    (lib/returned-columns query')
+              created-at (m/find-first #(and (= (:id %) (meta/id :orders :created-at))
                                              (= (lib.temporal-bucket/raw-temporal-bucket %) :quarter))
-                                       (lib/returned-columns query'))
+                                       columns)
               _          (assert created-at)
-              drill      (lib.drill-thru.zoom-in-timeseries/zoom-in-timeseries-drill query'
-                                                                                     -1
-                                                                                     {:column created-at
-                                                                                      :value  "2022-04-01T00:00:00"})]
+              drill      (lib.drill-thru.zoom-in-timeseries/zoom-in-timeseries-drill
+                           query' -1
+                           {:column     count-col
+                            :column-ref (lib/ref count-col)
+                            :value      19
+                            :dimensions [{:column     created-at
+                                          :column-ref (lib/ref created-at)
+                                          :value      "2022-04-01T00:00:00"}]})]
           (is (=? {:type         :drill-thru/zoom-in.timeseries
-                   :column       {:id                               (meta/id :orders :created-at)
-                                  :metabase.lib.field/temporal-unit :quarter}
-                   :value        "2022-04-01T00:00:00"
+                   :dimension    {:column     {:id                               (meta/id :orders :created-at)
+                                               :metabase.lib.field/temporal-unit :quarter}
+                                  :column-ref [:field {} (meta/id :orders :created-at)]
+                                  :value      "2022-04-01T00:00:00"}
                    :next-unit    :month
                    :display-name "See this quarter by month"}
                   drill))