Sort drill thrus should come back for aggregate columns (#34389)

src/metabase/lib/drill_thru/distribution.cljc

@@ -24,6 +24,7 @@
   (when (and (lib.drill-thru.common/mbql-stage? query stage-number)
              column
              (nil? value)
+             (not= (:lib/source column) :source/aggregations)
              (not (lib.types.isa/primary-key? column))
              (not (lib.types.isa/structured?  column))
              (not (lib.types.isa/comment?     column))

src/metabase/lib/drill_thru/sort.cljc

@@ -3,6 +3,7 @@
    [metabase.lib.drill-thru.common :as lib.drill-thru.common]
    [metabase.lib.equality :as lib.equality]
    [metabase.lib.order-by :as lib.order-by]
+   [metabase.lib.ref :as lib.ref]
    [metabase.lib.schema :as lib.schema]
    [metabase.lib.schema.drill-thru :as lib.schema.drill-thru]
    [metabase.lib.types.isa :as lib.types.isa]
@@ -10,26 +11,30 @@
 
 (mu/defn sort-drill :- [:maybe ::lib.schema.drill-thru/drill-thru.sort]
   "Sorting on a clicked column."
-  [query                  :- ::lib.schema/query
-   stage-number           :- :int
-   {:keys [column value]} :- ::lib.schema.drill-thru/context]
+  [query                                :- ::lib.schema/query
+   stage-number                         :- :int
+   {:keys [column value], :as _context} :- ::lib.schema.drill-thru/context]
+  ;; if we have a context with a `:column`, but no `:value`...
   (when (and (lib.drill-thru.common/mbql-stage? query stage-number)
              column
              (nil? value)
-             (not (lib.types.isa/structured? column))
-             (or (:lib/source column)
-                 (:id column)))
-    (let [orderable  (->> (lib.order-by/orderable-columns query stage-number)
-                          (map :id)
-                          set)
-          this-order (first (for [[dir _opts field] (lib.order-by/order-bys query stage-number)
-                                  :when (lib.equality/find-matching-column query stage-number field [column])]
-                              dir))]
-      (when (orderable (:id column))
+             (not (lib.types.isa/structured? column)))
+    ;; and the column is orderable (appears in [[lib.order-by/orderable-columns]]), we can return a sort drill thru.
+    (when (lib.equality/find-matching-column query
+                                             stage-number
+                                             (lib.ref/ref column)
+                                             (lib.order-by/orderable-columns query stage-number))
+      ;; check and see if there is already a sort on this column. If there is, we should only suggest flipping the
+      ;; direction to the opposite of what it is now. If there is no existing sort, then return both directions as
+      ;; options.
+      (let [existing-direction (some (fn [[dir _opts field]]
+                                       (when (lib.equality/find-matching-column query stage-number field [column])
+                                         dir))
+                                     (lib.order-by/order-bys query stage-number))]
         {:lib/type        :metabase.lib.drill-thru/drill-thru
          :type            :drill-thru/sort
          :column          column
-         :sort-directions (case this-order
+         :sort-directions (case existing-direction
                             :asc  [:desc]
                             :desc [:asc]
                             [:asc :desc])}))))

test/metabase/lib/drill_thru/distribution_test.cljc

+(ns metabase.lib.drill-thru.distribution-test
+  (:require
+   [clojure.test :refer [deftest is testing]]
+   [medley.core :as m]
+   [metabase.lib.core :as lib]
+   [metabase.lib.drill-thru.distribution :as lib.drill-thru.distribution]
+   [metabase.lib.test-metadata :as meta]
+   #?@(:cljs ([metabase.test-runner.assert-exprs.approximately-equal]))))
+
+#?(:cljs (comment metabase.test-runner.assert-exprs.approximately-equal/keep-me))
+
+(deftest ^:parallel aggregate-column-test
+  (testing "Don't suggest distribution drill thrus for aggregate columns like `count(*)`"
+    (let [query     (-> (lib/query meta/metadata-provider (meta/table-metadata :orders))
+                        (lib/aggregate (lib/count))
+                        (lib/breakout (meta/field-metadata :orders :product-id)))
+          count-col (m/find-first (fn [col]
+                                    (= (:display-name col) "Count"))
+                                  (lib/returned-columns query))
+          context   {:column 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

 (ns metabase.lib.drill-thru.sort-test
   (:require
    [clojure.test :refer [are deftest is testing]]
+   [medley.core :as m]
    [metabase.lib.convert :as lib.convert]
    [metabase.lib.core :as lib]
    [metabase.lib.drill-thru.sort :as lib.drill-thru.sort]
@@ -44,3 +45,28 @@
     (is (=? {:stages [{:lib/type :mbql.stage/mbql
                        :order-by [[:desc {} [:field {} (meta/id :orders :id)]]]}]}
             (lib/drill-thru query -1 drill :desc)))))
+
+(deftest ^:parallel aggregate-column-e2e-test
+  (testing "Sort drills should be suggested/work for aggregate columns like count (#34185)"
+    (let [query     (-> (lib/query meta/metadata-provider (meta/table-metadata :orders))
+                        (lib/aggregate (lib/count))
+                        (lib/breakout (meta/field-metadata :orders :product-id)))
+          count-col (m/find-first (fn [col]
+                                    (= (:display-name col) "Count"))
+                                  (lib/returned-columns query))
+          context   {:column 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
+                 :type            :drill-thru/sort
+                 :column          {:name "count"}
+                 :sort-directions [:asc :desc]}
+                drill))
+        (testing "Apply the drill"
+          (is (=? {:stages [{:aggregation [[:count {}]]
+                             :breakout    [[:field {} (meta/id :orders :product-id)]]
+                             :order-by    [[:desc
+                                            {}
+                                            [:aggregation {} string?]]]}]}
+                  (lib/drill-thru query -1 drill :desc))))))))

test/metabase/lib/drill_thru_test.cljc

@@ -70,25 +70,27 @@
    (test-drill-applications query context 0))
 
   ([query context depth]
-   (doseq [drill (lib/available-drill-thrus query -1 context)
-           args  (drill-thru-test-args drill)]
-     (condp = (:type drill)
-       :drill-thru/pivot
-       (log/warnf "drill-thru-method is not yet implemented for :drill-thru/pivot (#33559)")
+   (testing "\nTest drill applications"
+     (doseq [drill (lib/available-drill-thrus query -1 context)
+             args  (drill-thru-test-args drill)]
+       (condp = (:type drill)
+         :drill-thru/pivot
+         (log/warnf "drill-thru-method is not yet implemented for :drill-thru/pivot (#33559)")
 
-       :drill-thru/underlying-records
-       (log/warnf "drill-thru-method is not yet implemented for :drill-thru/underlying-records (#34233)")
+         :drill-thru/underlying-records
+         (log/warnf "drill-thru-method is not yet implemented for :drill-thru/underlying-records (#34233)")
 
-       (testing (str "\ndrill =\n" (u/pprint-to-str drill)
-                     "\nargs =\n" (u/pprint-to-str args))
-         (try
-           (let [query' (apply lib/drill-thru query -1 drill args)]
-             (is (not (me/humanize (mc/validate ::lib.schema/query query'))))
-             (when (< depth test-drill-applications-max-depth)
-               (testing (str "\n\nDEPTH = " (inc depth) "\n\nquery =\n" (u/pprint-to-str query'))
-                 (test-drill-applications query' context (inc depth)))))
-           (catch #?(:clj Throwable :cljs :default) e
-             (is (not e)))))))))
+         (testing (str "\nquery =\n" (u/pprint-to-str query)
+                       "\ndrill =\n" (u/pprint-to-str drill)
+                       "\nargs =\n" (u/pprint-to-str args))
+           (try
+             (let [query' (apply lib/drill-thru query -1 drill args)]
+               (is (not (me/humanize (mc/validate ::lib.schema/query query'))))
+               (when (< depth test-drill-applications-max-depth)
+                 (testing (str "\n\nDEPTH = " (inc depth) "\n\nquery =\n" (u/pprint-to-str query'))
+                   (test-drill-applications query' context (inc depth)))))
+             (catch #?(:clj Throwable :cljs :default) e
+               (is (not e))))))))))
 
 (deftest ^:parallel table-view-available-drill-thrus-headers-pk-test
   (testing "column headers: click on"
@@ -439,6 +441,43 @@
                     (lib/available-drill-thrus query -1 context)))
             (test-drill-applications query context)))))))
 
+(deftest ^:parallel table-view-available-drill-thrus-aggregate-column-header-test
+  (let [query (-> (lib/query meta/metadata-provider (meta/table-metadata :orders))
+                  (lib/aggregate (lib/count))
+                  (lib/aggregate (lib/sum (meta/field-metadata :orders :tax)))
+                  (lib/aggregate (lib/max (meta/field-metadata :orders :discount)))
+                  (lib/breakout (meta/field-metadata :orders :product-id))
+                  (lib/breakout (-> (meta/field-metadata :orders :created-at)
+                                    (lib/with-temporal-bucket :month))))]
+    (testing "Drills for count aggregation"
+      (let [count-col (m/find-first (fn [col]
+                                      (= (:display-name col) "Count"))
+                                    (lib/returned-columns query))]
+        (is (some? count-col))
+        (let [context {:column count-col
+                       :value  nil}]
+          (is (=? [{:type   :drill-thru/sort
+                    :column {:name "count"}}
+                   {:type         :drill-thru/summarize-column
+                    :column       {:name "count"}
+                    :aggregations [:distinct :sum :avg]}]
+                  (lib/available-drill-thrus query -1 context)))
+          (test-drill-applications query context))))
+    (testing "Drills for max(discount) aggregation"
+      (let [max-of-discount-col (m/find-first (fn [col]
+                                                (= (:display-name col) "Max of Discount"))
+                                              (lib/returned-columns query))]
+        (is (some? max-of-discount-col))
+        (let [context {:column max-of-discount-col
+                       :value  nil}]
+          (is (=? [{:type   :drill-thru/sort
+                    :column {:display-name "Max of Discount"}}
+                   {:type         :drill-thru/summarize-column
+                    :column       {:display-name "Max of Discount"}
+                    :aggregations [:distinct :sum :avg]}]
+                  (lib/available-drill-thrus query -1 context)))
+          (test-drill-applications query context))))))
+
 ;; TODO: Restore this test once zoom-in and underlying-records are checked properly.
 #_(deftest ^:parallel histogram-available-drill-thrus-test
   (testing "histogram breakout view"

test/metabase/lib/test_metadata/graph_provider.cljc

@@ -2,7 +2,9 @@
   (:require
    [clojure.core.protocols]
    [medley.core :as m]
-   [metabase.lib.metadata.protocols :as lib.metadata.protocols]))
+   [metabase.lib.metadata.protocols :as lib.metadata.protocols]
+   #?@(:clj
+       ([pretty.core :as pretty]))))
 
 (defn- graph-database [metadata-graph]
   (dissoc metadata-graph :tables))
@@ -65,4 +67,11 @@
 
   clojure.core.protocols/Datafiable
   (datafy [_this]
-    (list `->SimpleGraphMetadataProvider metadata-graph)))
+    (list `->SimpleGraphMetadataProvider metadata-graph))
+
+  #?@(:clj
+      [pretty/PrettyPrintable
+       (pretty [_this]
+         (if (identical? metadata-graph @(requiring-resolve 'metabase.lib.test-metadata/metadata))
+           'metabase.lib.test-metadata/metadata-provider
+           (list `->SimpleGraphMetadataProvider metadata-graph)))]))