Fix `mu/fn` and `mu/defn` support for map schemas on key-value args (#47153)

* Fix #46846

* Don't use `-v` to check if script has arg since it doesn't work on all bash versions

* Fix typo

.cljfmt.edn

@@ -35,6 +35,7 @@
  :extra-indents
  {;; clojure.core stuff
   fn*                                                                                  [[:inner 0]]
+  let*                                                                                 [[:block 1]]
   with-meta                                                                            [[:default]]
   ;; library stuff. Libraries that specify their own `:style/indent` specs are included in the section after this
   ;; since it is generated automatically.

bin/cljfmt-updated.sh

@@ -14,7 +14,7 @@ set -euo pipefail
 script_dir=`dirname "${BASH_SOURCE[0]}"`
 cd "$script_dir/.."
 
-if [ -v 1 ]; then
+if [ -n "${1:-}" ]; then
     diff_target="$1"
 else
     diff_target="HEAD"

bin/kondo-updated.sh

@@ -13,7 +13,7 @@ set -euo pipefail
 script_dir=`dirname "${BASH_SOURCE[0]}"`
 cd "$script_dir/.."
 
-if [ -v 1 ]; then
+if [ -n "${1:-}" ]; then
     diff_target="$1"
 else
     diff_target="HEAD"

src/metabase/util/malli/defn.clj

@@ -90,7 +90,7 @@
         {attr-map :meta} parsed
         attr-map         (merge
                           {:arglists (list 'quote (deparameterized-arglists parsed))
-                           :schema   (mu.fn/fn-schema parsed)}
+                           :schema   (mu.fn/fn-schema parsed {:target :target/metadata})}
                           attr-map)
         docstring        (annotated-docstring parsed)
         instrument?      (mu.fn/instrument-ns? *ns*)]

src/metabase/util/malli/fn.clj

@@ -52,10 +52,21 @@
 (defn- arity-schema
   "Given a `fn` arity as parsed by [[SchematizedParams]] an `return-schema`, return an appropriate `:=>` schema for
   the arity."
-  [{:keys [args], :as _arity} return-schema]
-  [:=>
-   (:schema (md/parse (add-default-schemas args)))
-   return-schema])
+  [{:keys [args], :as _arity} return-schema {:keys [target], :as _options}]
+  (let [parsed       (md/parse (add-default-schemas args))
+        varargs-info (get-in parsed [:parsed :rest :arg :arg])
+        varargs-type (cond
+                       (= (first varargs-info) :map) :varargs/map
+                       (seq varargs-info)            :varargs/sequential)
+        schema       (case target
+                       :target/metadata        (if (= varargs-type :varargs/map)
+                                                 (vec (concat (butlast (:schema parsed)) [[:* :any]]))
+                                                 (:schema parsed))
+                       :target/instrumentation (:schema parsed))]
+    [:=>
+     (cond-> schema
+       varargs-type (vary-meta assoc :varargs/type varargs-type))
+     return-schema]))
 
 (def ^:private SchematizedParams
   "This is exactly the same as [[malli.experimental/SchematizedParams]], but it preserves metadata from the arglists."
@@ -105,16 +116,29 @@
 
 (defn fn-schema
   "Implementation for [[fn]] and [[metabase.util.malli.defn/defn]]. Given an unparsed parametered fn tail, extract the
-  annotations and return a `:=>` or `:function` schema."
-  [parsed]
-  (let [{:keys [return arities]}     parsed
-        return-schema                (:schema return :any)
-        [arities-type arities-value] arities]
-    (case arities-type
-      :single   (arity-schema arities-value return-schema)
-      :multiple (into [:function]
-                      (for [arity (:arities arities-value)]
-                        (arity-schema arity return-schema))))))
+  annotations and return a `:=>` or `:function` schema.
+
+  `options` can contain `:target` which is either
+
+  * `:target/metadata`: generate the schema to attach to the metadata for a [[metabase.util.malli.defn/defn]]. For
+    key-value varargs like `& {:as kvs}` get a schema like `[:* :any]` in this case since the args aren't parsed to a
+    map yet
+
+  * `:target/instrumentation`: generate a schema for use in generating the instrumented `fn` form. `& {:as kvs}` can
+    have a real map schema here."
+
+  ([parsed]
+   (fn-schema parsed {:target :target/instrumentation}))
+
+  ([parsed options]
+   (let [{:keys [return arities]}     parsed
+         return-schema                (:schema return :any)
+         [arities-type arities-value] arities]
+     (case arities-type
+       :single   (arity-schema arities-value return-schema options)
+       :multiple (into [:function]
+                       (for [arity (:arities arities-value)]
+                         (arity-schema arity return-schema options)))))))
 
 (defn- deparameterized-arity [{:keys [body args prepost], :as _arity}]
   (concat
@@ -189,41 +213,43 @@
   (validate error-context schema value ::invalid-output)
   value)
 
-(defn- varargs-schema? [[_cat & args :as _input-schema]]
-  (letfn [(star-schema? [schema]
-            (and (sequential? schema)
-                 (= (first schema) :*)))]
-    (star-schema? (last args))))
+(defn- varargs-type [input-schema]
+  (-> input-schema meta :varargs/type))
 
 (defn- input-schema-arg-names [[_cat & args :as input-schema]]
-  (let [varargs?    (varargs-schema? input-schema)
-        normal-args (if varargs?
-                      (butlast args)
-                      args)]
+  (let [varargs-type (varargs-type input-schema)
+        normal-args  (if varargs-type
+                       (butlast args)
+                       args)]
     (concat
      (for [n (range (count normal-args))]
        (symbol (str (char (+ (int \a) n)))))
-     (when varargs?
-       ['more]))))
+     (case varargs-type
+       :varargs/sequential ['more]
+       :varargs/map        ['kvs]
+       nil))))
 
 (defn- input-schema->arglist [input-schema]
   (let [arg-names (input-schema-arg-names input-schema)]
-    (vec (if (varargs-schema? input-schema)
-           (concat (butlast arg-names) ['& (last arg-names)])
+    (vec (if-let [varargs-type (varargs-type input-schema)]
+           (concat (butlast arg-names) ['& (case varargs-type
+                                             :varargs/sequential (last arg-names)
+                                             :varargs/map        {:as (last arg-names)})])
            arg-names))))
 
 (defn- input-schema->validation-forms [error-context [_cat & schemas :as input-schema]]
   (let [arg-names (input-schema-arg-names input-schema)
-        schemas   (if (varargs-schema? input-schema)
+        schemas   (if (= (varargs-type input-schema) :varargs/sequential)
                     (concat (butlast schemas) [[:maybe (last schemas)]])
                     schemas)]
     (->> (map (core/fn [arg-name schema]
                 ;; 1. Skip checks against `:any` schema, there is no situation where it would fail.
                 ;;
-                ;; 2. Skip checks against the default varargs schema, there is no situation where [:maybe [:* :any]] is
+                ;; 2. Skip checks against the default varargs schemas, there is no situation where [:maybe [:* :any]] is
                 ;; going to fail.
-                (when-not (= schema (if (= arg-name 'more)
-                                      [:maybe [:* :any]]
+                (when-not (= schema (condp = arg-name
+                                      'more [:maybe [:* :any]]
+                                      'kvs  [:* :any]
                                       :any))
                   `(validate-input ~error-context ~schema ~arg-name)))
               arg-names
@@ -232,7 +258,7 @@
 
 (defn- input-schema->application-form [input-schema]
   (let [arg-names (input-schema-arg-names input-schema)]
-    (if (varargs-schema? input-schema)
+    (if (= (varargs-type input-schema) :varargs/sequential)
       (list* `apply '&f arg-names)
       (list* '&f arg-names))))
 
@@ -329,7 +355,8 @@
 
     {:fn-name 'whatever/my-multimethod, :dispatch-value :field}
 
-  If compiled in a namespace in [[namespaces-toskip]], during `config/is-prod?`, it will be emitted as a vanilla clojure.core/fn form.
+  If compiled in a namespace in [[namespaces-toskip]], during `config/is-prod?`, it will be emitted as a vanilla
+  clojure.core/fn form.
 
   Known issue: this version of `fn` does not capture the optional function name and make it available, e.g. you can't
   do

test/metabase/util/malli/fn_test.clj

@@ -74,40 +74,40 @@
                           (walk/macroexpand-all (mu.fn/instrumented-fn-form {} (mu.fn/parse-fn-tail form))))
     '([x :- :int y])
     '(let* [&f (fn* ([x y]))]
-           (fn* ([a b]
-                 (try
-                   (metabase.util.malli.fn/validate-input {} :int a)
-                   (&f a b)
-                   (catch java.lang.Exception error
-                     (throw (metabase.util.malli.fn/fixup-stacktrace error)))))))
+       (fn* ([a b]
+             (try
+               (metabase.util.malli.fn/validate-input {} :int a)
+               (&f a b)
+               (catch java.lang.Exception error
+                 (throw (metabase.util.malli.fn/fixup-stacktrace error)))))))
 
     '(:- :int [x :- :int y])
     '(let* [&f (fn* ([x y]))]
-           (fn* ([a b]
-                 (try
-                   (metabase.util.malli.fn/validate-input {} :int a)
-                   (metabase.util.malli.fn/validate-output {} :int (&f a b))
-                   (catch java.lang.Exception error
-                     (throw (metabase.util.malli.fn/fixup-stacktrace error)))))))
+       (fn* ([a b]
+             (try
+               (metabase.util.malli.fn/validate-input {} :int a)
+               (metabase.util.malli.fn/validate-output {} :int (&f a b))
+               (catch java.lang.Exception error
+                 (throw (metabase.util.malli.fn/fixup-stacktrace error)))))))
 
     '(:- :int [x :- :int y] (+ x y))
     '(let* [&f (fn* ([x y] (+ x y)))]
-           (fn* ([a b]
-                 (try
-                   (metabase.util.malli.fn/validate-input {} :int a)
-                   (metabase.util.malli.fn/validate-output {} :int (&f a b))
-                   (catch java.lang.Exception error
-                     (throw (metabase.util.malli.fn/fixup-stacktrace error)))))))
+       (fn* ([a b]
+             (try
+               (metabase.util.malli.fn/validate-input {} :int a)
+               (metabase.util.malli.fn/validate-output {} :int (&f a b))
+               (catch java.lang.Exception error
+                 (throw (metabase.util.malli.fn/fixup-stacktrace error)))))))
 
     '([x :- :int y] {:pre [(int? x)]})
     '(let* [&f (fn* ([x y]
                      {:pre [(int? x)]}))]
-           (fn* ([a b]
-                 (try
-                   (metabase.util.malli.fn/validate-input {} :int a)
-                   (&f a b)
-                   (catch java.lang.Exception error
-                     (throw (metabase.util.malli.fn/fixup-stacktrace error)))))))
+       (fn* ([a b]
+             (try
+               (metabase.util.malli.fn/validate-input {} :int a)
+               (&f a b)
+               (catch java.lang.Exception error
+                 (throw (metabase.util.malli.fn/fixup-stacktrace error)))))))
 
     '(:- :int
          ([x] (inc x))
@@ -116,18 +116,18 @@
                      (inc x))
                  ([x y]
                   (+ x y)))]
-           (fn*
-             ([a]
-              (try
-                (metabase.util.malli.fn/validate-output {} :int (&f a))
-                (catch java.lang.Exception error
-                  (throw (metabase.util.malli.fn/fixup-stacktrace error)))))
-             ([a b]
-              (try
-                (metabase.util.malli.fn/validate-input {} :int a)
-                (metabase.util.malli.fn/validate-output {} :int (&f a b))
-                (catch java.lang.Exception error
-                  (throw (metabase.util.malli.fn/fixup-stacktrace error)))))))))
+       (fn*
+         ([a]
+          (try
+            (metabase.util.malli.fn/validate-output {} :int (&f a))
+            (catch java.lang.Exception error
+              (throw (metabase.util.malli.fn/fixup-stacktrace error)))))
+         ([a b]
+          (try
+            (metabase.util.malli.fn/validate-input {} :int a)
+            (metabase.util.malli.fn/validate-output {} :int (&f a b))
+            (catch java.lang.Exception error
+              (throw (metabase.util.malli.fn/fixup-stacktrace error)))))))))
 
 (deftest ^:parallel fn-test
   (let [f (mu.fn/fn :- :int [y] y)]
@@ -167,13 +167,13 @@
     (is (= '(let* [&f (clojure.core/fn
                         [path opts & {:keys [token-check?], :or {token-check? true}}]
                         (merge {:path path, :token-check? token-check?} opts))]
-                  (clojure.core/fn
-                    ([a b & more]
-                     (try
-                       (metabase.util.malli.fn/validate-input {:fn-name 'my-fn} :map b)
-                       (clojure.core/apply &f a b more)
-                       (catch java.lang.Exception error
-                         (throw (metabase.util.malli.fn/fixup-stacktrace error)))))))
+              (clojure.core/fn
+                ([a b & {:as kvs}]
+                 (try
+                   (metabase.util.malli.fn/validate-input {:fn-name 'my-fn} :map b)
+                   (&f a b kvs)
+                   (catch java.lang.Exception error
+                     (throw (metabase.util.malli.fn/fixup-stacktrace error)))))))
            (macroexpand form)))
     (is (= [:=>
             [:cat :any :map [:* :any]]
@@ -190,6 +190,86 @@
     (is (= {:path "path", :token-check? false, :opts true}
            (f "path" {:opts true} :token-check? false)))))
 
+(deftest ^:parallel varargs-schema-test
+  (testing "Schemas on the varargs should work (#46864)"
+    (let [form '(metabase.util.malli.fn/fn my-plus :- :int
+                  [x :- :int
+                   y :- :int
+                   & more :- [:* :int]]
+                  (reduce + (list* x y more)))]
+      (is (= '(let* [&f (clojure.core/fn [x y & more]
+                          (reduce + (list* x y more)))]
+                (clojure.core/fn
+                  ([a b & more]
+                   (try
+                     (metabase.util.malli.fn/validate-input {:fn-name 'my-plus} :int a)
+                     (metabase.util.malli.fn/validate-input {:fn-name 'my-plus} :int b)
+                     (metabase.util.malli.fn/validate-input {:fn-name 'my-plus} [:maybe [:* :int]] more)
+                     (clojure.core/->>
+                      (clojure.core/apply &f a b more)
+                      (metabase.util.malli.fn/validate-output {:fn-name 'my-plus} :int))
+                     (catch java.lang.Exception error
+                       (throw (metabase.util.malli.fn/fixup-stacktrace error)))))))
+
+             (macroexpand form)))
+      (is (= [:=>
+              [:cat :int :int [:* :int]]
+              :int]
+             (mu.fn/fn-schema (mu.fn/parse-fn-tail (rest form))))))
+    (let [f (mu.fn/fn my-plus :- :int
+              [x :- :int
+               y :- :int
+               & more :- [:* :int]]
+              (reduce + (list* x y more)))]
+      (is (= 3
+             (f 1 2)))
+      (is (= 6
+             (f 1 2 3)))
+      (is (= 10
+             (f 1 2 3 4))))))
+
+(deftest ^:parallel varargs-map-schema-test
+  (testing "Schemas on the varargs map should work (#46864)"
+    (let [form '(metabase.util.malli.fn/fn my-plus :- :map
+                  [x :- :int
+                   y :- :int
+                   & {:as options} :- [:map [:integer? :boolean]]]
+                  {:options options, :output (+ x y)})]
+      (is (= '(let* [&f (clojure.core/fn [x y & {:as options}]
+                          {:options options, :output (+ x y)})]
+                (clojure.core/fn
+                  ([a b & {:as kvs}]
+                   (try
+                     (metabase.util.malli.fn/validate-input {:fn-name 'my-plus} :int a)
+                     (metabase.util.malli.fn/validate-input {:fn-name 'my-plus} :int b)
+                     (metabase.util.malli.fn/validate-input {:fn-name 'my-plus} [:map [:integer? :boolean]] kvs)
+                     (clojure.core/->>
+                      (&f a b kvs)
+                      (metabase.util.malli.fn/validate-output {:fn-name 'my-plus} :map))
+                     (catch java.lang.Exception error
+                       (throw (metabase.util.malli.fn/fixup-stacktrace error)))))))
+
+             (macroexpand form)))
+      (is (= [:=>
+              [:cat :int :int [:* :any]]
+              :map]
+             (mu.fn/fn-schema (mu.fn/parse-fn-tail (rest form)) {:target :target/metadata}))))
+    (let [f (mu.fn/fn my-plus :- :map
+              [x :- :int
+               y :- :int
+               & {:as options} :- [:map [:integer? :boolean]]]
+              {:options options, :output (+ x y)})]
+      (is (= {:options {:integer? true}, :output 3}
+             (f 1 2 :integer? true)))
+      (is (thrown-with-msg?
+           clojure.lang.ExceptionInfo
+           #"Invalid input:"
+           (f 1 2 :integer? 1)))
+      (is (thrown-with-msg?
+           clojure.lang.ExceptionInfo
+           #"Invalid input:"
+           (f 1 2))))))
+
 (deftest ^:parallel parse-fn-tail-preserve-metadata-test
   (is (= 'Integer
          (-> '(^{:private true} add-ints :- :int ^{:tag Integer} [x :- :int y :- :int] (+ x y))