Repository: TobikoData/sqlmesh (python)

Issue:
Fix Restatements in Airflow Rest API
Pydantic can't serialize other models as keys in dicts or values in sets. Therefore restatement will not work for Airflow REST API users (any who doesn't use MWAA). 

Example: https://github.com/TobikoData/sqlmesh/blob/28171c02de1ef271761daed89570bd2e4ec5eb0b/sqlmesh/schedulers/airflow/common.py#L45

Pull request description:
Fix: Support for plans with restatements when using the Airflow REST extension
Closes #2478

Patch:
diff --git a/sqlmesh/core/model/definition.py b/sqlmesh/core/model/definition.py
index 7db79d57..f247e84d 100644
--- a/sqlmesh/core/model/definition.py
+++ b/sqlmesh/core/model/definition.py
@@ -902,6 +902,7 @@ class _SqlBasedModel(_Model):
                 python_env=self.python_env,
                 only_execution_time=self.kind.only_execution_time,
                 default_catalog=self.default_catalog,
+                model_fqn=self.fqn,
             )
         return self.__statement_renderers[expression_key]
 
@@ -1669,7 +1670,7 @@ def create_seed_model(
         path=path,
         seed=seed,
         kind=seed_kind,
-        depends_on=kwargs.pop("depends_on", set()),
+        depends_on=kwargs.pop("depends_on", None),
         python_env=python_env,
         jinja_macros=jinja_macros,
         pre_statements=pre_statements,
diff --git a/sqlmesh/core/renderer.py b/sqlmesh/core/renderer.py
index c7d1ae95..e94f6918 100644
--- a/sqlmesh/core/renderer.py
+++ b/sqlmesh/core/renderer.py
@@ -3,7 +3,6 @@ from __future__ import annotations
 import logging
 import typing as t
 from contextlib import contextmanager
-from datetime import datetime
 from pathlib import Path
 
 from sqlglot import exp, parse
@@ -15,7 +14,7 @@ from sqlglot.optimizer.simplify import simplify
 from sqlmesh.core import constants as c
 from sqlmesh.core import dialect as d
 from sqlmesh.core.macros import MacroEvaluator, RuntimeStage
-from sqlmesh.utils.date import TimeLike, date_dict, make_inclusive, to_datetime
+from sqlmesh.utils.date import TimeLike, date_dict, make_inclusive_end, to_datetime
 from sqlmesh.utils.errors import (
     ConfigError,
     MacroEvalError,
@@ -31,8 +30,6 @@ if t.TYPE_CHECKING:
 
     from sqlmesh.core.snapshot import DeployabilityIndex, Snapshot
 
-CacheKey = t.Tuple[datetime, datetime, datetime, RuntimeStage]
-
 
 logger = logging.getLogger(__name__)
 
@@ -50,6 +47,7 @@ class BaseExpressionRenderer:
         schema: t.Optional[t.Dict[str, t.Any]] = None,
         default_catalog: t.Optional[str] = None,
         quote_identifiers: bool = True,
+        model_fqn: t.Optional[str] = None,
     ):
         self._expression = expression
         self._dialect = dialect
@@ -61,7 +59,8 @@ class BaseExpressionRenderer:
         self._default_catalog = default_catalog
         self._quote_identifiers = quote_identifiers
         self.update_schema({} if schema is None else schema)
-        self._cache: t.Dict[CacheKey, t.List[t.Optional[exp.Expression]]] = {}
+        self._cache: t.List[t.Optional[exp.Expression]] = []
+        self._model_fqn = model_fqn
 
     def update_schema(self, schema: t.Dict[str, t.Any]) -> None:
         self.schema = d.normalize_mapping_schema(schema, dialect=self._dialect)
@@ -93,119 +92,120 @@ class BaseExpressionRenderer:
             The rendered expressions.
         """
 
-        cache_key = self._cache_key(start, end, execution_time, runtime_stage)
+        should_cache = self._should_cache(
+            runtime_stage, start, end, execution_time, *kwargs.values()
+        )
+
+        if should_cache and self._cache:
+            return self._cache
 
-        if cache_key not in self._cache:
-            expressions = [self._expression]
+        if self._model_fqn and "this_model" not in kwargs:
+            if snapshots:
+                kwargs["this_model"] = self._to_table_mapping(
+                    [snapshots[self._model_fqn]], deployability_index
+                )
+            else:
+                kwargs["this_model"] = self._model_fqn
 
-            render_kwargs = {
-                **date_dict(
-                    cache_key[2],
-                    cache_key[0] if not self._only_execution_time else None,
-                    cache_key[1] if not self._only_execution_time else None,
-                ),
-                **kwargs,
-            }
+        expressions = [self._expression]
 
-            env = prepare_env(self._python_env)
-            jinja_env = self._jinja_macro_registry.build_environment(
-                **{**render_kwargs, **env},
-                snapshots=(snapshots or {}),
+        render_kwargs = {
+            **date_dict(
+                to_datetime(execution_time or c.EPOCH),
+                to_datetime(start or c.EPOCH) if not self._only_execution_time else None,
+                make_inclusive_end(end or c.EPOCH) if not self._only_execution_time else None,
+            ),
+            **kwargs,
+        }
+
+        jinja_env = self._jinja_macro_registry.build_environment(
+            **{**render_kwargs, **prepare_env(self._python_env)},
+            snapshots=(snapshots or {}),
+            table_mapping=table_mapping,
+            deployability_index=deployability_index,
+            default_catalog=self._default_catalog,
+        )
+
+        if isinstance(self._expression, d.Jinja):
+            try:
+                expressions = []
+                rendered_expression = jinja_env.from_string(self._expression.name).render()
+                if rendered_expression.strip():
+                    expressions = [e for e in parse(rendered_expression, read=self._dialect) if e]
+
+                    if not expressions:
+                        raise ConfigError(f"Failed to parse an expression:\n{self._expression}")
+            except ParsetimeAdapterCallError:
+                raise
+            except Exception as ex:
+                raise ConfigError(
+                    f"Could not render or parse jinja at '{self._path}'.\n{ex}"
+                ) from ex
+
+        macro_evaluator = MacroEvaluator(
+            self._dialect,
+            python_env=self._python_env,
+            jinja_env=jinja_env,
+            schema=self.schema,
+            runtime_stage=runtime_stage,
+            resolve_tables=lambda e: self._resolve_tables(
+                e,
+                snapshots=snapshots,
                 table_mapping=table_mapping,
                 deployability_index=deployability_index,
-                default_catalog=self._default_catalog,
-            )
-
-            if isinstance(self._expression, d.Jinja):
-                try:
-                    expressions = []
-                    rendered_expression = jinja_env.from_string(self._expression.name).render()
-                    if rendered_expression.strip():
-                        expressions = [
-                            e for e in parse(rendered_expression, read=self._dialect) if e
-                        ]
-
-                        if not expressions:
-                            raise ConfigError(f"Failed to parse an expression:\n{self._expression}")
-                except ParsetimeAdapterCallError:
-                    raise
-                except Exception as ex:
-                    raise ConfigError(
-                        f"Could not render or parse jinja at '{self._path}'.\n{ex}"
-                    ) from ex
-
-            macro_evaluator = MacroEvaluator(
-                self._dialect,
-                python_env=self._python_env,
-                jinja_env=jinja_env,
-                schema=self.schema,
+                start=start,
+                end=end,
+                execution_time=execution_time,
                 runtime_stage=runtime_stage,
-                resolve_tables=lambda e: self._resolve_tables(
-                    e,
-                    snapshots=snapshots,
-                    table_mapping=table_mapping,
-                    deployability_index=deployability_index,
-                    start=start,
-                    end=end,
-                    execution_time=execution_time,
-                    runtime_stage=runtime_stage,
-                ),
-                snapshots=snapshots,
-                default_catalog=self._default_catalog,
-            )
+            ),
+            snapshots=snapshots,
+            default_catalog=self._default_catalog,
+        )
 
-            for definition in self._macro_definitions:
-                try:
-                    macro_evaluator.evaluate(definition)
-                except MacroEvalError as ex:
-                    raise_config_error(f"Failed to evaluate macro '{definition}'. {ex}", self._path)
-
-            macro_evaluator.locals.update(render_kwargs)
-
-            resolved_expressions: t.List[t.Optional[exp.Expression]] = []
-            for expression in expressions:
-                try:
-                    expression = macro_evaluator.transform(expression)  # type: ignore
-                except MacroEvalError as ex:
-                    raise_config_error(f"Failed to resolve macro for expression. {ex}", self._path)
-
-                if expression:
-                    with self._normalize_and_quote(expression) as expression:
-                        if hasattr(expression, "selects"):
-                            for select in expression.selects:
-                                if not isinstance(select, exp.Alias) and select.output_name not in (
-                                    "*",
-                                    "",
-                                ):
-                                    alias = exp.alias_(
-                                        select, select.output_name, quoted=self._quote_identifiers
-                                    )
-                                    comments = alias.this.comments
-                                    if comments:
-                                        alias.add_comments(comments)
-                                        comments.clear()
-
-                                    select.replace(alias)
-                    resolved_expressions.append(expression)
-
-            # We dont cache here if columns_to_type was called in a macro.
-            # This allows the model's query to be re-rendered so that the
-            # MacroEvaluator can resolve columns_to_types calls and provide true schemas.
-            if not macro_evaluator.columns_to_types_called:
-                self._cache[cache_key] = resolved_expressions
-            return resolved_expressions
-
-        return self._cache[cache_key]
-
-    def update_cache(
-        self,
-        expression: t.Optional[exp.Expression],
-        start: t.Optional[TimeLike] = None,
-        end: t.Optional[TimeLike] = None,
-        execution_time: t.Optional[TimeLike] = None,
-        **kwargs: t.Any,
-    ) -> None:
-        self._cache[self._cache_key(start, end, execution_time)] = [expression]
+        for definition in self._macro_definitions:
+            try:
+                macro_evaluator.evaluate(definition)
+            except MacroEvalError as ex:
+                raise_config_error(f"Failed to evaluate macro '{definition}'. {ex}", self._path)
+
+        macro_evaluator.locals.update(render_kwargs)
+
+        resolved_expressions: t.List[t.Optional[exp.Expression]] = []
+
+        for expression in expressions:
+            try:
+                expression = macro_evaluator.transform(expression)  # type: ignore
+            except MacroEvalError as ex:
+                raise_config_error(f"Failed to resolve macro for expression. {ex}", self._path)
+
+            if expression:
+                with self._normalize_and_quote(expression) as expression:
+                    if hasattr(expression, "selects"):
+                        for select in expression.selects:
+                            if not isinstance(select, exp.Alias) and select.output_name not in (
+                                "*",
+                                "",
+                            ):
+                                alias = exp.alias_(
+                                    select, select.output_name, quoted=self._quote_identifiers
+                                )
+                                comments = alias.this.comments
+                                if comments:
+                                    alias.add_comments(comments)
+                                    comments.clear()
+
+                                select.replace(alias)
+                resolved_expressions.append(expression)
+
+        # We dont cache here if columns_to_type was called in a macro.
+        # This allows the model's query to be re-rendered so that the
+        # MacroEvaluator can resolve columns_to_types calls and provide true schemas.
+        if should_cache and (not self.schema.empty or not macro_evaluator.columns_to_types_called):
+            self._cache = resolved_expressions
+        return resolved_expressions
+
+    def update_cache(self, expression: t.Optional[exp.Expression]) -> None:
+        self._cache = [expression]
 
     def _resolve_tables(
         self,
@@ -223,13 +223,14 @@ class BaseExpressionRenderer:
         if not snapshots and not table_mapping and not expand:
             return expression
 
-        from sqlmesh.core.snapshot import to_table_mapping
-
         expression = expression.copy()
         with self._normalize_and_quote(expression) as expression:
             snapshots = snapshots or {}
             table_mapping = table_mapping or {}
-            mapping = {**to_table_mapping(snapshots.values(), deployability_index), **table_mapping}
+            mapping = {
+                **self._to_table_mapping(snapshots.values(), deployability_index),
+                **table_mapping,
+            }
             expand = set(expand) | {
                 name for name, snapshot in snapshots.items() if snapshot.is_embedded
             }
@@ -279,19 +280,6 @@ class BaseExpressionRenderer:
 
             return expression
 
-    def _cache_key(
-        self,
-        start: t.Optional[TimeLike] = None,
-        end: t.Optional[TimeLike] = None,
-        execution_time: t.Optional[TimeLike] = None,
-        runtime_stage: RuntimeStage = RuntimeStage.LOADING,
-    ) -> CacheKey:
-        return (
-            *make_inclusive(start or c.EPOCH, end or c.EPOCH),
-            to_datetime(execution_time or c.EPOCH),
-            runtime_stage,
-        )
-
     @contextmanager
     def _normalize_and_quote(self, query: E) -> t.Iterator[E]:
         with d.normalize_and_quote(
@@ -299,6 +287,16 @@ class BaseExpressionRenderer:
         ) as query:
             yield query
 
+    def _should_cache(self, runtime_stage: RuntimeStage, *args: t.Any) -> bool:
+        return runtime_stage == RuntimeStage.LOADING and not any(args)
+
+    def _to_table_mapping(
+        self, snapshots: t.Iterable[Snapshot], deployability_index: t.Optional[DeployabilityIndex]
+    ) -> t.Dict[str, str]:
+        from sqlmesh.core.snapshot import to_table_mapping
+
+        return to_table_mapping(snapshots, deployability_index)
+
 
 class ExpressionRenderer(BaseExpressionRenderer):
     def render(
@@ -342,40 +340,13 @@ class ExpressionRenderer(BaseExpressionRenderer):
 
 
 class QueryRenderer(BaseExpressionRenderer):
-    def __init__(
-        self,
-        query: exp.Expression,
-        dialect: str,
-        macro_definitions: t.List[d.MacroDef],
-        schema: t.Optional[t.Dict[str, t.Any]] = None,
-        model_fqn: t.Optional[str] = None,
-        path: Path = Path(),
-        jinja_macro_registry: t.Optional[JinjaMacroRegistry] = None,
-        python_env: t.Optional[t.Dict[str, Executable]] = None,
-        only_execution_time: bool = False,
-        default_catalog: t.Optional[str] = None,
-        quote_identifiers: bool = True,
-    ):
-        super().__init__(
-            expression=query,
-            dialect=dialect,
-            macro_definitions=macro_definitions,
-            path=path,
-            jinja_macro_registry=jinja_macro_registry,
-            python_env=python_env,
-            only_execution_time=only_execution_time,
-            schema=schema,
-            default_catalog=default_catalog,
-            quote_identifiers=quote_identifiers,
-        )
-
-        self._model_fqn = model_fqn
-
-        self._optimized_cache: t.Dict[CacheKey, exp.Query] = {}
+    def __init__(self, *args: t.Any, **kwargs: t.Any):
+        super().__init__(*args, **kwargs)
+        self._optimized_cache: t.Optional[exp.Query] = None
 
     def update_schema(self, schema: t.Dict[str, t.Any]) -> None:
         super().update_schema(schema)
-        self._optimized_cache = {}
+        self._optimized_cache = None
 
     def render(
         self,
@@ -410,9 +381,14 @@ class QueryRenderer(BaseExpressionRenderer):
         Returns:
             The rendered expression.
         """
-        cache_key = self._cache_key(start, end, execution_time, runtime_stage)
 
-        if not optimize or cache_key not in self._optimized_cache:
+        should_cache = self._should_cache(
+            runtime_stage, start, end, execution_time, *kwargs.values()
+        )
+
+        if should_cache and self._optimized_cache and optimize:
+            query = self._optimized_cache
+        else:
             try:
                 expressions = super()._render(
                     start=start,
@@ -433,7 +409,7 @@ class QueryRenderer(BaseExpressionRenderer):
             if len(expressions) > 1:
                 raise ConfigError(f"Too many statements in query:\n{self._expression}")
 
-            query = expressions[0]
+            query = expressions[0]  # type: ignore
 
             if not query:
                 return None
@@ -441,49 +417,39 @@ class QueryRenderer(BaseExpressionRenderer):
                 raise_config_error(f"Query needs to be a SELECT or a UNION {query}.", self._path)
                 raise
 
-            # we find deps here so that it is cached in the model cache
-            deps = d.find_tables(
-                query, default_catalog=self._default_catalog, dialect=self._dialect
-            )
-
             if optimize:
+                deps = d.find_tables(
+                    query, default_catalog=self._default_catalog, dialect=self._dialect
+                )
+
                 query = self._optimize_query(query, deps)
-                self._optimized_cache[cache_key] = query
-        else:
-            query = self._optimized_cache[cache_key]
 
-        # Table resolution MUST happen after optimization, otherwise the schema won't match the table names.
-        query = self._resolve_tables(
-            query,
-            snapshots=snapshots,
-            table_mapping=table_mapping,
-            expand=expand,
-            deployability_index=deployability_index,
-            start=start,
-            end=end,
-            execution_time=execution_time,
-            runtime_stage=runtime_stage,
-            **kwargs,
-        )
+                if should_cache:
+                    self._optimized_cache = query
+
+        if optimize:
+            query = self._resolve_tables(
+                query,
+                snapshots=snapshots,
+                table_mapping=table_mapping,
+                expand=expand,
+                deployability_index=deployability_index,
+                start=start,
+                end=end,
+                execution_time=execution_time,
+                runtime_stage=runtime_stage,
+                **kwargs,
+            )
+
         return query
 
-    def update_cache(
-        self,
-        expression: t.Optional[exp.Expression],
-        start: t.Optional[TimeLike] = None,
-        end: t.Optional[TimeLike] = None,
-        execution_time: t.Optional[TimeLike] = None,
-        optimized: bool = False,
-        **kwargs: t.Any,
-    ) -> None:
+    def update_cache(self, expression: t.Optional[exp.Expression], optimized: bool = False) -> None:
         if optimized:
             if not isinstance(expression, exp.Query):
                 raise SQLMeshError(f"Expected a Query but got: {expression}")
-            self._optimized_cache[self._cache_key(start, end, execution_time)] = expression
+            self._optimized_cache = expression
         else:
-            super().update_cache(
-                expression, start=start, end=end, execution_time=execution_time, **kwargs
-            )
+            super().update_cache(expression)
 
     def _optimize_query(self, query: exp.Query, all_deps: t.Set[str]) -> exp.Query:
         # We don't want to normalize names in the schema because that's handled by the optimizer
diff --git a/sqlmesh/magics.py b/sqlmesh/magics.py
index 04c4950d..e974e8b8 100644
--- a/sqlmesh/magics.py
+++ b/sqlmesh/magics.py
@@ -159,12 +159,14 @@ class SQLMeshMagics(Magics):
     def model(self, context: Context, line: str, sql: t.Optional[str] = None) -> None:
         """Renders the model and automatically fills in an editable cell with the model definition."""
         args = parse_argstring(self.model, line)
+
         model = context.get_model(args.model, raise_if_missing=True)
+        config = context.config_for_node(model)
 
         if sql:
-            config = context.config_for_node(model)
+            expressions = parse(sql, default_dialect=config.dialect)
             loaded = load_sql_based_model(
-                parse(sql, default_dialect=config.dialect),
+                expressions,
                 macros=context._macros,
                 jinja_macros=context._jinja_macros,
                 path=model._path,
@@ -176,11 +178,13 @@ class SQLMeshMagics(Magics):
 
             if loaded.name == args.model:
                 model = loaded
+        else:
+            with open(model._path, "r", encoding="utf-8") as file:
+                expressions = parse(file.read(), default_dialect=config.dialect)
 
-        context.upsert_model(model)
-        expressions = model.render_definition(include_python=False)
-
-        formatted = format_model_expressions(expressions, model.dialect)
+        formatted = format_model_expressions(
+            expressions, model.dialect, **config.format.generator_options
+        )
 
         self._shell.set_next_input(
             "\n".join(
diff --git a/sqlmesh/schedulers/airflow/client.py b/sqlmesh/schedulers/airflow/client.py
index 0060a1d0..8234de17 100644
--- a/sqlmesh/schedulers/airflow/client.py
+++ b/sqlmesh/schedulers/airflow/client.py
@@ -208,7 +208,7 @@ class AirflowClient(BaseAirflowClient):
             no_gaps=no_gaps,
             skip_backfill=skip_backfill,
             request_id=request_id,
-            restatements=restatements or {},
+            restatements={s.name: i for s, i in (restatements or {}).items()},
             notification_targets=notification_targets or [],
             backfill_concurrent_tasks=backfill_concurrent_tasks,
             ddl_concurrent_tasks=ddl_concurrent_tasks,
diff --git a/sqlmesh/schedulers/airflow/common.py b/sqlmesh/schedulers/airflow/common.py
index d789935c..d4685b31 100644
--- a/sqlmesh/schedulers/airflow/common.py
+++ b/sqlmesh/schedulers/airflow/common.py
@@ -42,7 +42,7 @@ class PlanApplicationRequest(PydanticModel):
     environment: Environment
     no_gaps: bool
     skip_backfill: bool
-    restatements: t.Dict[SnapshotId, SnapshotInterval]
+    restatements: t.Dict[str, SnapshotInterval]
     notification_targets: t.List[NotificationTarget]
     backfill_concurrent_tasks: int
     ddl_concurrent_tasks: int
diff --git a/sqlmesh/schedulers/airflow/plan.py b/sqlmesh/schedulers/airflow/plan.py
index 6db63469..1fbf03f7 100644
--- a/sqlmesh/schedulers/airflow/plan.py
+++ b/sqlmesh/schedulers/airflow/plan.py
@@ -87,6 +87,13 @@ def create_plan_dag_spec(
     stored_snapshots = state_sync.get_snapshots([*new_snapshots, *request.environment.snapshots])
     all_snapshots = {**new_snapshots, **stored_snapshots}
 
+    all_snaphots_by_name = {s.name: s for s in all_snapshots.values()}
+    restatements = {
+        all_snaphots_by_name[n].snapshot_id: i
+        for n, i in request.restatements.items()
+        if n in all_snaphots_by_name
+    }
+
     duplicated_snapshots = set(stored_snapshots).intersection(new_snapshots)
     if duplicated_snapshots:
         raise SQLMeshError(
@@ -102,9 +109,9 @@ def create_plan_dag_spec(
 
     if request.restatements:
         intervals_to_remove = [
-            (s, request.restatements[s.snapshot_id])
+            (s, restatements[s.snapshot_id])
             for s in all_snapshots.values()
-            if s.snapshot_id in request.restatements and s.snapshot_id not in new_snapshots
+            if s.snapshot_id in restatements and s.snapshot_id not in new_snapshots
         ]
         state_sync.remove_interval(
             intervals_to_remove,
@@ -125,7 +132,7 @@ def create_plan_dag_spec(
             end=end,
             execution_time=now(),
             deployability_index=deployability_index_for_evaluation,
-            restatements=request.restatements,
+            restatements=restatements,
             end_bounded=request.end_bounded,
         )
     else:


Test patch:
diff --git a/tests/core/test_model.py b/tests/core/test_model.py
index 0529d181..88e524ae 100644
--- a/tests/core/test_model.py
+++ b/tests/core/test_model.py
@@ -3731,3 +3731,31 @@ def test_render_quote_identifiers(dialect: str, kind: str) -> None:
         model.render_query_or_raise().sql(dialect="duckdb")
         == 'SELECT "a" AS "a", "b" AS "b", "c" AS "c" FROM "source_table" AS "source_table"'
     )
+
+
+def test_this_model() -> None:
+    expressions = d.parse(
+        """
+        MODEL (
+            name db.table,
+        );
+
+        JINJA_STATEMENT_BEGIN;
+        COPY {{ this_model }} TO 'a';
+        JINJA_END;
+
+        JINJA_QUERY_BEGIN;
+        SELECT '{{ this_model }}' as x
+        JINJA_END;
+
+        JINJA_STATEMENT_BEGIN;
+        COPY {{ this_model }} TO 'b';
+        JINJA_END;
+        """
+    )
+    model = load_sql_based_model(expressions)
+
+    assert model.render_query_or_raise().sql() == '''SELECT '"db"."table"' AS "x"'''
+
+    assert model.render_pre_statements()[0].sql() == """COPY "db"."table" TO 'a'"""
+    assert model.render_post_statements()[0].sql() == """COPY "db"."table" TO 'b'"""
diff --git a/tests/schedulers/airflow/test_client.py b/tests/schedulers/airflow/test_client.py
index 15db7fe3..fed35c9a 100644
--- a/tests/schedulers/airflow/test_client.py
+++ b/tests/schedulers/airflow/test_client.py
@@ -68,6 +68,9 @@ def test_apply_plan(mocker: MockerFixture, snapshot: Snapshot):
         request_id,
         models_to_backfill={'"test_model"'},
         directly_modified_snapshots=[snapshot.snapshot_id],
+        restatements={
+            snapshot.snapshot_id: (to_timestamp("2024-01-01"), to_timestamp("2024-01-02"))
+        },
     )
 
     apply_plan_mock.assert_called_once()
@@ -169,8 +172,11 @@ def test_apply_plan(mocker: MockerFixture, snapshot: Snapshot):
         "directly_modified_snapshots": [{"identifier": "2298286125", "name": '"test_model"'}],
         "indirectly_modified_snapshots": {},
         "removed_snapshots": [],
+        "restatements": {'"test_model"': [to_timestamp("2024-01-01"), to_timestamp("2024-01-02")]},
     }
 
+    common.PlanApplicationRequest.parse_raw(data["data"])
+
 
 def snapshot_url(snapshot_ids, key="ids") -> str:
     return urlencode({key: _list_to_json(snapshot_ids)[0]})
diff --git a/tests/schedulers/airflow/test_plan.py b/tests/schedulers/airflow/test_plan.py
index 76c58b76..01140dd1 100644
--- a/tests/schedulers/airflow/test_plan.py
+++ b/tests/schedulers/airflow/test_plan.py
@@ -239,7 +239,7 @@ def test_restatement(
         no_gaps=True,
         skip_backfill=False,
         restatements={
-            the_snapshot.snapshot_id: (
+            the_snapshot.name: (
                 to_timestamp("2022-01-02"),
                 to_timestamp("2022-01-04"),
             )