MLv2: TS wrappers for `fields` and `withFields`; minor tweaks (#31052)

* MLv2: TS wrappers for `fields` and `withFields`; minor tweaks

* Fix type declaration

frontend/src/metabase-lib/fields.ts

+import * as ML from "cljs/metabase.lib.js";
+import type { Clause, ColumnMetadata, Query } from "./types";
+
+const DEFAULT_STAGE_INDEX = -1;
+
+export function fields(
+  query: Query,
+  stageIndex = DEFAULT_STAGE_INDEX,
+): Clause[] {
+  return ML.fields(query, stageIndex);
+}
+
+export function withFields(
+  query: Query,
+  stageIndex: number,
+  newFields: ColumnMetadata[],
+): Query {
+  return ML.with_fields(query, stageIndex, newFields);
+}

frontend/src/metabase-lib/v2.ts

@@ -4,6 +4,7 @@ export * from "./column_types";
 export * from "./comparison";
 export * from "./metadata";
 export * from "./breakout";
+export * from "./fields";
 export * from "./limit";
 export * from "./order_by";
 export * from "./query";

src/metabase/lib/field.cljc

@@ -442,7 +442,7 @@
     (:name field-metadata)))
 
 (defn with-fields
-  "Specify the `:fields` for a query."
+  "Specify the `:fields` for a query. Pass `nil` or an empty sequence to remove `:fields`."
   ([xs]
    (fn [query stage-number]
      (with-fields query stage-number xs)))
@@ -456,10 +456,11 @@
                                    (x query stage-number)
                                    x)))
                   xs)]
-     (lib.util/update-query-stage query stage-number assoc :fields xs))))
+     (lib.util/update-query-stage query stage-number u/assoc-dissoc :fields (not-empty xs)))))
 
 (defn fields
-  "Fetches the `:fields` for a query."
+  "Fetches the `:fields` for a query. Returns `nil` if there are no `:fields`. `:fields` should never be empty; this is
+  enforced by the Malli schema."
   ([query]
    (fields query -1))
   ([query stage-number]

src/metabase/lib/js.cljs

@@ -338,3 +338,18 @@
   [[available-binning-strategies]] to get `available-aggregation`."
   [aggregation-operator]
   (to-array (lib.core/aggregation-operator-columns aggregation-operator)))
+
+(defn ^:export fields
+  "Get the current `:fields` in a query. Unlike the lib core version, this will return an empty sequence if `:fields` is
+  not specified rather than `nil` for JS-friendliness."
+  ([a-query]
+   (fields a-query -1))
+  ([a-query stage-number]
+   (to-array (lib.core/fields a-query stage-number))))
+
+(defn ^:export with-fields
+  "Specify the `:fields` for a query. Pass an empty sequence or `nil` to remove `:fields`."
+  ([a-query new-fields]
+   (with-fields a-query -1 new-fields))
+  ([a-query stage-number new-fields]
+   (lib.core/with-fields a-query stage-number new-fields)))

test/metabase/lib/field_test.cljc

@@ -458,3 +458,29 @@
                 :lib/source-column-alias  "avg_count"
                 :lib/desired-column-alias "avg_count"}]
               (lib.metadata.calculation/metadata query))))))
+
+(deftest ^:parallel with-fields-test
+  (let [query           (-> (lib/query-for-table-name meta/metadata-provider "VENUES")
+                            (lib/with-fields [(lib/field "VENUES" "ID") (lib/field "VENUES" "NAME")]))
+        fields-metadata (fn [query]
+                          (map (partial lib.metadata.calculation/metadata query)
+                               (lib/fields query)))
+        metadatas       (fields-metadata query)]
+    (is (=? [{:name "ID"}
+             {:name "NAME"}]
+            metadatas))
+    (testing "Set fields with metadatas"
+      (let [fields' [(last metadatas)]
+            query'  (lib/with-fields query fields')]
+        (is (=? [{:name "NAME"}]
+                (fields-metadata query')))))
+    (testing "remove fields by passing"
+      (doseq [new-fields [nil []]]
+        (testing (pr-str new-fields)
+          (let [query' (lib/with-fields query new-fields)]
+            (is (empty? (fields-metadata query')))
+            (letfn [(has-fields? [query]
+                      (get-in query [:stages 0 :fields]))]
+              (is (has-fields? query)
+                  "sanity check")
+              (is (not (has-fields? query'))))))))))