Applying sort-drill thru should REPLACE existing sort on same column (#34637)

src/metabase/lib/drill_thru/sort.cljc

 (ns metabase.lib.drill-thru.sort
   (:require
+   [medley.core :as m]
    [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.schema.order-by :as lib.schema.order-by]
    [metabase.lib.types.isa :as lib.types.isa]
    [metabase.util.malli :as mu]))
 
+(defn- orderable-column?
+  "Is `column` orderable? (Does it appear in [[lib.order-by/orderable-columns]]?)"
+  [query stage-number column]
+  (lib.equality/find-matching-column query
+                                     stage-number
+                                     (lib.ref/ref column)
+                                     (lib.order-by/orderable-columns query stage-number)))
+
+(mu/defn ^:private existing-order-by-clause :- [:maybe ::lib.schema.order-by/order-by]
+  [query stage-number column]
+  (m/find-first (fn [[_direction _opts expr, :as _asc-or-desc-clause]]
+                  (lib.equality/find-matching-column query stage-number expr [column]))
+                (lib.order-by/order-bys query stage-number)))
+
+(mu/defn ^:private existing-order-by-direction :- [:maybe ::lib.schema.order-by/direction]
+  [query stage-number column]
+  (when-let [[direction _opts _expr] (existing-order-by-clause query stage-number column)]
+    direction))
+
 (mu/defn sort-drill :- [:maybe ::lib.schema.drill-thru/drill-thru.sort]
   "Sorting on a clicked column."
   [query                                :- ::lib.schema/query
@@ -19,18 +40,12 @@
              column
              (nil? value)
              (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))
+    ;; ...and the column is orderable, we can return a sort drill-thru.
+    (when (orderable-column? query stage-number column)
       ;; 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))]
+      (let [existing-direction (existing-order-by-direction query stage-number column)]
         {:lib/type        :metabase.lib.drill-thru/drill-thru
          :type            :drill-thru/sort
          :column          column
@@ -46,8 +61,13 @@
   ([query                        :- ::lib.schema/query
     stage-number                 :- :int
     {:keys [column], :as _drill} :- ::lib.schema.drill-thru/drill-thru.sort
-    direction                    :- [:enum :asc :desc]]
-   (lib.order-by/order-by query stage-number column (keyword direction))))
+    direction                    :- ::lib.schema.order-by/direction]
+   ;; if you have an existing order by, the drill thru returned by [[sort-drill]] would only be one that would suggest
+   ;; changing it to the opposite direction, so we can safely assume we want to change the direction and
+   ;; use [[lib.order-by/change-direction]] here.
+   (if-let [existing-clause (existing-order-by-clause query stage-number column)]
+     (lib.order-by/change-direction query existing-clause)
+     (lib.order-by/order-by query stage-number column (keyword direction)))))
 
 (defmethod lib.drill-thru.common/drill-thru-info-method :drill-thru/sort
   [_query _stage-number {directions :sort-directions}]

test/metabase/lib/drill_thru/sort_test.cljc

@@ -70,3 +70,28 @@
                                             {}
                                             [:aggregation {} string?]]]}]}
                   (lib/drill-thru query -1 drill :desc))))))))
+
+(deftest ^:parallel remove-existing-sort-test
+  (testing "Applying sort to already sorted column should REPLACE original sort (#34497)"
+    ;; technically this query doesn't make sense, how are you supposed to have a max aggregation and then also order
+    ;; by a different column, but MBQL doesn't enforce that,
+    (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}
+          drill   (lib.drill-thru.sort/sort-drill query -1 context)]
+      (is (=? {:stages
+               [{:order-by [[:asc {} [:field {} (meta/id :orders :user-id)]]
+                            [:asc {} [:field {} (meta/id :orders :id)]]]}]}
+              query))
+      (is (=? {:lib/type        :metabase.lib.drill-thru/drill-thru
+               :type            :drill-thru/sort
+               :column          {:name "USER_ID"}
+               :sort-directions [:desc]}
+              drill))
+      (testing "We should REPLACE the original sort, as opposed to removing it and appending a new one"
+        (is (=? {:stages
+                 [{:order-by [[:desc {} [:field {} (meta/id :orders :user-id)]]
+                              [:asc {} [:field {} (meta/id :orders :id)]]]}]}
+                (lib/drill-thru query -1 drill :desc)))))))