Fix the zoom-in.timeseries drill for multi-stage queries (#50970) (#51118)

* Fix the zoom-in.timeseries drill for multi-stage queries

If no dimensions are provided to the zoom-in.timeseries drill, then attempt to find the matching dimension from the FE-provided row data instead, if present.

Related to https://github.com/metabase/metabase/issues/46932



* Only return zoom-in.timeseries drills for clicks on aggregated columns

Co-authored-by: appleby <86076+appleby@users.noreply.github.com>

src/metabase/lib/drill_thru/common.cljc

@@ -2,6 +2,7 @@
   (:require
    [metabase.lib.hierarchy :as lib.hierarchy]
    [metabase.lib.metadata.calculation :as lib.metadata.calculation]
+   [metabase.lib.underlying :as lib.underlying]
    [metabase.lib.util :as lib.util]))
 
 (defn mbql-stage?
@@ -46,3 +47,35 @@
   "Convert a drill value to a JS value"
   [value]
   (if (= value :null) nil value))
+
+(defn- has-source-or-underyling-source-fn
+  [source]
+  (fn has-source?
+    ([column]
+     (= (:lib/source column) source))
+    ([query column]
+     (and
+      (seq column)
+      (or (has-source? column)
+          (has-source? (lib.underlying/top-level-column query column)))))))
+
+(def aggregation-sourced?
+  "Does column or top-level-column have :source/aggregations?"
+  (has-source-or-underyling-source-fn :source/aggregations))
+
+(def breakout-sourced?
+  "Does column or top-level-column have :source/breakouts?"
+  (has-source-or-underyling-source-fn :source/breakouts))
+
+(defn strictly-underyling-aggregation?
+  "Does the top-level-column for `column` in `query` have :source/aggregations?"
+  [query column]
+  (and (not (aggregation-sourced? column))
+       (aggregation-sourced? query column)))
+
+(defn dimensions-from-breakout-columns
+  "Convert `row` data into dimensions for `column`s that come from an aggregation in a previous stage."
+  [query column row]
+  (when (strictly-underyling-aggregation? query column)
+    (not-empty (filterv #(breakout-sourced? query (:column %))
+                        row))))

src/metabase/lib/drill_thru/zoom_in_timeseries.cljc

@@ -14,7 +14,8 @@
   - `dimensions` have a date or datetime column with `year`, `quarter`, `month`, `week`, `day`, `hour` temporal unit.
     For other units, or when there is no temporal bucketing, this drill cannot be applied. Changing `hour` to `minute`
     ends the sequence for datetime columns (`week` to `day` for date columns). Only the first matching column would be
-    used in query transformation.
+    used in query transformation. If `dimensions` are not provided, `row` data will instead be used to try to find a
+    `matching-breakout-dimension`.
 
   - `displayInfo` returns `displayName` with `See this {0} by {1}` string using the current and the next available
     temporal unit.
@@ -45,6 +46,7 @@
    [metabase.lib.schema.drill-thru :as lib.schema.drill-thru]
    [metabase.lib.schema.temporal-bucketing :as lib.schema.temporal-bucketing]
    [metabase.lib.temporal-bucket :as lib.temporal-bucket]
+   [metabase.lib.underlying :as lib.underlying]
    [metabase.lib.util :as lib.util]
    [metabase.util.i18n :as i18n]
    [metabase.util.malli :as mu]))
@@ -62,14 +64,22 @@
   [query        :- ::lib.schema/query
    stage-number :- :int
    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))
+  (first (for [[breakout-ref breakout-col] (map vector
+                                                (lib.breakout/breakouts query stage-number)
+                                                (lib.breakout/breakouts-metadata query stage-number))
+               :when (and (lib.util/clause-of-type? breakout-ref :field)
+                          (lib.temporal-bucket/temporal-bucket breakout-ref))
                {: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))))
