Presto-JDBC Metabase test data extensions

presto_jdbc.clj

(ns metabase.test.data.presto-jdbc
  "Presto JDBC driver test extensions."
  (:require [clojure.string :as str]
            [metabase.config :as config]
            [metabase.connection-pool :as connection-pool]
            [metabase.driver :as driver]
            [metabase.driver.sql-jdbc.execute :as sql-jdbc.execute]
            [metabase.test.data.interface :as tx]
            [metabase.test.data.sql :as sql.tx]
            [metabase.test.data.sql-jdbc :as sql-jdbc.tx]
            [metabase.test.data.sql-jdbc.execute :as execute]
            [metabase.test.data.sql-jdbc.load-data :as load-data]
            [metabase.test.data.sql.ddl :as ddl]))

(sql-jdbc.tx/add-test-extensions! :presto-jdbc)

;; during unit tests don't treat presto as having FK support
(defmethod driver/supports? [:presto-jdbc :foreign-keys] [_ _] (not config/is-test?))

(doseq [[base-type db-type] {:type/BigInteger     "BIGINT"
                             :type/Boolean        "BOOLEAN"
                             :type/Date           "DATE"
                             :type/DateTime       "TIMESTAMP"
                             :type/DateTimeWithTZ "TIMESTAMP WITH TIME ZONE"
                             :type/Decimal        "DECIMAL"
                             :type/Float          "DOUBLE"
                             :type/Integer        "INTEGER"
                             :type/Text           "VARCHAR"
                             :type/Time           "TIME"
                             :type/TimeWithTZ     "TIME WITH TIME ZONE"}]
  (defmethod sql.tx/field-base-type->sql-type [:presto-jdbc base-type] [_ _] db-type))

;; in the past, we had to manually update our Docker image and add a new catalog for every new dataset definition we
;; added. That's insane. Just use the `test-data` catalog and put everything in that, and use
;; `db-qualified-table-name` like everyone else.
(def ^:private test-catalog-name "test_data")

(def ^:private test-schema-name "default")

(def ^:private connection-details
  (delay
    {:host                 (tx/db-test-env-var-or-throw :presto-jdbc :host "localhost")
     :port                 (tx/db-test-env-var-or-throw :presto-jdbc :port "8080" #_8843)
     :user                 (tx/db-test-env-var-or-throw :presto-jdbc :user "metabase")
     #_:additional-options #_ (tx/db-test-env-var-or-throw
                               :presto-jdbc
                               :additional-options
                               "SSLTrustStorePath=/tmp/cacerts-with-presto-ssl.jks&SSLTrustStorePassword=changeit")
     :ssl                  false
     :catalog              test-catalog-name}))

(defmethod tx/dbdef->connection-details :presto-jdbc
  [_ _ _]
  @connection-details)

(def ^:private jdbc-url
  (delay
    (format "jdbc:presto://%s:%s/%s/%s"
            (tx/db-test-env-var-or-throw :presto-jdbc :host "localhost")
            (tx/db-test-env-var-or-throw :presto-jdbc :port "8080" #_8843)
            test-catalog-name
            test-schema-name
            )))

(defn- connection ^java.sql.Connection []
  (java.sql.DriverManager/getConnection
   ^String @jdbc-url
   (connection-pool/map->properties
    (merge
     {:SSL  false
      :user (tx/db-test-env-var-or-throw :presto-jdbc :user "metabase")}
     ;; TODO -- additional SSQ options
     #_{:SSLTrustStorePath     "/tmp/cacerts-with-presto-ssl.jks"
        :SSLTrustStorePassword "changeit"}))))

(defmethod execute/execute-sql! :presto-jdbc
  [& args]
  (apply execute/sequentially-execute-sql! args))

(remove-method load-data/load-data! :presto-jdbc)
(defmethod load-data/load-data! :presto-jdbc
  [driver dbdef tabledef]
  (load-data/load-data-add-ids-chunked! driver dbdef tabledef))

(defmethod load-data/do-insert! :presto-jdbc
  [driver spec table-identifier row-or-rows]
  (let [statements (ddl/insert-rows-ddl-statements driver table-identifier row-or-rows)]
    (with-open [conn (connection)]
      (doseq [[^String sql & params] statements]
        (try
          (with-open [stmt (.prepareStatement conn sql)]
            (sql-jdbc.execute/set-parameters! driver stmt params)
            (let [rows-affected (.executeUpdate stmt)]
              (println (format "[%s] Inserted %d rows." driver rows-affected))))
          (catch Throwable e
            (throw (ex-info (format "[%s] Error executing SQL: %s" driver (ex-message e))
                            {:driver driver, :sql sql, :params params}
                            e))))))))

(defmethod sql.tx/drop-db-if-exists-sql :presto-jdbc [_ _] nil)
(defmethod sql.tx/create-db-sql         :presto-jdbc [_ _] nil)

(defmethod sql.tx/qualified-name-components :presto-jdbc
  ;; use the default schema from the in-memory connector
  ([_ db-name]                       ["default"])
  ([_ db-name table-name]            ["default" (tx/db-qualified-table-name db-name table-name)])
  ([_ db-name table-name field-name] ["default" (tx/db-qualified-table-name db-name table-name) field-name]))

(defmethod sql.tx/pk-sql-type :presto-jdbc
  [_]
  "INTEGER")

(defmethod sql.tx/create-table-sql :presto-jdbc
  [driver dbdef tabledef]
  ;; strip out the PRIMARY KEY stuff from the CREATE TABLE statement
  (let [sql ((get-method sql.tx/create-table-sql :sql/test-extensions) driver dbdef tabledef)]
    (str/replace sql #", PRIMARY KEY \([^)]+\)" "")))

;; Presto doesn't support FKs, at least not adding them via DDL
(defmethod sql.tx/add-fk-sql :presto-jdbc
  [_ _ _ _]
  nil)

;; FIXME Presto actually has very good timezone support
#_(defmethod tx/has-questionable-timezone-support? :presto-jdbc [_] true)