Allow a driver to customize the query remarks (#12422)

* Allow a driver to customize the query remarks

This allows driver-specific metadata to be added to each query. If the
driver can apply the default metadata, it should call
default-query->remark

Removed timezone test that isn't testing the driver itself, since the
driver doesn't specify a default timezone

Removed deprecated code from the redshift test

Leave the original parameters from the front-end as :user-parameters so
that they are accessible down in the drivers

Transform the original parameters into the field_values in the query
remark

backend/mbql/src/metabase/mbql/schema.clj

@@ -159,7 +159,8 @@
   {(s/optional-key :database_type) (s/maybe su/NonBlankString)
    (s/optional-key :base_type)     (s/maybe su/FieldType)
    (s/optional-key :special_type)  (s/maybe su/FieldType)
-   (s/optional-key :unit)          (s/maybe DatetimeFieldUnit)})
+   (s/optional-key :unit)          (s/maybe DatetimeFieldUnit)
+   (s/optional-key :name)          (s/maybe su/NonBlankString)})
 
 ;; Arguments to filter clauses are automatically replaced with [:value <value> <type-info>] clauses by the
 ;; `wrap-value-literals` middleware. This is done to make it easier to implement query processors, because most driver

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

@@ -228,9 +228,9 @@
   (let [database (qp.store/database)]
     (binding [bigquery.common/*bigquery-timezone-id* (effective-query-timezone-id database)]
       (log/tracef "Running BigQuery query in %s timezone" bigquery.common/*bigquery-timezone-id*)
-      (let [sql (str "-- " (qputil/query->remark outer-query) "\n" (if (seq params)
-                                                                     (unprepare/unprepare driver (cons sql params))
-                                                                     sql))]
+      (let [sql (str "-- " (qputil/query->remark :bigquery outer-query) "\n" (if (seq params)
+                                                                               (unprepare/unprepare driver (cons sql params))
+                                                                               sql))]
         (process-native* respond database sql)))))
 
 

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

@@ -305,7 +305,7 @@
    context
    respond]
   (let [sql     (str "-- "
-                     (qputil/query->remark outer-query) "\n"
+                     (qputil/query->remark :presto outer-query) "\n"
                      (unprepare/unprepare driver (cons sql params)))
         details (merge (:details (qp.store/database))
                        settings)]

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

 (ns metabase.driver.redshift
   "Amazon Redshift Driver."
-  (:require [clojure.java.jdbc :as jdbc]
+  (:require [cheshire.core :as json]
+            [clojure.java.jdbc :as jdbc]
             [clojure.string :as str]
             [clojure.tools.logging :as log]
             [honeysql.core :as hsql]
-            [metabase.driver :as driver]
+            [metabase
+             [driver :as driver]
+             [public-settings :as pubset]]
             [metabase.driver.common :as driver.common]
             [metabase.driver.sql-jdbc
              [connection :as sql-jdbc.conn]
              [execute :as sql-jdbc.execute]]
             [metabase.driver.sql-jdbc.execute.legacy-impl :as legacy]
             [metabase.driver.sql.query-processor :as sql.qp]
+            [metabase.query-processor
+             [store :as qp.store]
+             [util :as qputil]]
             [metabase.util
              [honeysql-extensions :as hx]
              [i18n :refer [trs]]])
@@ -155,3 +161,19 @@
  sql-jdbc.execute/set-parameter
  [::legacy/use-legacy-classes-for-read-and-set OffsetTime]
  [:postgres OffsetTime])
+
+(defn query->field-values
+  "Convert a MBQL query to a map of field->value"
+  [query]
+  (into {} (map (fn [param] [(:name (qp.store/field (get-in param [:target 1 1]))) (:value param)]) (:user-parameters query))))
+
+(defmethod qputil/query->remark :redshift
+  [_ {{:keys [executed-by query-hash card-id], :as info} :info, query-type :type :as query}]
+  (str "/* partner: \"metabase\", "
+       (json/generate-string {:dashboard_id nil ;; requires metabase/metabase#11909
+                              :chart_id card-id
+                              :optional_user_id executed-by
+                              :optional_account_id (pubset/site-uuid)
+                              :filter_values (query->field-values query)})
+       " */ "
+       (qputil/default-query->remark query)))

modules/drivers/redshift/test/metabase/driver/redshift_test.clj

 (ns metabase.driver.redshift-test
-  (:require [clojure.test :refer :all]
+  (:require [clojure
+             [string :as str]
+             [test :refer :all]]
+            [metabase
+             [public-settings :as pubset]
+             [query-processor :as qp]
+             [test :as mt]
+             [util :as u]]
+            [metabase.driver.sql-jdbc.execute :as execute]
             [metabase.plugins.jdbc-proxy :as jdbc-proxy]
-            [metabase.test
-             [fixtures :as fixtures]
-             [util :as tu]]
-            [metabase.test.data.datasets :refer [expect-with-driver]]))
+            [metabase.test.data.redshift :as rstest]
+            [metabase.test.fixtures :as fixtures]))
 
 (use-fixtures :once (fixtures/initialize :plugins))
