[MLv2] Native query functions for changing db and collection-name (#32051)

* [MLv2] Native query functions for changing db and collection-name

* Rename feature to be more specific

* Use native-extras based on feedback

* Fix missed rename

modules/drivers/mongo/src/metabase/driver/mongo.clj

@@ -302,6 +302,10 @@
       :semantic-version
       (driver.u/semantic-version-gte [4 2])))
 
+(defmethod driver/database-supports? [:mongo :native-requires-specified-collection]
+  [_driver _feature _db]
+  true)
+
 (defmethod driver/mbql->native :mongo
   [_ query]
   (mongo.qp/mbql->native query))

modules/drivers/mongo/test/metabase/driver/mongo_test.clj

@@ -11,6 +11,8 @@
    [metabase.driver.mongo.query-processor :as mongo.qp]
    [metabase.driver.mongo.util :as mongo.util]
    [metabase.driver.util :as driver.u]
+   [metabase.lib.core :as lib]
+   [metabase.lib.metadata.jvm :as lib.metadata.jvm]
    [metabase.mbql.util :as mbql.u]
    [metabase.models.card :refer [Card]]
    [metabase.models.database :refer [Database]]
@@ -93,7 +95,8 @@
       (testing (str "supports with " dbms_version)
         (is (= expected
                (let [db (first (t2/insert-returning-instances! Database {:name "dummy", :engine "mongo", :dbms_version dbms_version}))]
-                 (driver/database-supports? :mongo :expressions db))))))))
+                 (driver/database-supports? :mongo :expressions db))))))
+    (is (= #{:collection} (lib/required-native-extras (lib.metadata.jvm/application-database-metadata-provider (mt/id)))))))
 
 
 (def ^:private native-query

src/metabase/driver.clj

@@ -506,7 +506,10 @@
     :test/jvm-timezone-setting
 
     ;; Does the driver support connection impersonation (i.e. overriding the role used for individual queries)?
-    :connection-impersonation})
+    :connection-impersonation
+
+    ;; Does the driver require specifying a collection (table) for native queries? (mongo)
+    :native-requires-specified-collection})
 
 
 (defmulti supports?

src/metabase/lib/core.cljc

@@ -201,6 +201,10 @@
    with-native-query
    template-tags
    with-template-tags
+   required-native-extras
+   native-extras
+   with-native-extras
+   with-different-database
    extract-template-tags]
   [lib.order-by
    change-direction

src/metabase/lib/js.cljs

@@ -586,6 +586,30 @@
   [a-query]
   (lib.core/TemplateTags-> (lib.core/template-tags a-query)))
 
+(defn ^:export required-native-extras
+  "Returns whether the extra keys required by the database."
+  [database-id metadata]
+  (to-array (lib.core/required-native-extras (metadataProvider database-id metadata))))
+
+(defn ^:export with-different-database
+  "Changes the database for this query. The first stage must be a native type.
+   Native extras must be provided if the new database requires it."
+  ([a-query database-id metadata]
+   (with-different-database a-query database-id metadata nil))
+  ([a-query database-id metadata native-extras]
+   (lib.core/with-different-database a-query (metadataProvider database-id metadata) (js->clj native-extras :keywordize-keys true))))
+
+(defn ^:export with-native-extras
+  "Updates the extras required for the db to run this query.
+   The first stage must be a native type. Will ignore extras not in `required-native-extras`"
+  [a-query native-extras]
+  (lib.core/with-native-extras a-query (js->clj native-extras :keywordize-keys true)))
+
+(defn ^:export native-extras
+  "Returns the extra keys for native queries associated with this query."
+  [a-query]
+  (clj->js (lib.core/native-extras 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/native.cljc

@@ -13,7 +13,8 @@
    [metabase.lib.util :as lib.util]
    [metabase.shared.util.i18n :as i18n]
    [metabase.util.humanization :as u.humanization]
-   [metabase.util.malli :as mu]))
+   [metabase.util.malli :as mu]
+   [metabase.util.malli.registry :as mr]))
 
 (def ^:private TemplateTag
   [:map
@@ -136,24 +137,82 @@
        "Converter from a map of `TemplateTag`s keyed by their string names to vanilla JS."
        (converters/outgoing TemplateTags))))
 
