refactor(pinot): The python_date_format for a temporal column was not being passed to get_timestamp_expr (#24942)

tests/integration_tests/db_engine_specs/pinot_tests.py

@@ -27,61 +27,66 @@ class TestPinotDbEngineSpec(TestDbEngineSpec):
         col = column("tstamp")
         expr = PinotEngineSpec.get_timestamp_expr(col, "epoch_s", "P1D")
         result = str(expr.compile())
+        expected = (
+            "CAST(DATE_TRUNC('day', CAST("
+            + "DATETIMECONVERT(tstamp, '1:SECONDS:EPOCH', "
+            + "'1:SECONDS:EPOCH', '1:SECONDS') AS TIMESTAMP)) AS TIMESTAMP)"
+        )
         self.assertEqual(
             result,
-            "DATETIMECONVERT(tstamp, '1:SECONDS:EPOCH', '1:SECONDS:EPOCH', '1:DAYS')",
+            expected,
         )
 
     def test_pinot_time_expression_simple_date_format_1d_grain(self):
         col = column("tstamp")
         expr = PinotEngineSpec.get_timestamp_expr(col, "%Y-%m-%d %H:%M:%S", "P1D")
         result = str(expr.compile())
+        expected = "CAST(DATE_TRUNC('day', CAST(tstamp AS TIMESTAMP)) AS TIMESTAMP)"
         self.assertEqual(
             result,
-            (
-                "DATETIMECONVERT(tstamp, "
-                + "'1:SECONDS:SIMPLE_DATE_FORMAT:yyyy-MM-dd HH:mm:ss', "
-                + "'1:SECONDS:SIMPLE_DATE_FORMAT:yyyy-MM-dd HH:mm:ss', '1:DAYS')"
-            ),
+            expected,
         )
 
     def test_pinot_time_expression_simple_date_format_10m_grain(self):
         col = column("tstamp")
         expr = PinotEngineSpec.get_timestamp_expr(col, "%Y-%m-%d %H:%M:%S", "PT10M")
         result = str(expr.compile())
+        expected = (
+            "CAST(ROUND(DATE_TRUNC('minute', CAST(tstamp AS "
+            + "TIMESTAMP)), 600000) AS TIMESTAMP)"
+        )
         self.assertEqual(
             result,
-            (
-                "DATETIMECONVERT(tstamp, "
-                + "'1:SECONDS:SIMPLE_DATE_FORMAT:yyyy-MM-dd HH:mm:ss', "
-                + "'1:SECONDS:SIMPLE_DATE_FORMAT:yyyy-MM-dd HH:mm:ss', '10:MINUTES')"
-            ),
+            expected,
         )
 
     def test_pinot_time_expression_simple_date_format_1w_grain(self):
         col = column("tstamp")
         expr = PinotEngineSpec.get_timestamp_expr(col, "%Y-%m-%d %H:%M:%S", "P1W")
         result = str(expr.compile())
+        expected = "CAST(DATE_TRUNC('week', CAST(tstamp AS TIMESTAMP)) AS TIMESTAMP)"
         self.assertEqual(
             result,
-            (
-                "ToDateTime(DATETRUNC('week', FromDateTime(tstamp, "
-                + "'yyyy-MM-dd HH:mm:ss'), 'MILLISECONDS'), 'yyyy-MM-dd HH:mm:ss')"
-            ),
+            expected,
         )
 
     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())
+        expected = (
+            "CAST(DATE_TRUNC('month', CAST("
+            + "DATETIMECONVERT(tstamp, '1:SECONDS:EPOCH', "
+            + "'1:SECONDS:EPOCH', '1:SECONDS') AS TIMESTAMP)) AS TIMESTAMP)"
+        )
         self.assertEqual(
             result,
-            "DATETRUNC('month', tstamp, 'SECONDS')",
+            expected,
         )
 
     def test_invalid_get_time_expression_arguments(self):
         with self.assertRaises(NotImplementedError):
-            PinotEngineSpec.get_timestamp_expr(column("tstamp"), None, "P1M")
+            PinotEngineSpec.get_timestamp_expr(column("tstamp"), None, "P0.25Y")
 
         with self.assertRaises(NotImplementedError):
             PinotEngineSpec.get_timestamp_expr(

tests/unit_tests/db_engine_specs/test_pinot.py

@@ -0,0 +1,57 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from unittest import mock
+
+import pytest
+from sqlalchemy import column
+
+
+@pytest.mark.parametrize(
+    "time_grain,expected_result",
+    [
+        ("PT1S", "CAST(DATE_TRUNC('second', CAST(col AS TIMESTAMP)) AS TIMESTAMP)"),
+        (
+            "PT5M",
+            "CAST(ROUND(DATE_TRUNC('minute', CAST(col AS TIMESTAMP)), 300000) AS TIMESTAMP)",
+        ),
+        ("P1W", "CAST(DATE_TRUNC('week', CAST(col AS TIMESTAMP)) AS TIMESTAMP)"),
+        ("P1M", "CAST(DATE_TRUNC('month', CAST(col AS TIMESTAMP)) AS TIMESTAMP)"),
+        ("P3M", "CAST(DATE_TRUNC('quarter', CAST(col AS TIMESTAMP)) AS TIMESTAMP)"),
+        ("P1Y", "CAST(DATE_TRUNC('year', CAST(col AS TIMESTAMP)) AS TIMESTAMP)"),
+    ],
+)
+def test_timegrain_expressions(time_grain: str, expected_result: str) -> None:
+    """
+    DB Eng Specs (pinot): Test time grain expressions
+    """
+    from superset.db_engine_specs.pinot import PinotEngineSpec as spec
+
+    actual = str(
+        spec.get_timestamp_expr(col=column("col"), pdf=None, time_grain=time_grain)
+    )
+    assert actual == expected_result
+
+
+def test_extras_without_ssl() -> None:
+    from superset.db_engine_specs.pinot import PinotEngineSpec as spec
+    from tests.integration_tests.fixtures.database import default_db_extra
+
+    db = mock.Mock()
+    db.extra = default_db_extra
+    db.server_cert = None
+    extras = spec.get_extra_params(db)
+    assert "connect_args" not in extras["engine_params"]