[MLv2] Support segments (#32882)

* Support segments

Fixes #32409.

frontend/src/metabase-lib/segments.ts

+import * as ML from "cljs/metabase.lib.js";
+
+import type { SegmentMetadata, Query } from "./types";
+
+export function availableSegments(query: Query): SegmentMetadata[] {
+  return ML.available_segments(query);
+}

frontend/src/metabase-lib/types.ts

@@ -14,6 +14,9 @@ export type TableMetadata = unknown & { _opaque: typeof TableMetadata };
 declare const CardMetadata: unique symbol;
 export type CardMetadata = unknown & { _opaque: typeof CardMetadata };
 
+declare const SegmentMetadata: unique symbol;
+export type SegmentMetadata = unknown & { _opaque: typeof SegmentMetadata };
+
 declare const MetricMetadata: unique symbol;
 export type MetricMetadata = unknown & { _opaque: typeof MetricMetadata };
 

frontend/src/metabase-lib/v2.ts

@@ -15,6 +15,7 @@ export * from "./filter";
 export * from "./join";
 export * from "./limit";
 export * from "./metadata";
+export * from "./segments";
 export * from "./metrics";
 export * from "./order_by";
 export * from "./query";

src/metabase/lib/core.cljc

@@ -234,6 +234,8 @@
    rename-join
    replace-clause
    replace-join]
+  [lib.segment
+   available-segments]
   [lib.stage
    append-stage
    drop-stage]

src/metabase/lib/filter.cljc

@@ -6,6 +6,7 @@
    [clojure.string :as str]
    [medley.core :as m]
    [metabase.lib.common :as lib.common]
+   [metabase.lib.dispatch :as lib.dispatch]
    [metabase.lib.equality :as lib.equality]
    [metabase.lib.filter.operator :as lib.filter.operator]
    [metabase.lib.hierarchy :as lib.hierarchy]
@@ -158,9 +159,12 @@
   ([query :- :metabase.lib.schema/query
     stage-number :- [:maybe :int]
     boolean-expression]
-   (let [stage-number (clojure.core/or stage-number -1)
-         new-filter (lib.common/->op-arg boolean-expression)]
-     (lib.util/update-query-stage query stage-number update :filters (fnil conj []) new-filter))))
+   ;; if this is a Segment metadata, convert it to `:segment` MBQL clause before adding
+   (if (clojure.core/= (lib.dispatch/dispatch-value boolean-expression) :metadata/segment)
+     (recur query stage-number (lib.ref/ref boolean-expression))
+     (let [stage-number (clojure.core/or stage-number -1)
+           new-filter (lib.common/->op-arg boolean-expression)]
+       (lib.util/update-query-stage query stage-number update :filters (fnil conj []) new-filter)))))
 
 (mu/defn filters :- [:maybe [:ref ::lib.schema/filters]]
   "Returns the current filters in stage with `stage-number` of `query`.

src/metabase/lib/js.cljs

@@ -654,6 +654,12 @@
   [a-query]
   (clj->js (lib.core/native-extras a-query)))
 
+(defn ^:export available-segments
+  "Get a list of Segments that you may consider using as filters for a query. Returns JS array of opaque Segment
+  metadata objects."
+  [a-query]
+  (to-array (lib.core/available-segments a-query)))
+
 (defn ^:export available-metrics
   "Get a list of Metrics that you may consider using as aggregations for a query. Returns JS array of opaque Metric
   metadata objects."

src/metabase/lib/js/metadata.cljs

@@ -382,6 +382,12 @@
         :when              (and a-metric (= (:table-id a-metric) table-id))]
     a-metric))
 
