diff --git a/modules/drivers/druid/src/metabase/driver/druid.clj b/modules/drivers/druid/src/metabase/driver/druid.clj index 799d0887e6eac7174bf7cda6857cab7ffdefcc95..70581f759e61dbb8893b8d5feec6ce5aa981afe2 100644 --- a/modules/drivers/druid/src/metabase/driver/druid.clj +++ b/modules/drivers/druid/src/metabase/driver/druid.clj @@ -1,6 +1,7 @@ (ns metabase.driver.druid "Druid driver." - (:require [clj-http.client :as http] + (:require [cheshire.core :as json] + [clj-http.client :as http] [metabase.driver :as driver] [metabase.driver.druid.client :as druid.client] [metabase.driver.druid.execute :as druid.execute] @@ -29,11 +30,18 @@ [_ query] (druid.qp/mbql->native query)) +(defn- add-timeout-to-query [query timeout] + (let [parsed (if (string? query) + (json/parse-string query keyword) + query)] + (assoc-in parsed [:context :timeout] timeout))) + (defmethod driver/execute-reducible-query :druid [_ query context respond] (druid.execute/execute-reducible-query (partial druid.client/do-query-with-cancellation (qp.context/canceled-chan context)) - query respond)) + (update-in query [:native :query] add-timeout-to-query (qp.context/timeout context)) + respond)) (doseq [[feature supported?] {:set-timezone true :expression-aggregations true}] diff --git a/modules/drivers/druid/src/metabase/driver/druid/query_processor.clj b/modules/drivers/druid/src/metabase/driver/druid/query_processor.clj index 3474889fc5b9ed7638cc0878fe908d650f5ccdca..61b49a34e7b5c6b53128575cfe613ceb3f2a7374 100644 --- a/modules/drivers/druid/src/metabase/driver/druid/query_processor.clj +++ b/modules/drivers/druid/src/metabase/driver/druid/query_processor.clj @@ -136,8 +136,7 @@ (merge {:intervals ["1900-01-01/2100-01-01"] :granularity :all - :context {:timeout 60000 - :queryId (random-query-id)}} + :context {:queryId (random-query-id)}} (case query-type ::scan {:queryType :scan :limit qp.i/absolute-max-results} diff --git a/modules/drivers/druid/test/metabase/driver/druid/client_test.clj b/modules/drivers/druid/test/metabase/driver/druid/client_test.clj index 427fcee266f2179b8472c5bccffa93c693219313..1445dff603f8b0c0077408fd4d7ae2cf2f239e58 100644 --- a/modules/drivers/druid/test/metabase/driver/druid/client_test.clj +++ b/modules/drivers/druid/test/metabase/driver/druid/client_test.clj @@ -4,6 +4,7 @@ [metabase.driver.druid.client :as druid.client] [metabase.driver.util :as driver.u] [metabase.query-processor :as qp] + [metabase.query-processor.context.default :as default] [metabase.test :as mt] [metabase.test.util.log :as tu.log] [metabase.timeseries-query-processor-test.util :as tqpt])) @@ -29,6 +30,18 @@ (is (= ::cancel (mt/wait-for-result cancel-chan 2000))))))))) +(deftest query-timeout-test + (mt/test-driver :druid + (tqpt/with-flattened-dbdef + (let [query (mt/mbql-query checkins) + executed-query (atom nil)] + (with-redefs [druid.client/do-query-with-cancellation (fn [_chan _details query] + (reset! executed-query query) + [])] + (qp/process-query-sync query) + (is (partial= {:context {:timeout default/query-timeout-ms}} + @executed-query))))))) + (deftest ssh-tunnel-test (mt/test-driver :druid diff --git a/modules/drivers/druid/test/metabase/driver/druid/query_processor_test.clj b/modules/drivers/druid/test/metabase/driver/druid/query_processor_test.clj index 2d724c62e8de485280ada9d7d3e993b7a6c97335..d4ca47984d89077d0e02c51207097da975ebc1c8 100644 --- a/modules/drivers/druid/test/metabase/driver/druid/query_processor_test.clj +++ b/modules/drivers/druid/test/metabase/driver/druid/query_processor_test.clj @@ -106,7 +106,7 @@ :granularity :all :dataSource "checkins" :dimension "venue_price" - :context {:timeout 60000, :queryId "<Query ID>"} + :context {:queryId "<Query ID>"} :postAggregations [{:type :arithmetic :name "expression" :fn :* @@ -135,7 +135,7 @@ :granularity :all :dataSource "checkins" :dimension "venue_category_name" - :context {:timeout 60000, :queryId "<Query ID>"} + :context {:queryId "<Query ID>"} :intervals ["1900-01-01/2100-01-01"] :metric "__count_0" :aggregations [{:type :cardinality @@ -159,7 +159,7 @@ :granularity :all :dataSource "checkins" :dimensions ["venue_category_name", "user_name"] - :context {:timeout 60000, :queryId "<Query ID>"} + :context {:queryId "<Query ID>"} :intervals ["1900-01-01/2100-01-01"] :aggregations [{:type :cardinality :name "__count_0" @@ -186,7 +186,7 @@ :granularity :all :dataSource "checkins" :dimensions ["venue_category_name", "user_name"] - :context {:timeout 60000, :queryId "<Query ID>"} + :context {:queryId "<Query ID>"} :intervals ["1900-01-01/2100-01-01"] :aggregations [{:type :cardinality :name "__count_0" @@ -214,7 +214,7 @@ :query {:queryType :timeseries :granularity :all :dataSource "checkins" - :context {:timeout 60000, :queryId "<Query ID>"} + :context {:queryId "<Query ID>"} :intervals ["1900-01-01/2100-01-01"] :aggregations [{:type :cardinality :name "__distinct_0"