Feat: Add support for Restatements on SCD Type 2 models

Author: themisvaltinos

sqlmesh/core/engine_adapter/base.py

@@ -1514,6 +1514,7 @@ def _scd_type_2(
         unique_key: t.Sequence[exp.Expression],
         valid_from_col: exp.Column,
         valid_to_col: exp.Column,
+        start: TimeLike,
         execution_time: t.Union[TimeLike, exp.Column],
         invalidate_hard_deletes: bool = True,
         updated_at_col: t.Optional[exp.Column] = None,
@@ -1710,6 +1711,22 @@ def remove_managed_columns(
         )
         if truncate:
             existing_rows_query = existing_rows_query.limit(0)
+        else:
+            # If truncate is false it is not the first insert
+            # Determine the cleanup timestamp for restatement or a regular incremental run
+            cleanup_ts = to_time_column(start, time_data_type, self.dialect, nullable=True)
+
+            # Delete records that were created at or after cleanup point
+            self.delete_from(table_name=target_table, where=valid_from_col > cleanup_ts)
+
+            # "Re-open" records that were closed at or after cleanup point
+            self.update_table(
+                table_name=target_table,
+                properties={valid_to_col.name: exp.Null()},
+                where=exp.and_(
+                    valid_to_col > cleanup_ts,
+                ),
+            )
 
         with source_queries[0] as source_query:
             prefixed_columns_to_types = []

sqlmesh/core/engine_adapter/trino.py

@@ -256,6 +256,7 @@ def _scd_type_2(
         unique_key: t.Sequence[exp.Expression],
         valid_from_col: exp.Column,
         valid_to_col: exp.Column,
+        start: TimeLike,
         execution_time: t.Union[TimeLike, exp.Column],
         invalidate_hard_deletes: bool = True,
         updated_at_col: t.Optional[exp.Column] = None,
@@ -277,6 +278,7 @@ def _scd_type_2(
             unique_key,
             valid_from_col,
             valid_to_col,
+            start,
             execution_time,
             invalidate_hard_deletes,
             updated_at_col,

sqlmesh/core/model/kind.py

@@ -140,7 +140,6 @@ def full_history_restatement_only(self) -> bool:
             self.is_incremental_unmanaged
             or self.is_incremental_by_unique_key
             or self.is_incremental_by_partition
-            or self.is_scd_type_2
             or self.is_managed
             or self.is_full
             or self.is_view

sqlmesh/core/snapshot/definition.py

@@ -797,6 +797,21 @@ def get_removal_interval(
 
             removal_interval = expanded_removal_interval
 
+        # SCD Type 2 validation that end date is the latest interval if it was provided
+        if not is_preview and self.is_scd_type_2 and self.intervals:
+            requested_start, requested_end = removal_interval
+            latest_end = max(interval[1] for interval in self.intervals)
+            if requested_end != latest_end:
+                from sqlmesh.core.console import get_console
+
+                get_console().log_warning(
+                    f"SCD Type 2 model '{self.model.name}' does not support end date in restatements.\n"
+                    f"Requested end date [{to_ts(requested_end)}] doesn't match latest interval end date [{to_ts(latest_end)}].\n"
+                    f"You can set start date but end date must be the latest interval so ."
+                )
+
+                removal_interval = self.inclusive_exclusive(requested_start, latest_end, strict)
+
         return removal_interval
 
     @property

sqlmesh/core/snapshot/evaluator.py

@@ -1788,6 +1788,7 @@ def insert(
                 table_description=model.description,
                 column_descriptions=model.column_descriptions,
                 truncate=is_first_insert,
+                start=kwargs["start"],
             )
         elif isinstance(model.kind, SCDType2ByColumnKind):
             self.adapter.scd_type_2_by_column(
@@ -1805,6 +1806,7 @@ def insert(
                 table_description=model.description,
                 column_descriptions=model.column_descriptions,
                 truncate=is_first_insert,
+                start=kwargs["start"],
             )
         else:
             raise SQLMeshError(

tests/core/engine_adapter/integration/test_integration.py

@@ -744,6 +744,7 @@ def test_scd_type_2_by_time(ctx_query_and_df: TestContext):
         columns_to_types=input_schema,
         table_format=ctx.default_table_format,
         truncate=True,
+        start="2022-01-01 00:00:00",
     )
     results = ctx.get_metadata_results()
     assert len(results.views) == 0
@@ -807,6 +808,7 @@ def test_scd_type_2_by_time(ctx_query_and_df: TestContext):
         columns_to_types=input_schema,
         table_format=ctx.default_table_format,
         truncate=False,
+        start="2022-01-01 00:00:00",
     )
     results = ctx.get_metadata_results()
     assert len(results.views) == 0
@@ -899,6 +901,7 @@ def test_scd_type_2_by_column(ctx_query_and_df: TestContext):
         execution_time_as_valid_from=False,
         columns_to_types=ctx.columns_to_types,
         truncate=True,
+        start="2023-01-01",
     )
     results = ctx.get_metadata_results()
     assert len(results.views) == 0
@@ -970,6 +973,7 @@ def test_scd_type_2_by_column(ctx_query_and_df: TestContext):
         execution_time_as_valid_from=False,
         columns_to_types=ctx.columns_to_types,
         truncate=False,
+        start="2023-01-01",
     )
     results = ctx.get_metadata_results()
     assert len(results.views) == 0

tests/core/engine_adapter/test_base.py

@@ -1222,6 +1222,7 @@ def test_scd_type_2_by_time(make_mocked_engine_adapter: t.Callable):
             "test_valid_to": exp.DataType.build("TIMESTAMP"),
         },
         execution_time=datetime(2020, 1, 1, 0, 0, 0),
+        start=datetime(2020, 1, 1, 0, 0, 0),
     )
 
     assert (
@@ -1421,6 +1422,7 @@ def test_scd_type_2_by_time_no_invalidate_hard_deletes(make_mocked_engine_adapte
             "test_valid_to": exp.DataType.build("TIMESTAMP"),
         },
         execution_time=datetime(2020, 1, 1, 0, 0, 0),
+        start=datetime(2020, 1, 1, 0, 0, 0),
     )
 
     assert (
@@ -1609,6 +1611,7 @@ def test_merge_scd_type_2_pandas(make_mocked_engine_adapter: t.Callable):
             "test_valid_to": exp.DataType.build("TIMESTAMPTZ"),
         },
         execution_time=datetime(2020, 1, 1, 0, 0, 0),