+(defn- assert-native-query! [stage]
+  (assert (= (:lib/type stage) :mbql.stage/native) (i18n/tru "Must be a native query")))
+
+(def ^:private all-native-extra-keys
+  #{:collection})
+
+(mr/def ::native-extras
+  [:map
+   [:collection {:optional true} ::common/non-blank-string]])
+
+(mu/defn required-native-extras :- set?
+  "Returns the extra keys that are required for this database's native queries."
+  [metadata-provider :- lib.metadata/MetadataProviderable]
+  (let [db (lib.metadata/database metadata-provider)]
+   (cond-> #{}
+    (get-in db [:features :native-requires-specified-collection])
+    (conj :collection))))
+
+(mu/defn with-native-extras :- ::lib.schema/query
+  "Updates the extras required for the db to run this query.
+   The first stage must be a native type. Will ignore extras not in `required-native-extras`"
+  [query :- ::lib.schema/query
+   native-extras :- [:maybe ::native-extras]]
+  (let [required-extras (required-native-extras query)]
+    (lib.util/update-query-stage
+      query 0
+      (fn [stage]
+        (let [extras-to-remove (set/difference all-native-extra-keys required-extras)
+              stage-without-old-extras (apply dissoc stage extras-to-remove)
+              result (merge stage-without-old-extras (select-keys native-extras required-extras))
+              missing-keys (set/difference required-extras (set (keys native-extras)))]
+          (assert-native-query! (lib.util/query-stage query 0))
+          (assert (empty? missing-keys)
+                  (i18n/tru "Missing extra, required keys for native query: {0}"
+                            (pr-str missing-keys)))
+          result)))))
+
 (mu/defn native-query :- ::lib.schema/query
   "Create a new native query.
 
   Native in this sense means a pMBQL query with a first stage that is a native query."
   ([metadata-providerable :- lib.metadata/MetadataProviderable
     inner-query :- ::common/non-blank-string]
-   (native-query metadata-providerable nil inner-query))
+   (native-query metadata-providerable inner-query nil nil))
 
   ([metadata-providerable :- lib.metadata/MetadataProviderable
-    results-metadata      :- [:maybe lib.metadata/StageMetadata]
-    inner-query :- ::common/non-blank-string]
+    inner-query :- ::common/non-blank-string
+    results-metadata :- [:maybe lib.metadata/StageMetadata]
+    native-extras :- [:maybe ::native-extras]]
    (let [tags (extract-template-tags inner-query)]
-     (lib.query/query-with-stages metadata-providerable
-                                  [(-> {:lib/type           :mbql.stage/native
-                                        :lib/stage-metadata results-metadata
-                                        :template-tags      tags
-                                        :native             inner-query}
-                                       lib.options/ensure-uuid)]))))
+     (-> (lib.query/query-with-stages metadata-providerable
+                                      [(-> {:lib/type           :mbql.stage/native
+                                            :lib/stage-metadata results-metadata
+                                            :template-tags      tags
+                                            :native             inner-query}
+                                           lib.options/ensure-uuid)])
+         (with-native-extras native-extras)))))
+
+(mu/defn with-different-database :- ::lib.schema/query
+  "Changes the database for this query. The first stage must be a native type.
+   Native extras must be provided if the new database requires it."
+  ([query :- ::lib.schema/query
+    metadata-provider :- lib.metadata/MetadataProviderable]
+   (with-different-database query metadata-provider nil))
+  ([query :- ::lib.schema/query
+    metadata-provider :- lib.metadata/MetadataProviderable
+    native-extras :- [:maybe ::native-extras]]
+   (assert-native-query! (lib.util/query-stage query 0))
+   ;; Changing the database should also clean up template tags, see #31926
+   (-> (lib.query/query-with-stages metadata-provider (:stages query))
+       (with-native-extras native-extras))))
+
+(mu/defn native-extras :- [:maybe ::native-extras]
+  "Returns the extra keys for native queries associated with this query."
+  [query :- ::lib.schema/query]
+  (not-empty (select-keys (lib.util/query-stage query 0) (required-native-extras query))))
 
 (mu/defn with-native-query :- ::lib.schema/query
   "Update the raw native query, the first stage must already be a native type.
@@ -162,8 +221,8 @@
    inner-query :- ::common/non-blank-string]
   (lib.util/update-query-stage
     query 0
-    (fn [{existing-tags :template-tags stage-type :lib/type :as stage}]
-      (assert (= stage-type :mbql.stage/native) (i18n/tru "Must be a native query"))
+    (fn [{existing-tags :template-tags :as stage}]
+      (assert-native-query! stage)
       (assoc stage
         :native inner-query
         :template-tags (extract-template-tags inner-query existing-tags)))))
@@ -174,8 +233,8 @@
    tags :- TemplateTags]
   (lib.util/update-query-stage
     query 0
-    (fn [{existing-tags :template-tags stage-type :lib/type :as stage}]
-      (assert (= stage-type :mbql.stage/native) (i18n/tru "Must be a native query"))
+    (fn [{existing-tags :template-tags :as stage}]
+      (assert-native-query! stage)
       (let [valid-tags (keys existing-tags)]
         (assoc stage :template-tags
                (m/deep-merge existing-tags (select-keys tags valid-tags)))))))

test/metabase/lib/native_test.cljc

@@ -8,6 +8,7 @@
    [metabase.lib.metadata.calculation :as lib.metadata.calculation]
    [metabase.lib.native :as lib.native]
    [metabase.lib.test-metadata :as meta]
