Add an aggregate function to add aggregations to queries (#29225)

* Add an aggregate function to add aggregations to queries

* Call ->op-arg on aggregate input

* Allow any in aggregations for now because expressions on aggregates are allowed

* We have to recurse into clauses

.clj-kondo/macros/metabase/lib/common.clj

 (ns macros.metabase.lib.common)
 
 (defmacro defop
-  [op-name argv]
-  `(defn ~op-name "docstring"
-     ([~@argv]
-      ~(remove #{'&} argv))
-     ([~'query ~'stage ~@argv]
-      ~(into ['query 'stage] (remove #{'&} argv)))))
+  [op-name & argvecs]
+  `(do
+     (defn ~(symbol (str (name op-name) "-clause"))
+       ~(format "Create a standalone clause of type `%s`." (name op-name))
+       ~@(for [argvec argvecs
+               :let [arglist-expr (if (contains? (set argvec) '&)
+                                    (cons 'list* (remove #{'&} argvec))
+                                    argvec)]]
+           `([~'query ~'stage-number ~@argvec]
+             [~(keyword op-name) ~'query ~'stage-number ~@arglist-expr])))
+
+     (defn ~op-name
+       ~(format "Create a closure of clause of type `%s`." (name op-name))
+       ~@(for [argvec argvecs
+               :let [arglist-expr (if (contains? (set argvec) '&)
+                                    (cons 'list* (remove #{'&} argvec))
+                                    argvec)]]
+           `([~@argvec]
+             (fn ~(symbol (str (name op-name) "-closure"))
+               [~'query ~'stage-number]
+               (apply ~(symbol (str (name op-name) "-clause")) ~'query ~'stage-number ~@arglist-expr)))))))

src/metabase/lib/aggregation.cljc

@@ -2,26 +2,12 @@
   (:refer-clojure :exclude [count distinct max min])
   (:require
    [metabase.lib.common :as lib.common]
-   [metabase.lib.options :as lib.options]
+   [metabase.lib.schema :as lib.schema]
+   [metabase.lib.util :as lib.util]
    [metabase.util.malli :as mu])
   #?(:cljs (:require-macros [metabase.lib.aggregation])))
 
-(mu/defn count :- [:or
-                   fn?
-                   :mbql.clause/count]
-  "Create a `count` filter clause."
-  ([]
-   #_{:clj-kondo/ignore [:redundant-fn-wrapper]}
-   (fn [query stage-number]
-     (count query stage-number)))
-  ([x]
-   (fn [query stage-number]
-     (count query stage-number x)))
-  ([_query _stage-number]
-   (lib.options/ensure-uuid [:count]))
-  ([query stage-number x]
-   (lib.options/ensure-uuid [:count (lib.common/->op-arg query stage-number x)])))
-
+(lib.common/defop count [] [x])
 (lib.common/defop avg [x])
 (lib.common/defop count-where [x y])
 (lib.common/defop distinct [x])
@@ -33,3 +19,15 @@
 (lib.common/defop stddev [x])
 (lib.common/defop sum [x])
 (lib.common/defop sum-where [x y])
+
+(mu/defn aggregate :- ::lib.schema/query
+  "Adds an aggregation to query."
+  ([query an-aggregate-clause]
+   (aggregate query -1 an-aggregate-clause))
+  ([query stage-number an-aggregate-clause]
+   (let [stage-number (or stage-number -1)]
+     (lib.util/update-query-stage
+       query stage-number
+       update :aggregations
+       (fn [aggregations]
+         (conj (vec aggregations) (lib.common/->op-arg query stage-number an-aggregate-clause)))))))

src/metabase/lib/common.cljc

@@ -3,7 +3,9 @@
    [metabase.lib.dispatch :as lib.dispatch]
    [metabase.lib.field :as lib.field]
    #_{:clj-kondo/ignore [:unused-namespace]}
