chore(druid): Explicitly cast col to TIMESTAMP (#17101)

Co-authored-by: John Bodley <john.bodley@airbnb.com>
apache · Oct 18, 2021 · 565ee23 · 565ee23
1 parent ae4ced8
commit 565ee23
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 22 deletions.
diff --git a/superset/db_engine_specs/druid.py b/superset/db_engine_specs/druid.py
@@ -40,27 +40,27 @@ class DruidEngineSpec(BaseEngineSpec):
     allows_subqueries = True
 
     _time_grain_expressions = {
-        None: "{col}",
-        "PT1S": "TIME_FLOOR({col}, 'PT1S')",
-        "PT5S": "TIME_FLOOR({col}, 'PT5S')",
-        "PT30S": "TIME_FLOOR({col}, 'PT30S')",
-        "PT1M": "TIME_FLOOR({col}, 'PT1M')",
-        "PT5M": "TIME_FLOOR({col}, 'PT5M')",
-        "PT10M": "TIME_FLOOR({col}, 'PT10M')",
-        "PT15M": "TIME_FLOOR({col}, 'PT15M')",
-        "PT30M": "TIME_FLOOR({col}, 'PT30M')",
-        "PT1H": "TIME_FLOOR({col}, 'PT1H')",
-        "PT6H": "TIME_FLOOR({col}, 'PT6H')",
-        "P1D": "TIME_FLOOR({col}, 'P1D')",
-        "P1W": "TIME_FLOOR({col}, 'P1W')",
-        "P1M": "TIME_FLOOR({col}, 'P1M')",
-        "P3M": "TIME_FLOOR({col}, 'P3M')",
-        "P1Y": "TIME_FLOOR({col}, 'P1Y')",
+        None: "CAST({col} AS TIMESTAMP)",
+        "PT1S": "TIME_FLOOR(CAST({col} AS TIMESTAMP), 'PT1S')",
+        "PT5S": "TIME_FLOOR(CAST({col} AS TIMESTAMP), 'PT5S')",
+        "PT30S": "TIME_FLOOR(CAST({col} AS TIMESTAMP), 'PT30S')",
+        "PT1M": "TIME_FLOOR(CAST({col} AS TIMESTAMP), 'PT1M')",
+        "PT5M": "TIME_FLOOR(CAST({col} AS TIMESTAMP), 'PT5M')",
+        "PT10M": "TIME_FLOOR(CAST({col} AS TIMESTAMP), 'PT10M')",
+        "PT15M": "TIME_FLOOR(CAST({col} AS TIMESTAMP), 'PT15M')",
+        "PT30M": "TIME_FLOOR(CAST({col} AS TIMESTAMP), 'PT30M')",
+        "PT1H": "TIME_FLOOR(CAST({col} AS TIMESTAMP), 'PT1H')",
+        "PT6H": "TIME_FLOOR(CAST({col} AS TIMESTAMP), 'PT6H')",
+        "P1D": "TIME_FLOOR(CAST({col} AS TIMESTAMP), 'P1D')",
+        "P1W": "TIME_FLOOR(CAST({col} AS TIMESTAMP), 'P1W')",
+        "P1M": "TIME_FLOOR(CAST({col} AS TIMESTAMP), 'P1M')",
+        "P3M": "TIME_FLOOR(CAST({col} AS TIMESTAMP), 'P3M')",
+        "P1Y": "TIME_FLOOR(CAST({col} AS TIMESTAMP), 'P1Y')",
         "P1W/1970-01-03T00:00:00Z": (
-            "TIME_SHIFT(TIME_FLOOR(TIME_SHIFT({col}, 'P1D', 1), 'P1W'), 'P1D', 5)"
+            "TIME_SHIFT(TIME_FLOOR(TIME_SHIFT(CAST({col} AS TIMESTAMP), 'P1D', 1), 'P1W'), 'P1D', 5)"  # pylint: disable=line-too-long
         ),
         "1969-12-28T00:00:00Z/P1W": (
-            "TIME_SHIFT(TIME_FLOOR(TIME_SHIFT({col}, 'P1D', 1), 'P1W'), 'P1D', -1)"
+            "TIME_SHIFT(TIME_FLOOR(TIME_SHIFT(CAST({col} AS TIMESTAMP), 'P1D', 1), 'P1W'), 'P1D', -1)"  # pylint: disable=line-too-long
         ),
     }
 

diff --git a/tests/integration_tests/db_engine_specs/druid_tests.py b/tests/integration_tests/db_engine_specs/druid_tests.py
@@ -50,10 +50,10 @@ def test_timegrain_expressions(self):
         col = "__time"
         sqla_col = column(col)
         test_cases = {
-            "PT1S": f"TIME_FLOOR({col}, 'PT1S')",
-            "PT5M": f"TIME_FLOOR({col}, 'PT5M')",
-            "P1W/1970-01-03T00:00:00Z": f"TIME_SHIFT(TIME_FLOOR(TIME_SHIFT({col}, 'P1D', 1), 'P1W'), 'P1D', 5)",
-            "1969-12-28T00:00:00Z/P1W": f"TIME_SHIFT(TIME_FLOOR(TIME_SHIFT({col}, 'P1D', 1), 'P1W'), 'P1D', -1)",
+            "PT1S": f"TIME_FLOOR(CAST({col} AS TIMESTAMP), 'PT1S')",
+            "PT5M": f"TIME_FLOOR(CAST({col} AS TIMESTAMP), 'PT5M')",
+            "P1W/1970-01-03T00:00:00Z": f"TIME_SHIFT(TIME_FLOOR(TIME_SHIFT(CAST({col} AS TIMESTAMP), 'P1D', 1), 'P1W'), 'P1D', 5)",
+            "1969-12-28T00:00:00Z/P1W": f"TIME_SHIFT(TIME_FLOOR(TIME_SHIFT(CAST({col} AS TIMESTAMP), 'P1D', 1), 'P1W'), 'P1D', -1)",
         }
         for grain, expected in test_cases.items():
             actual = DruidEngineSpec.get_timestamp_expr(