Merge pull request #6015 from metabase/druid-format-timestamps-in-select

Added timestamp formatting for druid select queries [ci drivers]

src/metabase/driver/druid/query_processor.clj

 (ns metabase.driver.druid.query-processor
   (:require [cheshire.core :as json]
+            [clj-time.core :as time]
+            [clj-time.format :as tformat]
             [clojure.core.match :refer [match]]
             [clojure.math.numeric-tower :as math]
             [clojure.string :as s]
@@ -9,7 +11,9 @@
              [annotate :as annotate]
              [interface :as i]]
             [metabase.util :as u])
-  (:import [metabase.query_processor.interface AgFieldRef DateTimeField DateTimeValue Expression Field RelativeDateTimeValue Value]))
+  (:import java.util.TimeZone
+           [metabase.query_processor.interface AgFieldRef DateTimeField DateTimeValue Expression Field RelativeDateTimeValue Value]
+           org.joda.time.DateTimeZone))
 
 (def ^:private ^:const topN-max-results
   "Maximum number of rows the topN query in Druid should return. Huge values cause significant issues with the engine.
@@ -724,24 +728,34 @@
   {:projections projections
    :results     results})
 
-(defmethod post-process ::select  [_ projections results]
-  (->> results
-       first
-       :result
-       :events
-       (map :event)
-       (post-process-map projections)))
-
-(defmethod post-process ::total   [_ projections results]
+(def ^:private druid-ts-format (tformat/formatters :date-time))
+
+(defn- reformat-timestamp [timestamp target-formatter]
+  (->> timestamp
+       (tformat/parse druid-ts-format)
+       (tformat/unparse target-formatter)))
+
+(defmethod post-process ::select  [_ projections timezone results]
+  (let [update-ts-fn (if-let [target-formater (and timezone (tformat/with-zone druid-ts-format timezone))]
+                       #(update % :timestamp reformat-timestamp target-formater)
+                       identity)]
+    (->> results
+         first
+         :result
+         :events
+         (map (comp update-ts-fn :event))
+         (post-process-map projections))))
+
+(defmethod post-process ::total   [_ projections timezone results]
   (post-process-map projections (map :result results)))
 
-(defmethod post-process ::topN    [_ projections results]
+(defmethod post-process ::topN    [_ projections timezone results]
   (post-process-map projections (-> results first :result)))
 
-(defmethod post-process ::groupBy [_ projections results]
+(defmethod post-process ::groupBy [_ projections timezone results]
   (post-process-map projections (map :event results)))
 
-(defmethod post-process ::timeseries [_ projections results]
+(defmethod post-process ::timeseries [_ projections timezone results]
   (post-process-map (conj projections :timestamp)
                     (for [event results]
                       (conj {:timestamp (:timestamp event)} (:result event)))))
@@ -791,9 +805,26 @@
                                     k)
             k))))
 
