Some MLv2 joins conveniences (#33454)

src/metabase/lib/join.cljc

@@ -325,7 +325,7 @@
                   (lib.metadata.calculation/returned-columns query stage-number join options)))
         (:joins (lib.util/query-stage query stage-number))))
 
-(defmulti join-clause-method
+(defmulti ^:private join-clause-method
   "Convert something to a join clause."
   {:arglists '([joinable])}
   lib.dispatch/dispatch-value
@@ -359,12 +359,15 @@
                  :lib/type :mbql.stage/mbql}]}
       lib.options/ensure-uuid))
 
+(declare with-join-fields)
+
 (defmethod join-clause-method :metadata/table
-  [table]
-  (-> {:lib/type :mbql/join
-       :stages [{:source-table (:id table)
-                 :lib/type :mbql.stage/mbql}]}
-      lib.options/ensure-uuid))
+  [{::keys [join-alias join-fields], :as table-metadata}]
+  (cond-> (join-clause-method {:lib/type     :mbql.stage/mbql
+                               :lib/options  {:lib/uuid (str (random-uuid))}
+                               :source-table (:id table-metadata)})
+    join-alias  (with-join-alias join-alias)
+    join-fields (with-join-fields join-fields)))
 
 (defn- with-join-conditions-add-alias-to-rhses
   "Add `join-alias` to the RHS of all [[standard-join-condition?]] `conditions` that don't already have a `:join-alias`.
@@ -544,6 +547,10 @@
                           conditions)))
           (with-join-alias join-alias)))))
 
+(declare join-conditions
+         joined-thing
+         suggested-join-condition)
+
 (mu/defn join :- ::lib.schema/query
   "Add a join clause to a `query`."
   ([query a-join]
@@ -551,10 +558,15 @@
 
   ([query        :- ::lib.schema/query
     stage-number :- :int
-    a-join       :- PartialJoin]
-   (let [a-join (add-default-alias query stage-number a-join)]
-     (lib.util/update-query-stage query stage-number update :joins (fn [joins]
-                                                                     (conj (vec joins) a-join))))))
+    a-join       :- [:or PartialJoin Joinable]]
+   (let [a-join              (join-clause a-join)
+         suggested-condition (when (empty? (join-conditions a-join))
+                               (suggested-join-condition query stage-number (joined-thing query a-join)))
+         a-join              (cond-> a-join
+                               suggested-condition (with-join-conditions [suggested-condition]))
+         a-join              (add-default-alias query stage-number a-join)]
+     (lib.util/update-query-stage query stage-number update :joins (fn [existing-joins]
+                                                                     (conj (vec existing-joins) a-join))))))
 
 (mu/defn joins :- [:maybe ::lib.schema.join/joins]
   "Get all joins in a specific `stage` of a `query`. If `stage` is unspecified, returns joins in the final stage of the
@@ -632,7 +644,7 @@
 (mu/defn joined-thing :- [:maybe Joinable]
   "Return metadata about the origin of `a-join` using `metadata-providerable` as the source of information."
   [metadata-providerable :- lib.metadata/MetadataProviderable
-   a-join                :- ::lib.schema.join/join]
+   a-join                :- PartialJoin]
   (let [origin (-> a-join :stages first)]
     (cond
       (:source-card origin)  (lib.metadata/card metadata-providerable (:source-card origin))
@@ -808,16 +820,23 @@
   (m/find-first lib.types.isa/primary-key?
                 (lib.metadata.calculation/visible-columns query stage-number x)))
 
+(mu/defn ^:private fk-column-for-pk-in :- [:maybe lib.metadata/ColumnMetadata]
+  [pk-col          :- [:maybe lib.metadata/ColumnMetadata]
+   visible-columns :- [:maybe [:sequential lib.metadata/ColumnMetadata]]]
+  (when-let [pk-id (:id pk-col)]
+    (m/find-first (fn [{:keys [fk-target-field-id], :as col}]
+                    (and (lib.types.isa/foreign-key? col)
+                         (= fk-target-field-id pk-id)))
+                  visible-columns)))
+
 (mu/defn ^:private fk-column-for :- [:maybe lib.metadata/ColumnMetadata]
   "Given a query stage find an FK column that points to the PK `pk-col`."
   [query        :- ::lib.schema/query
    stage-number :- :int
    pk-col       :- [:maybe lib.metadata/ColumnMetadata]]
-  (when-let [pk-id (:id pk-col)]
-    (m/find-first (fn [{:keys [fk-target-field-id], :as col}]
-                    (and (lib.types.isa/foreign-key? col)
-                         (= fk-target-field-id pk-id)))
-                  (lib.metadata.calculation/visible-columns query stage-number (lib.util/query-stage query stage-number)))))
+  (when pk-col
+    (let [visible-columns (lib.metadata.calculation/visible-columns query stage-number (lib.util/query-stage query stage-number))]
+      (fk-column-for-pk-in pk-col visible-columns))))
 
 (mu/defn suggested-join-condition :- [:maybe ::lib.schema.expression/boolean] ; i.e., a filter clause
   "Return a suggested default join condition when constructing a join against `joinable`, e.g. a Table, Saved
@@ -830,9 +849,14 @@
   ([query         :- ::lib.schema/query
     stage-number  :- :int
     joinable]
-   (when-let [pk-col (pk-column query stage-number joinable)]
-     (when-let [fk-col (fk-column-for query stage-number pk-col)]
-       (lib.filter/filter-clause (lib.filter.operator/operator-def :=) fk-col pk-col)))))
+   (letfn [(filter-clause [x y]
+             (lib.filter/filter-clause (lib.filter.operator/operator-def :=) x y))]
+     (or (when-let [pk-col (pk-column query stage-number joinable)]
+           (when-let [fk-col (fk-column-for query stage-number pk-col)]
+             (filter-clause fk-col pk-col)))
+         (when-let [pk-col (pk-column query stage-number (lib.util/query-stage query stage-number))]
+           (when-let [fk-col (fk-column-for-pk-in pk-col (lib.metadata.calculation/visible-columns query stage-number joinable))]
+             (filter-clause pk-col fk-col)))))))
 
 (defn- add-join-alias-to-joinable-columns [cols a-join]
   (let [join-alias     (current-join-alias a-join)

src/metabase/lib/table.cljc

 (ns metabase.lib.table
   (:require
-   [metabase.lib.join :as lib.join]
    [metabase.lib.metadata :as lib.metadata]
    [metabase.lib.metadata.calculation :as lib.metadata.calculation]
    [metabase.lib.util :as lib.util]
@@ -55,11 +54,3 @@
                        :lib/source               :source/table-defaults
                        :lib/source-column-alias  (:name col)
                        :lib/desired-column-alias (unique-name-fn (or (:name col) ""))))))))
