From 9044c1efda70ea28c88eca54f98b1f9d7c688c17 Mon Sep 17 00:00:00 2001
From: Xiang Fu <fx19880617@gmail.com>
Date: Thu, 14 Jan 2021 19:44:42 -0800
Subject: [PATCH] Fixing Pinot queries for time granularities:
 WEEKS/MONTHS/QUARTERS/YEARS

---
 superset/db_engine_specs/pinot.py    | 26 +++++++++++++++++++++-----
 tests/db_engine_specs/pinot_tests.py | 12 ++++++++++--
 2 files changed, 31 insertions(+), 7 deletions(-)

diff --git a/superset/db_engine_specs/pinot.py b/superset/db_engine_specs/pinot.py
index 7a11d08fb547..1f504b840a2f 100644
--- a/superset/db_engine_specs/pinot.py
+++ b/superset/db_engine_specs/pinot.py
@@ -35,10 +35,10 @@ class PinotEngineSpec(BaseEngineSpec):  # pylint: disable=abstract-method
         "PT1M": "1:MINUTES",
         "PT1H": "1:HOURS",
         "P1D": "1:DAYS",
-        "P1W": "1:WEEKS",
-        "P1M": "1:MONTHS",
-        "P0.25Y": "3:MONTHS",
-        "P1Y": "1:YEARS",
+        "P1W": "week",
+        "P1M": "month",
+        "P0.25Y": "quarter",
+        "P1Y": "year",
     }
 
     _python_to_java_time_patterns: Dict[str, str] = {
@@ -50,6 +50,17 @@ class PinotEngineSpec(BaseEngineSpec):  # pylint: disable=abstract-method
         "%S": "ss",
     }
 
+    _use_date_trunc_function: Dict[str, bool] = {
+        "PT1S": False,
+        "PT1M": False,
+        "PT1H": False,
+        "P1D": False,
+        "P1W": True,
+        "P1M": True,
+        "P0.25Y": True,
+        "P1Y": True,
+    }
+
     @classmethod
     def get_timestamp_expr(
         cls,
@@ -86,8 +97,13 @@ def get_timestamp_expr(
                 raise NotImplementedError("No pinot grain spec for " + str(time_grain))
         else:
             return TimestampExpression("{{col}}", col)
+
         # In pinot the output is a string since there is no timestamp column like pg
-        time_expr = f"DATETIMECONVERT({{col}}, '{tf}', '{tf}', '{granularity}')"
+        if cls._use_date_trunc_function.get(time_grain):
+            time_expr = f"DATETRUNC('{granularity}', {{col}}, '{seconds_or_ms}')"
+        else:
+            time_expr = f"DATETIMECONVERT({{col}}, '{tf}', '{tf}', '{granularity}')"
+
         return TimestampExpression(time_expr, col)
 
     @classmethod
diff --git a/tests/db_engine_specs/pinot_tests.py b/tests/db_engine_specs/pinot_tests.py
index 973282737b2d..4dcca0ced118 100644
--- a/tests/db_engine_specs/pinot_tests.py
+++ b/tests/db_engine_specs/pinot_tests.py
@@ -23,11 +23,19 @@
 class TestPinotDbEngineSpec(TestDbEngineSpec):
     """ Tests pertaining to our Pinot database support """
 
+    def test_pinot_time_expression_sec_one_1d_grain(self):
+        col = column("tstamp")
+        expr = PinotEngineSpec.get_timestamp_expr(col, "epoch_s", "P1D")
+        result = str(expr.compile())
+        self.assertEqual(
+            result,
+            "DATETIMECONVERT(tstamp, '1:SECONDS:EPOCH', '1:SECONDS:EPOCH', '1:DAYS')",
+        )  # noqa
+
     def test_pinot_time_expression_sec_one_1m_grain(self):
         col = column("tstamp")
         expr = PinotEngineSpec.get_timestamp_expr(col, "epoch_s", "P1M")
         result = str(expr.compile())
         self.assertEqual(
-            result,
-            "DATETIMECONVERT(tstamp, '1:SECONDS:EPOCH', '1:SECONDS:EPOCH', '1:MONTHS')",
+            result, "DATETRUNC('month', tstamp, 'SECONDS')",
         )  # noqa