Fix query remarks on Snowflake (#29888)

* Fix query remarks on Snowflake

* Move docstrings out of defmethod

* Add test to Snowflake query remarks

* Improve test on Snowflake query remark

* Remove unnecessary comment

docs/developers-guide/driver-changelog.md

@@ -75,6 +75,9 @@ title: Driver interface changelog
   will be used to call the corresponding `metabase.driver.sql.query-processor/date` implementation to convert the `field`.
   Returns `nil` if the conversion is not necessary for this `field` and `param-type` combination.
 
+- The multimethod `metabase.driver.sql-jdbc.execute/inject-remark` has been added. It allows JDBC-based drivers to
+  override the default behavior of how SQL query remarks are added to queries (prepending them as a comment).
+
 ## Metabase 0.47.0
 
 - A new driver feature has been added: `:schemas`. This feature signals whether the database organizes tables in

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

 (ns metabase.driver.snowflake
   "Snowflake Driver."
   (:require
+   [buddy.core.codecs :as codecs]
+   [cheshire.core :as json]
    [clojure.java.jdbc :as jdbc]
    [clojure.set :as set]
    [clojure.string :as str]
@@ -24,9 +26,11 @@
    [metabase.driver.sync :as driver.s]
    [metabase.lib.metadata :as lib.metadata]
    [metabase.models.secret :as secret]
+   [metabase.public-settings :as public-settings]
    [metabase.query-processor.error-type :as qp.error-type]
    [metabase.query-processor.store :as qp.store]
    [metabase.query-processor.timezone :as qp.timezone]
+   [metabase.query-processor.util :as qp.util]
    [metabase.query-processor.util.add-alias-info :as add]
    [metabase.util :as u]
    [metabase.util.date-2 :as u.date]
@@ -587,6 +591,28 @@
   [driver ps i t]
   (sql-jdbc.execute/set-parameter driver ps i (t/sql-timestamp (t/with-zone-same-instant t (t/zone-id "UTC")))))
 
+;;; --------------------------------------------------- Query remarks ---------------------------------------------------
+
+;  Snowflake strips comments prepended to the SQL statement (default remark injection behavior). We should append the
+;  remark instead.
+(defmethod sql-jdbc.execute/inject-remark :snowflake
+  [_ sql remark]
+  (str sql "\n\n-- " remark))
+
+(defmethod qp.util/query->remark :snowflake
+  [_ {{:keys [context executed-by card-id pulse-id dashboard-id query-hash]} :info,
+      query-type :type,
+      database-id :database}]
+  (json/generate-string {:client      "Metabase"
+                         :context     context
+                         :queryType   query-type
+                         :userId      executed-by
+                         :pulseId     pulse-id
+                         :cardId      card-id
+                         :dashboardId dashboard-id
+                         :databaseId  database-id
+                         :queryHash   (when (bytes? query-hash) (codecs/bytes->hex query-hash))
+                         :serverId    (public-settings/site-uuid)}))
 
 ;;; ------------------------------------------------- User Impersonation --------------------------------------------------
 

modules/drivers/snowflake/test/metabase/driver/snowflake_test.clj

 (ns metabase.driver.snowflake-test
   (:require
+   [clojure.data.json :as json]
    [clojure.java.jdbc :as jdbc]
    [clojure.set :as set]
    [clojure.string :as str]
@@ -7,8 +8,10 @@
    [metabase.driver :as driver]
    [metabase.driver.sql :as driver.sql]
    [metabase.driver.sql-jdbc.connection :as sql-jdbc.conn]
+   [metabase.driver.sql-jdbc.execute :as sql-jdbc.execute]
    [metabase.models :refer [Table]]
    [metabase.models.database :refer [Database]]
+   [metabase.public-settings :as public-settings]
    [metabase.query-processor :as qp]
    [metabase.sync :as sync]
    [metabase.sync.util :as sync-util]
@@ -26,6 +29,17 @@
 
 (set! *warn-on-reflection* true)
 
+(defn- query->native [query]
+  (let [check-sql-fn (fn [_ _ sql & _]
+                       (throw (ex-info "done" {::native-query sql})))]
+    (with-redefs [sql-jdbc.execute/prepared-statement check-sql-fn
+                  sql-jdbc.execute/execute-statement! check-sql-fn]
+      (try
+        (qp/process-query query)
+        (is false "no statement created")
+        (catch Exception e
+          (-> e u/all-ex-data ::native-query))))))
+
 (use-fixtures :each (fn [thunk]
                       ;; 1. If sync fails when loading a test dataset, don't swallow the error; throw an Exception so we
                       ;;    can debug it. This is much less confusing when trying to fix broken tests.
@@ -298,3 +312,33 @@
     ;; Special characters
     (is (= "USE ROLE \"Role.123\";"   (driver.sql/set-role-statement :snowflake "Role.123")))
     (is (= "USE ROLE \"$role\";"      (driver.sql/set-role-statement :snowflake "$role")))))
+
+(deftest remark-test
+  (testing "Queries should have a remark formatted as JSON appended to them with additional metadata"
+    (mt/test-driver :snowflake
+      (let [expected-map {"pulseId" nil
+                          "serverId" (public-settings/site-uuid)
+                          "client" "Metabase"
+                          "queryHash" "cb83d4f6eedc250edb0f2c16f8d9a21e5d42f322ccece1494c8ef3d634581fe2"
+                          "queryType" "query"
+                          "cardId" 1234
+                          "dashboardId" 5678
+                          "context" "ad-hoc"
+                          "userId" 1000
+                          "databaseId" (mt/id)}
+            result-query (driver/prettify-native-form :snowflake
+                           (query->native
+                             (assoc
+                               (mt/mbql-query users {:limit 2000})
+                               :parameters [{:type   "id"
+                                             :target [:dimension [:field (mt/id :users :id) nil]]
+                                             :value  ["1" "2" "3"]}]
+                               :info {:executed-by  1000
+                                      :card-id      1234
+                                      :dashboard-id 5678
+                                      :context      :ad-hoc
+                                      :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])})))
+            result-comment (second (re-find #"-- (\{.*\})" result-query))
+            result-map (json/read-str result-comment)]
+        (is (= expected-map result-map))))))

src/metabase/driver/sql_jdbc/execute.clj

@@ -662,13 +662,26 @@
   (let [row-thunk (row-thunk driver rs rsmeta)]
     (qp.reducible/reducible-rows row-thunk canceled-chan)))
 
+(defmulti inject-remark
+  "Injects the remark into the SQL query text."
+  {:added "0.48.0", :arglists '([driver sql remark])}
+  driver/dispatch-on-initialized-driver
+  :hierarchy #'driver/hierarchy)
+
+;  Combines the original SQL query with query remarks. Most databases using sql-jdbc based drivers support prepending the
+;  remark to the SQL statement, so we have it as a default. However, some drivers do not support it, so we allow it to
+;  be overriden.
+(defmethod inject-remark :default
+  [_ sql remark]
+  (str "-- " remark "\n" sql))
+
 (defn execute-reducible-query
   "Default impl of `execute-reducible-query` for sql-jdbc drivers."
   {:added "0.35.0", :arglists '([driver query context respond] [driver sql params max-rows context respond])}
   ([driver {{sql :query, params :params} :native, :as outer-query} context respond]
    {:pre [(string? sql) (seq sql)]}
    (let [remark   (qp.util/query->remark driver outer-query)
-         sql      (str "-- " remark "\n" sql)
+         sql      (inject-remark driver sql remark)
          max-rows (limit/determine-query-max-rows outer-query)]
      (execute-reducible-query driver sql params max-rows context respond)))