Merge pull request #8583 from metabase/qp-refactor-part-6

QP refactor part 6 - rework add-dimension-projections

src/metabase/models/dimension.clj

 (ns metabase.models.dimension
+  "Dimensions are used to define remappings for Fields handled automatically when those Fields are encountered by the
+  Query Processor. For a more detailed explanation, refer to the documentation in
+  `metabase.query-processor.middleware.add-dimension-projections`."
   (:require [toucan.models :as models]
             [metabase.util :as u]))
 

src/metabase/query_processor.clj

@@ -103,8 +103,8 @@
       format-rows/format-rows
       binning/update-binning-strategy
       resolve/resolve-middleware
-      add-dim/add-remapping
       expand/expand-middleware                         ; ▲▲▲ QUERY EXPANSION POINT  ▲▲▲ All functions *above* will see EXPANDED query during PRE-PROCESSING
+      add-dim/add-remapping
       implicit-clauses/add-implicit-clauses
       bucket-datetime/auto-bucket-datetime-breakouts
       source-table/resolve-source-table-middleware
@@ -118,7 +118,7 @@
       fetch-source-query/fetch-source-query
       store/initialize-store
       log-query/log-initial-query
-      ;; TODO - bind *query* here ?
+      ;; TODO - bind `*query*` here ?
       cache/maybe-return-cached-results
       log-query/log-results-metadata
       validate/validate-query

