Support JDBC explicitPrepare flag

Support JDBC explicitPrepare flag

CHANGELOG.md

@@ -1,5 +1,8 @@
 ## [Unreleased]
 
+- Add a new "Optimized prepared statement" option when defining a database
+- Add further detail in the query comments such as account ID, dashboard ID or card ID when available
+
 ## [4.0.0] - 2023-09-27
 
 -   Dump Metabase version to v1.47.2

Makefile

@@ -80,6 +80,10 @@ test: start_trino_if_missing link_to_driver update_deps_files
 	@echo "Testing Starburst driver..."
 	cd $(makefile_dir)/metabase/; DRIVERS=starburst MB_STARBURST_TEST_PORT=$(trino_port) clojure -X:dev:drivers:drivers-dev:test
 
+testOptimized: start_trino_if_missing link_to_driver update_deps_files
+	@echo "Testing Starburst driver (explicitPrepare=true)..."
+	cd $(makefile_dir)/metabase/; DRIVERS=starburst MB_STARBURST_TEST_PORT=$(trino_port) clojure -J-DexplicitPrepare=false -X:dev:drivers:drivers-dev:test
+
 build: clone_metabase_if_missing update_deps_files link_to_driver front_end driver
 
 docker-image:

README.md

@@ -39,13 +39,14 @@ This command starts a local Metabase server on port `3000`. If you want to build
 Once you have built all required resources with the `make build` command, run `make test`. This command builds your local driver changes and then starts Starburst driver tests.
 
 #### Staring a Trino Server in Docker 
-Running `make test` will start a Trino server for you on port 8082 when needed, but if you want to start one, you can run `make start_trino_if_missing`.
+Running `make test` will start a Trino server for you on port 8082 when needed, but if you want to start one, you can run `make start_trino_if_missing`. Run `make testOptimized` to test Metabase with the "Optimized prepared statements" flag on.
 
 *Note:* Running `make test` will populate the Trino catalogs with mock data that is used for testing. You can then connect Metabase to this Trino server to view that data. This is useful for manual testing.
 
 #### Executing Specific Tests
 You can cd into the metabase repo and run commands like:
 `DRIVERS=starburst clojure -X:dev:drivers:drivers-dev:test :only metabase.query-processor-test.timezones-test/filter-test`
+`DRIVERS=starburst clojure -J-DexplicitPrepare=false -X:dev:drivers:drivers-dev:test :only metabase.query-processor-test.timezones-test/filter-test`
 
 or even
 

app_versions.json

@@ -1,5 +1,5 @@
 {
-    "trino": "426",
+    "trino": "431",
     "clojure": "1.11.1.1262",
     "metabase": "v1.47.2"
 }

drivers/starburst/deps.edn

@@ -15,7 +15,7 @@
  ["src" "resources"]
 
  :deps
- {io.trino/trino-jdbc {:mvn/version "426"}}
+ {io.trino/trino-jdbc {:mvn/version "431"}}
 
  ;; build the driver with clojure -X:build
  :aliases

drivers/starburst/resources/metabase-plugin.yaml

@@ -42,6 +42,12 @@ driver:
       type: boolean
       required: false
       description: Impersonate end user when running queries
+    - name: prepared-optimized
+      display-name: Optimized prepared statements
+      type: boolean
+      required: false
+      default: false
+      description: Requires Starburst Galaxy, Starburst Enterprise (version 420-e or higher), or Trino (version 418 or higher)
     - advanced-options-start
     - merge:
         - additional-options

drivers/starburst/src/metabase/driver/implementation/connectivity.clj

