Feat: expose columns to types of parent models in MacroEvaluator (#1488)

* Feat: expose explicit model dependencies in the macro context

* Temporarily add debug statement

* Add another temporary debug statement

* Add another temporary debug statement

* Set 'columns' instead of 'columns_to_types_'

* Formatting

* mocked star

* formatting

* Start refactoring - backtrack to previous codebase state

* Refactor

* Fixups

* Simplify

* Fix lineage test

* Ensure sqlglot's diff doesn't fail when when prev/new queries are identical

* Don't actually need to call diff if queries are identical

* Rename to excluded

* Yet another refactor

* Add render test for sushi.waiters

* Remove unnecessary depends_on field

Author: georgesittas

examples/sushi/models/waiters.py

@@ -30,8 +30,14 @@ def entrypoint(evaluator: MacroEvaluator) -> exp.Select:
     FROM sushi.orders AS o
     WHERE @incremental_by_ds(ds)
     """
+    excluded = {"id", "customer_id", "start_ts", "end_ts"}
+    projections = []
+    for column, dtype in evaluator.columns_to_types("sushi.orders").items():
+        if column not in excluded:
+            projections.append(f"{column}::{dtype}")
+
     return (
-        exp.select("waiter_id::int as waiter_id", "ds::text as ds")
+        exp.select(*projections)
         .from_("sushi.orders AS o")
         .where(incremental_by_ds(evaluator, exp.to_column("ds")))
         .distinct()

sqlmesh/core/loader.py

@@ -43,6 +43,7 @@
 
 logger = logging.getLogger(__name__)
 
+
 # TODO: consider moving this to context
 def update_model_schemas(
     dag: DAG[str],
@@ -300,6 +301,7 @@ def _load() -> Model:
                             raise ConfigError(
                                 f"Failed to parse a model definition at '{path}': {ex}."
                             )
+
                     return load_sql_based_model(
                         expressions,
                         defaults=config.model_defaults.dict(),

sqlmesh/core/macros.py

@@ -10,6 +10,7 @@
 from sqlglot.executor.env import ENV
 from sqlglot.executor.python import Python
 from sqlglot.helper import csv, ensure_collection
+from sqlglot.schema import MappingSchema
 
 from sqlmesh.core.dialect import (
     SQLMESH_MACRO_PREFIX,
@@ -25,6 +26,8 @@
 from sqlmesh.utils.jinja import JinjaMacroRegistry, has_jinja
 from sqlmesh.utils.metaprogramming import Executable, prepare_env, print_exception
 
+SQLMESH_MOCKED_STAR = "__SQLMESH_MOCKED_STAR__"
+
 
 class MacroStrTemplate(Template):
     delimiter = SQLMESH_MACRO_PREFIX
@@ -99,6 +102,7 @@ def __init__(
         dialect: str = "",
         python_env: t.Optional[t.Dict[str, Executable]] = None,
         jinja_env: t.Optional[Environment] = None,
+        schema: t.Optional[t.Dict[str, t.Any]] = None,
     ):
         self.dialect = dialect
         self.generator = MacroDialect().generator()
@@ -107,6 +111,7 @@ def __init__(
         self.python_env = python_env or {}
         self._jinja_env: t.Optional[Environment] = jinja_env
         self.macros = {normalize_macro_name(k): v.func for k, v in macro.get_registry().items()}
+        self._schema = MappingSchema(schema, dialect=dialect, normalize=False) if schema else {}
 
         prepare_env(self.python_env, self.env)
         for k, v in self.python_env.items():
@@ -261,6 +266,17 @@ def jinja_env(self) -> Environment:
             self._jinja_env = JinjaMacroRegistry().build_environment(**jinja_env_methods)
         return self._jinja_env
 
+    def columns_to_types(self, model_name: str) -> t.Dict[str, exp.DataType]:
+        """Returns the columns-to-types mapping corresponding to the specified model."""
+        if not isinstance(self._schema, MappingSchema):
+            return {SQLMESH_MOCKED_STAR: exp.DataType.build("unknown")}
+
+        columns_to_types = self._schema.find(exp.to_table(model_name))
+        if columns_to_types is None:
+            raise SQLMeshError(f"Model '{model_name}' not found in the macro evaluator's context.")
+
+        return columns_to_types  # type: ignore
+
 
 class macro(registry_decorator):
     """Specifies a function is a macro and registers it the global MACROS registry.

sqlmesh/core/model/definition.py

@@ -17,12 +17,13 @@
 from sqlglot import diff, exp
 from sqlglot.diff import Insert, Keep
 from sqlglot.helper import ensure_list
+from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
 from sqlglot.schema import MappingSchema, nested_set
 from sqlglot.time import format_time
 
 from sqlmesh.core import constants as c
 from sqlmesh.core import dialect as d
-from sqlmesh.core.macros import MacroRegistry, macro
+from sqlmesh.core.macros import SQLMESH_MOCKED_STAR, MacroRegistry, macro
 from sqlmesh.core.model.common import expression_validator
 from sqlmesh.core.model.kind import (
     IncrementalByTimeRangeKind,
@@ -1026,6 +1027,14 @@ def update_schema(
         super().update_schema(
             schema, default_schema=default_schema, default_catalog=default_catalog
         )
+
+        mocked_star = normalize_identifiers(SQLMESH_MOCKED_STAR, dialect=self.dialect)
+        if mocked_star.name in (self.columns_to_types or {}):
+            # We reset the unoptimized query cache here as well to allow the model's query
+            # to be re-rendered so that the MacroEvaluator can resolve columns_to_types calls
+            # and get rid of the mocked star column
+            self._query_renderer._cache = {}
+
         self._columns_to_types = None
         self._query_renderer._optimized_cache = {}
 

sqlmesh/core/renderer.py

@@ -38,17 +38,6 @@
 logger = logging.getLogger(__name__)
 
 
-def _dates(
-    start: t.Optional[TimeLike] = None,
-    end: t.Optional[TimeLike] = None,
-    execution_time: t.Optional[TimeLike] = None,
-) -> t.Tuple[datetime, datetime, datetime]:
-    return (
-        *make_inclusive(start or c.EPOCH, end or c.EPOCH),
-        to_datetime(execution_time or c.EPOCH),
-    )
-
-
 class BaseExpressionRenderer:
     def __init__(
         self,
@@ -59,6 +48,7 @@ def __init__(
         jinja_macro_registry: t.Optional[JinjaMacroRegistry] = None,
         python_env: t.Optional[t.Dict[str, Executable]] = None,
         only_execution_time: bool = False,
+        schema: t.Optional[t.Dict[str, t.Any]] = None,
     ):
         self._expression = expression
         self._dialect = dialect
@@ -67,6 +57,7 @@ def __init__(
         self._jinja_macro_registry = jinja_macro_registry or JinjaMacroRegistry()
         self._python_env = python_env or {}
         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]] = {}
 
@@ -96,7 +87,7 @@ def _render(
             The rendered expressions.
         """
 
-        cache_key = _dates(start, end, execution_time)
+        cache_key = self._cache_key(start, end, execution_time)
         start_dt, end_dt, execution_dt = cache_key
         if cache_key not in self._cache:
             expressions = [self._expression]
@@ -140,6 +131,7 @@ def _render(
                 self._dialect,
                 python_env=self._python_env,
                 jinja_env=jinja_env,
+                schema=self.schema,
             )
 
             for definition in self._macro_definitions:
@@ -174,7 +166,7 @@ def update_cache(
         execution_time: t.Optional[TimeLike] = None,
         **kwargs: t.Any,
     ) -> None:
-        self._cache[_dates(start, end, execution_time)] = [expression]
+        self._cache[self._cache_key(start, end, execution_time)] = [expression]
 
     def _resolve_tables(
         self,
@@ -206,6 +198,17 @@ def _resolve_tables(
                 **kwargs,
             )
 
+    def _cache_key(
+        self,
+        start: t.Optional[TimeLike] = None,
+        end: t.Optional[TimeLike] = None,
+        execution_time: t.Optional[TimeLike] = None,
+    ) -> t.Tuple[datetime, datetime, datetime]:
+        return (
+            *make_inclusive(start or c.EPOCH, end or c.EPOCH),
+            to_datetime(execution_time or c.EPOCH),
+        )
+
 
 class ExpressionRenderer(BaseExpressionRenderer):
     def render(
@@ -265,14 +268,13 @@ def __init__(
             jinja_macro_registry=jinja_macro_registry,
             python_env=python_env,
             only_execution_time=only_execution_time,
+            schema=schema,
         )
 
         self._model_name = model_name
 
         self._optimized_cache: t.Dict[t.Tuple[datetime, datetime, datetime], exp.Expression] = {}
 
-        self.schema = {} if schema is None else schema
-
     def render(
         self,
         start: t.Optional[TimeLike] = None,
@@ -305,7 +307,7 @@ def render(
         Returns:
             The rendered expression.
         """
-        cache_key = _dates(start, end, execution_time)
+        cache_key = self._cache_key(start, end, execution_time)
 
         if not optimize or cache_key not in self._optimized_cache:
             try:
@@ -367,7 +369,7 @@ def update_cache(
                 expression, start=start, end=end, execution_time=execution_time, **kwargs
             )
         else:
-            self._optimized_cache[_dates(start, end, execution_time)] = expression
+            self._optimized_cache[self._cache_key(start, end, execution_time)] = expression
 
     def _optimize_query(self, query: exp.Subqueryable) -> exp.Subqueryable:
         # We don't want to normalize names in the schema because that's handled by the optimizer

tests/core/test_model.py

@@ -985,7 +985,7 @@ def test_lookback():
     assert to_timestamp(model.lookback_start("Jan 1 2020")) == to_timestamp("Jan 1 2018")
 
 
-def test_render_query(assert_exp_eq):
+def test_render_query(assert_exp_eq, sushi_context):
     model = SqlModel(
         name="test",
         cron="1 0 * * *",
@@ -1040,6 +1040,18 @@ def test_render_query(assert_exp_eq):
         'SELECT COUNT(DISTINCT "a") FILTER (WHERE "b" > 0) AS "c" FROM "x" AS "x"',
     )
 
+    assert_exp_eq(
+        sushi_context.models["sushi.waiters"].render_query().sql(),
+        """
+        SELECT DISTINCT
+          CAST("o"."waiter_id" AS INT) AS "waiter_id",
+          CAST("o"."ds" AS TEXT) AS "ds"
+        FROM "sushi"."orders" AS "o"
+        WHERE
+          "o"."ds" <= '1970-01-01' AND "o"."ds" >= '1970-01-01'
+        """,
+    )
+
 
 def test_time_column():
     expressions = d.parse(