Feat: expose a macro var to check if we're evaluating snapshots (#1562)

* Feat: expose a macro var to check if we're evaluating snapshots

* PR feedback

* Fixup

* PR comments

* Make evaluating an optional in _cache_key again

* Refactor to use enum, address PR feedback, add unit test

* Rename unit test

Author: georgesittas

sqlmesh/core/macros.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import typing as t
+from enum import Enum
 from functools import reduce
 from string import Template
 
@@ -29,6 +30,12 @@
 SQLMESH_MOCKED_STAR = "__SQLMESH_MOCKED_STAR__"
 
 
+class RuntimeStage(Enum):
+    LOADING = "loading"
+    CREATING = "creating"
+    EVALUATING = "evaluating"
+
+
 class MacroStrTemplate(Template):
     delimiter = SQLMESH_MACRO_PREFIX
 
@@ -103,10 +110,11 @@ def __init__(
         python_env: t.Optional[t.Dict[str, Executable]] = None,
         jinja_env: t.Optional[Environment] = None,
         schema: t.Optional[t.Dict[str, t.Any]] = None,
+        runtime_stage: RuntimeStage = RuntimeStage.LOADING,
     ):
         self.dialect = dialect
         self.generator = MacroDialect().generator()
-        self.locals: t.Dict[str, t.Any] = {}
+        self.locals: t.Dict[str, t.Any] = {"runtime_stage": runtime_stage}
         self.env = {**ENV, "self": self}
         self.python_env = python_env or {}
         self._jinja_env: t.Optional[Environment] = jinja_env

sqlmesh/core/renderer.py

@@ -18,7 +18,7 @@
 
 from sqlmesh.core import constants as c
 from sqlmesh.core import dialect as d