+(defn- utc?
+  "There are several timezone ids that mean UTC. This will create a
+  TimeZone object from `TIMEZONE` and check to see if it's a UTC
+  timezone"
+  [^DateTimeZone timezone]
+  (.hasSameRules (TimeZone/getTimeZone "UTC")
+                 (.toTimeZone timezone)))
+
+(defn- resolve-timezone
+  "Returns the timezone object (either report-timezone or JVM
+  timezone). Returns nil if the timezone is UTC as the timestamps from
+  Druid are already in UTC and don't need to be converted"
+  [{:keys [settings]}]
+  (let [tz (time/time-zone-for-id (:report-timezone settings (System/getProperty "user.timezone")))]
+    (when-not (utc? tz)
+      tz)))
+
 (defn execute-query
   "Execute a query for a Druid DB."
-  [do-query {database :database, {:keys [query query-type mbql? projections]} :native}]
+  [do-query {database :database, {:keys [query query-type mbql? projections]} :native :as query-ctx}]
   {:pre [database query]}
   (let [details       (:details database)
         query         (if (string? query)
@@ -802,7 +833,7 @@
         query-type    (or query-type (keyword "metabase.driver.druid.query-processor" (name (:queryType query))))
         post-proc-map (->> query
                            (do-query details)
-                           (post-process query-type projections))
+                           (post-process query-type projections (resolve-timezone query-ctx)))
         columns       (if mbql?
                         (->> post-proc-map
                              :projections

test/metabase/driver/druid_test.clj

 (ns metabase.driver.druid-test
   (:require [cheshire.core :as json]
+            [clj-time.core :as time]
             [expectations :refer [expect]]
             [medley.core :as m]
             [metabase
@@ -11,10 +12,12 @@
             metabase.driver.druid
             [metabase.models
              [field :refer [Field]]
-             [table :refer [Table]]
-             [metric :refer [Metric]]]
+             [metric :refer [Metric]]
+             [table :refer [Table]]]
             [metabase.query-processor.middleware.expand :as ql]
-            [metabase.test.data :as data]
+            [metabase.test
+             [data :as data]
+             [util :as tu]]
             [metabase.test.data.datasets :as datasets :refer [expect-with-engine]]
             [toucan.util.test :as tt])
   (:import metabase.driver.druid.DruidDriver))
@@ -28,11 +31,39 @@
    ["1000" "Tito's Tacos"                "2014-06-03T07:00:00.000Z"]
    ["101"  "Golden Road Brewing"         "2015-09-04T07:00:00.000Z"]]
   (->> (driver/table-rows-sample (Table (data/id :checkins))
-         [(Field (data/id :checkins :id))
-          (Field (data/id :checkins :venue_name))])
+                                 [(Field (data/id :checkins :id))
+                                  (Field (data/id :checkins :venue_name))])
        (sort-by first)
        (take 5)))
 
+(datasets/expect-with-engine :druid
+  ;; druid returns a timestamp along with the query, but that shouldn't really matter here :D
+  [["1"    "The Misfit Restaurant + Bar" "2014-04-07T00:00:00.000-07:00"]
+   ["10"   "Dal Rae Restaurant"          "2015-08-22T00:00:00.000-07:00"]
+   ["100"  "PizzaHacker"                 "2014-07-26T00:00:00.000-07:00"]
+   ["1000" "Tito's Tacos"                "2014-06-03T00:00:00.000-07:00"]
+   ["101"  "Golden Road Brewing"         "2015-09-04T00:00:00.000-07:00"]]
+  (tu/with-temporary-setting-values [report-timezone "America/Los_Angeles"]
+    (->> (driver/table-rows-sample (Table (data/id :checkins))
+                                   [(Field (data/id :checkins :id))
+                                    (Field (data/id :checkins :venue_name))])
+         (sort-by first)
+         (take 5))))
+
+(datasets/expect-with-engine :druid
+  ;; druid returns a timestamp along with the query, but that shouldn't really matter here :D
+  [["1"    "The Misfit Restaurant + Bar" "2014-04-07T02:00:00.000-05:00"]
+   ["10"   "Dal Rae Restaurant"          "2015-08-22T02:00:00.000-05:00"]
+   ["100"  "PizzaHacker"                 "2014-07-26T02:00:00.000-05:00"]
+   ["1000" "Tito's Tacos"                "2014-06-03T02:00:00.000-05:00"]
+   ["101"  "Golden Road Brewing"         "2015-09-04T02:00:00.000-05:00"]]
+  (tu/with-jvm-tz (time/time-zone-for-id "America/Chicago")
+    (->> (driver/table-rows-sample (Table (data/id :checkins))
+                                   [(Field (data/id :checkins :id))
+                                    (Field (data/id :checkins :venue_name))])
+         (sort-by first)
+         (take 5))))
+
 (def ^:const ^:private ^String native-query-1
   (json/generate-string
     {:intervals   ["1900-01-01/2100-01-01"]

test/metabase/query_processor_test/date_bucketing_test.clj

@@ -58,35 +58,6 @@
   [engine]
   (contains? #{:oracle} engine))
 
-(defn- call-with-jvm-tz
-  "Invokes the thunk `F` with the JVM timezone set to `DTZ`, puts the
-  various timezone settings back the way it found it when it exits."
-  [^DateTimeZone dtz f]
-  (let [orig-tz (TimeZone/getDefault)
-        orig-dtz (time/default-time-zone)
-        orig-tz-prop (System/getProperty "user.timezone")]
-    (try
-      ;; It looks like some DB drivers cache the timezone information
-      ;; when instantiated, this clears those to force them to reread
-      ;; that timezone value
-      (reset! @#'metabase.driver.generic-sql/database-id->connection-pool {})
-      ;; Used by JDBC, and most JVM things
-      (TimeZone/setDefault (.toTimeZone dtz))
-      ;; Needed as Joda time has a different default TZ
-      (DateTimeZone/setDefault dtz)
-      ;; We read the system property directly when formatting results, so this needs to be changed
-      (System/setProperty "user.timezone" (.getID dtz))
-      (f)
-      (finally
-        ;; We need to ensure we always put the timezones back the way
-        ;; we found them as it will cause test failures
-        (TimeZone/setDefault orig-tz)
-        (DateTimeZone/setDefault orig-dtz)
-        (System/setProperty "user.timezone" orig-tz-prop)))))
-
-(defmacro ^:private with-jvm-tz [dtz & body]
-  `(call-with-jvm-tz ~dtz (fn [] ~@body)))
-
 (defn- sad-toucan-incidents-with-bucketing
   "Returns 10 sad toucan incidents grouped by `UNIT`"
   ([unit]
@@ -233,7 +204,7 @@
     :else
     (sad-toucan-result (source-date-formatter utc-tz) (result-date-formatter pacific-tz)))
 
-  (with-jvm-tz pacific-tz
+  (tu/with-jvm-tz pacific-tz
     (sad-toucan-incidents-with-bucketing :default eastern-tz)))
 
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
@@ -528,7 +499,7 @@
                     (result-date-formatter pacific-tz)
                     [6 10 4 9 9 8 8 9 7 9]))
 
-  (with-jvm-tz pacific-tz
+  (tu/with-jvm-tz pacific-tz
     (sad-toucan-incidents-with-bucketing :day pacific-tz)))
 
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
@@ -724,7 +695,7 @@
     (results-by-week date-formatter-without-time
                      (result-date-formatter pacific-tz)
                      [46 47 40 60 7]))
-  (with-jvm-tz pacific-tz
+  (tu/with-jvm-tz pacific-tz
     (sad-toucan-incidents-with-bucketing :week pacific-tz)))
 
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

test/metabase/test/util.clj

 (ns metabase.test.util
   "Helper functions and macros for writing unit tests."
   (:require [cheshire.core :as json]
+            [clj-time.core :as time]
             [clojure.tools.logging :as log]
             [clojure.walk :as walk]
             [clojurewerkz.quartzite.scheduler :as qs]
@@ -28,7 +29,8 @@
             [metabase.test.data :as data]
             [metabase.test.data.datasets :refer [*driver*]]
             [toucan.util.test :as test])
-  (:import org.joda.time.DateTime
+  (:import java.util.TimeZone
+           [org.joda.time DateTime DateTimeZone]
            [org.quartz CronTrigger JobDetail JobKey Scheduler Trigger]))
 
 ;; ## match-$
@@ -418,3 +420,34 @@
         .getChronology
         .getZone
         .getID)))
+
+(defn call-with-jvm-tz
+  "Invokes the thunk `F` with the JVM timezone set to `DTZ`, puts the
+  various timezone settings back the way it found it when it exits."
+  [^DateTimeZone dtz f]
+  (let [orig-tz (TimeZone/getDefault)
+        orig-dtz (time/default-time-zone)
+        orig-tz-prop (System/getProperty "user.timezone")]
+    (try
+      ;; It looks like some DB drivers cache the timezone information
+      ;; when instantiated, this clears those to force them to reread
+      ;; that timezone value
+      (reset! @#'metabase.driver.generic-sql/database-id->connection-pool {})
+      ;; Used by JDBC, and most JVM things
+      (TimeZone/setDefault (.toTimeZone dtz))
+      ;; Needed as Joda time has a different default TZ
+      (DateTimeZone/setDefault dtz)
+      ;; We read the system property directly when formatting results, so this needs to be changed
+      (System/setProperty "user.timezone" (.getID dtz))
+      (f)
+      (finally
+        ;; We need to ensure we always put the timezones back the way
+        ;; we found them as it will cause test failures
+        (TimeZone/setDefault orig-tz)
+        (DateTimeZone/setDefault orig-dtz)
+        (System/setProperty "user.timezone" orig-tz-prop)))))
+
+(defmacro with-jvm-tz
+  "Invokes `BODY` with the JVM timezone set to `DTZ`"
+  [dtz & body]
+  `(call-with-jvm-tz ~dtz (fn [] ~@body)))