-
-(expect-with-driver :redshift
-  "UTC"
-  (tu/db-timezone-id))
+(use-fixtures :once (fixtures/initialize :db))
 
 (deftest correct-driver-test
-  (is (= "com.amazon.redshift.jdbc.Driver"
-         (.getName (class (jdbc-proxy/wrapped-driver (java.sql.DriverManager/getDriver "jdbc:redshift://host:5432/testdb")))))
-      "Make sure we're using the correct driver for Redshift"))
+  (mt/test-driver
+    :redshift
+    (is (= "com.amazon.redshift.jdbc.Driver"
+           (.getName (class (jdbc-proxy/wrapped-driver (java.sql.DriverManager/getDriver "jdbc:redshift://host:5432/testdb")))))
+        "Make sure we're using the correct driver for Redshift")))
+
+(defn- query->native [query]
+  (let [native-query (atom nil)]
+    (with-redefs [execute/prepared-statement (fn [_ _ sql _]
+                                               (reset! native-query sql)
+                                               (throw (Exception. "done")))]
+      (u/ignore-exceptions
+        (qp/process-query query))
+      @native-query)))
+
+(deftest remark-test
+  (let [expected (str/replace
+                  (str
+                   "-- /* partner: \"metabase\", {\"dashboard_id\":null,\"chart_id\":1234,\"optional_user_id\":1000,"
+                   "\"optional_account_id\":\"" (pubset/site-uuid) "\","
+                   "\"filter_values\":{\"id\":[\"1\",\"2\",\"3\"]}} */"
+                   " Metabase:: userID: 1000 queryType: MBQL queryHash: cb83d4f6eedc250edb0f2c16f8d9a21e5d42f322ccece1494c8ef3d634581fe2\n"
+                   "SELECT \"%schema%\".\"test_data_users\".\"id\" AS \"id\","
+                   " \"%schema%\".\"test_data_users\".\"name\" AS \"name\","
+                   " \"%schema%\".\"test_data_users\".\"last_login\" AS \"last_login\""
+                   " FROM \"%schema%\".\"test_data_users\""
+                   " WHERE (\"%schema%\".\"test_data_users\".\"id\" = 1 OR \"%schema%\".\"test_data_users\".\"id\" = 2"
+                   " OR \"%schema%\".\"test_data_users\".\"id\" = 3)"
+                   " LIMIT 2000")
+                  "%schema%" rstest/session-schema-name)]
+   (mt/test-driver
+    :redshift
+    (is (= expected
+           (query->native
+            (assoc
+             (mt/mbql-query users {:limit 2000})
+             :parameters [{:type "id", :target ["dimension" ["field-id" (mt/id :users :id)]], :value ["1" "2" "3"]}]
+             :info {:executed-by 1000
+                    :card-id 1234
+                    :context :ad-hoc
+                    :nested? false
+                    :query-hash (byte-array [-53, -125, -44, -10, -18, -36, 37, 14, -37, 15, 44, 22, -8, -39, -94, 30, 93, 66, -13, 34, -52, -20, -31, 73, 76, -114, -13, -42, 52, 88, 31, -30])})))
+        "if I run a Redshift query, does it get a remark added to it?"))))

modules/drivers/redshift/test/metabase/test/data/redshift.clj

@@ -46,7 +46,7 @@
 (defonce ^:private session-schema-number
   (rand-int 240)) ; there's a maximum of 256 schemas per DB so make sure we don't go over that limit
 