+        start=datetime(2020, 1, 1, 0, 0, 0),
     )
 
     assert (
@@ -1790,6 +1793,7 @@ def test_scd_type_2_by_column(make_mocked_engine_adapter: t.Callable):
             "test_valid_to": exp.DataType.build("TIMESTAMP"),
         },
         execution_time=datetime(2020, 1, 1, 0, 0, 0),
+        start=datetime(2020, 1, 1, 0, 0, 0),
         extra_col_ignore="testing",
     )
 
@@ -1972,6 +1976,7 @@ def test_scd_type_2_by_column_composite_key(make_mocked_engine_adapter: t.Callab
             "test_valid_to": exp.DataType.build("TIMESTAMP"),
         },
         execution_time=datetime(2020, 1, 1, 0, 0, 0),
+        start=datetime(2020, 1, 1, 0, 0, 0),
     )
 
     assert (
@@ -2164,6 +2169,7 @@ def test_scd_type_2_truncate(make_mocked_engine_adapter: t.Callable):
         },
         execution_time=datetime(2020, 1, 1, 0, 0, 0),
         truncate=True,
+        start=datetime(2020, 1, 1, 0, 0, 0),
     )
 
     assert (
@@ -2346,6 +2352,7 @@ def test_scd_type_2_by_column_star_check(make_mocked_engine_adapter: t.Callable)
             "test_valid_to": exp.DataType.build("TIMESTAMP"),
         },
         execution_time=datetime(2020, 1, 1, 0, 0, 0),
