refactor: refine timestamp expr function (#21510)

apache · Sep 20, 2022 · 737d4dc · 737d4dc
1 parent 4200082
commit 737d4dc
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 9 deletions.
diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py
@@ -375,9 +375,7 @@ def get_timestamp_expression(
             col = literal_column(expression, type_=type_)
         else:
             col = column(self.column_name, type_=type_)
-        time_expr = self.db_engine_spec.get_timestamp_expr(
-            col, pdf, time_grain, self.type
-        )
+        time_expr = self.db_engine_spec.get_timestamp_expr(col, pdf, time_grain)
         return self.table.make_sqla_column_compatible(time_expr, label)
 
     def dttm_sql_literal(self, dttm: DateTime) -> str:
@@ -1147,7 +1145,6 @@ def adhoc_column_to_sqla(
                 col=sqla_column,
                 pdf=None,
                 time_grain=time_grain,
-                type_=str(getattr(sqla_column, "type", "")),
             )
         return self.make_sqla_column_compatible(sqla_column, label)
 

diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py
@@ -476,18 +476,17 @@ def get_timestamp_expr(
         col: ColumnClause,
         pdf: Optional[str],
         time_grain: Optional[str],
-        type_: Optional[str] = None,
     ) -> TimestampExpression:
         """
         Construct a TimestampExpression to be used in a SQLAlchemy query.
 
         :param col: Target column for the TimestampExpression
         :param pdf: date format (seconds or milliseconds)
         :param time_grain: time grain, e.g. P1Y for 1 year
-        :param type_: the source column type
         :return: TimestampExpression object
         """
         if time_grain:
+            type_ = str(getattr(col, "type", ""))
             time_expr = cls.get_time_grain_expressions().get(time_grain)
             if not time_expr:
                 raise NotImplementedError(

diff --git a/superset/db_engine_specs/pinot.py b/superset/db_engine_specs/pinot.py
@@ -75,7 +75,6 @@ def get_timestamp_expr(
         col: ColumnClause,
         pdf: Optional[str],
         time_grain: Optional[str],
-        type_: Optional[str] = None,
     ) -> TimestampExpression:
         if not pdf:
             raise NotImplementedError(f"Empty date format for '{col}'")

diff --git a/tests/integration_tests/db_engine_specs/bigquery_tests.py b/tests/integration_tests/db_engine_specs/bigquery_tests.py
@@ -77,8 +77,9 @@ def test_timegrain_expressions(self):
             "TIMESTAMP": "TIMESTAMP_TRUNC(temporal, HOUR)",
         }
         for type_, expected in test_cases.items():
+            col.type = type_
             actual = BigQueryEngineSpec.get_timestamp_expr(
-                col=col, pdf=None, time_grain="PT1H", type_=type_
+                col=col, pdf=None, time_grain="PT1H"
             )
             self.assertEqual(str(actual), expected)
 
@@ -99,8 +100,9 @@ def test_custom_minute_timegrain_expressions(self):
             ") AS TIMESTAMP)",
         }
         for type_, expected in test_cases.items():
+            col.type = type_
             actual = BigQueryEngineSpec.get_timestamp_expr(
-                col=col, pdf=None, time_grain="PT5M", type_=type_
+                col=col, pdf=None, time_grain="PT5M"
             )
             assert str(actual) == expected