feat: add support for set timezone (#68)

src/metabase/driver/firebolt.clj

@@ -14,9 +14,12 @@
             [metabase.driver.sql-jdbc.execute.legacy-impl :as legacy]
             [metabase.driver.sql.util.unprepare :as unprepare]
             [metabase.driver.sql.query-processor :as sql.qp]
+            [metabase.driver.sql-jdbc.execute :as sql-jdbc.execute]
             [metabase.util
+             [i18n :as i18n]
              [date-2 :as u.date]
-             [ssh :as ssh]]
+             [ssh :as ssh]
+             [log :as log]]
              [metabase.util.honey-sql-2 :as h2x])
   (:import [java.sql Types Connection ResultSet]
            [java.time LocalTime OffsetDateTime OffsetTime ZonedDateTime]))
@@ -48,6 +51,26 @@
         (sql-jdbc.common/handle-additional-options  (select-keys details [:password :classname :subprotocol :user :subname :additional-options :account :engine_name]))
         )))
 
+(defn set-connection-timezone [conn timezone]
+  (try
+    (with-open [stmt (.createStatement conn)]
+      (.execute stmt (str "SET time_zone = '" timezone "'")))
+    (catch Throwable e
+      (log/error e (i18n/trs "Failed to set timezone ''{0}'' for {1} database" timezone :firebolt)))))
+
+; Set timezone if provided in connection options
+(defmethod sql-jdbc.execute/do-with-connection-with-options :firebolt
+  [driver db-or-id-or-spec options f]
+  (sql-jdbc.execute/do-with-resolved-connection
+    driver
+    db-or-id-or-spec
+    options
+    (fn [^Connection conn]
+      (let [timezone (get options :session-timezone)]
+        (when timezone
+          (set-connection-timezone conn timezone)))
+      (f conn))))
+
 ; Testing the firebolt database connection
 (defmethod driver/can-connect? :firebolt [driver details]
    (let [connection (sql-jdbc.conn/connection-details->spec driver (ssh/include-ssh-tunnel! details))]
@@ -164,6 +187,19 @@
 (defmethod sql.qp/cast-temporal-string [:firebolt :Coercion/ISO8601->Time]
  [_driver _semantic_type expr]
   (h2x/maybe-cast :string expr))
+
+;; insert datetime with timezone in UTC
+(defmethod sql-jdbc.execute/set-parameter [::use-legacy-classes-for-read-and-set java.time.OffsetTime]
+  [driver preparedStmt idx timestamp]
+  (sql-jdbc.execute/set-parameter driver preparedStmt idx (t/sql-time (t/with-offset-same-instant timestamp (t/zone-offset 0)))))
+
+(defmethod sql-jdbc.execute/set-parameter [::use-legacy-classes-for-read-and-set java.time.OffsetDateTime]
+  [driver preparedStmt idx timestamp]
+  (sql-jdbc.execute/set-parameter driver preparedStmt idx (t/sql-timestamp (t/with-offset-same-instant timestamp (t/zone-offset 0)))))
+
+(defmethod sql-jdbc.execute/set-parameter [:firebolt java.time.ZonedDateTime]
+  [driver preparedStmt idx timestamp]
+  (sql-jdbc.execute/set-parameter driver preparedStmt idx (t/sql-timestamp (t/with-zone-same-instant timestamp (t/zone-id "UTC")))))
 ;;; ------------------------------------------------- query handling -------------------------------------------------
 
 ;(models/defmodel Table :metabase_table)
@@ -266,7 +302,7 @@
 
 (defmethod driver/database-supports? [:firebolt :case-sensitivity-string-filter-options]  [_ _ _] false)
 
-(defmethod driver/database-supports? [:firebolt :set-timezone]  [_ _ _] false)
+(defmethod driver/database-supports? [:firebolt :set-timezone]  [_ _ _] true)
 
 (defmethod driver/database-supports? [:firebolt :nested-fields]  [_ _ _] false)
 

test/metabase/test/data/firebolt.clj

@@ -9,19 +9,14 @@
              [driver :as driver]]
             [metabase.driver.ddl.interface :as ddl.i]
             [metabase.driver.sql-jdbc.connection :as sql-jdbc.conn]
-            [metabase.driver.sql.util.unprepare :as unprepare]
             [metabase.test.data.sql-jdbc
              [execute :as execute]
              [load-data :as load-data]]
-            [metabase.test.data.sql.ddl :as ddl]
             [metabase.driver.sql.util :as sql.u]
             [clojure.java.jdbc :as jdbc]
-            [metabase.connection-pool :as connection-pool]
             [metabase.driver.sql-jdbc.execute :as sql-jdbc.execute]
-            [clojure.tools.logging :as log]
-            [metabase.util :as u]
     )
-  (:import [java.sql ResultSet Connection DriverManager PreparedStatement SQLException]))
+  (:import [java.sql ResultSet]))
 
 
 (sql-jdbc.tx/add-test-extensions! :firebolt)
@@ -41,15 +36,16 @@
 ;;; ----------------------------------------------- Sync -----------------------------------------------
 
 ; Map firebolt data types to base type
-(doseq [[base-type sql-type] {:type/Array          "Array"
-                              :type/DateTime       "DateTime"
-                              :type/Date           "Date"
-                              :type/Time           "String"
-                              :type/Float          "Float"
-                              :type/Integer        "Int"
-                              :type/Text           "String"
-                              :type/Boolean        "Boolean"
-                              :type/BigInteger     "BIGINT"}]
+(doseq [[base-type sql-type] {:type/Array          "array"
+                              :type/DateTime       "timestamp"
+                              :type/DateTimeWithTZ "timestamptz"
+                              :type/Date           "date"
+                              :type/Time           "text"
+                              :type/Float          "double precision"
+                              :type/Integer        "int"
+                              :type/Text           "text"
+                              :type/Boolean        "boolean"
+                              :type/BigInteger     "bigint"}]
   (defmethod sql.tx/field-base-type->sql-type [:firebolt base-type] [_ _] sql-type))
 
 ;;; ----------------------------------------------- Query handling -----------------------------------------------
@@ -108,22 +104,6 @@
 (defmethod load-data/load-data! :firebolt [& args]
   (apply load-data/load-data-add-ids! args))
 
-; by default, firebolt allows to run a query of 50000 characters.
-; So setting the max_ast_elements flag to huge number to make the bulk insert queries work
-(defmethod load-data/do-insert! :firebolt
-  [driver spec table-identifier row-or-rows]
-
-  (let [statements (ddl/insert-rows-ddl-statements driver table-identifier row-or-rows)]
-      (try
-        (doseq [sql+args statements]
-          (log/tracef "[insert] %s" (pr-str sql+args))
-          (jdbc/execute! spec sql+args {:set-parameters (fn [stmt params]
-                                                          (sql-jdbc.execute/set-parameters! driver stmt params))}))
-        (catch SQLException e
-          (println (u/format-color 'red "INSERT FAILED: \n%s\n" statements))
-          (jdbc/print-sql-exception-chain e)
-          (throw e)))))
-
 ; Modified the table name to be in the format of db_name_table_name.
 ; So get the table and view names to make all test cases to use this format while forming the query to run tests
 (defmethod driver/describe-database :firebolt