@@ -15,6 +15,7 @@
   "Connectivity implementation for Starburst driver."
   (:require [clojure.set :as set]
             [clojure.string :as str]
+            [clojure.tools.logging :as log]
             [metabase.api.common :as api]
             [metabase.db.spec :as mdb.spec]
             [metabase.driver.sql-jdbc.common :as sql-jdbc.common]
@@ -109,6 +110,9 @@
     (:impersonation details-map)
     (not= (get (deref api/*current-user*) :email) (:user details-map))))
 
+(defn assoc-if [coll key value condition]
+  (if condition (assoc coll key value) coll))
+
 (defmethod sql-jdbc.conn/connection-details->spec :starburst
   [_ details-map]
   (let [props (-> details-map
@@ -121,7 +125,8 @@
                   (update :kerberos-delegation bool->str)
                   (assoc :SSL (:ssl details-map))
                   (assoc :source "Starburst Metabase 4.0.0")
-                  (assoc :clientInfo (if (:impersonation details-map) "impersonate:true" ""))
+                  (assoc-if :clientInfo "impersonate:true" (:impersonation details-map))
+                  (assoc-if :explicitPrepare "false" (:prepared-optimized details-map))
                   (dissoc (if (remove-role? details-map) :roles :test))
 
                 ;; remove any Metabase specific properties that are not recognized by the Trino JDBC driver, which is
@@ -136,6 +141,6 @@
                                  :source :applicationNamePrefix ::accessToken :SSL :SSLVerification :SSLKeyStorePath
                                  :SSLKeyStorePassword :SSLKeyStoreType :SSLTrustStorePath :SSLTrustStorePassword :SSLTrustStoreType :SSLUseSystemTrustStore
                                  :extraCredentials :roles :sessionProperties :externalAuthentication :externalAuthenticationTokenCache :disableCompression 
-                                 :assumeLiteralNamesInMetadataCallsForNonConformingClients]
+                                 :explicitPrepare :assumeLiteralNamesInMetadataCallsForNonConformingClients]
                                 (keys kerb-props->url-param-names))))]
     (jdbc-spec props)))

drivers/starburst/src/metabase/driver/implementation/execute.clj

@@ -17,6 +17,7 @@
             [clojure.tools.logging :as log]
             [java-time :as t]
             [metabase.driver.implementation.sync :refer [starburst-type->base-type]]
+            [metabase.driver.implementation.messages :as msg]
             [metabase.driver.sql-jdbc.execute :as sql-jdbc.execute]
             [metabase.driver.sql.parameters.substitution :as sql.params.substitution]
             [metabase.query-processor.timezone :as qp.timezone]
@@ -168,10 +169,75 @@
 ; Metabase tests require a specific error when an invalid number of parameters are passed
 (defn handle-execution-error
   [e]
-  (log/fatalf (.getMessage e))
-  (if (clojure.string/includes? (.getMessage e) "Incorrect number of parameters")
-    (throw (Exception. "It looks like we got more parameters than we can handle, remember that parameters cannot be used in comments or as identifiers."))
-    (throw e)))
+  (let [message (.getMessage e)]
+    (cond
+      (clojure.string/includes? message "Expecting: 'USING'")
+      (throw (Exception. (str message msg/STARBURST_MAYBE_INCOMPATIBLE)))
+      (clojure.string/includes? message "Incorrect number of parameters")
+      (throw (Exception. msg/TOO_MANY_PARAMETERS))
+      :else (throw e))))
+
+; Optimized prepared statement where a proxy is generated and set-parameters! called on that proxy.
+; Metabase is sometimes calling getParametersMetaData() on the prepared statement in order to count
+; the number of parameters and verify they are correct with what is expected
+; This unfortunately defeats the purpose of using EXECUTE IMMEDIATE as it forces the JDBC driver
+; to call an explicit PREPARE. However this call is optional, so the solution is to create a proxy
+; which defines all methods used by Metabase *except* for metadata methods. When Metabase tries
+; to count the number of parameters:
+; - The proxy issues an exception as getParametersMetaData() is not defined
+; - Metabase catches the exception and does not perform the check
+; - An invalid query is sent to Trino, which fails with a "Incorrect number of parameters" message
+; - This message is caught by the driver and replaced with the exact same Metabase message
+; In the end, the exact same message is presented to the user when the number of arguments is
+; incorrect except we now execute the query to display the error message
+(defn proxy-optimized-prepared-statement
+  [driver conn stmt params]
+  (let [ps (proxy [java.sql.PreparedStatement] []
+          (executeQuery []
+            (try
+              (let [rs (.executeQuery stmt)]
+                (remove-impersonation conn)
+                rs)
+                (catch Throwable e (handle-execution-error e))))
+          (setMaxRows [nb] (.setMaxRows stmt nb))
+          (setObject
+            ([index obj] (.setObject stmt index obj))
+            ([index obj type] (.setObject stmt index obj type)))
+          (setTime
+            ([index val] (.setTime stmt index val))
+            ([index val cal] (.setTime stmt index val cal)))
+          (setTimestamp
+            ([index val] (.setTimestamp stmt index val))
+            ([index val cal] (.setTimestamp stmt index val cal)))
+          (setDate
+            ([index val] (.setDate stmt index val))
+            ([index val cal] (.setDate stmt index val cal)))
+          (setArray [index val] (.setArray stmt index val))
+          (setBoolean [index val] (.setBoolean stmt index val))
+          (setByte [index val] (.setByte stmt index val))
+          (setBytes [index val] (.setBytes stmt index val))
+          (setInt [index val] (.setInt stmt index val))
+          (setShort [index val] (.setShort stmt index val))
+          (setLong [index val] (.setLong stmt index val))
+          (setFloat [index val] (.setFloat stmt index val))
+          (setDouble [index val] (.setDouble stmt index val))
+          (close [] (.close stmt)))]
+    (sql-jdbc.execute/set-parameters! driver ps params)
+    ps))
+
+; Default prepared statement where set-parameters! is called before generating the proxy
+(defn proxy-prepared-statement
+  [driver conn stmt params]
+  (sql-jdbc.execute/set-parameters! driver stmt params)
+  (proxy [java.sql.PreparedStatement] []
+    (executeQuery []
+      (try
+        (let [rs (.executeQuery stmt)]
+          (remove-impersonation conn)
+          rs)
+          (catch Throwable e (handle-execution-error e))))
+    (setMaxRows [nb] (.setMaxRows stmt nb))
+    (close [] (.close stmt))))
 
 (defmethod sql-jdbc.execute/prepared-statement :starburst
   [driver ^Connection conn ^String sql params]
@@ -187,16 +253,10 @@
         (.setFetchDirection stmt ResultSet/FETCH_FORWARD)
         (catch Throwable e
           (log/debug e (trs "Error setting prepared statement fetch direction to FETCH_FORWARD"))))
-      (sql-jdbc.execute/set-parameters! driver stmt params)
-      (proxy [java.sql.PreparedStatement] []
-        (executeQuery []
-          (try
-            (let [rs (.executeQuery stmt)]
-              (remove-impersonation conn)
-              rs)
-              (catch Throwable e (handle-execution-error e))))
-        (setMaxRows [nb] (.setMaxRows stmt nb))
-        (close [] (.close stmt)))
+      (if
+        (.useExplicitPrepare (.unwrap conn TrinoConnection))
+        (proxy-prepared-statement driver conn stmt params)
+        (proxy-optimized-prepared-statement driver conn stmt params))
       (catch Throwable e
         (remove-impersonation conn)
         (.close stmt)

drivers/starburst/src/metabase/driver/implementation/messages.clj

@@ -0,0 +1,18 @@
+;;
+;; Licensed under the Apache License, Version 2.0 (the "License");
+;; you may not use this file except in compliance with the License.
+;; You may obtain a copy of the License at
+
+;;     http://www.apache.org/licenses/LICENSE-2.0
+
+;; Unless required by applicable law or agreed to in writing, software
+;; distributed under the License is distributed on an "AS IS" BASIS,
+;; WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+;; See the License for the specific language governing permissions and
+;; limitations under the License.
+;;
+(ns metabase.driver.implementation.messages "Starburst messages")
+
+(def STARBURST_INCOMPATIBLE_WITH_OPTIMIZED_PREPARED "\"Optimized prepared statements\" require Starburst Galaxy, Starburst Enterprise (version 420-e or higher), or Trino (version 418 or higher)")
+(def STARBURST_MAYBE_INCOMPATIBLE ". If the database has the \"Optimized prepared statements\" option on, it require Starburst Galaxy, Starburst Enterprise (version 420-e or higher), or Trino (version 418 or higher)")
+(def TOO_MANY_PARAMETERS "It looks like we got more parameters than we can handle, remember that parameters cannot be used in comments or as identifiers.")

drivers/starburst/src/metabase/driver/starburst.clj

@@ -14,8 +14,13 @@
 (ns metabase.driver.starburst
   "Starburst driver."
   (:require [metabase.driver :as driver]
+            [clojure.java.jdbc :as jdbc]
+            [clojure.tools.logging :as log]
+            [metabase.util.ssh :as ssh]
+            [metabase.driver.sql-jdbc.connection :as connection]
             [metabase.query-processor.util :as qp.util]
             [metabase.public-settings :as public-settings]
+            [metabase.driver.implementation.messages :as msg]
             [metabase.driver.sql-jdbc.execute.legacy-impl :as sql-jdbc.legacy]))
 (driver/register! :starburst, :parent #{::sql-jdbc.legacy/use-legacy-classes-for-read-and-set})
 
@@ -35,6 +40,22 @@
     (format-field "dashboardID" dashboard-id)
     (format-field "cardID" card-id)))
 
+(defn handle-execution-error
+  [e details]
+  (let [message (.getMessage e)
+        execute-immediate (get details :prepared-optimized false)]
+    (cond
+      (and (clojure.string/includes? message "Expecting: 'USING'") execute-immediate)
+      (throw (Exception. msg/STARBURST_INCOMPATIBLE_WITH_OPTIMIZED_PREPARED))
+      :else (throw e))))
+
+(defmethod driver/can-connect? :starburst
+  [driver details]
+  (try
+    (connection/with-connection-spec-for-testing-connection [jdbc-spec [driver details]]
+      (connection/can-connect-with-spec? jdbc-spec))
+    (catch Throwable e (handle-execution-error e details))))
+
 ;;; +----------------------------------------------------------------------------------------------------------------+
 ;;; |                                                  Load implemetation files                                      |
 ;;; +----------------------------------------------------------------------------------------------------------------+

drivers/starburst/test/metabase/driver/starburst_test.clj

@@ -284,6 +284,57 @@
       (is (= (str "impersonate:true")
              (:clientInfo jdbc-spec))))))
 
+(defn prepared-statements-helper
+  [prepared-optimized]
+        (is (= [["2023-11-06T00:00:00Z" 42 "It was created"]]
+          (mt/rows
+          (let [details {
+            :host                         "localhost"
+            :port                         8082
+            :catalog                      "catalog"
+            :user                         "admin"
+            :ssl                          false
+            :prepared-optimized           prepared-optimized}]
+              (t2.with-temp/with-temp [Database db {:engine :starburst, :name "Temp Trino JDBC Schema DB", :details details}]
+              (mt/with-db db
+              (qp/process-query
+                {:database     (mt/id)
+                :type         :native
+                :native       {
+                  :query         "SELECT {{created_at}}, {{nb_created}}, {{detail}}"
+                  :template-tags {:created_at {:name         "created_at"
+                                               :display_name "created_at"
+                                               :type         :date
+                                               :required     true}
+                                  :nb_created {:name         "nb_created"
+                                               :display_name "nb_created"
+                                               :type         :number
+                                               :required     true}
+                                  :detail {:name         "detail"
+                                               :display_name "detail"
+                                               :type         :text
+                                               :required     true}}}
+                :parameters    [{:type   :date
+                                 :name   "created_at"
+                                 :target [:variable [:template-tag "created_at"]]
+                                 :value  "2023-11-06"}
+                                {:type   :number
+                                 :name   "nb_created"
+                                 :target [:variable [:template-tag "nb_created"]]
+                                 :value  "42"}
+                                 {:type   :text
+                                 :name   "detail"
+                                 :target [:variable [:template-tag "detail"]]
+                                 :value  "It was created"}]}))))))))
+
+(deftest prepared-statements
+  (mt/test-driver :starburst
+    (testing "Make sure prepared statements work"
+        ;; If impersonation is set, then the Trino user should be the current Metabase user, i.e. [email protected]
+        ;; The role is ignored as Metabase users may not have the role defined in the database connection
+        (prepared-statements-helper true)
+        (prepared-statements-helper false))))
+
 (deftest impersonation-query
   (mt/test-driver :starburst
     (testing "Make sure the right credentials are used depending on the impersonation checkbox"

drivers/starburst/test/metabase/test/data/starburst.clj

@@ -61,6 +61,7 @@
    :port                               (tx/db-test-env-var :starburst :port "8082")
    :user                               (tx/db-test-env-var-or-throw :starburst :user "metabase")
    :additional-options                 (tx/db-test-env-var :starburst :additional-options nil)
+   :prepared-optimized                 (tx/db-test-env-var :starburst :prepared-optimized (not= (System/getProperty "explicitPrepare" "true") "true"))
    :ssl                                (tx/db-test-env-var :starburst :ssl "false")
    :kerberos                           (tx/db-test-env-var :starburst :kerberos "false")
    :kerberos-principal                 (tx/db-test-env-var :starburst :kerberos-principal nil)