Add check query tests

dbt-labs · Jan 22, 2025 · c7a436f · c7a436f
1 parent bc7fe7d
commit c7a436f
Show file tree

Hide file tree

Showing 2 changed files with 147 additions and 0 deletions.
diff --git a/tests_metricflow/integration/test_cases/itest_granularity.yaml b/tests_metricflow/integration/test_cases/itest_granularity.yaml
@@ -961,3 +961,133 @@ integration_test:
       GROUP BY subq_2.martian_day
     ) subq_5
     ON subq_6.metric_time__martian_day = subq_5.metric_time__martian_day
+---
+integration_test:
+  name: custom_offset_window
+  description: Test querying a metric with a custom offset window
+  model: SIMPLE_MODEL
+  metrics: ["bookings_offset_one_martian_day"]
+  group_bys: ["metric_time__day"]
+  check_query: |
+    WITH cte AS (
+      SELECT
+        ds AS ds__day
+        , martian_day AS ds__martian_day
+        , FIRST_VALUE(ds) OVER (
+          PARTITION BY martian_day
+          ORDER BY ds
+          ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING
+        ) AS ds__martian_day__first_value
+        , LAST_VALUE(ds) OVER (
+          PARTITION BY martian_day
+          ORDER BY ds
+          ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING
+        ) AS ds__martian_day__last_value
+        , ROW_NUMBER() OVER (
+          PARTITION BY martian_day
+          ORDER BY ds
+        ) AS ds__day__row_number
+      FROM {{ source_schema }}.mf_time_spine ts
+    )
+
+    SELECT
+      subq_8.ds__day__lead AS metric_time__day
+      , SUM(1) AS bookings_offset_one_martian_day
+    FROM (
+      SELECT
+        cte.ds__day AS ds__day
+        , CASE
+          WHEN {{ render_date_add("subq_7.ds__martian_day__first_value__offset", "(cte.ds__day__row_number - 1)", TimeGranularity.DAY) }} <= subq_7.ds__martian_day__last_value__offset
+            THEN {{ render_date_add("subq_7.ds__martian_day__first_value__offset", "(cte.ds__day__row_number - 1)", TimeGranularity.DAY) }}
+          ELSE subq_7.ds__martian_day__last_value__offset
+        END AS ds__day__lead
+      FROM cte
+      INNER JOIN (
+        SELECT
+          ds__martian_day
+          , LEAD(ds__martian_day__first_value, 1) OVER (ORDER BY ds__martian_day) AS ds__martian_day__first_value__offset
+          , LEAD(ds__martian_day__last_value, 1) OVER (ORDER BY ds__martian_day) AS ds__martian_day__last_value__offset
+        FROM (
+          SELECT
+            ds__martian_day__first_value
+            , ds__martian_day__last_value
+            , ds__martian_day
+          FROM cte
+          GROUP BY
+            ds__martian_day__first_value
+            , ds__martian_day__last_value
+            , ds__martian_day
+        ) subq_5
+      ) subq_7
+      ON cte.ds__martian_day = subq_7.ds__martian_day
+    ) subq_8
+    INNER JOIN {{ source_schema }}.fct_bookings b ON subq_8.ds__day = {{ render_date_trunc("b.ds", TimeGranularity.DAY) }}
+    GROUP BY subq_8.ds__day__lead
+---
+integration_test:
+  name: custom_offset_window_with_grain_and_date_part
+  description: Test querying a metric with a custom offset window
+  model: SIMPLE_MODEL
+  metrics: ["bookings_offset_one_martian_day"]
+  group_by_objs: [{"name": "booking__ds", "grain": "week"}, {"name": "metric_time", "date_part": "month"}, {"name": "booking__ds", "grain": "martian_day"}]
+  check_query: |
+    WITH cte AS (
+      SELECT
+        ds AS ds__day
+        , martian_day AS ds__martian_day
+        , FIRST_VALUE(ds) OVER (
+          PARTITION BY martian_day
+          ORDER BY ds
+          ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING
+        ) AS ds__martian_day__first_value
+        , LAST_VALUE(ds) OVER (
+          PARTITION BY martian_day
+          ORDER BY ds
+          ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING
+        ) AS ds__martian_day__last_value
+        , ROW_NUMBER() OVER (
+          PARTITION BY martian_day
+          ORDER BY ds
+        ) AS ds__day__row_number
+      FROM {{ source_schema }}.mf_time_spine ts
+    )
+
+    SELECT
+      subq_11.martian_day AS booking__ds__martian_day
+      , {{ render_date_trunc("subq_8.ds__day__lead", TimeGranularity.WEEK) }} AS booking__ds__week
+      , {{ render_extract("subq_8.ds__day__lead", DatePart.MONTH) }} AS metric_time__extract_month
+      , SUM(1) AS bookings_offset_one_martian_day
+    FROM (
+      SELECT
+        cte.ds__day
+        , CASE
+          WHEN {{ render_date_add("subq_7.ds__martian_day__first_value__offset", "(cte.ds__day__row_number - 1)", TimeGranularity.DAY) }} <= subq_7.ds__martian_day__last_value__offset
+            THEN {{ render_date_add("subq_7.ds__martian_day__first_value__offset", "(cte.ds__day__row_number - 1)", TimeGranularity.DAY) }}
+          ELSE subq_7.ds__martian_day__last_value__offset
+        END AS ds__day__lead
+      FROM cte
+      INNER JOIN (
+        SELECT
+          ds__martian_day
+          , LEAD(ds__martian_day__first_value, 1) OVER (ORDER BY ds__martian_day) AS ds__martian_day__first_value__offset
+          , LEAD(ds__martian_day__last_value, 1) OVER (ORDER BY ds__martian_day) AS ds__martian_day__last_value__offset
+        FROM (
+          SELECT
+            ds__martian_day__first_value
+            , ds__martian_day__last_value
+            , ds__martian_day
+          FROM cte
+          GROUP BY
+            ds__martian_day__first_value
+            , ds__martian_day__last_value
+            , ds__martian_day
+        ) subq_5
+      ) subq_7
+      ON cte.ds__martian_day = subq_7.ds__martian_day
+    ) subq_8
+    INNER JOIN {{ source_schema }}.fct_bookings b ON subq_8.ds__day = {{ render_date_trunc("b.ds", TimeGranularity.DAY) }}
+    LEFT OUTER JOIN {{ source_schema }}.mf_time_spine subq_11 ON subq_8.ds__day__lead = subq_11.ds
+    GROUP BY
+      subq_11.martian_day
+      , {{ render_date_trunc("subq_8.ds__day__lead", TimeGranularity.WEEK) }}
+      , {{ render_extract("subq_8.ds__day__lead", DatePart.MONTH) }}
diff --git a/tests_metricflow/integration/test_configured_cases.py b/tests_metricflow/integration/test_configured_cases.py
@@ -16,6 +16,7 @@
 from metricflow_semantics.protocols.query_parameter import DimensionOrEntityQueryParameter
 from metricflow_semantics.specs.query_param_implementations import DimensionOrEntityParameter, TimeDimensionParameter
 from metricflow_semantics.sql.sql_exprs import (
+    SqlAddTimeExpression,
     SqlCastToTimestampExpression,
     SqlColumnReference,
     SqlColumnReferenceExpression,
@@ -98,6 +99,20 @@ def render_date_sub(
         )
         return self._sql_client.sql_plan_renderer.expr_renderer.render_sql_expr(expr).sql
 
+    def render_date_add(
+        self,
+        date_column: str,
+        count_column: str,
+        granularity: TimeGranularity,
+    ) -> str:
+        """Renders a date add expression."""
+        expr = SqlAddTimeExpression.create(
+            arg=SqlStringExpression.create(sql_expr=date_column, requires_parenthesis=False),
+            count_expr=SqlStringExpression.create(sql_expr=count_column, requires_parenthesis=False),
+            granularity=granularity,
+        )
+        return self._sql_client.sql_query_plan_renderer.expr_renderer.render_sql_expr(expr).sql
+
     def render_date_trunc(self, expr: str, granularity: TimeGranularity) -> str:
         """Return the DATE_TRUNC() call that can be used for converting the given expr to the granularity."""
         renderable_expr = SqlDateTruncExpression.create(
@@ -291,6 +306,7 @@ def test_case(
                         TimeGranularity=TimeGranularity,
                         DatePart=DatePart,
                         render_date_sub=check_query_helpers.render_date_sub,
+                        render_date_add=check_query_helpers.render_date_add,
                         render_date_trunc=check_query_helpers.render_date_trunc,
                         render_extract=check_query_helpers.render_extract,
                         render_percentile_expr=check_query_helpers.render_percentile_expr,
@@ -324,6 +340,7 @@ def test_case(
             TimeGranularity=TimeGranularity,
             DatePart=DatePart,
             render_date_sub=check_query_helpers.render_date_sub,
+            render_date_add=check_query_helpers.render_date_add,
             render_date_trunc=check_query_helpers.render_date_trunc,
             render_extract=check_query_helpers.render_extract,
             render_percentile_expr=check_query_helpers.render_percentile_expr,