+(defn- segments [metadata table-id]
+  (for [[_id segment-delay] (some-> metadata :segments deref)
+        :let               [a-segment (some-> segment-delay deref)]
+        :when              (and a-segment (= (:table-id a-segment) table-id))]
+    a-segment))
+
 (defn metadata-provider
   "Use a `metabase-lib/metadata/Metadata` as a [[metabase.lib.metadata.protocols/MetadataProvider]]."
   [database-id unparsed-metadata]
@@ -397,6 +403,7 @@
       (tables   [_this]            (tables   metadata database-id))
       (fields   [_this table-id]   (fields   metadata table-id))
       (metrics  [_this table-id]   (metrics  metadata table-id))
+      (segments [_this table-id]   (segments metadata table-id))
 
       ;; for debugging: call [[clojure.datafy/datafy]] on one of these to parse all of our metadata and see the whole
       ;; thing at once.

src/metabase/lib/metadata.cljc

@@ -135,7 +135,12 @@
    {:error/message "Valid Segment metadata"}
    [:lib/type [:= :metadata/segment]]
    [:id       ::lib.schema.id/segment]
-   [:name     ::lib.schema.common/non-blank-string]])
+   [:name     ::lib.schema.common/non-blank-string]
+   [:table-id   ::lib.schema.id/table]
+   ;; the MBQL snippet defining this Segment; this may still be in legacy
+   ;; format. [[metabase.lib.segment/segment-definition]] handles conversion to pMBQL if needed.
+   [:definition :map]
+   [:description {:optional true} [:maybe ::lib.schema.common/non-blank-string]]])
 
 (def MetricMetadata
   "Malli schema for a legacy v1 [[metabase.models.metric]], but with kebab-case keys. A Metric defines an MBQL snippet

src/metabase/lib/metadata/protocols.cljc

@@ -64,8 +64,12 @@
   the [[metabase.lib.metadata/ColumnMetadata]] schema. If no such Table exists, this should error.")
 
   (metrics [metadata-provider table-id]
-    "Return a sequence of legacy v1 Metrics associated with a Table with the given `table-id`. Metrics should satisfy
-  the [[metabase.lib.metadata/MetricMetadata]] schema. If no such Table exists, this should error."))
+    "Return a sequence of legacy Metrics associated with a Table with the given `table-id`. Metrics should satisfy
+  the [[metabase.lib.metadata/MetricMetadata]] schema. If no such Table exists, this should error.")
+
+  (segments [metadata-provider table-id]
+    "Return a sequence of legacy Segments associated with a Table with the given `table-id`. Segments should satisfy
+  the [[metabase.lib.metadata/SegmentMetadata]] schema. If no Table with ID `table-id` exists, this should error."))
 
 (defn metadata-provider?
   "Whether `x` is a valid [[MetadataProvider]]."

src/metabase/lib/schema.cljc

@@ -58,8 +58,14 @@
 (mr/def ::fields
   [:sequential {:min 1} [:ref ::ref/ref]])
 
+;; this is just for enabling round-tripping filters with named segment references
+(mr/def ::filterable
+  [:or
+   [:ref ::expression/boolean]
+   [:tuple [:= :segment] :map :string]])
+
 (mr/def ::filters
-  [:sequential {:min 1} [:ref ::expression/boolean]])
+  [:sequential {:min 1} ::filterable])
 
 (defn- bad-ref-clause? [ref-type valid-ids x]
   (and (vector? x)

src/metabase/lib/schema/ref.cljc

@@ -94,6 +94,11 @@
 
 (lib.hierarchy/derive :aggregation ::ref)
 
+(mbql-clause/define-tuple-mbql-clause :segment :- :type/Boolean
+  #_segment-id [:schema [:ref ::id/segment]])
+
+(lib.hierarchy/derive :segment ::ref)
+
 (mbql-clause/define-tuple-mbql-clause :metric :- ::expression/type.unknown
   #_metric-id [:schema [:ref ::id/metric]])
 

src/metabase/lib/segment.cljc

 (ns metabase.lib.segment
-  "A Metric is a saved MBQL query stage snippet with `:filter`. Segments are always boolean expressions."
+  "A Segment is a saved MBQL query stage snippet with `:filter`. Segments are always boolean expressions."
   (:require
    [metabase.lib.metadata :as lib.metadata]
    [metabase.lib.metadata.calculation :as lib.metadata.calculation]
-   [metabase.shared.util.i18n :as i18n]))
+   [metabase.lib.metadata.protocols :as lib.metadata.protocols]
+   [metabase.lib.options :as lib.options]
+   [metabase.lib.ref :as lib.ref]
+   [metabase.lib.schema :as lib.schema]
+   [metabase.lib.schema.id :as lib.schema.id]
+   [metabase.lib.util :as lib.util]
+   [metabase.shared.util.i18n :as i18n]
+   [metabase.util.malli :as mu]))
+
+(defn- resolve-segment [query segment-id]
+  (when (integer? segment-id)
+    (lib.metadata/segment query segment-id)))
+
+(defmethod lib.ref/ref-method :metadata/segment
+  [{:keys [id]}]
+  (lib.options/ensure-uuid [:segment {} id]))
+
+(defmethod lib.metadata.calculation/type-of-method :metadata/segment
+  [_query _stage-number _metric-metadata]
+  :type/Boolean)
+
+(defmethod lib.metadata.calculation/type-of-method :segment
+  [_query _stage-number _segment-clause]
+  :type/Boolean)
+
+(defn- fallback-display-name []
+  (i18n/tru "[Unknown Segment]"))
 
 (defmethod lib.metadata.calculation/display-name-method :metadata/segment
   [_query _stage-number segment-metadata _style]
-  (or ((some-fn :display-name :name) segment-metadata)
-      (i18n/tru "[Unknown Segment]")))
+  (or (:display-name segment-metadata)
+      (:name segment-metadata)
+      (fallback-display-name)))
 
 (defmethod lib.metadata.calculation/display-name-method :segment
   [query stage-number [_tag _opts segment-id-or-name] style]
   (or (when (integer? segment-id-or-name)
         (when-let [segment-metadata (lib.metadata/segment query segment-id-or-name)]
           (lib.metadata.calculation/display-name query stage-number segment-metadata style)))
-      (i18n/tru "[Unknown Segment]")))
+      (fallback-display-name)))
+
+(mu/defn ^:private filtering-by-segment? :- :boolean
+  "Whether a given stage of a query currently includes a `:segment` ref clause in its filters."
+  [query        :- ::lib.schema/query
+   stage-number :- :int
+   segment-id    :- ::lib.schema.id/segment]
+  (boolean
+   (some (fn [[tag _opts id]]
+           (and (= tag :segment)
+                (= id segment-id)))
+         (:filters (lib.util/query-stage query stage-number)))))
+
+(defmethod lib.metadata.calculation/display-info-method :metadata/segment
+  [query stage-number {:keys [id description], :as segment-metadata}]
+  (let [default-display-info-method (get-method lib.metadata.calculation/display-info-method :default)
+        default-display-info        (default-display-info-method query stage-number segment-metadata)]
+    (cond-> default-display-info
+      description                                   (assoc :description description)
+      (filtering-by-segment? query stage-number id) (assoc :selected true))))
+
+(defmethod lib.metadata.calculation/display-info-method :segment
+  [query stage-number [_tag _opts segment-id-or-name]]
+  (if-let [segment-metadata (resolve-segment query segment-id-or-name)]
+    (lib.metadata.calculation/display-info query stage-number segment-metadata)
+    {:effective-type    :type/Boolean
+     :display-name      (fallback-display-name)
+     :long-display-name (fallback-display-name)}))
+
+(mu/defn available-segments :- [:maybe [:sequential {:min 1} lib.metadata/SegmentMetadata]]
+  "Get a list of Segments that you may consider using as filter for a query. Only Segments that have the same
+  `table-id` as the `source-table` for this query will be suggested."
+  [query :- ::lib.schema/query]
+  (when-let [source-table-id (lib.util/source-table-id query)]
+    (not-empty (lib.metadata.protocols/segments (lib.metadata/->metadata-provider query) source-table-id))))

test/metabase/lib/convert_test.cljc

@@ -238,6 +238,19 @@
 (deftest ^:parallel round-trip-test
   ;; Miscellaneous queries that have caused test failures in the past, captured here for quick feedback.
   (are [query] (test-round-trip query)
+    ;; query with segment
+    {:database 282,
+     :type :query,
+     :query {:source-table 661,
+             :aggregation [[:metric "ga:totalEvents"]],
+             :filter [:and
+                      [:segment "gaid::-4"]
+                      [:segment 42]
+                      [:= [:field 1972 nil] "Run Query"]
+                      [:time-interval [:field 1974 nil] -30 :day]
+                      [:!= [:field 1973 nil] "(not set)" "url"]],
+             :breakout [[:field 1973 nil]]}}
+
     ;; :aggregation-options on a non-aggregate expression with an inner aggregate.
     {:database 194
      :query {:aggregation [[:aggregation-options
@@ -387,20 +400,20 @@
     {:type :query
      :database 5
      :query {:source-table 5822,
-     :expressions {"Additional Information Capture" [:coalesce
-                                                     [:field 519195 nil]
-                                                     [:field 519196 nil]
-                                                     [:field 519194 nil]
-                                                     [:field 519193 nil]
-                                                     "None"],
-                   "Additional Info Present" [:case
-                                              [[[:= [:expression "Additional Information Capture"] "None"]
-                                                "No"]]
-                                              {:default "Yes"}]},
-      :filter [:= [:field 518086 nil] "active"],
-      :aggregation [[:aggregation-options
-                     [:share [:= [:expression "Additional Info Present"] "Yes"]]
-                     {:name "Additional Information", :display-name "Additional Information"}]]}}))
+             :expressions {"Additional Information Capture" [:coalesce
+                                                             [:field 519195 nil]
+                                                             [:field 519196 nil]
+                                                             [:field 519194 nil]
+                                                             [:field 519193 nil]
+                                                             "None"],
+                           "Additional Info Present" [:case
+                                                      [[[:= [:expression "Additional Information Capture"] "None"]
+                                                        "No"]]
+                                                      {:default "Yes"}]},
+             :filter [:= [:field 518086 nil] "active"],
+             :aggregation [[:aggregation-options
+                            [:share [:= [:expression "Additional Info Present"] "Yes"]]
+                            {:name "Additional Information", :display-name "Additional Information"}]]}}))
 
 (deftest ^:parallel round-trip-aggregation-with-case-test
   (test-round-trip

test/metabase/lib/filter_test.cljc

@@ -99,7 +99,7 @@
        :day))
 
     (testing "segment"
-      (doseq [id [7 "6"]]
+      (let [id 7]
         (test-clause
          [:segment {:lib/uuid string?} id]
          lib/segment

test/metabase/lib/schema/filter_test.cljc

@@ -71,8 +71,7 @@
            [:time-interval field :last :hour]
            [:time-interval field 4 :hour]
            [:time-interval {:include-current true} field :next :day]
-           [:segment 1]
-           [:segment "segment-id"]]]
+           [:segment 1]]]
       (doseq [op (filter-ops filter-expr)]
           (is (not (identical? (get-method expression/type-of-method op)
         (testing (str op " is a registered MBQL clause (a type-of-method method is registered for it)")

test/metabase/lib/segment_test.cljc

+(ns metabase.lib.segment-test
+  (:require
+   [clojure.test :refer [are deftest is testing]]
+   [metabase.lib.core :as lib]
+   [metabase.lib.metadata :as lib.metadata]
+   [metabase.lib.metadata.calculation :as lib.metadata.calculation]
+   [metabase.lib.test-metadata :as meta]
+   [metabase.lib.test-util :as lib.tu]
+   #?@(:cljs ([metabase.test-runner.assert-exprs.approximately-equal]))))
+
+#?(:cljs (comment metabase.test-runner.assert-exprs.approximately-equal/keep-me))
+
+(def ^:private segment-id 100)
+
+(def ^:private segment-definition
+  {:source-table (meta/id :venues)
+   :aggregation  [[:count]]
+   :filter       [:and
+                  [:> [:field (meta/id :venues :id) nil] [:* [:field (meta/id :venues :price) nil] 11]]
+                  [:contains [:field (meta/id :venues :name) nil] "BBQ" {:case-sensitive true}]]})
+
+(def ^:private metadata-provider
+  (lib.tu/mock-metadata-provider
+   {:database meta/metadata
+    :tables   [(meta/table-metadata :venues)]
+    :fields   (mapv #(meta/field-metadata :venues %) [:id :price])
+    :segments [{:id          segment-id
+                :name        "PriceID-BBQ"
+                :table-id    (meta/id :venues)
+                :definition  segment-definition
+                :description "The ID is greater than 11 times the price and the name contains \"BBQ\"."}]}))
+
+(def ^:private segment-clause
+  [:segment {:lib/uuid (str (random-uuid))} segment-id])
+
+(def ^:private query-with-segment
+  (-> (lib/query metadata-provider (meta/table-metadata :venues))
+      (lib/filter (lib/= (meta/field-metadata :venues :id) 5))
+      (lib/filter segment-clause)))
+
+(def ^:private segment-metadata
+  (lib.metadata/segment query-with-segment segment-id))
+
+(deftest ^:parallel query-suggested-name-test
+  (is (= "Venues, Filtered by ID equals 5 and PriceID-BBQ"
+         (lib.metadata.calculation/suggested-name query-with-segment))))
+
+(deftest ^:parallel display-info-test
+  (are [segment] (=? {:name              "priceid_bbq",
+                      :display-name      "PriceID-BBQ",
+                      :long-display-name "PriceID-BBQ",
+                      :effective-type    :type/Boolean,
+                      :description       "The ID is greater than 11 times the price and the name contains \"BBQ\".",
+                      :selected          true}
+                     (lib.metadata.calculation/display-info query-with-segment segment))
+    segment-clause
+    segment-metadata))
+
+(deftest ^:parallel display-info-unselected-segment-test
+  (testing "Include `:selected false` in display info for Segments not in aggregations"
+    (are [segment] (not (:selected (lib.metadata.calculation/display-info lib.tu/venues-query segment)))
+      segment-clause
+      segment-metadata)))
+
+(deftest ^:parallel unknown-display-info-test
+  (is (=? {:effective-type    :type/Boolean
+           :display-name      "[Unknown Segment]"
+           :long-display-name "[Unknown Segment]"}
+          (lib.metadata.calculation/display-info query-with-segment [:segment {} (inc segment-id)]))))
+
+(deftest ^:parallel available-segments-test
+  (testing "Should return Segments with the same Table ID as query's `:source-table`"
+    (is (=? [{:lib/type    :metadata/segment
+              :id          segment-id
+              :name        "PriceID-BBQ"
+              :table-id    (meta/id :venues)
+              :definition  segment-definition
+              :description "The ID is greater than 11 times the price and the name contains \"BBQ\"."}]
+            (lib/available-segments (lib/query metadata-provider (meta/table-metadata :venues))))))
+  (testing "query with different Table -- don't return Segments"
+    (is (nil? (lib/available-segments (lib/query metadata-provider (meta/table-metadata :orders)))))))
+
+(deftest ^:parallel filter-with-segment-test
+  (testing "Should be able to pass a Segment metadata to `filter`"
+    (let [query    (lib/query metadata-provider (meta/table-metadata :venues))
+          segments (lib/available-segments query)]
+      (is (= 1
+             (count segments)))
+      ;; test with both `:metadata/segment` and with a `:segment` ref clause
+      (doseq [segment [(first segments)
+                      [:segment {:lib/uuid (str (random-uuid))} segment-id]]]
+        (testing (pr-str (list 'lib/filter 'query segment))
+          (let [query' (lib/filter query segment)]
+            (is (=? {:lib/type :mbql/query
+                     :stages   [{:lib/type     :mbql.stage/mbql
+                                 :source-table (meta/id :venues)
+                                 :filters  [[:segment {:lib/uuid string?} segment-id]]}]}
+                    query'))
+            (is (=? [[:segment {:lib/uuid string?} segment-id]]
+                    (lib/filters query')))
+            (is (=? [{:name              "priceid_bbq",
+                      :display-name      "PriceID-BBQ",
+                      :long-display-name "PriceID-BBQ",
+                      :effective-type    :type/Boolean,
+                      :description       "The ID is greater than 11 times the price and the name contains \"BBQ\".",
+                      :selected          true}]
+                    (map (partial lib/display-info query')
+                         (lib/filters query'))))))))))

test/metabase/lib/test_util.cljc

@@ -89,6 +89,9 @@
     (metrics  [_this table-id]   (for [metric metrics
                                        :when (= (:table-id metric) table-id)]
                                    (assoc metric :lib/type :metadata/metric)))
+    (segments [_this table-id]   (for [segment segments
+                                       :when (= (:table-id segment) table-id)]
+                                   (assoc segment :lib/type :metadata/segment)))
 
     clojure.core.protocols/Datafiable
     (datafy [_this]