feat: Make Jinja template applied in timestamp columns (#17237)

* Update models.py * Add optional template processor to get_timestamp_expression() * Update models.py
apache · Oct 28, 2021 · 5aaa333 · 5aaa333
1 parent 93f59e0
commit 5aaa333
Showing 1 changed file with 19 additions and 5 deletions.
diff --git 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)