Fix: Improve tracking of var dependencies in dbt models

Author: izeigerman

sqlmesh/dbt/adapter.py

@@ -245,6 +245,56 @@ def _raise_parsetime_adapter_call_error(action: str) -> None:
         raise ParsetimeAdapterCallError(f"Can't {action} at parse time.")
 
 
+class StubParsetimeAdapter(BaseAdapter):
+    """Same as ParsetimeAdapter, but returns stub / empty values instead of raising an error."""
+
+    def get_relation(self, database: str, schema: str, identifier: str) -> t.Optional[BaseRelation]:
+        return None
+
+    def load_relation(self, relation: BaseRelation) -> t.Optional[BaseRelation]:
+        return None
+
+    def list_relations(self, database: t.Optional[str], schema: str) -> t.List[BaseRelation]:
+        return []
+
+    def list_relations_without_caching(self, schema_relation: BaseRelation) -> t.List[BaseRelation]:
+        return []
+
+    def get_columns_in_relation(self, relation: BaseRelation) -> t.List[Column]:
+        return []
+
+    def get_missing_columns(
+        self, from_relation: BaseRelation, to_relation: BaseRelation
+    ) -> t.List[Column]:
+        return []
+
+    def create_schema(self, relation: BaseRelation) -> None:
+        pass
+
+    def drop_schema(self, relation: BaseRelation) -> None:
+        pass
+
+    def drop_relation(self, relation: BaseRelation) -> None:
+        pass
+
+    def rename_relation(self, from_relation: BaseRelation, to_relation: BaseRelation) -> None:
+        pass
+
+    def execute(
+        self, sql: str, auto_begin: bool = False, fetch: bool = False
+    ) -> t.Tuple[AdapterResponse, agate.Table]:
+        from dbt.adapters.base.impl import AdapterResponse
+        from sqlmesh.dbt.util import empty_table
+
+        return AdapterResponse(""), empty_table()
+
+    def resolve_schema(self, relation: BaseRelation) -> t.Optional[str]:
+        return relation.schema
+
+    def resolve_identifier(self, relation: BaseRelation) -> t.Optional[str]:
+        return relation.identifier
+
+
 class RuntimeAdapter(BaseAdapter):
     def __init__(
         self,

sqlmesh/dbt/basemodel.py

@@ -309,37 +309,24 @@ def sqlmesh_model_kwargs(
         self,
         context: DbtContext,
         column_types_override: t.Optional[t.Dict[str, ColumnConfig]] = None,
+        extra_dependencies: t.Optional[Dependencies] = None,
     ) -> t.Dict[str, t.Any]:
         """Get common sqlmesh model parameters"""
         self.remove_tests_with_invalid_refs(context)
         self.check_for_circular_test_refs(context)
+
+        dependencies = self.dependencies
+        if extra_dependencies:
+            dependencies = dependencies.union(extra_dependencies)
+
         model_dialect = self.dialect(context)
         model_context = context.context_for_dependencies(
-            self.dependencies.union(self.tests_ref_source_dependencies)
+            dependencies.union(self.tests_ref_source_dependencies)
         )
         jinja_macros = model_context.jinja_macros.trim(
-            self.dependencies.macros, package=self.package_name
-        )
-
-        model_node: AttributeDict[str, t.Any] = AttributeDict(
-            {
-                k: v
-                for k, v in context._manifest._manifest.nodes[self.node_name].to_dict().items()
-                if k in self.dependencies.model_attrs
-            }
-            if context._manifest and self.node_name in context._manifest._manifest.nodes
-            else {}
-        )
-
-        jinja_macros.add_globals(
-            {
-                "this": self.relation_info,
-                "model": model_node,
-                "schema": self.table_schema,
-                "config": self.config_attribute_dict,
-                **model_context.jinja_globals,  # type: ignore
-            }
+            dependencies.macros, package=self.package_name
         )
+        jinja_macros.add_globals(self._model_jinja_context(model_context, dependencies))
         return {
             "audits": [(test.name, {}) for test in self.tests],
             "columns": column_types_to_sqlmesh(
@@ -369,3 +356,28 @@ def to_sqlmesh(
         virtual_environment_mode: VirtualEnvironmentMode = VirtualEnvironmentMode.default,
     ) -> Model:
         """Convert DBT model into sqlmesh Model"""
+
+    def _model_jinja_context(
+        self, context: DbtContext, dependencies: Dependencies
+    ) -> t.Dict[str, t.Any]:
+        model_node: AttributeDict[str, t.Any] = AttributeDict(
+            {
+                k: v
+                for k, v in context._manifest._manifest.nodes[self.node_name].to_dict().items()
+                if k in dependencies.model_attrs
+            }
+            if context._manifest and self.node_name in context._manifest._manifest.nodes
+            else {}
+        )
+        return {
+            "this": self.relation_info,
+            "model": model_node,
+            "schema": self.table_schema,
+            "config": self.config_attribute_dict,
+            **context.jinja_globals,
+        }
+
+    def _track_dependencies_on_render(self, input: str, context: DbtContext) -> Dependencies:
+        return context.track_dependencies_on_render(
+            input, self._model_jinja_context(context, self.dependencies), self.package_name
+        )

sqlmesh/dbt/builtin.py

@@ -17,7 +17,7 @@
 from sqlmesh.core.console import get_console
 from sqlmesh.core.engine_adapter import EngineAdapter
 from sqlmesh.core.snapshot.definition import DeployabilityIndex
-from sqlmesh.dbt.adapter import BaseAdapter, ParsetimeAdapter, RuntimeAdapter
+from sqlmesh.dbt.adapter import BaseAdapter, ParsetimeAdapter, RuntimeAdapter, StubParsetimeAdapter
 from sqlmesh.dbt.relation import Policy
 from sqlmesh.dbt.target import TARGET_TYPE_TO_CONFIG_CLASS
 from sqlmesh.dbt.util import DBT_VERSION
@@ -384,15 +384,15 @@ def create_builtin_globals(
             builtin_globals["this"] = this
 
     sources = jinja_globals.pop("sources", None)
-    if sources is not None:
+    if sources is not None and "source" not in jinja_globals:
         builtin_globals["source"] = generate_source(sources, api)
 
     refs = jinja_globals.pop("refs", None)
-    if refs is not None:
+    if refs is not None and "ref" not in jinja_globals:
         builtin_globals["ref"] = generate_ref(refs, api)
 
     variables = jinja_globals.pop("vars", None)
-    if variables is not None:
+    if variables is not None and "var" not in jinja_globals:
         builtin_globals["var"] = Var(variables)
 
     deployability_index = (
@@ -415,6 +415,7 @@ def create_builtin_globals(
         {k: builtin_globals.get(k) for k in ("ref", "source", "config", "var")}
     )
 
+    execute = True
     if engine_adapter is not None:
         builtin_globals["flags"] = Flags(which="run")
         adapter: BaseAdapter = RuntimeAdapter(
@@ -435,7 +436,11 @@ def create_builtin_globals(
         )
     else:
         builtin_globals["flags"] = Flags(which="parse")
-        adapter = ParsetimeAdapter(
+        adapter_class: t.Type[BaseAdapter] = ParsetimeAdapter
+        if jinja_globals.get("use_stub_adapter", False):
+            adapter_class = StubParsetimeAdapter
+            execute = False
+        adapter = adapter_class(
             jinja_macros,
             jinja_globals={**builtin_globals, **jinja_globals},
             project_dialect=project_dialect,
@@ -446,7 +451,7 @@ def create_builtin_globals(
     builtin_globals.update(
         {
             "adapter": adapter,
-            "execute": True,
+            "execute": execute,
             "load_relation": adapter.load_relation,
             "store_result": sql_execution.store_result,
             "load_result": sql_execution.load_result,

sqlmesh/dbt/common.py

@@ -184,13 +184,16 @@ class Dependencies(PydanticModel):
     variables: t.Set[str] = set()
     model_attrs: t.Set[str] = set()
 
+    has_dynamic_var_names: bool = False
+
     def union(self, other: Dependencies) -> Dependencies:
         return Dependencies(
             macros=list(set(self.macros) | set(other.macros)),
             sources=self.sources | other.sources,
             refs=self.refs | other.refs,
             variables=self.variables | other.variables,
             model_attrs=self.model_attrs | other.model_attrs,
+            has_dynamic_var_names=self.has_dynamic_var_names or other.has_dynamic_var_names,
         )
 
     @field_validator("macros", mode="after")

sqlmesh/dbt/context.py

@@ -7,7 +7,8 @@
 from dbt.adapters.base import BaseRelation
 
 from sqlmesh.core.config import Config as SQLMeshConfig
-from sqlmesh.dbt.builtin import _relation_info_to_relation
+from sqlmesh.dbt.builtin import _relation_info_to_relation, Var
+from sqlmesh.dbt.common import Dependencies
 from sqlmesh.dbt.manifest import ManifestHelper
 from sqlmesh.dbt.target import TargetConfig
 from sqlmesh.utils import AttributeDict
@@ -22,7 +23,6 @@
 if t.TYPE_CHECKING:
     from jinja2 import Environment
 
-    from sqlmesh.dbt.basemodel import Dependencies
     from sqlmesh.dbt.model import ModelConfig
     from sqlmesh.dbt.relation import Policy
     from sqlmesh.dbt.seed import SeedConfig
@@ -212,6 +212,38 @@ def target(self, value: TargetConfig) -> None:
     def render(self, source: str, **kwargs: t.Any) -> str:
         return self.jinja_environment.from_string(source).render(**kwargs)
 
+    def track_dependencies_on_render(
+        self, input: str, jinja_context: t.Dict[str, t.Any], package_name: t.Optional[str] = None
+    ) -> Dependencies:
+        dependencies_on_render = Dependencies()
+
+        class TrackingVar(Var):
+            def __call__(
+                self, name: str, default: t.Optional[t.Any] = None, **kwargs: t.Any
+            ) -> t.Any:
+                dependencies_on_render.variables.add(name)
+                return super().__call__(name, default, **kwargs)
+
+            def has_var(self, name: str) -> bool:
+                dependencies_on_render.variables.add(name)
+                return super().has_var(name)
+
+        if package_name:
+            top_level_packages = [*self.jinja_macros.top_level_packages, package_name]
+            jinja_macros = self.jinja_macros.copy(update={"top_level_packages": top_level_packages})
+        else:
+            jinja_macros = self.jinja_macros
+
+        jinja_environment = jinja_macros.build_environment(
+            **{
+                **jinja_context,
+                "var": TrackingVar(self.variables),
+                "use_stub_adapter": True,
+            }
+        )
+        jinja_environment.from_string(input).render()
+        return dependencies_on_render
+
     def get_callable_macro(
         self, name: str, package: t.Optional[str] = None
     ) -> t.Optional[t.Callable]:

sqlmesh/dbt/manifest.py

@@ -554,6 +554,9 @@ def _extra_dependencies(self, target: str, package: str) -> Dependencies:
                 args = [jinja_call_arg_name(arg) for arg in node.args]
                 if args and args[0]:
                     dependencies.variables.add(args[0])
+                else:
+                    # We couldn't determine the var name statically
+                    dependencies.has_dynamic_var_names = True
                 dependencies.macros.append(MacroReference(name="var"))
             elif len(call_name) == 1:
                 macro_name = call_name[0]

sqlmesh/dbt/model.py

@@ -25,7 +25,8 @@
 )
 from sqlmesh.core.model.kind import SCDType2ByTimeKind, OnDestructiveChange, OnAdditiveChange
 from sqlmesh.dbt.basemodel import BaseModelConfig, Materialization, SnapshotStrategy
-from sqlmesh.dbt.common import SqlStr, extract_jinja_config, sql_str_validator
+from sqlmesh.dbt.column import ColumnConfig
+from sqlmesh.dbt.common import SqlStr, extract_jinja_config, sql_str_validator, Dependencies
 from sqlmesh.utils.errors import ConfigError
 from sqlmesh.utils.pydantic import field_validator
 
@@ -436,6 +437,30 @@ def sqlmesh_config_fields(self) -> t.Set[str]:
             "physical_version",
         }
 
+    def sqlmesh_model_kwargs(
+        self,
+        context: DbtContext,
+        column_types_override: t.Optional[t.Dict[str, ColumnConfig]] = None,
+        extra_dependencies: t.Optional[Dependencies] = None,
+    ) -> t.Dict[str, t.Any]:
+        if not self.dependencies.has_dynamic_var_names:
+            return super().sqlmesh_model_kwargs(context, column_types_override, extra_dependencies)
+
+        extra_dependencies = extra_dependencies or Dependencies()
+        extra_dependencies = extra_dependencies.union(
+            self._track_dependencies_on_render(self.sql_no_config, context)
+        )
+        for pre_hook in self.pre_hook:
+            extra_dependencies = extra_dependencies.union(
+                self._track_dependencies_on_render(pre_hook.sql, context)
+            )
+        for post_hook in self.post_hook:
+            extra_dependencies = extra_dependencies.union(
+                self._track_dependencies_on_render(post_hook.sql, context)
+            )
+
+        return super().sqlmesh_model_kwargs(context, column_types_override, extra_dependencies)
+
     def to_sqlmesh(
         self,
         context: DbtContext,

tests/dbt/test_config.py

@@ -362,6 +362,7 @@ def test_variables(assert_exp_eq, sushi_test_project):
         "nested_vars": {
             "some_nested_var": 2,
         },
+        "dynamic_test_var": 3,
         "list_var": [
             {"name": "item1", "value": 1},
             {"name": "item2", "value": 2},
@@ -385,6 +386,7 @@ def test_variables(assert_exp_eq, sushi_test_project):
         "nested_vars": {
             "some_nested_var": 2,
         },
+        "dynamic_test_var": 3,
         "list_var": [
             {"name": "item1", "value": 1},
             {"name": "item2", "value": 2},

tests/dbt/test_manifest.py

@@ -79,6 +79,7 @@ def test_manifest_helper(caplog):
     waiter_revenue_by_day_config = models["waiter_revenue_by_day_v2"]
     assert waiter_revenue_by_day_config.dependencies == Dependencies(
         macros={
+            MacroReference(name="dynamic_var_name_dependency"),
             MacroReference(name="log_value"),
             MacroReference(name="test_dependencies"),
             MacroReference(package="customers", name="duckdb__current_engine"),
@@ -87,6 +88,7 @@ def test_manifest_helper(caplog):
         },
         sources={"streaming.items", "streaming.orders", "streaming.order_items"},
         variables={"yet_another_var", "nested_vars"},
+        has_dynamic_var_names=True,
     )
     assert waiter_revenue_by_day_config.materialized == "incremental"
     assert waiter_revenue_by_day_config.incremental_strategy == "delete+insert"