Handle METRICs inside expression aggregation clauses [ci drivers]

src/metabase/query_processor/macros.clj

 (ns metabase.query-processor.macros
+  "TODO - this namespace is ancient and written with MBQL '95 in mind, e.g. it is case-sensitive.
+   At some point this ought to be reworked to be case-insensitive and cleaned up."
   (:require [clojure.core.match :refer [match]]
+            [clojure.walk :as walk]
             [metabase.db :as db]
             [metabase.util :as u]))
 
@@ -46,56 +49,39 @@
 
 ;;; ------------------------------------------------------------ Metrics ------------------------------------------------------------
 
-(defn- merge-aggregation [aggregations new-ag]
-  (if (map? aggregations)
-    (recur [aggregations] new-ag)
-    (conj aggregations new-ag)))
-
-(defn- is-metric? [aggregation]
+(defn- metric? [aggregation]
   (match aggregation
     ["METRIC" (_ :guard integer?)] true
     _                              false))
 
 (defn- metric-id [metric]
-  (when (is-metric? metric)
+  (when (metric? metric)
     (second metric)))
 
-(defn- merge-aggregations {:style/indent 0} [query-dict [aggregation & more]]
-  (if-not aggregation
-    ;; no more aggregations? we're done
-    query-dict
-    ;; otherwise determine if this aggregation is a METRIC and recur
-    (let [metric-def (when (is-metric? aggregation)
-                       (db/select-one-field :definition 'Metric, :id (metric-id aggregation)))]
-      (recur (if-not metric-def
-               ;; not a metric, move to next aggregation
-               query-dict
-               ;; it *is* a metric, insert it into the query appropriately
-               (-> query-dict
-                   (update-in [:query :aggregation] merge-aggregation (:aggregation metric-def))
-                   (update-in [:query :filter] merge-filter-clauses (:filter metric-def))))
-             more))))
-
-(defn- remove-metrics [aggregations]
-  (if-not (and (sequential? aggregations)
-               (every? coll? aggregations))
-    (recur [aggregations])
-    (vec (for [ag    aggregations
-               :when (not (is-metric? ag))]
-           ag))))
+(defn- expand-metric [metric-clause filter-clauses-atom]
+  (let [{filter-clause :filter, ag-clause :aggregation} (db/select-one-field :definition 'Metric, :id (metric-id metric-clause))]
+    (when filter-clause
+      (swap! filter-clauses-atom conj filter-clause))
+    ag-clause))
+
+(defn- expand-metrics-in-ag-clause [query-dict filter-clauses-atom]
+  (walk/postwalk (fn [form]
+                   (if-not (metric? form)
+                     form
+                     (expand-metric form filter-clauses-atom)))
+                 query-dict))
+
+(defn- expand-metrics [query-dict]
+  (let [filter-clauses-atom (atom [])
+        query-dict          (expand-metrics-in-ag-clause query-dict filter-clauses-atom)]
+    (update-in query-dict [:query :filter] merge-filter-clauses @filter-clauses-atom)))
 
 (defn- macroexpand-metric [{{aggregations :aggregation} :query, :as query-dict}]
   (if-not (seq aggregations)
     ;; :aggregation is empty, so no METRIC to expand
     query-dict
-    ;; we have an aggregation clause, so lets see if we are using a METRIC
-    ;; (since `:aggregation` can be either single or multiple, wrap single ones so `merge-aggregations` can always assume input is multiple)
-    (merge-aggregations
-      (update-in query-dict [:query :aggregation] remove-metrics)
-      (if (and (sequential? aggregations)
-               (every? coll? aggregations))
-        aggregations
-        [aggregations]))))
+    ;; otherwise walk the query dict and expand METRIC clauses
+    (expand-metrics query-dict)))
 
 
 ;;; ------------------------------------------------------------ Middleware ------------------------------------------------------------

test/metabase/driver/druid_test.clj

 (ns metabase.driver.druid-test
   (:require [cheshire.core :as json]
             [expectations :refer :all]
+            [metabase.models.metric :refer [Metric]]
             [metabase.query-processor :as qp]
             [metabase.query-processor.expand :as ql]
             [metabase.query-processor-test :refer [rows rows+column-names]]
             [metabase.test.data :as data]
             [metabase.test.data.datasets :as datasets, :refer [expect-with-engine]]
+            [metabase.test.util :as tu]
             [metabase.timeseries-query-processor-test :as timeseries-qp-test]
-            [metabase.query :as q]))
+            [metabase.util :as u]))
 
 (def ^:const ^:private ^String native-query-1
   (json/generate-string
@@ -230,3 +232,18 @@
     (druid-query
       (ql/aggregation (ql/named (ql/- (ql/sum $venue_price) 41) "Sum-41"))
       (ql/breakout $venue_price))))
+
+;; check that we can handle METRICS (ick) inside expression aggregation clauses
+(expect-with-engine :druid
+  [["2" 1231.0]
+   ["3"  346.0]
+   ["4" 197.0]]
+  (timeseries-qp-test/with-flattened-dbdef
+    (tu/with-temp Metric [metric {:definition {:aggregation [:sum [:field-id (data/id :checkins :venue_price)]]
+                                               :filter      [:> [:field-id (data/id :checkins :venue_price)] 1]}}]
+      (rows (qp/process-query
+              {:database (data/id)
+               :type     :query
+               :query    {:source-table (data/id :checkins)
+                          :aggregation  [:+ ["METRIC" (u/get-id metric)] 1]
+                          :breakout     [(ql/breakout (ql/field-id (data/id :checkins :venue_price)))]}})))))

test/metabase/query_processor_test/expression_aggregations_test.clj

 (ns metabase.query-processor-test.expression-aggregations-test
   "Tests for expression aggregations."
   (:require [expectations :refer :all]
+            [metabase.models.metric :refer [Metric]]
+            [metabase.query-processor :as qp]
             [metabase.query-processor.expand :as ql]
             [metabase.query-processor-test :refer :all]
             [metabase.test.data :as data]
             [metabase.test.data.datasets :as datasets, :refer [*engine*]]
+            [metabase.test.util :as tu]
             [metabase.util :as u]))
 
 ;; sum, *
@@ -182,3 +185,19 @@
     (rows+column-names (data/run-query venues
                          (ql/aggregation (ql/named (ql/- (ql/sum $price) 41) "Sum-41"))
                          (ql/breakout $price)))))
+
+;; check that we can handle METRICS (ick) inside expression aggregation clauses
+(datasets/expect-with-engines (engines-that-support :expression-aggregations)
+  [[2 119]
+   [3  40]
+   [4  25]]
+  (tu/with-temp Metric [metric {:table_id   (data/id :venues)
+                                :definition {:aggregation [:sum [:field-id (data/id :venues :price)]]
+                                             :filter      [:> [:field-id (data/id :venues :price)] 1]}}]
+    (format-rows-by [int int]
+      (rows (qp/process-query
+              {:database (data/id)
+               :type     :query
+               :query    {:source-table (data/id :venues)
+                          :aggregation  [:+ ["METRIC" (u/get-id metric)] 1]
+                          :breakout     [(ql/breakout (ql/field-id (data/id :venues :price)))]}})))))

test/metabase/test/data.clj

@@ -105,6 +105,7 @@
   "Call `driver/process-query` on expanded inner QUERY, looking up the `Database` ID for the `source-table.`
 
      (run-query* (query (source-table 5) ...))"
+  {:style/indent 0}
   [query :- qi/Query]
   (qp/process-query (wrap-inner-query query)))
 

test/metabase/test/util.clj

@@ -124,7 +124,7 @@
   {:with-temp-defaults (fn [_] {:base_type :type/Text
                                 :name      (random-name)
                                 :position  1
-                                :table_id  (data/id :venues)})})
+                                :table_id  (data/id :checkins)})})
 
 (u/strict-extend (class Metric)
   WithTempDefaults
@@ -132,7 +132,7 @@
                                 :definition  {}
                                 :description "Lookin' for a blueberry"
                                 :name        "Toucans in the rainforest"
-                                :table_id    (data/id :venues)})})
+                                :table_id    (data/id :checkins)})})
 
 (u/strict-extend (class PermissionsGroup)
   WithTempDefaults
@@ -172,7 +172,7 @@
                                 :definition  {}
                                 :description "Lookin' for a blueberry"
                                 :name        "Toucans in the rainforest"
-                                :table_id    (data/id :venues)})})
+                                :table_id    (data/id :checkins)})})
 
 ;; TODO - `with-temp` doesn't return `Sessions`, probably because their ID is a string?