From a0fd7930ad7fe742d11469a1e22eba5a614c5479 Mon Sep 17 00:00:00 2001
From: Courtney Holcomb <courtneyeholcomb@gmail.com>
Date: Mon, 11 Sep 2023 20:14:06 -0700
Subject: [PATCH] Tests for all DateParts

---
 metricflow/plan_conversion/dataflow_to_sql.py |  4 +++
 .../integration/test_cases/itest_metrics.yaml | 31 ++++++++++++++-----
 metricflow/test/sql/test_sql_expr_render.py   |  9 ++++++
 3 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/metricflow/plan_conversion/dataflow_to_sql.py b/metricflow/plan_conversion/dataflow_to_sql.py
index 84b2e0d261..39d4597e42 100644
--- a/metricflow/plan_conversion/dataflow_to_sql.py
+++ b/metricflow/plan_conversion/dataflow_to_sql.py
@@ -552,6 +552,10 @@ def visit_aggregate_measures_node(self, node: AggregateMeasuresNode) -> SqlDataS
         the measure names as references.
 
         """
+        # TODO: would it make more sense to always apply granularity / date_part in the AggregateMeasuresNode?
+        # Keep that logic all in once place; currently happens in too many places for special cases
+        # Standard: source node; special: cumulative metrics, time offset metrics
+
         # Get the data from the parent, and change measure instances to the aggregated state.
         from_data_set: SqlDataSet = node.parent_node.accept(self)
         aggregated_instance_set = from_data_set.instance_set.transform(
diff --git a/metricflow/test/integration/test_cases/itest_metrics.yaml b/metricflow/test/integration/test_cases/itest_metrics.yaml
index f08254aaa1..e4ba24d7e4 100644
--- a/metricflow/test/integration/test_cases/itest_metrics.yaml
+++ b/metricflow/test/integration/test_cases/itest_metrics.yaml
@@ -1049,13 +1049,35 @@ integration_test:
   description: Test query using date_part
   model: SIMPLE_MODEL
   metrics: ["bookings"]
-  group_bys: ["metric_time__extract_dayofweek"]
+  group_bys: ["metric_time__year"]
   check_query: |
     SELECT
       SUM(1) AS bookings
+      , {{ render_extract("ds", DatePart.YEAR) }} AS metric_time__year
+    FROM {{ source_schema }}.fct_bookings
+    GROUP BY {{ render_extract("ds", DatePart.YEAR) }};
+---
+integration_test:
+  name: simple_query_with_multiple_date_parts
+  description: Test query using multiple date_parts
+  model: SIMPLE_MODEL
+  metrics: ["bookings"]
+  group_bys: ["metric_time__extract_quarter", "metric_time__extract_dayofweek", "metric_time__extract_dayofyear", "metric_time__extract_day", "metric_time__extract_week"]
+  check_query: |
+    SELECT
+      SUM(1) AS bookings
+      , {{ render_extract("ds", DatePart.QUARTER) }} AS metric_time__extract_quarter
       , {{ render_extract("ds", DatePart.DAYOFWEEK) }} AS metric_time__extract_dayofweek
+      , {{ render_extract("ds", DatePart.DAYOFYEAR) }} AS metric_time__extract_dayofyear
+      , {{ render_extract("ds", DatePart.DAY) }} AS metric_time__extract_day
+      , {{ render_extract("ds", DatePart.WEEK) }} AS metric_time__extract_week
     FROM {{ source_schema }}.fct_bookings
-    GROUP BY {{ render_extract("ds", DatePart.DAYOFWEEK) }};
+    GROUP BY
+      {{ render_extract("ds", DatePart.QUARTER) }}
+      , {{ render_extract("ds", DatePart.DAYOFWEEK) }}
+      , {{ render_extract("ds", DatePart.DAYOFYEAR) }}
+      , {{ render_extract("ds", DatePart.DAY) }}
+      , {{ render_extract("ds", DatePart.WEEK) }};
 ---
 integration_test:
   name: derived_metric_offset_window_and_date_part
@@ -1076,8 +1098,3 @@ integration_test:
     ) b
     ON {{ render_date_sub("a", "ds", 5, TimeGranularity.DAY) }} = b.metric_time__day
     GROUP BY metric_time__extract_month
-
-# TODO:
-# test each date part syntax with each engine
-# dataflow plan tests?
-# dataflow to sql tests?
diff --git a/metricflow/test/sql/test_sql_expr_render.py b/metricflow/test/sql/test_sql_expr_render.py
index c123843ae2..035dca2f2e 100644
--- a/metricflow/test/sql/test_sql_expr_render.py
+++ b/metricflow/test/sql/test_sql_expr_render.py
@@ -17,6 +17,7 @@
     SqlComparison,
     SqlComparisonExpression,
     SqlDateTruncExpression,
+    SqlExtractExpression,
     SqlFunction,
     SqlIsNullExpression,
     SqlLogicalExpression,
@@ -29,6 +30,7 @@
     SqlWindowFunctionExpression,
     SqlWindowOrderByArgument,
 )
+from metricflow.time.date_part import DatePart
 
 logger = logging.getLogger(__name__)
 
@@ -193,6 +195,13 @@ def test_date_trunc_expr(default_expr_renderer: DefaultSqlExpressionRenderer) ->
     assert actual == "DATE_TRUNC('month', ds)"
 
 
+def test_extract_expr(default_expr_renderer: DefaultSqlExpressionRenderer) -> None:  # noqa: D
+    actual = default_expr_renderer.render_sql_expr(
+        SqlExtractExpression(date_part=DatePart.DAYOFYEAR, arg=SqlStringExpression("ds"))
+    ).sql
+    assert actual == "EXTRACT(DAYOFYEAR FROM ds)"
+
+
 def test_ratio_computation_expr(default_expr_renderer: DefaultSqlExpressionRenderer) -> None:  # noqa: D
     actual = default_expr_renderer.render_sql_expr(
         SqlRatioComputationExpression(