+               :when (and (lib.equality/find-matching-column breakout-ref [column])
+                          (= (lib.temporal-bucket/raw-temporal-bucket breakout-ref)
+                             (or (lib.temporal-bucket/raw-temporal-bucket column)
+                                 ;; If query is multi-stage and column comes from a call
+                                 ;; to [[lib.calculation/returned-columns]], then it may have an
+                                 ;; :inherited-temporal-unit instead of a :temporal-unit.
+                                 (:inherited-temporal-unit column))))]
+           ;; If stage-number is not -1, then the column from the input dimension will be from the last stage,
+           ;; whereas breakout-col will be the corresponding breakout column from [[lib.underlying/top-level-stage]].
+           (assoc dimension :column breakout-col :column-ref breakout-ref))))
 
 (mu/defn- next-breakout-unit :- [:maybe ::lib.schema.drill-thru/drill-thru.zoom-in.timeseries.next-unit]
   [query :- ::lib.schema/query
@@ -96,27 +106,34 @@
   For example: The month of a year, days or weeks of a quarter, smaller lat/long regions, etc.
 
   This is different from the `:drill-thru/zoom` type, which is for showing the details of a single object."
-  [query                              :- ::lib.schema/query
-   stage-number                       :- :int
-   {:keys [dimensions], :as _context} :- ::lib.schema.drill-thru/context]
-  (when (and (lib.drill-thru.common/mbql-stage? query stage-number)
-             (not-empty dimensions))
-    (when-let [{:keys [value column-ref], :as dimension} (matching-breakout-dimension query stage-number dimensions)]
-      (when value
-        (when-let [next-unit (next-breakout-unit query stage-number column-ref)]
-          {: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})))))
+  [query                                         :- ::lib.schema/query
+   _stage-number                                 :- :int
+   {:keys [column dimensions row], :as _context} :- ::lib.schema.drill-thru/context]
+  ;; For multi-stage queries, we want the stage-number of the underlying stage with breakouts or aggregations.
+  ;; In such cases, the FE will not pass dimensions, so use the row data instead, if available.
+  (let [stage-number (lib.underlying/top-level-stage-number query)
+        dimensions   (or (not-empty dimensions)
+                         (lib.drill-thru.common/dimensions-from-breakout-columns query column row))]
+    (when (and (lib.drill-thru.common/mbql-stage? query stage-number)
+               dimensions)
+      (when-let [{:keys [value column-ref], :as dimension}
+                 (matching-breakout-dimension query stage-number dimensions)]
+        (when value
+          (when-let [next-unit (next-breakout-unit query stage-number column-ref)]
+            {: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
+   _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)]
+        new-breakout           (lib.temporal-bucket/with-temporal-bucket old-breakout next-unit)
+        stage-number           (lib.underlying/top-level-stage-number query)]
     (-> query
         (lib.filter/filter stage-number (lib.filter/= column value))
         (lib.remove-replace/replace-clause stage-number old-breakout new-breakout))))

src/metabase/lib/field.cljc

@@ -100,8 +100,8 @@
                   ;; the [[lib.metadata.calculation/*propagate-inherited-temoral-unit*]] is thruthy, ie. bound. Intent
                   ;; is to pass it from ref to column only during [[returned-columns]] call. Otherwise eg.
                   ;; [[orderable-columns]] would contain that too. That could be problematic, because original ref that
-                  ;; contained `:temporal-unit` contains no `:inherent-temporal-unit`. If the column like this was used
-                  ;; to generate ref for eg. order by it would contain the `:inherent-temporal-unit`, while
+                  ;; contained `:temporal-unit` contains no `:inherited-temporal-unit`. If the column like this was used
+                  ;; to generate ref for eg. order by it would contain the `:inherited-temporal-unit`, while
                   ;; the original column (eg. in breakout) would not.
                   (let [inherited-temporal-unit-keys (cond-> (list :inherited-temporal-unit)
                                                        lib.metadata.calculation/*propagate-inherited-temoral-unit*

src/metabase/lib/underlying.cljc

@@ -15,31 +15,51 @@
    [metabase.lib.util :as lib.util]
    [metabase.util.malli :as mu]))
 
-(mu/defn- pop-until-aggregation-or-breakout :- [:maybe ::lib.schema/query]
-  "Strips off any trailing stages that do not contain aggregations.
+(mu/defn- pop-until-aggregation-or-breakout :- [:tuple [:maybe ::lib.schema/query] [:int {:max -1}]]
+  "Strips off any trailing stages that do not contain aggregations or breakouts.
 
-  If there are no such stages, returns nil."
+  Returns a tuple of [query, stage-number] where `query` is the first stage with aggregations or breakouts and
+  `stage-number` is the (negative) index of that stage relative to the original query.
+
+  If there are no such stages, returns [nil, -1]."
   [query :- ::lib.schema/query]
-  (if (and (empty? (lib.aggregation/aggregations query -1))
-           (empty? (lib.breakout/breakouts query -1)))
-    ;; No aggregations or breakouts in the last stage, so pop it off and recur.
-    (let [popped (update query :stages pop)]
-      (when (seq (:stages popped))
-        (recur popped)))
-    query))
+  (loop [query query
+         i -1]
+    (if (and (empty? (lib.aggregation/aggregations query -1))
+             (empty? (lib.breakout/breakouts query -1)))
+      ;; No aggregations or breakouts in the last stage, so pop it off and recur.
+      (let [popped (update query :stages pop)]
+        (if (seq (:stages popped))
+          (recur popped (dec i))
+          [nil, -1]))
+      [query, i])))
+
+(mu/defn- query-database-supports? :- :boolean
+  [query   :- ::lib.schema/query
+   feature :- :keyword]
+  (contains? (-> query lib.metadata/database :features) feature))
 
 (mu/defn top-level-query :- ::lib.schema/query
   "Returns the \"top-level\" query for the given query.
 
-  That means dropping any trailing filters, fields, etc. to get back to the last stage that has an aggregation. If there
-  are no stages with aggregations, the original query is returned.
+  That means dropping any trailing filters, fields, etc. to get back to the last stage that has an aggregation or
+  breakout. If there are no such stages, the original query is returned.
 
-  If the database does not support nested queries, this also returns the original."
+  If the database does not support nested queries, this also returns the original query."
   [query :- ::lib.schema/query]
-  (or (when ((-> query lib.metadata/database :features) :nested-queries)
-        (pop-until-aggregation-or-breakout query))
+  (or (when (query-database-supports? query :nested-queries)
+        (first (pop-until-aggregation-or-breakout query)))
       query))
 
+(mu/defn top-level-stage-number :- :int
+  "Returns the stage-number of the [[top-level-query]] for the given query.
+
+  If there are no such stages, or if the database does not supported nested queries, returns -1."
+  [query :- ::lib.schema/query]
+  (or (when (query-database-supports? query :nested-queries)
+        (second (pop-until-aggregation-or-breakout query)))
+      -1))
+
 (mu/defn top-level-column :- ::lib.schema.metadata/column
   "Given a column, returns the \"top-level\" equivalent.
 
@@ -67,4 +87,4 @@
 (mu/defn has-aggregation-or-breakout?
   "Whether the `query` has an aggregation or breakout clause in some query stage."
   [query :- ::lib.schema/query]
-  (some? (pop-until-aggregation-or-breakout query)))
+  (some? (first (pop-until-aggregation-or-breakout query))))

test/metabase/lib/drill_thru/zoom_in_timeseries_test.cljc

@@ -189,33 +189,77 @@
           bucketing
           {"CREATED_AT" "2022-12-01"}))))))
 
+(deftest ^:parallel returns-zoom-in-timeseries-for-multi-stage-queries
+  (lib.drill-thru.tu/test-returns-drill
+   {:drill-type   :drill-thru/zoom-in.timeseries
+    :click-type   :cell
+    :query-type   :aggregated
+    :column-name  "count"
+    :custom-query (let [base-query (-> (lib/query meta/metadata-provider (meta/table-metadata :orders))
+                                       (lib/aggregate (lib/count))
+                                       (lib/breakout (meta/field-metadata :orders :product-id))
+                                       (lib/breakout (-> (meta/field-metadata :orders :created-at)
+                                                         (lib/with-temporal-bucket :month)))
+                                       lib/append-stage)
+                        count-col  (m/find-first #(= (:name %) "count")
+                                                 (lib/returned-columns base-query))
+                        _          (is (some? count-col))]
+                    (lib/filter base-query (lib/> count-col 0)))
+    :custom-row   {"PRODUCT_ID" 3
+                   "CREATED_AT" "2023-12-01"
+                   "count"      77}
+    :expected     {:type      :drill-thru/zoom-in.timeseries
+                   :next-unit :week
+                   ;; the "underlying" dimension is reconstructed from the row.
+                   :dimension {:column     {:name       "CREATED_AT"
+                                            :lib/source :source/breakouts}
+                               :column-ref [:field {} (meta/id :orders :created-at)]
+                               :value      "2023-12-01"}}}))
+
 (deftest ^:parallel returns-zoom-in-timeseries-e2e-test-2
   (testing "zoom-in.timeseries should be returned for a"
-    (let [query          (-> (lib/query meta/metadata-provider (meta/table-metadata :orders))
-                             (lib/aggregate (lib/count))
-                             (lib/breakout (-> (meta/field-metadata :orders :created-at)
-                                               (lib/with-temporal-bucket :year))))
-          created-at-col (m/find-first #(= (:name %) "CREATED_AT")
-                                       (lib/returned-columns query))
-          _              (is (some? created-at-col))]
-      (doseq [[message context] {"pivot cell (no column, value = NULL) (#36173)"
-                                 {:value      :null
-                                  :column     nil
-                                  :column-ref nil
-                                  :dimensions [{:column     created-at-col
-                                                :column-ref (lib/ref created-at-col)
-                                                :value      "2022-12-01T00:00:00+02:00"}]}
+    (let [base-query        (-> (lib/query meta/metadata-provider (meta/table-metadata :orders))
+                                (lib/aggregate (lib/count))
+                                (lib/breakout (-> (meta/field-metadata :orders :created-at)
+                                                  (lib/with-temporal-bucket :year))))
+          created-at-col    (m/find-first #(= (:name %) "CREATED_AT")
+                                          (lib/returned-columns base-query))
+          _                 (is (some? created-at-col))
+          multi-stage-query (lib/append-stage base-query)
+          count-col         (m/find-first #(= (:name %) "count")
+                                          (lib/returned-columns multi-stage-query))
+          _                 (is (some? count-col))
+          multi-stage-query (lib/filter multi-stage-query (lib/> count-col 0))]
+      (doseq [[query-type query] {"single-stage" base-query
+                                  "multi-stage" multi-stage-query}
+              [message context]  {"pivot cell (no column, value = NULL) (#36173)"
+                                  {:value      :null
+                                   :column     nil
+                                   :column-ref nil
+                                   :dimensions [{:column     created-at-col
+                                                 :column-ref (lib/ref created-at-col)
+                                                 :value      "2022-12-01T00:00:00+02:00"}]}
 
-                                 "legend item (no column, no value) (#36173)"
-                                 {:value      nil
-                                  :column     nil
-                                  :column-ref nil
-                                  :dimensions [{:column     created-at-col
-                                                :column-ref (lib/ref created-at-col)
-                                                :value      "2022-12-01T00:00:00+02:00"}]}}]
-        (testing message
+                                  "legend item (no column, no value) (#36173)"
+                                  {:value      nil
+                                   :column     nil
+                                   :column-ref nil
+                                   :dimensions [{:column     created-at-col
+                                                 :column-ref (lib/ref created-at-col)
+                                                 :value      "2022-12-01T00:00:00+02:00"}]}}]
+        (testing (str query-type " query: " message)
           (let [[drill :as drills] (filter #(= (:type %) :drill-thru/zoom-in.timeseries)
-                                           (lib/available-drill-thrus query -1 context))]
+                                           (lib/available-drill-thrus query -1 context))
+                ;; both queries have the base-query stage where the drill filter should be added
+                expected-stages (cond-> [{:aggregation [[:count {}]]
+                                          :breakout    [[:field {:temporal-unit :quarter} (meta/id :orders :created-at)]],
+                                          :filters     [[:=
+                                                         {}
+                                                         [:field {:temporal-unit :year} (meta/id :orders :created-at)]
+                                                         "2022-12-01T00:00:00+02:00"]]}]
+                                  ;; the multi-stage-query has an additional filter stage
+                                  (= query multi-stage-query) (conj {:filters
+                                                                     [[:> {} [:field {} "count"] 0]]}))]
             (is (= 1
                    (count drills)))
             (is (=? {:lib/type     :metabase.lib.drill-thru/drill-thru
@@ -225,43 +269,49 @@
                                     :value  "2022-12-01T00:00:00+02:00"}
                      :next-unit    :quarter}
                     drill))
-            (is (=? {:stages [{:aggregation [[:count {}]]
-                               :breakout    [[:field {:temporal-unit :quarter} (meta/id :orders :created-at)]],
-                               :filters     [[:=
-                                              {}
-                                              [:field {:temporal-unit :year} (meta/id :orders :created-at)]
-                                              "2022-12-01T00:00:00+02:00"]]}]}
+            (is (=? {:stages expected-stages}
                     (lib/drill-thru query -1 nil drill)))))))))
 
 (deftest ^:parallel zoom-in-timeseries-unit-tower-test
   (doseq [[unit1 unit2] datetime-unit-pairs]
-    (testing (str "zoom-in.timeseries for a DateTime column should zoom from " unit1 " to " unit2)
-      (let [query (-> (lib/query meta/metadata-provider (meta/table-metadata :orders))
-                      (lib/aggregate (lib/count))
-                      (lib/breakout (-> (meta/field-metadata :orders :created-at)
-                                        (lib/with-temporal-bucket unit1))))]
-        (lib.drill-thru.tu/test-drill-application
-         {:click-type     :cell
-          :query-type     :aggregated
-          :custom-query   query
-          :custom-row     {"count"      100
-                           "CREATED_AT" "2022-12-09T11:22:33+02:00"}
-          :column-name    "count"
-          :drill-type     :drill-thru/zoom-in.timeseries
-          :expected       {:type         :drill-thru/zoom-in.timeseries
-                           :display-name (str "See this " (name unit1) " by " (name unit2))
-                           :dimension    {:column     {:name "CREATED_AT"}
-                                          :column-ref [:field {:temporal-unit unit1} (meta/id :orders :created-at)]
-                                          :value      "2022-12-09T11:22:33+02:00"}
-                           :next-unit    unit2}
-          :expected-query {:stages [{:source-table (meta/id :orders)
-                                     :aggregation  [[:count {}]]
-                                     :breakout     [[:field
-                                                     {:temporal-unit unit2}
-                                                     (meta/id :orders :created-at)]]
-                                     :filters      [[:= {}
-                                                     [:field {:temporal-unit unit1} (meta/id :orders :created-at)]
-                                                     "2022-12-09T11:22:33+02:00"]]}]}})))))
+    (let [base-query        (-> (lib/query meta/metadata-provider (meta/table-metadata :orders))
+                                (lib/aggregate (lib/count))
+                                (lib/breakout (-> (meta/field-metadata :orders :created-at)
+                                                  (lib/with-temporal-bucket unit1))))
+          multi-stage-query (lib/append-stage base-query)
+          count-col         (m/find-first #(= (:name %) "count")
+                                          (lib/returned-columns multi-stage-query))
+          _                 (is (some? count-col))
+          multi-stage-query (lib/filter multi-stage-query (lib/> count-col 0))]
+      (doseq [[query-type query] {"single-stage" base-query
+                                  "multi-stage"  multi-stage-query}]
+        (testing (str "zoom-in.timeseries for a DateTime column in a " query-type " query should zoom from " unit1 " to " unit2)
+          (let [expected-stages (cond-> [{:source-table (meta/id :orders)
+                                          :aggregation  [[:count {}]]
+                                          :breakout     [[:field
+                                                          {:temporal-unit unit2}
+                                                          (meta/id :orders :created-at)]]
+                                          :filters      [[:= {}
+                                                          [:field {:temporal-unit unit1} (meta/id :orders :created-at)]
+                                                          "2022-12-09T11:22:33+02:00"]]}]
+                                  ;; the multi-stage-query has an additional filter stage
+                                  (= query multi-stage-query) (conj {:filters
+                                                                     [[:> {} [:field {} "count"] 0]]}))]
+            (lib.drill-thru.tu/test-drill-application
+             {:click-type     :cell
+              :query-type     :aggregated
+              :custom-query   query
+              :custom-row     {"count"      100
+                               "CREATED_AT" "2022-12-09T11:22:33+02:00"}
+              :column-name    "count"
+              :drill-type     :drill-thru/zoom-in.timeseries
+              :expected       {:type         :drill-thru/zoom-in.timeseries
+                               :display-name (str "See this " (name unit1) " by " (name unit2))
+                               :dimension    {:column     {:name "CREATED_AT"}
+                                              :column-ref [:field {:temporal-unit unit1} (meta/id :orders :created-at)]
+                                              :value      "2022-12-09T11:22:33+02:00"}
+                               :next-unit    unit2}
+              :expected-query {:stages expected-stages}})))))))
 
 (deftest ^:parallel zoom-in-timeseries-unit-tower-test-2
   (doseq [[unit1 unit2] date-unit-pairs]

test/metabase/lib/drill_thru_test.cljc

@@ -554,11 +554,10 @@
                                              :month)))
             columns      (lib/returned-columns query)
             sum          (by-name columns "sum")
-            breakout     (by-name columns "CREATED_AT")
             sum-dim      {:column     sum
                           :column-ref (lib/ref sum)
                           :value      42295.12}
-            breakout-dim {:column     breakout
+            breakout-dim {:column     (first (lib/breakouts-metadata query))
                           :column-ref (first (lib/breakouts query))
                           :value      "2024-11-01T00:00:00Z"}
             context      (merge sum-dim