-from sqlmesh.core.macros import MacroEvaluator
+from sqlmesh.core.macros import MacroEvaluator, RuntimeStage
 from sqlmesh.utils.date import TimeLike, date_dict, make_inclusive, to_datetime
 from sqlmesh.utils.errors import (
     ConfigError,
@@ -34,6 +34,8 @@
 
     from sqlmesh.core.snapshot import Snapshot
 
+CacheKey = t.Tuple[datetime, datetime, datetime, RuntimeStage]
+
 
 logger = logging.getLogger(__name__)
 
@@ -59,7 +61,7 @@ def __init__(
         self._only_execution_time = only_execution_time
         self.schema = {} if schema is None else schema
 
-        self._cache: t.Dict[t.Tuple[datetime, datetime, datetime], t.List[exp.Expression]] = {}
+        self._cache: t.Dict[CacheKey, t.List[exp.Expression]] = {}
 
     def _render(
         self,
@@ -69,6 +71,7 @@ def _render(
         snapshots: t.Optional[t.Dict[str, Snapshot]] = None,
         table_mapping: t.Optional[t.Dict[str, str]] = None,
         is_dev: bool = False,
+        runtime_stage: t.Optional[RuntimeStage] = None,
         **kwargs: t.Any,
     ) -> t.List[exp.Expression]:
         """Renders a expression, expanding macros with provided kwargs
@@ -77,18 +80,19 @@ def _render(
             start: The start datetime to render. Defaults to epoch start.
             end: The end datetime to render. Defaults to epoch start.
             execution_time: The date/time time reference to use for execution time.
-            kwargs: Additional kwargs to pass to the renderer.
             snapshots: All upstream snapshots (by model name) to use for expansion and mapping of physical locations.
             table_mapping: Table mapping of physical locations. Takes precedence over snapshot mappings.
             is_dev: Indicates whether the rendering happens in the development mode and temporary
                 tables / table clones should be used where applicable.
+            runtime_stage: Indicates the current runtime stage, for example if we're still loading the project, etc.
+            kwargs: Additional kwargs to pass to the renderer.
 
         Returns:
             The rendered expressions.
         """
 
-        cache_key = self._cache_key(start, end, execution_time)
-        start_dt, end_dt, execution_dt = cache_key
+        cache_key = self._cache_key(start, end, execution_time, runtime_stage)
+        start_dt, end_dt, execution_dt, runtime_stage = cache_key
         if cache_key not in self._cache:
             expressions = [self._expression]
 
@@ -132,6 +136,7 @@ def _render(
                 python_env=self._python_env,
                 jinja_env=jinja_env,
                 schema=self.schema,
+                runtime_stage=runtime_stage,
             )
 
             for definition in self._macro_definitions:
@@ -203,10 +208,12 @@ def _cache_key(
         start: t.Optional[TimeLike] = None,
         end: t.Optional[TimeLike] = None,
         execution_time: t.Optional[TimeLike] = None,
-    ) -> t.Tuple[datetime, datetime, datetime]:
+        runtime_stage: t.Optional[RuntimeStage] = None,
+    ) -> CacheKey:
         return (
             *make_inclusive(start or c.EPOCH, end or c.EPOCH),
             to_datetime(execution_time or c.EPOCH),
+            runtime_stage or RuntimeStage.LOADING,
         )
 
 
@@ -273,7 +280,7 @@ def __init__(
 
         self._model_name = model_name
 
-        self._optimized_cache: t.Dict[t.Tuple[datetime, datetime, datetime], exp.Expression] = {}
+        self._optimized_cache: t.Dict[CacheKey, exp.Expression] = {}
 
     def render(
         self,
@@ -285,6 +292,7 @@ def render(
         is_dev: bool = False,
         expand: t.Iterable[str] = tuple(),
         optimize: bool = True,
+        runtime_stage: t.Optional[RuntimeStage] = None,
         **kwargs: t.Any,
     ) -> t.Optional[exp.Subqueryable]:
         """Renders a query, expanding macros with provided kwargs, and optionally expanding referenced models.
@@ -302,12 +310,13 @@ def render(
                 that depend on materialized tables.  Model definitions are inlined and can thus be run end to
                 end on the fly.
             optimize: Whether to optimize the query.
+            runtime_stage: Indicates the current runtime stage, for example if we're still loading the project, etc.
             kwargs: Additional kwargs to pass to the renderer.
 
         Returns:
             The rendered expression.
         """
-        cache_key = self._cache_key(start, end, execution_time)
+        cache_key = self._cache_key(start, end, execution_time, runtime_stage)
 
         if not optimize or cache_key not in self._optimized_cache:
             try:

sqlmesh/core/snapshot/evaluator.py

@@ -35,6 +35,7 @@
 from sqlmesh.core.dialect import schema_
 from sqlmesh.core.engine_adapter import EngineAdapter
 from sqlmesh.core.engine_adapter.base import InsertOverwriteStrategy
+from sqlmesh.core.macros import RuntimeStage
 from sqlmesh.core.model import IncrementalUnmanagedKind, Model, SCDType2Kind, ViewKind
 from sqlmesh.core.snapshot import (
     QualifiedViewName,
@@ -158,6 +159,7 @@ def apply(query_or_df: QueryOrDF, index: int = 0) -> None:
             engine_adapter=self.adapter,
             snapshots=snapshots,
             is_dev=is_dev,
+            runtime_stage=RuntimeStage.EVALUATING,
             **common_render_kwargs,
         )
 
@@ -411,6 +413,7 @@ def _create_snapshot(
             engine_adapter=self.adapter,
             snapshots=parent_snapshots_by_name,
             is_dev=is_dev,
+            runtime_stage=RuntimeStage.CREATING,
         )
 
         evaluation_strategy = _evaluation_strategy(snapshot, self.adapter)

tests/core/test_snapshot_evaluator.py

@@ -11,7 +11,7 @@
 from sqlmesh.core.engine_adapter import EngineAdapter, create_engine_adapter
 from sqlmesh.core.engine_adapter.base import InsertOverwriteStrategy
 from sqlmesh.core.environment import EnvironmentNamingInfo
-from sqlmesh.core.macros import macro
+from sqlmesh.core.macros import RuntimeStage, macro
 from sqlmesh.core.model import (
     FullKind,
     IncrementalByTimeRangeKind,
@@ -148,6 +148,43 @@ def x(evaluator, y=None) -> None:
     )
 
 
+def test_runtime_stages(capsys, mocker, adapter_mock, make_snapshot):
+    evaluator = SnapshotEvaluator(adapter_mock)
+
+    @macro()
+    def increment_stage_counter(evaluator) -> None:
+        # Hack which allows us to intercept the different runtime stage values
+        print(f"RuntimeStage value: {evaluator.locals['runtime_stage'].value}")
+
+    model = load_sql_based_model(
+        parse(  # type: ignore
+            """
+            MODEL (
+                name test_schema.test_model,
+                kind FULL,
+            );
+
+            @increment_stage_counter();
+
+            SELECT 1 AS a;
+            """
+        ),
+        macros=macro.get_registry(),
+    )
+
+    capsys.readouterr()
+
+    snapshot = make_snapshot(model)
+    snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
+    assert f"RuntimeStage value: {RuntimeStage.LOADING.value}" in capsys.readouterr().out
+
+    evaluator.create([snapshot], {})
+    assert f"RuntimeStage value: {RuntimeStage.CREATING.value}" in capsys.readouterr().out
+
+    evaluator.evaluate(snapshot, "2020-01-01", "2020-01-02", "2020-01-02", snapshots={})
+    assert f"RuntimeStage value: {RuntimeStage.EVALUATING.value}" in capsys.readouterr().out
+
+
 def test_evaluate_paused_forward_only_upstream(mocker: MockerFixture, make_snapshot):
     model = SqlModel(name="test_schema.test_model", query=parse_one("SELECT a, ds"))
     snapshot = make_snapshot(model)