-
-(defmethod lib.join/join-clause-method :metadata/table
-  [{::keys [join-alias join-fields], :as table-metadata}]
-  (cond-> (lib.join/join-clause {:lib/type     :mbql.stage/mbql
-                                 :lib/options  {:lib/uuid (str (random-uuid))}
-                                 :source-table (:id table-metadata)})
-    join-alias  (lib.join/with-join-alias join-alias)
-    join-fields (lib.join/with-join-fields join-fields)))

test/metabase/lib/join_test.cljc

@@ -604,7 +604,7 @@
            {:input :none, :expected {:fields :none}}
            ;; (with-join-fields ... []) should set :fields to :none
            {:input [], :expected {:fields :none}}
-           {:input nil, :expected nil}]]
+           {:input nil, :expected {:fields :all}}]]
     (test-with-join-fields input expected)))
 
 (deftest ^:parallel with-join-fields-explicit-fields-test
@@ -858,11 +858,14 @@
   ;; this is to preserve the existing behavior from MLv1, it doesn't necessarily make sense, but we don't want to have
   ;; to update a million tests, right? Once v1-compatible joins lands then maybe we can go in and make this work,
   ;; since it seems like it SHOULD work.
-  (testing "Don't suggest join conditions for a PK -> FK relationship"
-    (is (nil?
-         (lib/suggested-join-condition
-          (lib/query meta/metadata-provider (meta/table-metadata :categories))
-          (meta/table-metadata :venues))))))
+  (testing "DO suggest join conditions for a PK -> FK relationship"
+    (is (=? [:=
+             {}
+             [:field {} (meta/id :categories :id)]
+             [:field {} (meta/id :venues :category-id)]]
+            (lib/suggested-join-condition
+             (lib/query meta/metadata-provider (meta/table-metadata :categories))
+             (meta/table-metadata :venues))))))
 
 (deftest ^:parallel suggested-join-condition-fk-from-join-test
   (testing "DO suggest join conditions for a FK -> PK relationship if the FK comes from a join"
@@ -1131,3 +1134,27 @@
           (is (= {:lib/type :option/temporal-bucketing, :unit :month} (lib/temporal-bucket lhs-column))))
         (let [[rhs-column] (filter :selected? (lib/join-condition-rhs-columns query 0 join lhs rhs))]
           (is (= {:lib/type :option/temporal-bucketing, :unit :month} (lib/temporal-bucket rhs-column))))))))
+
+(deftest ^:parallel join-a-table-test
+  (testing "As a convenience, we should support calling `join` with a Table metadata and do the right thing automatically"
+    (is (=? {:stages [{:source-table (meta/id :orders)
+                       :joins        [{:stages     [{:source-table (meta/id :products)}]
+                                       :fields     :all
+                                       :alias      "Products"
+                                       :conditions [[:=
+                                                     {}
+                                                     [:field {} (meta/id :orders :product-id)]
+                                                     [:field {:join-alias "Products"} (meta/id :products :id)]]]}]}]}
+            (-> (lib/query meta/metadata-provider (meta/table-metadata :orders))
+                (lib/join (meta/table-metadata :products)))))
+    (testing "with reverse PK <- FK relationship"
+      (is (=? {:stages [{:source-table (meta/id :products)
+                         :joins        [{:stages     [{:source-table (meta/id :orders)}]
+                                         :fields     :all
+                                         :alias      "Orders"
+                                         :conditions [[:=
+                                                       {}
+                                                       [:field {} (meta/id :products :id)]
+                                                       [:field {:join-alias "Orders"} (meta/id :orders :product-id)]]]}]}]}
+              (-> (lib/query meta/metadata-provider (meta/table-metadata :products))
+                  (lib/join (meta/table-metadata :orders))))))))