src/metabase/query_processor/middleware/add_dimension_projections.clj

 (ns metabase.query-processor.middleware.add-dimension-projections
-  "Middleware for adding remapping and other dimension related projections"
-  (:require [metabase.query-processor
-             [interface :as i]
-             [util :as qputil]]
-            [metabase.query-processor.middleware.resolve :as resolve]))
+  "Middleware for adding remapping and other dimension related projections. This remaps Fields that have a corresponding
+  Dimension object (which defines a remapping) in two different ways, depending on the `:type` attribute of the
+  Dimension:
+
+  `external` type Dimensions mean the Field's values will be replaced with corresponding values from a column on a
+  different table, joined via a foreign key. A common use-case would be to replace FK IDs with the name of whatever it
+  references, for example replacing a values of `venue.category_id` with values of `category.name`. Actual replacement
+  of values happens on the frontend, so this middleware simply adds the column to be used for replacement (e.g.
+  `category.name`) to the `:fields` clause in pre-processing, so the Field will be fetched. Recall that Fields
+  referenced via with `:fk->` clauses imply that JOINs will take place, which are automatically handled later in the
+  Query Processor pipeline. Additionally, this middleware will swap out and `:order-by` clauses referencing the
+  original Field with ones referencing the remapped Field (for example, so we would sort by `category.name` instead of
+  `category_id`).
+
+  `internal` type Dimensions mean the Field's values are replaced by a user-defined map of values, stored in the
+  `human_readable_values` column of a corresponding `FieldValues` object. A common use-case for this scenario would be
+  to replace integer enum values with something more descriptive, for example replacing values of an enum `can_type`
+  -- `0` becomes `Toucan`, `1` becomes `Pelican`, and so forth. This is handled exclusively in post-processing by
+  adding extra columns and values to the results.
+
+  In both cases, to accomplish values replacement on the frontend, the post-processing part of this middleware adds
+  appropriate `:remapped_from` and `:remapped_to` attributes in the result `:cols` in post-processing.
+  `:remapped_from` and `:remapped_to` are the names of the columns, e.g. `category_id` is `:remapped_to` `name`, and
+  `name` is `:remapped_from` `:category_id`."
+  (:require [metabase.mbql
+             [schema :as mbql.s]
+             [util :as mbql.u]]
+            [metabase.models.dimension :refer [Dimension]]
+            [metabase.util :as u]
+            [metabase.util.schema :as su]
+            [schema.core :as s]
+            [toucan.db :as db]))
+
+(def ^:private ExternalRemappingDimension
+  "Schema for the info we fetch about `external` type Dimensions that will be used for remappings in this Query. Fetched
+  by the pre-processing portion of the middleware, and passed along to the post-processing portion."
+  {:name                    su/NonBlankString       ; display name for the remapping
+   :field_id                su/IntGreaterThanZero   ; ID of the Field being remapped
+   :human_readable_field_id su/IntGreaterThanZero}) ; ID of the FK Field to remap values to
+
+
+;;; ----------------------------------------- add-fk-remaps (pre-processing) -----------------------------------------
+
+(s/defn ^:private fields->field-id->remapping-dimension :- (s/maybe {su/IntGreaterThanZero ExternalRemappingDimension})
+  "Given a sequence of field clauses (from the `:fields` clause), return a map of `:field-id` clause (other clauses
+  are ineligable) to a remapping dimension information for any Fields that have an `external` type dimension remapping."
+  [fields :- [mbql.s/Field]]
+  (when-let [field-ids (seq (map second (filter (partial mbql.u/is-clause? :field-id) fields)))]
+    (u/key-by :field_id (db/select [Dimension :field_id :name :human_readable_field_id]
+                          :field_id [:in (set field-ids)]
+                          :type     "external"))))
+
+(s/defn ^:private create-remap-col-tuples :- [[(s/one mbql.s/field-id            "Field")
+                                               (s/one mbql.s/fk->                "remapped FK Field")
+                                               (s/one ExternalRemappingDimension "remapping Dimension info")]]
+  "Return tuples of `:field-id` clauses, the new remapped column `:fk->` clauses that the Field should be remapped to,
+  and the Dimension that suggested the remapping, which is used later in this middleware for post-processing. Order is
+  important here, because the results are added to the `:fields` column in order. (TODO - why is it important, if they
+  get hidden when displayed anyway?)"
+  [fields :- [mbql.s/Field]]
+  (when-let [field-id->remapping-dimension (fields->field-id->remapping-dimension fields)]
+    (vec (for [field fields
+               :when (mbql.u/is-clause? :field-id field)
+               :let  [dimension (field-id->remapping-dimension (second field))]
+               :when dimension]
+           [field
+            [:fk-> field [:field-id (:human_readable_field_id dimension)]]
+            dimension]))))
+
+(s/defn ^:private update-remapped-order-by :- [mbql.s/OrderBy]
+  "Order by clauses that include an external remapped column should be replace that original column in the order by with
+  the newly remapped column. This should order by the text of the remapped column vs. the id of the source column
+  before the remapping"
+  [field->remapped-col :- {mbql.s/field-id, mbql.s/fk->}, order-by-clauses :- [mbql.s/OrderBy]]
+  (vec
+   (for [[direction field, :as order-by-clause] order-by-clauses]
+     (if-let [remapped-col (get field->remapped-col field)]
+       [direction remapped-col]
+       order-by-clause))))
+
+(s/defn ^:private add-fk-remaps :- [(s/one (s/maybe [ExternalRemappingDimension]) "external remapping dimensions")
+                                    (s/one mbql.s/Query "query")]
+  "Add any Fields needed for `:external` remappings to the `:fields` clause of the query, and update `:order-by`
+  clause as needed. Returns a pair like `[external-remapping-dimensions updated-query]`."
+  [{{:keys [fields order-by]} :query, :as query} :- mbql.s/Query]
+  ;; fetch remapping column pairs if any exist...
+  (if-let [remap-col-tuples (seq (create-remap-col-tuples fields))]
+    ;; if they do, update `:fields` and `:order-by` clauses accordingly and add to the query
+    (let [new-fields   (vec (concat fields (map second remap-col-tuples)))
+          ;; make a map of field-id-clause -> fk-clause from the tuples
+          new-order-by (update-remapped-order-by (into {} (for [[field-clause fk-clause] remap-col-tuples]
+                                                            [field-clause fk-clause]))
+                                                 order-by)]
+      ;; return the Dimensions we are using and the query
+      [(map last remap-col-tuples)
+       (cond-> (assoc-in query [:query :fields] new-fields)
+         (seq new-order-by) (assoc-in [:query :order-by] new-order-by))])
+    ;; otherwise return query as-is
+    [nil query]))
+
+
+;;; ---------------------------------------- remap-results (post-processing) -----------------------------------------
+
+(s/defn ^:private add-remapping-info :- [su/Map]
+  "Add `:display_name`, `:remapped_to`, and `:remapped_from` keys to columns for the results, needed by the frontend.
+  To get this critical information, this uses the `remapping-dimensions` info saved by the pre-processing portion of
+  this middleware for external remappings, and the internal-only remapped columns handled by post-processing
+  middleware below for internal columns."
+  [remapping-dimensions   :- (s/maybe [ExternalRemappingDimension])
+   columns                :- [su/Map]
+   internal-remap-columns :- (s/maybe [su/Map])]
+  (let [column-id->column              (u/key-by :id columns)
+        name->internal-remapped-to-col (u/key-by :remapped_from internal-remap-columns)
+        id->remapped-to-dimension      (u/key-by :field_id                remapping-dimensions)
+        id->remapped-from-dimension    (u/key-by :human_readable_field_id remapping-dimensions)]
+    (for [{:keys [id], column-name :name, :as column} columns]
+      (merge
+       column
+       ;; if one of the internal remapped columns says it's remapped from this column, add a matching `:remapped_to`
+       ;; entry
+       (when-let [{remapped-to-name :name} (get name->internal-remapped-to-col column-name)]
+         {:remapped_to remapped-to-name})
+       ;; if the pre-processing remapping Dimension info contains an entry where this Field's ID is `:field_id`, add
+       ;; an entry noting the name of the Field it gets remapped to
+       (when-let [{remapped-to-id :human_readable_field_id} (get id->remapped-to-dimension id)]
+         {:remapped_to (:name (get column-id->column remapped-to-id))})
+       ;; if the pre-processing remapping Dimension info contains an entry where this Field's ID is
+       ;; `:human_readable_field_id`, add an entry noting the name of the Field it gets remapped from, and use the
+       ;; `:display_name` of the Dimension
+       (when-let [{dimension-name :name, remapped-from-id :field_id} (get id->remapped-from-dimension id)]
+         {:display_name  dimension-name
+          :remapped_from (:name (get column-id->column remapped-from-id))})))))
 
 (defn- create-remapped-col [col-name remapped-from]
   {:description     nil
@@ -18,47 +145,18 @@
    :remapped_from   remapped-from
    :remapped_to     nil})
 
