Skip to content

Commit

Permalink
Merge pull request #8 from dacort/feature/date-handling
Browse files Browse the repository at this point in the history 
Add date handling for aggregation/filtering support
  • Loading branch information
@dacort
dacort authored May 11, 2019
2 parents 41eac38 + 84b7780 commit 0c77c4e
Showing 1 changed file with 72 additions and 26 deletions.
98 changes: 72 additions & 26 deletions src/metabase/driver/athena.clj
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
[date :as du]
[honeysql-extensions :as hx]
[i18n :refer [trs]]])
(:import [java.sql DatabaseMetaData Timestamp]))
(:import [java.sql DatabaseMetaData Timestamp] java.util.Date java.sql.Time))

(driver/register! :athena, :parent :sql-jdbc)

Expand All @@ -28,6 +28,7 @@
;;; | metabase.driver method impls |
;;; +----------------------------------------------------------------------------------------------------------------+


(defmethod driver/supports? [:athena :foreign-keys] [_ _] false)

;;; +----------------------------------------------------------------------------------------------------------------+
Expand Down Expand Up @@ -59,6 +60,7 @@
:varbinary :type/*
:boolean :type/Boolean
:char :type/Text
:date :type/Date
:decimal :type/Decimal
:double :type/Float
:float :type/Float
Expand All @@ -70,7 +72,51 @@
:struct :type/*
:timestamp :type/DateTime
:tinyint :type/Integer
:varchar :type/Text} database-type ))
:varchar :type/Text} database-type))

;;; ------------------------------------------------- date functions -------------------------------------------------
(defmethod unprepare/unprepare-value [:athena Date] [_ value]
(unprepare/unprepare-date-with-iso-8601-fn :from_iso8601_timestamp value))

(prefer-method unprepare/unprepare-value [:sql Time] [:athena Date])

; Helper function for truncating dates - currently unused
(defn- date-trunc [unit expr] (hsql/call :date_trunc (hx/literal unit) expr))

; Example of handling report timezone
; (defn- date-trunc
; "date_trunc('interval', timezone, timestamp): truncates a timestamp to a given interval"
; [unit expr]
; (let [timezone (get-in sql.qp/*query* [:settings :report-timezone])]
; (if (nil? timezone)
; (hsql/call :date_trunc (hx/literal unit) expr)
; (hsql/call :date_trunc (hx/literal unit) timezone expr))))

; Helper function to cast `expr` to a timestamp if necessary
(defn- expr->literal [expr]
(if (instance? Timestamp expr)
expr
(hx/cast :timestamp expr)))

; If `expr` is a date, we need to cast it to a timestamp before we can truncate to a finer granularity
; Ideally, we should make this conditional. There's a generic approach above, but different use cases should b tested.
(defmethod sql.qp/date [:athena :minute] [_ _ expr] (hsql/call :date_trunc (hx/literal :minute) (expr->literal expr)))
(defmethod sql.qp/date [:athena :hour] [_ _ expr] (hsql/call :date_trunc (hx/literal :hour) (expr->literal expr)))
(defmethod sql.qp/date [:athena :day] [_ _ expr] (hsql/call :date_trunc (hx/literal :day) expr))
(defmethod sql.qp/date [:athena :week] [_ _ expr] (hsql/call :date_trunc (hx/literal :week) expr))
(defmethod sql.qp/date [:athena :month] [_ _ expr] (hsql/call :date_trunc (hx/literal :month) expr))
(defmethod sql.qp/date [:athena :quarter] [_ _ expr] (hsql/call :date_trunc (hx/literal :quarter) expr))
(defmethod sql.qp/date [:athena :year] [_ _ expr] (hsql/call :date_trunc (hx/literal :year) expr))

; Extraction functions
(defmethod sql.qp/date [:athena :minute-of-hour] [_ _ expr] (hsql/call :minute expr))
(defmethod sql.qp/date [:athena :hour-of-day] [_ _ expr] (hsql/call :hour expr))
(defmethod sql.qp/date [:athena :day-of-week] [_ _ expr] (hsql/call :day_of_week expr))
(defmethod sql.qp/date [:athena :day-of-month] [_ _ expr] (hsql/call :day_of_month expr))
(defmethod sql.qp/date [:athena :day-of-year] [_ _ expr] (hsql/call :day_of_year expr))
(defmethod sql.qp/date [:athena :week-of-year] [_ _ expr] (hsql/call :week_of_year expr))
(defmethod sql.qp/date [:athena :month-of-year] [_ _ expr] (hsql/call :month expr))
(defmethod sql.qp/date [:athena :quarter-of-year] [_ _ expr] (hsql/call :quarter expr))

;; keyword function converts database-type variable to a symbol, so we use symbols above to map the types
(defn- database-type->base-type-or-warn
Expand All @@ -86,37 +132,37 @@
(defn describe-table-fields
"Returns a set of column metadata for `schema` and `table-name` using `metadata`. "
[^DatabaseMetaData metadata, driver, {^String schema :schema, ^String table-name :name}, & [^String db-name-or-nil]]
(try
(try
(with-open [rs (.getColumns metadata db-name-or-nil schema table-name nil)]
(set
(for [{database-type :type_name
column-name :column_name
remarks :remarks} (jdbc/metadata-result rs)]
(merge
{:name column-name
:database-type database-type
:base-type (database-type->base-type-or-warn driver database-type)}
(when (not (str/blank? remarks))
{:field-comment remarks})
))))
(set
(for [{database-type :type_name
column-name :column_name
remarks :remarks} (jdbc/metadata-result rs)]
(merge
{:name column-name
:database-type database-type
:base-type (database-type->base-type-or-warn driver database-type)}
(when (not (str/blank? remarks))
{:field-comment remarks})))))
(catch Throwable e
(log/error e (trs "Error retreiving fields for DB {0}.{1}" schema table-name))
(throw e))
)
)
(throw e))))


;; Becuse describe-table-fields might fail, we catch the error here and return an empty set of columns


(defmethod driver/describe-table :athena [driver database table]
(jdbc/with-db-metadata [metadata (sql-jdbc.conn/db->pooled-connection-spec database)]
(->> (assoc (select-keys table [:name :schema])
:fields (try
(describe-table-fields metadata driver table)
(catch Throwable e (set nil)))
))))
:fields (try
(describe-table-fields metadata driver table)
(catch Throwable e (set nil)))))))


;; EXTERNAL_TABLE is required for Athena


(defn- get-tables [^DatabaseMetaData metadata, ^String schema-or-nil, ^String db-name-or-nil]
;; tablePattern "%" = match all tables
(with-open [rs (.getTables metadata db-name-or-nil schema-or-nil "%"
Expand Down Expand Up @@ -147,10 +193,10 @@

; Unsure if this is the right way to approach building the parameterized query...but it works
(defn- prepare-query [driver {:keys [database settings], query :native, :as outer-query}]
(cond-> outer-query
(seq (:params query))
(merge {:native {:params nil
:query (unprepare/unprepare driver (cons (:query query) (:params query)))}})))
(cond-> outer-query
(seq (:params query))
(merge {:native {:params nil
:query (unprepare/unprepare driver (cons (:query query) (:params query)))}})))

(defmethod driver/execute-query :athena [driver query]
(sql-jdbc.execute/execute-query driver (prepare-query driver, query)))
(sql-jdbc.execute/execute-query driver (prepare-query driver, query)))

0 comments on commit 0c77c4e

Please sign in to comment.