-(defonce ^:private session-schema-name
+(defonce session-schema-name
   (str "schema_" session-schema-number))
 
 ;; When we test against Redshift we use a session-unique schema so we can run simultaneous tests

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

@@ -128,7 +128,7 @@
 (defmethod driver/execute-reducible-query :sparksql
   [driver {:keys [database settings], {sql :query, :keys [params], :as inner-query} :native, :as outer-query} context respond]
   (let [inner-query (-> (assoc inner-query
-                               :remark (qputil/query->remark outer-query)
+                               :remark (qputil/query->remark :sparksql outer-query)
                                :query  (if (seq params)
                                          (unprepare/unprepare driver (cons sql params))
                                          sql)

src/metabase/driver/sql_jdbc/execute.clj

@@ -379,7 +379,7 @@
   {:added "0.35.0", :arglists '([driver query context respond])}
   [driver {{sql :query, params :params} :native, :as outer-query} context respond]
   {:pre [(string? sql) (seq sql)]}
-  (let [remark   (qputil/query->remark outer-query)
+  (let [remark   (qputil/query->remark driver outer-query)
         sql      (str "-- " remark "\n" sql)
         max-rows (or (mbql.u/query->max-rows-limit outer-query)
                      qp.i/absolute-max-results)]

src/metabase/query_processor/middleware/parameters.clj

 (ns metabase.query-processor.middleware.parameters
   "Middleware for substituting parameters in queries."
-  (:require [clojure.data :as data]
+  (:require [clojure
+             [data :as data]
+             [set :as set]]
             [clojure.tools.logging :as log]
             [metabase.mbql
              [normalize :as normalize]
@@ -66,7 +68,7 @@
   "Move any top-level parameters to the same level (i.e., 'inner query') as the query the affect."
   [{:keys [parameters], query-type :type, :as outer-query}]
   {:pre [(#{:query :native} query-type)]}
-  (cond-> (dissoc outer-query :parameters)
+  (cond-> (set/rename-keys outer-query {:parameters :user-parameters})
     (seq parameters)
     (assoc-in [query-type :parameters] parameters)))
 

src/metabase/query_processor/middleware/wrap_value_literals.clj

@@ -24,7 +24,7 @@
 (defmethod type-info :default [_] nil)
 
 (defmethod type-info (class Field) [this]
-  (let [field-info (select-keys this [:base_type :special_type :database_type])]
+  (let [field-info (select-keys this [:base_type :special_type :database_type :name])]
     (merge
      field-info
      ;; add in a default unit for this Field so we know to wrap datetime strings in `absolute-datetime` below based on

src/metabase/query_processor/util.clj

@@ -5,6 +5,7 @@
              [hash :as hash]]
             [cheshire.core :as json]
             [clojure.string :as str]
+            [metabase.driver :as driver]
             [metabase.util.schema :as su]
             [schema.core :as s]))
 
@@ -23,11 +24,10 @@
        (not page)
        (nil? aggregations)))
 
-(defn query->remark
-  "Generate an approparite REMARK to be prepended to a query to give DBAs additional information about the query being
-  executed. See documentation for `mbql->native` and [issue #2386](https://github.com/metabase/metabase/issues/2386)
-  for more information."
-  ^String [{{:keys [executed-by query-hash], :as info} :info, query-type :type}]
+(defn default-query->remark
+  "Generates the default query remark. Exists as a separate function so that overrides of the query->remark multimethod
+   can access the default value."
+  [{{:keys [executed-by query-hash], :as info} :info, query-type :type}]
   (str "Metabase" (when executed-by
                     (assert (instance? (Class/forName "[B") query-hash))
                     (format ":: userID: %s queryType: %s queryHash: %s"
@@ -37,6 +37,18 @@
                               :native "native")
                             (codecs/bytes->hex query-hash)))))
 
+(defmulti query->remark
+  "Generate an appropriate remark `^String` to be prepended to a query to give DBAs additional information about the query
+  being executed. See documentation for `mbql->native` and [issue #2386](https://github.com/metabase/metabase/issues/2386)
+  for more information."
+  {:arglists '(^String [driver data])}
+  driver/dispatch-on-initialized-driver
+  :hierarchy #'driver/hierarchy)
+
+(defmethod query->remark :default
+  [_ params]
+  (default-query->remark params))
+
 
 ;;; ------------------------------------------------- Normalization --------------------------------------------------
 

test/metabase/query_processor/middleware/parameters_test.clj

@@ -14,7 +14,8 @@
 
 (deftest move-top-level-params-to-inner-query-test
   (is (= {:type   :native
-          :native {:query "WOW", :parameters ["My Param"]}}
+          :native {:query "WOW", :parameters ["My Param"]}
+          :user-parameters ["My Param"]}
          (#'parameters/move-top-level-params-to-inner-query
           {:type :native, :native {:query "WOW"}, :parameters ["My Param"]}))))
 
@@ -36,7 +37,8 @@
   (testing "can we expand native params if they are specified at the top level?"
     (is (= (data/query nil
              {:type   :native
-              :native {:query "SELECT * FROM venues WHERE price = 1;", :params []}})
+              :native {:query "SELECT * FROM venues WHERE price = 1;", :params []}
+              :user-parameters [{:type :category, :target [:variable [:template-tag "price"]], :value "1"}]})
            (substitute-params
             (data/query nil
               {:type       :native

test/metabase/query_processor/middleware/wrap_value_literals_test.clj

@@ -27,7 +27,8 @@
                      $id
                      [:value 50 {:base_type     :type/BigInteger
                                  :special_type  :type/PK
-                                 :database_type "BIGINT"}]]})
+                                 :database_type "BIGINT"
+                                 :name "ID"}]]})
          (wrap-value-literals
            (mt/mbql-query venues
              {:filter [:> $id 50]}))))
@@ -35,10 +36,12 @@
            {:filter [:and
                      [:> $id [:value 50 {:base_type     :type/BigInteger
                                          :special_type  :type/PK
-                                         :database_type "BIGINT"}]]
+                                         :database_type "BIGINT"
+                                         :name "ID"}]]
                      [:< $price [:value 5 {:base_type     :type/Integer
                                            :special_type  :type/Category
-                                           :database_type "INTEGER"}]]]})
+                                           :database_type "INTEGER"
+                                           :name "PRICE"}]]]})
          (wrap-value-literals
            (mt/mbql-query venues
              {:filter [:and
@@ -137,7 +140,8 @@
                        [:value "2018-10-01" {:base_type     :type/Date
                                              :special_type  nil
                                              :database_type "DATE"
-                                             :unit          :month}]]})
+                                             :unit          :month
+                                             :name          "DATE"}]]})
            (wrap-value-literals
              (mt/mbql-query checkins
                {:filter [:starts-with !month.date "2018-10-01"]}))))))