Fix sorting by `Offset()` custom expression aggregation (#42737)

* Fix sorting by `Offset()` custom expression aggregation

* Unskip test

* Revert breaking change to `mt/id` and related

e2e/test/scenarios/question/offset.cy.spec.ts

@@ -106,7 +106,7 @@ describe("scenarios > question > offset", () => {
     ]);
   });
 
-  it.skip("works with a single breakout and sorting by aggregation (metabase#42554)", () => {
+  it("works with a single breakout and sorting by aggregation (metabase#42554)", () => {
     const query: StructuredQuery = {
       "source-table": ORDERS_ID,
       aggregation: [OFFSET_SUM_TOTAL_AGGREGATION],

src/metabase/driver.clj

@@ -16,6 +16,7 @@
    [metabase.util :as u]
    [metabase.util.i18n :refer [deferred-tru tru]]
    [metabase.util.log :as log]
+   [metabase.util.malli :as mu]
    [potemkin :as p]
    [toucan2.core :as t2]))
 
@@ -682,8 +683,8 @@
   dispatch-on-initialized-driver
   :hierarchy #'hierarchy)
 
-(defmethod escape-alias ::driver
-  [_driver alias-name]
+(mu/defmethod escape-alias ::driver :- :string
+  [_driver alias-name :- :string]
   (driver.impl/truncate-alias alias-name))
 
 (defmulti humanize-connection-error-message

src/metabase/driver/sql/query_processor.clj

@@ -982,6 +982,9 @@
     #{:+ :- :* :/}
     (->honeysql driver &match)
 
+    [:offset (options :guard :name) _expr _n]
+    (->honeysql driver (h2x/identifier :field-alias (:name options)))
+
     ;; for everything else just use the name of the aggregation as an identifer, e.g. `:sum`
     ;;
     ;; TODO -- I don't think we will ever actually get to this anymore because everything should have been given a name
@@ -1630,9 +1633,10 @@
       add/add-alias-info
       nest-query/nest-expressions))
 
-(defn mbql->honeysql
+(mu/defn mbql->honeysql :- :map
   "Build the HoneySQL form we will compile to SQL and execute."
-  [driver {inner-query :query}]
+  [driver               :- :keyword
+   {inner-query :query} :- :map]
   (binding [driver/*driver* driver]
     (let [inner-query (preprocess driver inner-query)]
       (log/tracef "Compiling MBQL query\n%s" (u/pprint-to-str 'magenta inner-query))
@@ -1641,10 +1645,13 @@
 
 ;;;; MBQL -> Native
 
-(defn mbql->native
+(mu/defn mbql->native :- [:map
+                          [:query  :string]
+                          [:params [:maybe [:sequential :any]]]]
   "Transpile MBQL query into a native SQL statement. This is the `:sql` driver implementation
   of [[driver/mbql->native]] (actual multimethod definition is in [[metabase.driver.sql]]."
-  [driver outer-query]
+  [driver      :- :keyword
+   outer-query :- :map]
   (let [honeysql-form (mbql->honeysql driver outer-query)
         [sql & args]  (format-honeysql driver honeysql-form)]
     {:query sql, :params args}))

src/metabase/query_processor/util/add_alias_info.clj

@@ -79,7 +79,8 @@
   []
   (let [unique-name-fn (lib.util/unique-name-generator (qp.store/metadata-provider))]
     (fn unique-alias-fn [position original-alias]
-      {:pre (string? original-alias)}
+      (assert (string? original-alias)
+              (format "unique-alias-fn expected string, got: %s" (pr-str original-alias)))
       (unique-name-fn position (driver/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
@@ -405,6 +406,29 @@
              {::desired-alias (unique-alias-fn position (field-desired-alias inner-query field-clause expensive-info))
               ::position      position}))))
 
+(defmulti ^:private aggregation-name
+  {:arglists '([mbql-clause])}
+  (fn [x]
+    (when (mbql.u/mbql-clause? x)
+      (first x))))
+
+;;; make sure we have an `:aggregation-options` or other fully-preprocessed aggregation clause (i.e., `:offset`) 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
+(defmethod aggregation-name :default
+  [ag-clause]
+  (throw (ex-info (format "Expected :aggregation-options or other fully-preprocessed aggregation, got %s."
+                          (pr-str ag-clause))
+                  {:clause ag-clause})))
+
+(mu/defmethod aggregation-name :aggregation-options :- :string
+  [[_aggregation-options _wrapped-ag opts]]
+  (:name opts))
+
+(mu/defmethod aggregation-name :offset :- :string
+  [[_offset opts _expr _n]]
+  (:name opts))
+
 (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)]
@@ -415,13 +439,7 @@
                       {:type   qp.error-type/invalid-query
                        :clause ag-ref-clause
                        :query  inner-query})))
-    (let [[_ _ {ag-name :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})))
+    (let [ag-name (aggregation-name (nth aggregations index))]
       {::desired-alias (unique-alias-fn position ag-name)
        ::position      position})))
 

test/metabase/query_processor_test/offset_test.clj

@@ -4,12 +4,14 @@
    [clojure.test :refer :all]
    [java-time.api :as t]
    [medley.core :as m]
+   [metabase.driver :as driver]
    [metabase.lib.core :as lib]
    [metabase.lib.metadata :as lib.metadata]
    [metabase.lib.metadata.jvm :as lib.metadata.jvm]
    [metabase.query-processor :as qp]
    [metabase.query-processor.test-util :as qp.test-util]
-   [metabase.test :as mt]))
+   [metabase.test :as mt]
+   [metabase.test.data.interface :as tx]))
 
 (defn- ->local-date [t]
   (t/local-date
@@ -283,3 +285,27 @@
                (mt/formatted-rows
                 [->local-date 2.0]
                 (qp/process-query (mt/obj->json->obj query)))))))))
+
+(deftest ^:parallel sort-by-offset-aggregation-test
+  (testing "Should be able to sort by an Offset() expression in an aggregation (#42554)"
+    (mt/test-drivers (mt/normal-drivers-with-feature :window-functions/offset)
+      (let [query (-> (mt/mbql-query orders
+                        {:breakout    [[:field %created_at {:base-type :type/DateTime, :temporal-unit :month}]],
+                         :aggregation [[:offset
+                                        {:display-name "X", :name "X", :lib/uuid "59590ea6-b853-4c2f-99dd-a3b0f5662fa7"}
+                                        [:sum [:field %total {:base-type :type/Float}]]
+                                        -1]]
+                         :order-by [[:asc [:aggregation 0]]]
+                         :limit    3})
+                      (assoc-in [:middleware :format-rows?] false))]
+        (mt/with-native-query-testing-context query
+          (is (= (if (tx/sorts-nil-first? driver/*driver* :type/Float)
+                   [[#t "2016-04-01" nil]
+                    [#t "2016-05-01" 52.76]
+                    [#t "2016-06-01" 1265.73]]
+                   [[#t "2016-05-01" 52.76]
+                    [#t "2016-06-01" 1265.73]
+                    [#t "2016-07-01" 2072.92]])
+                 (mt/formatted-rows
+                  [->local-date 2.0]
+                  (qp/process-query query)))))))))