+        start=datetime(2020, 1, 1, 0, 0, 0),
     )
 
     assert (
@@ -2541,6 +2548,7 @@ def test_scd_type_2_by_column_no_invalidate_hard_deletes(make_mocked_engine_adap
             "test_valid_to": exp.DataType.build("TIMESTAMP"),
         },
         execution_time=datetime(2020, 1, 1, 0, 0, 0),
+        start=datetime(2020, 1, 1, 0, 0, 0),
     )
 
     assert (

tests/core/engine_adapter/test_clickhouse.py

@@ -606,6 +606,8 @@ def test_scd_type_2_by_time(
             "test_valid_to": exp.DataType.build("TIMESTAMP"),
         },
         execution_time=datetime(2020, 1, 1, 0, 0, 0),
+        start=datetime(2020, 1, 1, 0, 0, 0),
+        truncate=True,
     )
 
     assert to_sql_calls(adapter)[4] == parse_one(
@@ -637,7 +639,7 @@ def test_scd_type_2_by_time(
     TRUE AS "_exists"
   FROM ""__temp_target_efgh""
   WHERE
-    NOT "test_valid_to" IS NULL
+    NOT "test_valid_to" IS NULL LIMIT 0
 ), "latest" AS (
   SELECT
     "id",
@@ -649,7 +651,7 @@ def test_scd_type_2_by_time(
     TRUE AS "_exists"
   FROM ""__temp_target_efgh""
   WHERE
-    "test_valid_to" IS NULL
+    "test_valid_to" IS NULL LIMIT 0
 ), "deleted" AS (
   SELECT
     "static"."id",
@@ -811,6 +813,8 @@ def test_scd_type_2_by_column(
             "test_valid_to": exp.DataType.build("TIMESTAMP"),
         },
         execution_time=datetime(2020, 1, 1, 0, 0, 0),
+        start=datetime(2020, 1, 1, 0, 0, 0),
+        truncate=True,
     )
 
     assert to_sql_calls(adapter)[4] == parse_one(
@@ -840,7 +844,7 @@ def test_scd_type_2_by_column(
     TRUE AS "_exists"
   FROM "__temp_target_efgh"
   WHERE
-    NOT "test_valid_to" IS NULL
+    NOT ("test_valid_to" IS NULL) LIMIT 0
 ), "latest" AS (
   SELECT
     "id",
@@ -851,7 +855,7 @@ def test_scd_type_2_by_column(
     TRUE AS "_exists"
   FROM "__temp_target_efgh"
   WHERE
-    "test_valid_to" IS NULL
+    "test_valid_to" IS NULL LIMIT 0
 ), "deleted" AS (
   SELECT
     "static"."id",
@@ -907,7 +911,7 @@ def test_scd_type_2_by_column(
     COALESCE("joined"."t_id", "joined"."id") AS "id",
     COALESCE("joined"."t_name", "joined"."name") AS "name",
     COALESCE("joined"."t_price", "joined"."price") AS "price",
-    COALESCE("t_test_VALID_from", CAST('2020-01-01 00:00:00' AS Nullable(DateTime64(6)))) AS "test_VALID_from",
+    COALESCE("t_test_VALID_from", CAST('1970-01-01 00:00:00' AS Nullable(DateTime64(6)))) AS "test_VALID_from",
     CASE
       WHEN "joined"."_exists" IS NULL
       OR (

tests/core/engine_adapter/test_spark.py

@@ -569,6 +569,8 @@ def check_table_exists(table_name: exp.Table) -> bool:
             "test_valid_to": exp.DataType.build("TIMESTAMP"),
         },
         execution_time=datetime(2020, 1, 1, 0, 0, 0),
+        start=datetime(2020, 1, 1, 0, 0, 0),
+        truncate=True,
     )
 
     assert to_sql_calls(adapter) == [
@@ -613,7 +615,7 @@ def check_table_exists(table_name: exp.Table) -> bool:
     TRUE AS `_exists`
   FROM `db`.`temp_target_abcdefgh`
   WHERE
-    NOT `test_valid_to` IS NULL
+    NOT `test_valid_to` IS NULL LIMIT 0
 ), `latest` AS (
   SELECT
     `id`,
@@ -625,7 +627,7 @@ def check_table_exists(table_name: exp.Table) -> bool:
     TRUE AS `_exists`
   FROM `db`.`temp_target_abcdefgh`
   WHERE
-    `test_valid_to` IS NULL
+    `test_valid_to` IS NULL LIMIT 0
 ), `deleted` AS (
   SELECT
     `static`.`id`,