[MLv2] Add tests and fix a special case for fk-filter drills (#38297)

These were not adding an extra stage in order to filter queryies with
aggregations.

Progress towards #36253.

src/metabase/lib/drill_thru/fk_filter.cljc

@@ -21,6 +21,7 @@
   Question transformation:
   - None"
   (:require
+   [metabase.lib.aggregation :as lib.aggregation]
    [metabase.lib.drill-thru.common :as lib.drill-thru.common]
    [metabase.lib.filter :as lib.filter]
    [metabase.lib.metadata :as lib.metadata]
@@ -28,6 +29,7 @@
    [metabase.lib.options :as lib.options]
    [metabase.lib.schema :as lib.schema]
    [metabase.lib.schema.drill-thru :as lib.schema.drill-thru]
+   [metabase.lib.stage :as lib.stage]
    [metabase.lib.types.isa :as lib.types.isa]
    [metabase.lib.util :as lib.util]
    [metabase.util.malli :as mu]))
@@ -65,12 +67,16 @@
    stage-number :- :int
    drill-thru   :- ::lib.schema.drill-thru/drill-thru.fk-filter
    & _args]
-  ;; if the stage in question is an MBQL stage, we can simply add a `=` filter to it. If it's a native stage, we have
-  ;; to apply the drill to the stage after that stage, which will be an MBQL stage, adding it if needed (native stages
-  ;; are currently only allowed to be the first stage.)
+  ;; If the stage in question is an MBQL stage, we can simply add a `=` filter to it.
+  ;; If it's a native stage, we have to apply the drill to the stage after that stage, which will be an MBQL stage,
+  ;; adding it if needed (native stages are currently only allowed to be the first stage.)
+  ;; Similarly if the query contains aggregations we will have to add a new stage to do the filtering.
   (let [[query stage-number] (if (lib.drill-thru.common/mbql-stage? query stage-number)
-                               [query stage-number]
-                               ;; native stage
+                               ;; MBQL stage - append a stage if there are aggregations
+                               (if (seq (lib.aggregation/aggregations query stage-number))
+                                 (lib.stage/ensure-extra-stage query stage-number)
+                                 [query stage-number])
+                               ;; native stage - append an MBQL stage
                                (let [;; convert the stage number e.g. `-1` to the canonical non-relative stage number
                                      stage-number      (lib.util/canonical-stage-index query stage-number)
                                      ;; make sure the query has at least one MBQL stage after the native stage, which we

src/metabase/lib/stage.cljc

@@ -352,3 +352,18 @@
   (if-not (has-clauses? query -1)
     (drop-stage query)
     query))
+
+(mu/defn ensure-extra-stage :- [:tuple ::lib.schema/query :int]
+  "Given a query and current stage, returns a tuple of `[query next-stage-number]`.
+
+  If that stage already exists, the query is unchanged. If it does not, a new (MBQL) stage is appended and its index
+  is returned."
+  [query        :- ::lib.schema/query
+   stage-number :- :int]
+  (let [stage-number (lib.util/canonical-stage-index query stage-number)]
+    (if-let [next-number (lib.util/next-stage-number query stage-number)]
+      ;; There is already a next stage, so just return it.
+      [query next-number]
+      ;; Otherwise append a stage and return the new query and updated stage number.
+      (let [query (append-stage query)]
+        [query (lib.util/next-stage-number query stage-number)]))))

test/metabase/lib/drill_thru/fk_filter_test.cljc

@@ -4,12 +4,22 @@
    [medley.core :as m]
    [metabase.lib.core :as lib]
    [metabase.lib.drill-thru.test-util :as lib.drill-thru.tu]
+   [metabase.lib.drill-thru.test-util.canned :as canned]
    [metabase.lib.test-metadata :as meta]
    [metabase.util :as u]
    #?@(:cljs ([metabase.test-runner.assert-exprs.approximately-equal]))))
 
 #?(:cljs (comment metabase.test-runner.assert-exprs.approximately-equal/keep-me))
 
+(deftest ^:parallel fk-filter-availability-test
+  (testing "fk-filter is available for cell clicks on FKs with non-NULL values"
+    (doseq [[test-case context {:keys [click column-type]}] (canned/canned-clicks)]
+      (if (and (= click :cell)
+               (= column-type :fk)
+               (not= (:value context) :null))
+        (is (canned/returned test-case context :drill-thru/fk-filter))
+        (is (not (canned/returned test-case context :drill-thru/fk-filter)))))))
+
 (deftest ^:parallel returns-fk-filter-test-1
   (lib.drill-thru.tu/test-returns-drill
    {:drill-type  :drill-thru/fk-filter
@@ -134,3 +144,35 @@
             :column-name "Venues → Category ID"
             :table-name "Checkins"}
            (lib/display-info query fk-filter-drill)))))
+
+(deftest ^:parallel fk-filter-application-test
+  (testing "adds an = filter for the selected column and value"
+    (testing "in the same stage for a plain query"
+      (lib.drill-thru.tu/test-drill-application
+        {:click-type     :cell
+         :query-type     :unaggregated
+         :column-name    "USER_ID"
+         :drill-type     :drill-thru/fk-filter
+         :expected       {:lib/type  :metabase.lib.drill-thru/drill-thru
+                          :type      :drill-thru/fk-filter
+                          :column-name "User ID"
+                          :table-name  "Orders"}
+         :expected-query {:stages [{:source-table (meta/id :orders)
+                                    :filters      [[:= {}
+                                                    [:field {} (meta/id :orders :user-id)]
+                                                    (get-in lib.drill-thru.tu/test-queries
+                                                            ["ORDERS" :unaggregated :row "USER_ID"])]]}]}}))
+    (testing "in a new stage for an aggregated query"
+      (lib.drill-thru.tu/test-drill-application
+        {:click-type     :cell
+         :query-type     :aggregated
+         :column-name    "PRODUCT_ID"
+         :drill-type     :drill-thru/fk-filter
+         :expected       {:lib/type  :metabase.lib.drill-thru/drill-thru
+                          :type      :drill-thru/fk-filter
+                          :column-name "Product ID"
+                          :table-name  "Orders"}
+         :expected-query {:stages [(-> (get lib.drill-thru.tu/test-queries "ORDERS") :aggregated :query :stages first)
+                                   {:filters [[:= {} [:field {} (meta/id :orders :product-id)]
+                                               (get-in lib.drill-thru.tu/test-queries
+                                                       ["ORDERS" :aggregated :row "PRODUCT_ID"])]]}]}}))))