-   [metabase.lib.options :as lib.options])
+   [metabase.lib.options :as lib.options]
+   #_{:clj-kondo/ignore [:unused-namespace]}
+   [metabase.util.malli :as mu])
   #?(:cljs (:require-macros [metabase.lib.common])))
 
 (defmulti ->op-arg
@@ -13,8 +15,10 @@
     (lib.dispatch/dispatch-value x)))
 
 (defmethod ->op-arg :default
-  [_query _stage-number x]
-  x)
+  [query stage-number x]
+  (if (vector? x)
+    (mapv #(->op-arg query stage-number %) x)
+    x))
 
 (defmethod ->op-arg :metadata/field
   [query stage-number field-metadata]
@@ -28,21 +32,31 @@
    (defmacro defop
      "Defines a clause creating function with given args.
       Calling the clause without query and stage produces a fn that can be resolved later."
-     [op-name argvec]
+     [op-name & argvecs]
      {:pre [(symbol? op-name)
-            (vector? argvec) (every? symbol? argvec)
-            (not-any? #{'query 'stage-number} argvec)]}
-     (let [arglist-expr (if (contains? (set argvec) '&)
-                          (cons 'list* (remove #{'&} argvec))
-                          argvec)]
-       `(mu/defn ~op-name :- [:or fn? ~(keyword "mbql.clause" (name op-name))]
-          ~(format "Create a clause of type `%s`." (name op-name))
-          ([~@argvec]
-           (fn [~'query ~'stage-number]
-             (~op-name ~'query ~'stage-number ~@argvec)))
-          ([~'query ~'stage-number ~@argvec]
-           (-> (into [~(keyword op-name)]
-                     (map (fn [~'arg]
-                            (->op-arg ~'query ~'stage-number ~'arg)))
-                     ~arglist-expr)
-               lib.options/ensure-uuid))))))
+            (every? vector? argvecs) (every? #(every? symbol? %) argvecs)
+            (every? #(not-any? #{'query 'stage-number} %) argvecs)]}
+     `(do
+          (mu/defn ~(symbol (str (name op-name) "-clause")) :- ~(keyword "mbql.clause" (name op-name))
+            ~(format "Create a standalone clause of type `%s`." (name op-name))
+            ~@(for [argvec argvecs
+                    :let [arglist-expr (if (contains? (set argvec) '&)
+                                         (cons `list* (remove #{'&} argvec))
+                                         argvec)]]
+                `([~'query ~'stage-number ~@argvec]
+                  (-> (into [~(keyword op-name)]
+                            (map (fn [~'arg]
+                                   (->op-arg ~'query ~'stage-number ~'arg)))
+                            ~arglist-expr)
+                      lib.options/ensure-uuid))))
+
+          (mu/defn ~op-name :- fn?
+            ~(format "Create a closure of clause of type `%s`." (name op-name))
+            ~@(for [argvec argvecs
+                    :let [arglist-expr (if (contains? (set argvec) '&)
+                                         (filterv (complement #{'&}) argvec)
+                                         (conj argvec []))]]
+               `([~@argvec]
+                 (fn ~(symbol (str (name op-name) "-closure"))
+                   [~'query ~'stage-number]
+                   (apply ~(symbol (str (name op-name) "-clause")) ~'query ~'stage-number ~@arglist-expr))))))))

src/metabase/lib/core.cljc

@@ -23,6 +23,7 @@
 
 (shared.ns/import-fns
   [lib.aggregation
+   aggregate
    count
    avg
    count-where

src/metabase/lib/schema/expression.cljc

@@ -77,7 +77,7 @@
      "Registers [[tag]] with [[type-of*]] in terms of its first incoming [[expr]].
       Useful for clauses that are polymorphic on their argument."
      [tag]
-     `(defmethod type-of* ~tag [[_tag# _opts# expr# & _#]] (type-of expr#))))
+     `(defmethod type-of* ~tag [[_tag# _opts# expr#]] (type-of expr#))))
 
 (defn- expression-schema
   "Schema that matches the following rules:

test/metabase/lib/aggregation_test.cljc

@@ -13,12 +13,6 @@
                   expected-args)
             (f {:lib/metadata meta/metadata} -1)))))
 
-(defn- is-clause? [op tag args expected-args]
-  (is (=? (into [tag
-                 {:lib/uuid string?}]
-                 expected-args)
-          (apply op {:lib/metadata meta/metadata} -1 args))))
-
 (deftest ^:parallel aggregation-test
   (let [q1 (lib/query-for-table-name meta/metadata-provider "CATEGORIES")
         venues-category-id-metadata (lib.metadata/field q1 nil "VENUES" "CATEGORY_ID")
@@ -26,10 +20,7 @@
     (testing "count"
       (testing "without query/stage-number, return a function for later resolution"
         (is-fn? lib/count :count [] [])
-        (is-fn? lib/count :count [venues-category-id-metadata] [venue-field-check]))
-      (testing "with query/stage-number, return clause right away"
-        (is-clause? lib/count :count [] [])
-        (is-clause? lib/count :count [venues-category-id-metadata] [venue-field-check])))
+        (is-fn? lib/count :count [venues-category-id-metadata] [venue-field-check])))
     (testing "single arg aggregations"
       (doseq [[op tag] [[lib/avg :avg]
                         [lib/max :max]
@@ -39,6 +30,17 @@
                         [lib/stddev :stddev]
                         [lib/distinct :distinct]]]
         (testing "without query/stage-number, return a function for later resolution"
-          (is-fn? op tag [venues-category-id-metadata] [venue-field-check]))
-        (testing "with query/stage-number, return clause right away"
-          (is-clause? op tag [venues-category-id-metadata] [venue-field-check]))))))
+          (is-fn? op tag [venues-category-id-metadata] [venue-field-check]))))))
+
+(deftest ^:parallel join-test
+  (is (=? {:lib/type :mbql/query,
+           :database (meta/id) ,
+           :type :pipeline,
+           :stages [{:lib/type :mbql.stage/mbql,
+                     :source-table (meta/id :venues) ,
+                     :lib/options {:lib/uuid string?},
+                     :aggregations [[:sum {:lib/uuid string?}
+                                     [:field {:base-type :type/Integer, :lib/uuid string?} (meta/id :venues :category-id)]]]}]}
+          (-> (lib/query-for-table-name meta/metadata-provider "VENUES")
+              (lib/aggregate (lib/sum (lib/field "VENUES" "CATEGORY_ID")))
+              (dissoc :lib/metadata)))))