-(defn- create-fk-remap-col [fk-field-id dest-field-id remapped-from field-display-name]
-  (i/map->FieldPlaceholder {:fk-field-id        fk-field-id
-                            :field-id           dest-field-id
-                            :remapped-from      remapped-from
-                            :remapped-to        nil
-                            :field-display-name field-display-name}))
-
-(defn- row-map-fn [dim-seq]
-  (fn [row]
-    (concat row (map (fn [{:keys [col-index xform-fn]}]
-                          (xform-fn (nth row col-index)))
-                        dim-seq))))
-
 (defn- transform-values-for-col
   "Converts `values` to a type compatible with the base_type found for `col`. These values should be directly comparable
   with the values returned from the database for the given `col`."
   [{:keys [base_type] :as col} values]
-  (cond
-    (isa? base_type :type/Decimal)
-    (map bigdec values)
-
-    (isa? base_type :type/Float)
-    (map double values)
-
-    (isa? base_type :type/BigInteger)
-    (map bigint values)
-
-    (isa? base_type :type/Integer)
-    (map int values)
-
-    (isa? base_type :type/Text)
-    (map str values)
-
-    :else
-    values))
-
-(defn- assoc-remapped-to [from->to]
-  (fn [col]
-    (-> col
-        (update :remapped_to #(or % (from->to (:name col))))
-        (update :remapped_from #(or % nil)))))
+  (map (condp #(isa? %2 %1) base_type
+         :type/Decimal    bigdec
+         :type/Float      double
+         :type/BigInteger bigint
+         :type/Integer    int
+         :type/Text       str
+         identity)
+       values))
 
 (defn- col->dim-map
   [idx {{remap-to :dimension-name, remap-type :dimension-type, field-id :field-id} :dimensions, :as col}]
@@ -72,71 +170,38 @@
        :new-column     (create-remapped-col remap-to remap-from)
        :dimension-type remap-type})))
 
-(defn- create-remap-col-pairs
-  "Return pairs of field id and the new remapped column that the field should be remapped to. This is a list of pairs as
-  we want to preserve order"
-  [fields]
-  (for [{{:keys [field-id human-readable-field-id dimension-type dimension-name]} :dimensions,
-         field-name :field-name, source-field-id :field-id} fields
-        :when (= :external dimension-type)]
-    [source-field-id (create-fk-remap-col field-id
-                                          human-readable-field-id
-                                          field-name
-                                          dimension-name)]))
-
-(defn- update-remapped-order-by
-  "Order by clauses that include an external remapped column should be replace that original column in the order by with
-  the newly remapped column. This should order by the text of the remapped column vs. the id of the source column
-  before the remapping"
-  [remap-cols-by-id order-by-seq]
-  (when (seq order-by-seq)
-    (mapv (fn [{{:keys [field-id]} :field :as order-by-clause}]
-            (if-let [remapped-col (get remap-cols-by-id field-id)]
-              (assoc order-by-clause :field remapped-col)
-              order-by-clause))
-          order-by-seq)))
-
-(defn- add-fk-remaps
-  "Function that will include FK references needed for external remappings. This will then flow through to the resolver
-  to get the new tables included in the join."
-  [query]
-  (let [resolved-fields (resolve/resolve-fields-if-needed (qputil/get-in-query query [:fields]))
-        remap-col-pairs (create-remap-col-pairs resolved-fields)]
-    (if (seq remap-col-pairs)
-      (let [order-by (qputil/get-in-query query [:order-by])]
-        (-> query
-            (qputil/assoc-in-query [:order-by] (update-remapped-order-by (into {} remap-col-pairs) order-by))
-            (qputil/assoc-in-query [:fields] (concat resolved-fields (map second remap-col-pairs)))))
-      query)))
-
-(defn- remap-results
+(defn- row-map-fn [dim-seq]
+  (fn [row]
+    (concat row (map (fn [{:keys [col-index xform-fn]}]
+                          (xform-fn (nth row col-index)))
+                        dim-seq))))
+
+(s/defn ^:private remap-results
   "Munges results for remapping after the query has been executed. For internal remappings, a new column needs to be
   added and each row flowing through needs to include the remapped data for the new column. For external remappings,
   the column information needs to be updated with what it's being remapped from and the user specified name for the
   remapped column."
-  [results]
+  [remapping-dimensions :- (s/maybe [ExternalRemappingDimension]), results]
   (let [indexed-dims       (keep-indexed col->dim-map (:cols results))
         internal-only-dims (filter #(= :internal (:dimension-type %)) indexed-dims)
         remap-fn           (row-map-fn internal-only-dims)
-        columns            (concat (:cols results)
-                                   (map :new-column internal-only-dims))
-        from->to           (reduce (fn [acc {:keys [remapped_from name]}]
-                                     (if remapped_from
-                                       (assoc acc remapped_from name)
-                                       acc))
-                                   {} columns)]
+        internal-only-cols (map :new-column internal-only-dims)]
     (-> results
         (update :columns into (map :to internal-only-dims))
-        ;; TODO - this code doesn't look right... why use `update` if we're not using the value we're updating?
-        (update :cols (fn [_]
-                        (mapv (comp #(dissoc % :dimensions :values)
-                                    (assoc-remapped-to from->to))
-                              columns)))
-        (update :rows #(map remap-fn %)))))
+        (assoc  :cols    (map #(dissoc % :dimensions :values)
+                              (concat (add-remapping-info remapping-dimensions (:cols results) internal-only-cols)
+                                      internal-only-cols)))
+        (update :rows    #(map remap-fn %)))))
+
+
+;;; --------------------------------------------------- middleware ---------------------------------------------------
 
 (defn add-remapping
-  "Query processor middleware. `QP` is the query processor, returns a function that works on a `QUERY` map. Delgates to
+  "Query processor middleware. `qp` is the query processor, returns a function that works on a `query` map. Delgates to
   `add-fk-remaps` for making remapping changes to the query (before executing the query). Then delegates to
   `remap-results` to munge the results after query execution."
   [qp]
-  (comp remap-results qp add-fk-remaps))
+  (fn [query]
+    (let [[remapping-dimensions query] (add-fk-remaps query)
+          results                      (qp query)]
+      (remap-results remapping-dimensions results))))

src/metabase/query_processor/middleware/auto_bucket_datetime_breakouts.clj

@@ -10,15 +10,17 @@
             [toucan.db :as db]))
 
 (def ^:private FieldTypeInfo
-  {:id           su/IntGreaterThanZero
-   :base_type    (s/maybe su/FieldType)
-   :special_type (s/maybe su/FieldType)})
+  {:base_type    (s/maybe su/FieldType)
+   :special_type (s/maybe su/FieldType)
+   s/Keyword     s/Any})
 
 (s/defn ^:private is-datetime-field?
   [{base-type :base_type, special-type :special_type} :- (s/maybe FieldTypeInfo)]
   (or (isa? base-type :type/DateTime)
       (isa? special-type :type/DateTime)))
 
+;; TODO - we should check the store to see if these Fields have already been resolved! And if they haven't, we should
+;; go ahead and resolve them, and save them in the store...
 (s/defn ^:private unbucketed-breakouts->field-id->type-info :- {su/IntGreaterThanZero (s/maybe FieldTypeInfo)}
   "Fetch a map of Field ID -> type information for the Fields referred to by the `unbucketed-breakouts`."
   [unbucketed-breakouts :- (su/non-empty [mbql.s/field-id])]

test/metabase/api/dataset_test.clj

@@ -52,8 +52,9 @@
   [ ;; API call response
    {:data                   {:rows    [[1000]]
                              :columns ["count"]
-                             :cols    [{:base_type "type/Integer", :special_type "type/Number", :name "count", :display_name "count", :id nil, :table_id nil,
-                                        :description nil, :target nil, :extra_info {}, :source "aggregation", :remapped_from nil, :remapped_to nil}]
+                             :cols    [{:base_type "type/Integer", :special_type "type/Number", :name "count",
+                                        :display_name "count", :id nil, :table_id nil, :description nil, :target nil,
+                                        :extra_info {}, :source "aggregation"}]
                              :native_form true}
     :row_count              1
     :status                 "completed"

test/metabase/api/embed_test.clj

@@ -66,7 +66,7 @@
    {:data       {:columns ["count"]
                  :cols    [{:description nil, :table_id nil, :special_type "type/Number", :name "count",
                             :source "aggregation", :extra_info {}, :id nil, :target nil, :display_name "count",
-                            :base_type "type/Integer", :remapped_from nil, :remapped_to nil}]
+                            :base_type "type/Integer"}]
                  :rows    [[100]]}
     :json_query {:parameters nil}
     :status     "completed"})

test/metabase/driver/bigquery_test.clj

@@ -19,9 +19,6 @@
             [metabase.test.data.datasets :refer [expect-with-engine]]
             [toucan.util.test :as tt]))
 
-(def ^:private col-defaults
-  {:remapped_to nil, :remapped_from nil})
-
 ;; Test native queries
 (expect-with-engine :bigquery
   [[100]
@@ -52,10 +49,9 @@
 ;; ordering shouldn't apply (Issue #2821)
 (expect-with-engine :bigquery
   {:columns ["venue_id" "user_id" "checkins_id"],
-   :cols    (mapv #(merge col-defaults %)
-                  [{:name "venue_id",    :display_name "Venue ID",    :base_type :type/Integer}
-                   {:name "user_id",     :display_name  "User ID",    :base_type :type/Integer}
-                   {:name "checkins_id", :display_name "Checkins ID", :base_type :type/Integer}])}
+   :cols    [{:name "venue_id",    :display_name "Venue ID",    :base_type :type/Integer}
+             {:name "user_id",     :display_name  "User ID",    :base_type :type/Integer}
+             {:name "checkins_id", :display_name "Checkins ID", :base_type :type/Integer}]}
 
   (select-keys (:data (qp/process-query
                         {:native   {:query (str "SELECT `test_data.checkins`.`venue_id` AS `venue_id`, "

test/metabase/driver/druid_test.clj

@@ -84,7 +84,7 @@
           (m/dissoc-in [:data :results_metadata])))))
 
 (def ^:private col-defaults
-  {:base_type :type/Text, :remapped_from nil, :remapped_to nil})
+  {:base_type :type/Text})
 
 ;; test druid native queries
 (expect-with-engine :druid

test/metabase/driver/generic_sql/native_test.clj

 (ns metabase.driver.generic-sql.native-test
+  "Tests for running native queries against SQL databases."
   (:require [expectations :refer :all]
             [medley.core :as m]
-            [metabase.models.database :refer [Database]]
             [metabase.query-processor :as qp]
-            [metabase.test.data :as data]
-            [toucan.db :as db]))
-
-(def ^:private col-defaults
-  {:remapped_from nil, :remapped_to nil})
+            [metabase.test.data :as data]))
 
 ;; Just check that a basic query works
 (expect
@@ -16,7 +12,7 @@
    :data      {:rows        [[100]
                              [99]]
                :columns     ["ID"]
-               :cols        [(merge col-defaults {:name "ID", :display_name "ID", :base_type :type/Integer})]
+               :cols        [{:name "ID", :display_name "ID", :base_type :type/Integer}]
                :native_form {:query "SELECT ID FROM VENUES ORDER BY ID DESC LIMIT 2", :params []}}}
   (-> (qp/process-query {:native   {:query "SELECT ID FROM VENUES ORDER BY ID DESC LIMIT 2"}
                          :type     :native
@@ -30,10 +26,9 @@
    :data      {:rows        [[100 "Mohawk Bend" 46]
                              [99 "Golden Road Brewing" 10]]
                :columns     ["ID" "NAME" "CATEGORY_ID"]
-               :cols        (mapv #(merge col-defaults %)
-                                  [{:name "ID",          :display_name "ID",          :base_type :type/Integer}
-                                   {:name "NAME",        :display_name "Name",        :base_type :type/Text}
-                                   {:name "CATEGORY_ID", :display_name "Category ID", :base_type :type/Integer}])
+               :cols        [{:name "ID",          :display_name "ID",          :base_type :type/Integer}
+                             {:name "NAME",        :display_name "Name",        :base_type :type/Text}
+                             {:name "CATEGORY_ID", :display_name "Category ID", :base_type :type/Integer}]
                :native_form {:query "SELECT ID, NAME, CATEGORY_ID FROM VENUES ORDER BY ID DESC LIMIT 2", :params []}}}
   (-> (qp/process-query {:native   {:query "SELECT ID, NAME, CATEGORY_ID FROM VENUES ORDER BY ID DESC LIMIT 2"}
                          :type     :native
@@ -50,13 +45,3 @@
           :stacktrace
           :query
           :expanded-query))
-
-;; Check that we're not allowed to run SQL against an H2 database with a non-admin account
-(expect "Running SQL queries against H2 databases using the default (admin) database user is forbidden."
-  ;; Insert a fake Database. It doesn't matter that it doesn't actually exist since query processing should
-  ;; fail immediately when it realizes this DB doesn't have a USER
-  (let [db (db/insert! Database, :name "Fake-H2-DB", :engine "h2", :details {:db "mem:fake-h2-db"})]
-    (try (:error (qp/process-query {:database (:id db)
-                                    :type     :native
-                                    :native   {:query "SELECT 1"}}))
-         (finally (db/delete! Database :name "Fake-H2-DB")))))

test/metabase/driver/h2_test.clj

@@ -2,10 +2,13 @@
   (:require [expectations :refer :all]
             [metabase
              [db :as mdb]
-             [driver :as driver]]
+             [driver :as driver]
+             [query-processor :as qp]]
             [metabase.driver.h2 :as h2]
+            [metabase.models.database :refer [Database]]
             [metabase.test.data.datasets :refer [expect-with-engine]]
-            [metabase.test.util :as tu])
+            [metabase.test.util :as tu]
+            [toucan.db :as db])
   (:import metabase.driver.h2.H2Driver))
 
 ;; Check that the functions for exploding a connection string's options work as expected
@@ -42,3 +45,13 @@
 (expect-with-engine :h2
   "UTC"
   (tu/db-timezone-id))
+
+;; Check that we're not allowed to run SQL against an H2 database with a non-admin account
+(expect "Running SQL queries against H2 databases using the default (admin) database user is forbidden."
+  ;; Insert a fake Database. It doesn't matter that it doesn't actually exist since query processing should
+  ;; fail immediately when it realizes this DB doesn't have a USER
+  (let [db (db/insert! Database, :name "Fake-H2-DB", :engine "h2", :details {:db "mem:fake-h2-db"})]
+    (try (:error (qp/process-query {:database (:id db)
+                                    :type     :native
+                                    :native   {:query "SELECT 1"}}))
+         (finally (db/delete! Database :name "Fake-H2-DB")))))

test/metabase/query_processor/middleware/add_dimension_projections_test.clj

 (ns metabase.query-processor.middleware.add-dimension-projections-test
-  "Tests for the Query Processor cache."
   (:require [expectations :refer :all]
-            [metabase.query-processor.interface :as i]
-            [metabase.query-processor.middleware.add-dimension-projections :as add-dim-projections :refer :all]))
+            [metabase.query-processor.middleware.add-dimension-projections :as add-dim-projections]))
+
+;;; ----------------------------------------- add-fk-remaps (pre-processing) -----------------------------------------
+
+(def ^:private example-query
+  {:database 1
+   :type     :query
+   :query    {:source-table 1
+              :fields       [[:field-id 1]
+                             [:field-id 2]
+                             [:field-id 3]]}})
+
+(defn- do-with-fake-remappings-for-field-3 [f]
+  (with-redefs [add-dim-projections/fields->field-id->remapping-dimension
+                (constantly
+                 {3 {:name "Product", :field_id 3, :human_readable_field_id 4}})]
+    (f)))
+
+;; make sure FK remaps add an entry for the FK field to `:fields`, and returns a pair of [dimension-info updated-query]
+(expect
+  [[{:name "Product", :field_id 3, :human_readable_field_id 4}]
+   (update-in example-query [:query :fields]
+              conj [:fk-> [:field-id 3] [:field-id 4]])]
+  (do-with-fake-remappings-for-field-3
+   (fn []
+     (#'add-dim-projections/add-fk-remaps example-query))))
+
+;; adding FK remaps should replace any existing order-bys for a field with order bys for the FK remapping Field
+(expect
+  [[{:name "Product", :field_id 3, :human_readable_field_id 4}]
+   (-> example-query
+       (assoc-in [:query :order-by] [[:asc [:fk-> [:field-id 3] [:field-id 4]]]])
+       (update-in [:query :fields]
+                  conj [:fk-> [:field-id 3] [:field-id 4]]))]
+  (do-with-fake-remappings-for-field-3
+   (fn []
+     (#'add-dim-projections/add-fk-remaps (assoc-in example-query [:query :order-by] [[:asc [:field-id 3]]])))))
+
+
+;;; ---------------------------------------- remap-results (post-processing) -----------------------------------------
 
 (def ^:private col-defaults
   {:description     nil
    :source          :fields
    :extra_info      {}
    :fk_field_id     nil
-   :values          []
-   :dimensions      []
    :visibility_type :normal
    :target          nil
    :remapped_from   nil
    :remapped_to     nil})
 
-(def ^:private example-resultset
-  {:rows
-   [[1 "Red Medicine" 4 3]
-    [2 "Stout Burgers & Beers" 11 2]
-    [3 "The Apple Pan" 11 2]
-    [4 "Wurstküche" 29 2]
-    [5 "Brite Spot Family Restaurant" 20 2]],
-   :columns ["ID" "NAME" "CATEGORY_ID" "PRICE"],
-   :cols
-   (mapv #(merge col-defaults %)
-         [
-          ;; 0
-          {:table_id 4,
-           :schema_name "PUBLIC",
-           :special_type :type/PK,
-           :name "ID",
-           :id 12,
-           :display_name "ID",
-           :base_type :type/BigInteger}
-          ;; 1
-          {:table_id 4,
-           :schema_name "PUBLIC",
-           :special_type :type/Name,
-           :name "NAME",
-           :id 15,
-           :display_name "Name",
-           :base_type :type/Text}
-          ;; 2
-          {:table_id 4,
-           :schema_name "PUBLIC",
-           :special_type :type/FK,
-           :name "CATEGORY_ID",
-           :extra_info {:target_table_id 1},
-           :id 11,
-           :values {:field-value-id 1, :human-readable-values ["Foo" "Bar" "Baz" "Qux"],
-                    :values [4 11 29 20], :field-id 33}
-           :dimensions {:dimension-id 1 :dimension-type :internal :dimension-name "Foo" :field-id 10}
-           :display_name "Category ID",
-           :base_type :type/Integer}
-          ;; 3
-          {:table_id 4,
-           :schema_name "PUBLIC",
-           :special_type :type/Category,
-           :name "PRICE",
-           :id 16,
-           :display_name "Price",
-           :base_type :type/Integer}])})
+(def ^:private example-result-cols-id
+  (merge
+   col-defaults
+   {:table_id     4
+    :schema_name  "PUBLIC"
+    :special_type :type/PK
+    :name         "ID"
+    :id           12
+    :display_name "ID"
+    :base_type    :type/BigInteger}))
 
-(expect
-  (-> example-resultset
-      (assoc :rows [[1 "Red Medicine" 4 3 "Foo"]
-                    [2 "Stout Burgers & Beers" 11 2 "Bar"]
-                    [3 "The Apple Pan" 11 2 "Bar"]
-                    [4 "Wurstküche" 29 2 "Baz"]
-                    [5 "Brite Spot Family Restaurant" 20 2 "Qux"]])
-      (update :columns conj "Foo")
-      (update :cols (fn [cols]
-                      (conj
-                       (mapv (fn [col]
-                               (let [new-col (dissoc col :dimensions :values)]
-                                 (if (= "CATEGORY_ID" (:name new-col))
-                                   (assoc new-col
-                                     :remapped_to "Foo"
-                                     :remapped_from nil)
-                                   new-col)))
-                             cols)
-                       {:description nil,
-                        :id nil,
-                        :table_id nil,
-                        :expression-name "Foo",
-                        :source :fields,
-                        :name "Foo",
-                        :display_name "Foo",
-                        :target nil,
-                        :extra_info {}
-                        :remapped_from "CATEGORY_ID"
-                        :remapped_to nil}))))
-  (#'add-dim-projections/remap-results example-resultset))
-
-(def ^:private field-defaults
-  {:dimensions [],
-   :values [],
-   :visibility-type :normal})
+(def ^:private example-result-cols-name
+  (merge
+   col-defaults
+   {:table_id     4
+    :schema_name  "PUBLIC"
+    :special_type :type/Name
+    :name         "NAME"
+    :id           15
+    :display_name "Name"
+    :base_type    :type/Text}))
 
-(def ^:private example-query
-  {:query
-   {:order-by nil
-    :fields
-    (mapv #(merge field-defaults %)
-          [{:description "A unique internal identifier for the review. Should not be used externally.",
-            :base-type :type/BigInteger,
-            :table-id 4,
-            :special-type :type/PK,
-            :field-name "ID",
-            :field-display-name "ID",
-            :position 0,
-            :field-id 31,
-            :table-name "REVIEWS",
-            :schema-name "PUBLIC"}
-           {:description "The review the user left. Limited to 2000 characters.",
-            :base-type :type/Text,
-            :table-id 4,
-            :special-type :type/Description,
-            :field-name "BODY",
-            :field-display-name "Body",
-            :position 0,
-            :field-id 29,
-            :table-name "REVIEWS",
-            :schema-name "PUBLIC"}
-           {:field-id 32,
-            :field-name "PRODUCT_ID",
-            :field-display-name "Product ID",
-            :base-type :type/Integer,
-            :special-type :type/FK,
-            :table-id 4,
-            :schema-name "PUBLIC",
-            :table-name "REVIEWS",
-            :position 0,
-            :fk-field-id nil,
-            :description "The product the review was for",
-            :parent-id nil,
-            :parent nil,
-            :remapped-from nil,
-            :remapped-to nil,
-            :dimensions {:dimension-id 2, :dimension-name "Product", :field-id 32, :human-readable-field-id 27, :dimension-type :external}}])}})
+(def ^:private example-result-cols-category-id
+  (merge
+   col-defaults
+   {:table_id     4
+    :schema_name  "PUBLIC"
+    :special_type :type/FK
+    :name         "CATEGORY_ID"
+    :extra_info   {:target_table_id 1}
+    :id           11
+    :display_name "Category ID"
+    :base_type    :type/Integer}))
 
-(expect
-  (update-in example-query [:query :fields]
-             conj (i/map->FieldPlaceholder {:fk-field-id 32
-                                            :field-id 27
-                                            :remapped-from "PRODUCT_ID"
-                                            :remapped-to nil
-                                            :field-display-name "Product"}))
-  (#'add-dim-projections/add-fk-remaps example-query))
+(def ^:private example-result-cols-price
+  (merge
+   col-defaults
+   {:table_id     4
+    :schema_name  "PUBLIC"
+    :special_type :type/Category
+    :name         "PRICE"
+    :id           16
+    :display_name "Price"
+    :base_type    :type/Integer}))
+
+;; test that internal get the appropriate values and columns injected in, and the `:remapped_from`/`:remapped_to` info
+(def ^:private example-result-cols-foo
+  {:description     nil
+   :table_id        nil
+   :name            "Foo"
+   :expression-name "Foo"
+   :source          :fields
+   :remapped_from   "CATEGORY_ID"
+   :extra_info      {}
+   :remapped_to     nil
+   :id              nil
+   :target          nil
+   :display_name    "Foo"})
 
 (expect
-  (-> example-query
-      (assoc-in [:query :order-by] [{:direction :ascending
-                                     :field (i/map->FieldPlaceholder {:fk-field-id 32
-                                                                      :field-id 27
-                                                                      :remapped-from "PRODUCT_ID"
-                                                                      :remapped-to nil
-                                                                      :field-display-name "Product"})}])
-      (update-in [:query :fields]
-                 conj (i/map->FieldPlaceholder {:fk-field-id 32
-                                                :field-id 27
-                                                :remapped-from "PRODUCT_ID"
-                                                :remapped-to nil
-                                                :field-display-name "Product"})))
-  (-> example-query
-      (assoc-in [:query :order-by] [{:direction :ascending :field {:field-id 32}}])
-      (#'add-dim-projections/add-fk-remaps)))
-
-(def ^:private external-remapped-result
-  (-> example-resultset
-      (update :cols conj {:description "The name of the product as it should be displayed to customers.",
-                          :table_id 3,
-                          :schema_name nil,
-                          :special_type :type/Category,
-                          :name "CATEGORY",
-                          :source :fields,
-                          :remapped_from "CATEGORY_ID",
-                          :extra_info {},
-                          :fk_field_id 32,
-                          :remapped_to nil,
-                          :id 27,
-                          :visibility_type :normal,
-                          :target nil,
-                          :display_name "Category",
-                          :base_type :type/Text})
-      (update-in [:cols 2]
-                 (fn [col]
-                   (-> col
-                       (update :values merge {:human-readable-values []})
-                       (update :dimensions merge {:dimension-type :external :human-readable_field-id 27}))))))
+  {:rows    [[1 "Red Medicine"                  4 3 "Foo"]
+             [2 "Stout Burgers & Beers"        11 2 "Bar"]
+             [3 "The Apple Pan"                11 2 "Bar"]
+             [4 "Wurstküche"                   29 2 "Baz"]
+             [5 "Brite Spot Family Restaurant" 20 2 "Qux"]]
+   :columns ["ID" "NAME" "CATEGORY_ID" "PRICE" "Foo"]
+   :cols    [example-result-cols-id
+             example-result-cols-name
+             (assoc example-result-cols-category-id
+               :remapped_to "Foo")
+             example-result-cols-price
+             example-result-cols-foo]}
+  (#'add-dim-projections/remap-results
+   nil
+   {:rows    [[1 "Red Medicine"                  4 3]
+              [2 "Stout Burgers & Beers"        11 2]
+              [3 "The Apple Pan"                11 2]
+              [4 "Wurstküche"                   29 2]
+              [5 "Brite Spot Family Restaurant" 20 2]]
+    :columns ["ID" "NAME" "CATEGORY_ID" "PRICE"]
+    :cols    [example-result-cols-id
+              example-result-cols-name
+              (assoc example-result-cols-category-id
+                :dimensions {:dimension-id 1, :dimension-type :internal, :dimension-name "Foo", :field-id 10}
+                :values     {:field-value-id        1
+                             :human-readable-values ["Foo" "Bar" "Baz" "Qux"]
+                             :values                [4 11 29 20]
+                             :field-id              33})
+              example-result-cols-price]}))
+
+;; test that external remappings get the appropriate `:remapped_from`/`:remapped_to` info
+(def ^:private example-result-cols-category
+  (merge
+   col-defaults
+   {:description     "The name of the product as it should be displayed to customers."
+    :table_id        3
+    :schema_name     nil
+    :special_type    :type/Category
+    :name            "CATEGORY"
+    :source          :fields
+    :extra_info      {}
+    :fk_field_id     32
+    :id              27
+    :visibility_type :normal
+    :display_name    "Category"
+    :base_type       :type/Text}))
 
 (expect
-  (-> external-remapped-result
-      (update :cols (fn [col] (mapv #(dissoc % :dimensions :values) col)))
-      (update-in [:cols 2] assoc :remapped_to "CATEGORY"))
-  (#'add-dim-projections/remap-results external-remapped-result))
+  {:rows    []
+   :columns ["ID" "NAME" "CATEGORY_ID" "PRICE" "CATEGORY"]
+   :cols    [example-result-cols-id
+             example-result-cols-name
+             (assoc example-result-cols-category-id
+               :remapped_to "CATEGORY")
+             example-result-cols-price
+             (assoc example-result-cols-category
+               :remapped_from "CATEGORY_ID"
+               :display_name  "My Venue Category")]}
+  (#'add-dim-projections/remap-results
+   [{:name "My Venue Category", :field_id 11, :human_readable_field_id 27}]
+   {:rows    []
+    :columns ["ID" "NAME" "CATEGORY_ID" "PRICE" "CATEGORY"]
+    :cols    [example-result-cols-id
+              example-result-cols-name
+              example-result-cols-category-id
+              example-result-cols-price
+              example-result-cols-category]}))

test/metabase/query_processor_test.clj

@@ -270,32 +270,28 @@
   {:arglists '([ag-col-kw] [ag-col-kw field])}
   ([ag-col-kw]
    (case ag-col-kw
-     :count  {:base_type       :type/Integer
-              :special_type    :type/Number
-              :name            "count"
-              :display_name    "count"
-              :id              nil
-              :table_id        nil
-              :description     nil
-              :source          :aggregation
-              :extra_info      {}
-              :target          nil
-              :remapped_from   nil
-              :remapped_to     nil}))
+     :count {:base_type    :type/Integer
+             :special_type :type/Number
+             :name         "count"
+             :display_name "count"
+             :id           nil
+             :table_id     nil
+             :description  nil
+             :source       :aggregation
+             :extra_info   {}
+             :target       nil}))
   ([ag-col-kw {:keys [base_type special_type]}]
    {:pre [base_type special_type]}
    {:base_type    base_type
-    :special_type  special_type
-    :id            nil
-    :table_id      nil
-    :description   nil
-    :source        :aggregation
-    :extra_info    {}
-    :target        nil
-    :name          (name ag-col-kw)
-    :display_name  (name ag-col-kw)
-    :remapped_from nil
-    :remapped_to   nil}))
+    :special_type special_type
+    :id           nil
+    :table_id     nil
+    :description  nil
+    :source       :aggregation
+    :extra_info   {}
+    :target       nil
+    :name         (name ag-col-kw)
+    :display_name (name ag-col-kw)}))
 
 (defn breakout-col [column]
   (assoc column :source :breakout))