+   [metabase.lib.test-metadata.graph-provider :as meta.graph-provider]
    [metabase.util.humanization :as u.humanization]))
 
 (deftest ^:parallel variable-tag-test
@@ -169,10 +170,10 @@
            :stages   [{:lib/type    :mbql.stage/native
                        :lib/options {:lib/uuid string?}
                        :native      "SELECT * FROM VENUES;"}]}
-          (lib/native-query meta/metadata-provider meta/qp-results-metadata "SELECT * FROM VENUES;"))))
+          (lib/native-query meta/metadata-provider "SELECT * FROM VENUES;" meta/qp-results-metadata nil))))
 
 (deftest ^:parallel native-query-suggested-name-test
-  (let [query (lib/native-query meta/metadata-provider meta/qp-results-metadata "SELECT * FROM VENUES;")]
+  (let [query (lib/native-query meta/metadata-provider "SELECT * FROM VENUES;" meta/qp-results-metadata nil)]
     (is (= "Native query"
            (lib.metadata.calculation/describe-query query)))
     (is (nil? (lib.metadata.calculation/suggested-name query)))))
@@ -199,7 +200,12 @@
       (is (empty?
             (-> query
                 (lib/with-native-query "select * from venues")
-                lib/template-tags))))))
+                lib/template-tags))))
+    (is (thrown-with-msg?
+          #?(:clj Throwable :cljs :default)
+          #"Must be a native query"
+          (-> (lib/query meta/metadata-provider (meta/table-metadata :venues))
+              (lib/with-native-query "foobar"))))))
 
 (deftest ^:parallel with-template-tags-test
   (let [query (lib/native-query meta/metadata-provider "select * from venues where id = {{myid}}")
@@ -218,4 +224,63 @@
       (is (= original-tags
              (-> query
                  (lib/with-template-tags {"garbage" (assoc (get original-tags "myid") :display-name "Foobar")})
-                 lib/template-tags))))))
+                 lib/template-tags))))
+    (is (thrown-with-msg?
+          #?(:clj Throwable :cljs :default)
+          #"Must be a native query"
+          (-> (lib/query meta/metadata-provider (meta/table-metadata :venues))
+              (lib/with-template-tags {"myid" (assoc (get original-tags "myid") :display-name "My ID")}))))))
+
+(defn ^:private metadata-provider-requiring-collection []
+  (meta.graph-provider/->SimpleGraphMetadataProvider (-> meta/metadata
+                                                         (update :features conj :native-requires-specified-collection))))
+
+(deftest ^:parallel native-query+collection-test
+  (testing "building when collection is not required"
+    (is (=? {:stages [(complement :collection)]}
+            (lib/native-query meta/metadata-provider "myquery")))
+    (is (=? {:stages [(complement :collection)]}
+            (lib/native-query meta/metadata-provider "myquery" nil {:collection "mycollection"}))))
+  (testing "building when requires collection"
+    (is (=? {:stages [{:collection "mycollection"}]}
+            (lib/native-query (metadata-provider-requiring-collection) "myquery" nil {:collection "mycollection"})))
+    (is (thrown-with-msg?
+          #?(:clj Throwable :cljs :default)
+          #"Missing extra, required keys for native query: .*:collection.*"
+          (lib/native-query (metadata-provider-requiring-collection) "myquery")))))
+
+(deftest ^:parallel with-different-database-test
+  (let [query (lib/native-query meta/metadata-provider "myquery")]
+    (testing "database id should change"
+      (is (=? {:database 9999}
+              (-> query
+                  (lib/with-different-database
+                    (meta.graph-provider/->SimpleGraphMetadataProvider
+                      (assoc meta/metadata :id 9999)))))))
+    (testing "Checks collection requirement"
+      (is (=? {:stages [(complement :collection)]}
+              (-> query
+                  (lib/with-different-database meta/metadata-provider {:collection "mycollection"}))))
+      (is (thrown-with-msg?
+            #?(:clj Throwable :cljs :default)
+            #"Missing extra, required keys for native query: .*:collection.*"
+            (-> query
+                (lib/with-different-database (metadata-provider-requiring-collection))))))
+    (is (thrown-with-msg?
+          #?(:clj Throwable :cljs :default)
+          #"Must be a native query"
+          (-> (lib/query meta/metadata-provider (meta/table-metadata :venues))
+              (lib/with-different-database meta/metadata-provider))))))
+
+(deftest ^:parallel with-native-collection-test
+  (is (=? {:stages [{:collection "mynewcollection"}]}
+          (-> (lib/native-query (metadata-provider-requiring-collection) "myquery" nil {:collection "mycollection"})
+              (lib/with-native-extras {:collection "mynewcollection"}))))
+  (is (=? {:stages [(complement :collection)]}
+        (-> (lib/native-query meta/metadata-provider "myquery")
+            (lib/with-native-extras {:collection "mycollection"}))))
+  (is (thrown-with-msg?
+        #?(:clj Throwable :cljs :default)
+        #"Must be a native query"
+        (-> (lib/query (metadata-provider-requiring-collection) (meta/table-metadata :venues))
+            (lib/with-native-extras {:collection "mycollection"})))))