[MLv2] Add `lib.field/add-field` and `remove-field` (#32679)

These are "porcelain" APIs to make FE management of the set of selected
columns for a query much easier.

`add-field` and `remove-field` will handle regular table fields,
explicit join fields, and custom expressions correctly.

One known quirk: if the set of `:fields` on a stage happens to match the
default set, this is *not* noticed or fixed. Once an explicit list is
set on a stage, it will forever have a list.

src/metabase/lib/core.cljc

@@ -142,10 +142,12 @@
    upper
    lower]
   [lib.field
+   add-field
    field-id
+   fieldable-columns
    fields
-   with-fields
-   fieldable-columns]
+   remove-field
+   with-fields]
   [lib.filter
    filter
    filters

src/metabase/lib/field.cljc

@@ -13,6 +13,7 @@
    [metabase.lib.normalize :as lib.normalize]
    [metabase.lib.options :as lib.options]
    [metabase.lib.ref :as lib.ref]
+   [metabase.lib.remove-replace :as lib.remove-replace]
    [metabase.lib.schema :as lib.schema]
    [metabase.lib.schema.common :as lib.schema.common]
    [metabase.lib.schema.id :as lib.schema.id]
@@ -222,19 +223,19 @@
                                       ;;    Products → Products → Category
                                       (not (str/includes? field-display-name " → ")))
                              (or
-                               (when fk-field-id
+                              (when fk-field-id
                                  ;; Implicitly joined column pickers don't use the target table's name, they use the FK field's name with
                                  ;; "ID" dropped instead.
                                  ;; This is very intentional: one table might have several FKs to one foreign table, each with different
                                  ;; meaning (eg. ORDERS.customer_id vs. ORDERS.supplier_id both linking to a PEOPLE table).
                                  ;; See #30109 for more details.
-                                 (if-let [field (lib.metadata/field query fk-field-id)]
-                                   (-> (lib.metadata.calculation/display-info query stage-number field)
-                                       :display-name
-                                       lib.util/strip-id)
-                                   (let [table (table-metadata query table-id)]
-                                     (lib.metadata.calculation/display-name query stage-number table style))))
-                               (or join-alias (lib.join/current-join-alias field-metadata))))
+                                (if-let [field (lib.metadata/field query fk-field-id)]
+                                  (-> (lib.metadata.calculation/display-info query stage-number field)
+                                      :display-name
+                                      lib.util/strip-id)
+                                  (let [table (table-metadata query table-id)]
+                                    (lib.metadata.calculation/display-name query stage-number table style))))
+                              (or join-alias (lib.join/current-join-alias field-metadata))))
         display-name       (if join-display-name
                              (str join-display-name " → " field-display-name)
                              field-display-name)]
@@ -534,3 +535,154 @@
   "Find the field id for something or nil."
   [field-metadata :- lib.metadata/ColumnMetadata]
   (:id field-metadata))
