From 5aaa3330a66c90047da3dcc557970cf9383d4933 Mon Sep 17 00:00:00 2001
From: XcYyCx <33602700+xingyc15@users.noreply.github.com>
Date: Thu, 28 Oct 2021 01:18:03 -0700
Subject: [PATCH] feat: Make Jinja template applied in timestamp columns
 (#17237)

* Update models.py

* Add optional template processor to get_timestamp_expression()

* Update models.py
---
 superset/connectors/sqla/models.py | 24 +++++++++++++++++++-----
 1 file changed, 19 insertions(+), 5 deletions(-)

diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py
index 83abc27c2759..c19f01fb2cfb 100644
--- a/superset/connectors/sqla/models.py
+++ b/superset/connectors/sqla/models.py
@@ -303,7 +303,10 @@ def get_time_filter(
         return and_(*l)
 
     def get_timestamp_expression(
-        self, time_grain: Optional[str], label: Optional[str] = None
+        self,
+        time_grain: Optional[str],
+        label: Optional[str] = None,
+        template_processor: Optional[BaseTemplateProcessor] = None,
     ) -> Union[TimestampExpression, Label]:
         """
         Return a SQLAlchemy Core element representation of self to be used in a query.
@@ -322,7 +325,10 @@ def get_timestamp_expression(
             sqla_col = column(self.column_name, type_=type_)
             return self.table.make_sqla_column_compatible(sqla_col, label)
         if self.expression:
-            col = literal_column(self.expression, type_=type_)
+            expression = self.expression
+            if template_processor:
+                expression = template_processor.process_template(self.expression)
+            col = literal_column(expression, type_=type_)
         else:
             col = column(self.column_name, type_=type_)
         time_expr = self.db_engine_spec.get_timestamp_expr(
@@ -1088,7 +1094,11 @@ def get_sqla_query(  # pylint: disable=too-many-arguments,too-many-locals,too-ma
                 # if groupby field/expr equals granularity field/expr
                 table_col = columns_by_name.get(selected)
                 if table_col and table_col.type_generic == GenericDataType.TEMPORAL:
-                    outer = table_col.get_timestamp_expression(time_grain, selected)
+                    outer = table_col.get_timestamp_expression(
+                        time_grain=time_grain,
+                        label=selected,
+                        template_processor=template_processor,
+                    )
                 # if groupby field equals a selected column
                 elif table_col:
                     outer = table_col.get_sqla_col()
@@ -1121,7 +1131,9 @@ def get_sqla_query(  # pylint: disable=too-many-arguments,too-many-locals,too-ma
             time_filters = []
 
             if is_timeseries:
-                timestamp = dttm_col.get_timestamp_expression(time_grain)
+                timestamp = dttm_col.get_timestamp_expression(
+                    time_grain=time_grain, template_processor=template_processor
+                )
                 # always put timestamp as the first column
                 select_exprs.insert(0, timestamp)
                 groupby_all_columns[timestamp.name] = timestamp
@@ -1186,7 +1198,9 @@ def get_sqla_query(  # pylint: disable=too-many-arguments,too-many-locals,too-ma
 
             if col_obj:
                 if filter_grain:
-                    sqla_col = col_obj.get_timestamp_expression(filter_grain)
+                    sqla_col = col_obj.get_timestamp_expression(
+                        time_grain=filter_grain, template_processor=template_processor
+                    )
                 else:
                     sqla_col = col_obj.get_sqla_col()
                 col_spec = db_engine_spec.get_column_spec(col_obj.type)