+
+(defn- populate-fields-for-stage
+  "Given a query and stage, sets the `:fields` list to be the fields which would be selected by default.
+  This is exactly [[lib.metadata.calculation/returned-columns]] filtered by the `:lib/source`.
+  Fields from explicit joins are listed on the join itself; custom expressions are always included and should not be
+  listed in `:fields`."
+  [query stage-number]
+  (lib.util/update-query-stage query stage-number
+                               (fn [stage]
+                                 (assoc stage :fields
+                                        (into [] (comp (remove (comp #{:source/joins :source/expressions}
+                                                                     :lib/source))
+                                                       (map lib.ref/ref))
+                                              (lib.metadata.calculation/returned-columns query stage-number stage))))))
+
+(defn- query-with-fields
+  "If the given stage already has a `:fields` clause, do nothing. If it doesn't, populate the `:fields` clause with the
+  full set of `returned-columns`. (See [[populate-fields-for-stage]] for the details.)"
+  [query stage-number]
+  (cond-> query
+    (not (:fields (lib.util/query-stage query stage-number))) (populate-fields-for-stage stage-number)))
+
+(defn- include-field [query stage-number column]
+  (let [populated  (query-with-fields query stage-number)
+        column-ref (lib.ref/ref column)]
+    (if (lib.equality/find-closest-matching-ref populated column-ref (fields populated stage-number))
+      ;; If the column is already found, do nothing and return the original query.
+      query
+      (lib.util/update-query-stage populated stage-number update :fields conj column-ref))))
+
+(defn- add-field-to-join [query stage-number column]
+  (let [column-ref   (lib.ref/ref column)
+        [join field] (first (for [join  (lib.join/joins query stage-number)
+                                  field (lib.join/joinable-columns query stage-number join)
+                                  :when (lib.equality/find-closest-matching-ref query column-ref
+                                                                                [(lib.ref/ref field)])]
+                              [join field]))
+        join-fields  (lib.join/join-fields join)]
+
+    ;; Nothing to do if it's already selected, or if this join already has :fields :all.
+    ;; Otherwise, append it to the list of fields.
+    (if (or (and field (:selected? field))
+            (= join-fields :all))
+      query
+      (lib.remove-replace/replace-join query stage-number join
+                                       (lib.join/with-join-fields join
+                                         (if (= join-fields :none)
+                                           [column]
+                                           (conj join-fields column)))))))
+
+(defn- native-query-fields-edit-error []
+  (i18n/tru "Fields cannot be adjusted on native queries. Either edit the native query, or save this question and edit the fields in a GUI question based on this one."))
+
+(mu/defn add-field :- ::lib.schema/query
+  "Adds a given field (`ColumnMetadata`, as returned from eg. [[visible-columns]]) to the fields returned by the query.
+  Exactly what this means depends on the source of the field:
+  - Source table/card, previous stage of the query, aggregation or breakout:
+      - Add it to the `:fields` list
+      - If `:fields` is missing, it's implicitly `:all`, so do nothing.
+  - Implicit join: add it to the `:fields` list; query processor will do the right thing with it.
+  - Explicit join: add it to that join's `:fields` list.
+  - Custom expression: Do nothing - expressions are always included."
+  [query      :- ::lib.schema/query
+   stage-number :- :int
+   column       :- lib.metadata.calculation/ColumnMetadataWithSource]
+  (let [stage (lib.util/query-stage query stage-number)]
+    (case (:lib/source column)
+      (:source/table-defaults
+       :source/card
+       :source/previous-stage
+       :source/aggregations
+       :source/breakouts)         (cond-> query
+                                    (contains? stage :fields) (include-field stage-number column))
+      :source/joins               (add-field-to-join query stage-number column)
+      :source/implicitly-joinable (include-field query stage-number column)
+      :source/native              (throw (ex-info (native-query-fields-edit-error) {:query query :stage stage-number}))
+      ;; Default case - for columns from a native query or a custom expression - these are always returned and cannot be
+      ;; selected off and on.
+      query)))
+
+(defn- exclude-field
+  "This is called only for fields that plausibly need removing. If the stage has no `:fields`, this will populate it.
+  It shouldn't happen that we can't find the target field, but if that does happen, this will return the original query
+  unchanged. (In particular, if `:fields` did not exist before it will still be omitted.)"
+  [query stage-number column]
+  (let [old-fields (-> query
+                       (query-with-fields stage-number)
+                       (lib.util/query-stage stage-number)
+                       :fields)
+        column-ref (lib.ref/ref column)
+        index      (first (keep-indexed (fn [index field]
+                                          (when (lib.equality/find-closest-matching-ref query column-ref [field])
+                                            index))
+                                        old-fields))
+        new-fields (when index
+                     (let [[pre [_ & post]] (split-at index old-fields)]
+                       (vec (concat pre post))))]
+    ;; If we couldn't find the field, return the original query unchanged too.
+    (cond-> query
+      new-fields (lib.util/update-query-stage stage-number assoc :fields new-fields))))
+
+(defn- remove-field-from-join [query stage-number column]
+  (let [field-ref    (lib.ref/ref column)
+        join         (lib.join/resolve-join query stage-number (::lib.join/join-alias column))
+        join-fields  (lib.join/join-fields join)]
+    (if (or (nil? join-fields)
+            (= join-fields :none))
+      ;; Nothing to do if there's already no join fields.
+      query
+      (let [join-fields (if (= join-fields :all)
+                          (map lib.ref/ref (lib.metadata.calculation/returned-columns query stage-number join))
+                          join-fields)
+            removed     (remove #(lib.equality/find-closest-matching-ref query field-ref [%]) join-fields)]
+        (cond-> query
+          ;; If we actually removed a field, replace the join. Otherwise return the query unchanged.
+          (< (count removed) (count join-fields))
+          (lib.remove-replace/replace-join stage-number join (lib.join/with-join-fields join removed)))))))
+
+(mu/defn remove-field :- ::lib.schema/query
+  "Removes the field (a `ColumnMetadata`, as returned from eg. [[visible-columns]]) from those fields returned by the
+  query. Exactly what this means depends on the source of the field:
+  - Source table/card, previous stage, aggregations or breakouts:
+      - If `:fields` is missing, it's implicitly `:all` - populate it with all the columns except the removed one.
+      - Remove the target column from the `:fields` list
+  - Implicit join: remove it from the `:fields` list; do nothing if it's not there.
+      - (An implicit join only exists in the `:fields` clause, so if it's not there then it's not anywhere.)
+  - Explicit join: remove it from that join's `:fields` list (handle `:fields :all` like for source tables).
+  - Custom expression: Throw! Custom expressions are always returned. To remove a custom expression, the expression
+    itself should be removed from the query."
+  [query      :- ::lib.schema/query
+   stage-number :- :int
+   column       :- lib.metadata.calculation/ColumnMetadataWithSource]
+  (let [stage (lib.util/query-stage query stage-number)]
+    (case (:lib/source column)
+      (:source/table-defaults
+       :source/breakouts
+       :source/aggregations
+       :source/card
+       :source/previous-stage)    (exclude-field query stage-number column)
+      :source/implicitly-joinable (cond-> query
+                                    ;; If there are fields, exclude this column from it.
+                                    ;; If :fields is implied, then there can't be any implicitly joined fields in it.
+                                    (:fields stage) (exclude-field stage-number column))
+      :source/joins               (remove-field-from-join query stage-number column)
+      :source/expressions         (throw (ex-info (i18n/tru "Custom expressions cannot be de-selected. Delete the expression instead.")
+                                                  {:query      query
+                                                   :stage      stage-number
+                                                   :expression column}))
+      :source/native              (throw (ex-info (native-query-fields-edit-error) {:query query :stage stage-number}))
+      ;; Default case: do nothing and return the query unchaged.
+      query)))

src/metabase/lib/js.cljs

@@ -401,6 +401,32 @@
   [a-query stage-number]
   (to-array (lib.core/fieldable-columns a-query stage-number)))
 
+(defn ^:export add-field
+  "Adds a given field (`ColumnMetadata`, as returned from eg. [[visible-columns]]) to the fields returned by the query.
+  Exactly what this means depends on the source of the field:
+  - Source table/card, previous stage of the query, aggregation or breakout:
+      - Add it to the `:fields` list
+      - If `:fields` is missing, it's implicitly `:all`, so do nothing.
+  - Implicit join: add it to the `:fields` list; query processor will do the right thing with it.
+  - Explicit join: add it to that join's `:fields` list.
+  - Custom expression: Do nothing - expressions are always included."
+  [a-query stage-number column]
+  (lib.core/add-field a-query stage-number column))
+
+(defn ^:export remove-field
+  "Removes the field (a `ColumnMetadata`, as returned from eg. [[visible-columns]]) from those fields returned by the
+  query. Exactly what this means depends on the source of the field:
+  - Source table/card, previous stage, aggregations or breakouts:
+      - If `:fields` is missing, it's implicitly `:all` - populate it with all the columns except the removed one.
+      - Remove the target column from the `:fields` list
+  - Implicit join: remove it from the `:fields` list; do nothing if it's not there.
+      - (An implicit join only exists in the `:fields` clause, so if it's not there then it's not anywhere.)
+  - Explicit join: remove it from that join's `:fields` list (handle `:fields :all` like for source tables).
+  - Custom expression: Throw! Custom expressions are always returned. To remove a custom expression, the expression
+    itself should be removed from the query."
+  [a-query stage-number column]
+  (lib.core/remove-field a-query stage-number column))
+
 (defn ^:export join-strategy
   "Get the strategy (type) of a given join as an opaque JoinStrategy object."
   [a-join]