PR feedback: remove DBT_ALL_MODEL_ATTRS placeholder

Author: georgesittas

sqlmesh/dbt/basemodel.py

@@ -19,7 +19,6 @@
     column_types_to_sqlmesh,
 )
 from sqlmesh.dbt.common import (
-    DBT_ALL_MODEL_ATTRS,
     DbtConfig,
     Dependencies,
     GeneralConfig,
@@ -379,11 +378,11 @@ def _model_jinja_context(
     ) -> t.Dict[str, t.Any]:
         if context._manifest and self.node_name in context._manifest._manifest.nodes:
             attributes = context._manifest._manifest.nodes[self.node_name].to_dict()
-            if DBT_ALL_MODEL_ATTRS in dependencies.model_attrs:
+            if dependencies.model_attrs.all_attrs:
                 model_node: AttributeDict[str, t.Any] = AttributeDict(attributes)
             else:
                 model_node = AttributeDict(
-                    filter(lambda kv: kv[0] in dependencies.model_attrs, attributes.items())
+                    filter(lambda kv: kv[0] in dependencies.model_attrs.attrs, attributes.items())
                 )
 
             raw_code_key = "raw_code" if DBT_VERSION >= (1, 3, 0) else "raw_sql"  # type: ignore

sqlmesh/dbt/common.py

@@ -2,6 +2,7 @@
 
 import re
 import typing as t
+from dataclasses import dataclass
 from pathlib import Path
 
 from ruamel.yaml.constructor import DuplicateKeyError
@@ -19,7 +20,6 @@
 T = t.TypeVar("T", bound="GeneralConfig")
 
 PROJECT_FILENAME = DBT_PROJECT_FILENAME
-DBT_ALL_MODEL_ATTRS = "__DBT_ALL_MODEL_ATTRS__"
 
 JINJA_ONLY = {
     "adapter",
@@ -173,6 +173,12 @@ def sqlmesh_config_fields(self) -> t.Set[str]:
         return set()
 
 
+@dataclass
+class ModelAttrs:
+    attrs: t.Set[str]
+    all_attrs: bool = False
+
+
 class Dependencies(PydanticModel):
     """
     DBT dependencies for a model, macro, etc.
@@ -187,7 +193,7 @@ class Dependencies(PydanticModel):
     sources: t.Set[str] = set()
     refs: t.Set[str] = set()
     variables: t.Set[str] = set()
-    model_attrs: t.Set[str] = set()
+    model_attrs: ModelAttrs = ModelAttrs(attrs=set())
 
     has_dynamic_var_names: bool = False
 
@@ -197,7 +203,10 @@ def union(self, other: Dependencies) -> Dependencies:
             sources=self.sources | other.sources,
             refs=self.refs | other.refs,
             variables=self.variables | other.variables,
-            model_attrs=self.model_attrs | other.model_attrs,
+            model_attrs=ModelAttrs(
+                attrs=self.model_attrs.attrs | other.model_attrs.attrs,
+                all_attrs=self.model_attrs.all_attrs or other.model_attrs.all_attrs,
+            ),
             has_dynamic_var_names=self.has_dynamic_var_names or other.has_dynamic_var_names,
         )
 

sqlmesh/dbt/manifest.py

@@ -45,7 +45,7 @@
 from sqlmesh.utils.errors import SQLMeshError
 from sqlmesh.core.config import ModelDefaultsConfig
 from sqlmesh.dbt.builtin import BUILTIN_FILTERS, BUILTIN_GLOBALS, OVERRIDDEN_MACROS
-from sqlmesh.dbt.common import DBT_ALL_MODEL_ATTRS, Dependencies
+from sqlmesh.dbt.common import Dependencies
 from sqlmesh.dbt.model import ModelConfig
 from sqlmesh.dbt.package import HookConfig, MacroConfig
 from sqlmesh.dbt.seed import SeedConfig
@@ -584,7 +584,7 @@ def _extra_dependencies(
 
             if isinstance(node, jinja2.nodes.Getattr):
                 if call_name[0] == "model":
-                    dependencies.model_attrs.add(call_name[1])
+                    dependencies.model_attrs.attrs.add(call_name[1])
             elif call_name[0] == "source":
                 args = [jinja_call_arg_name(arg) for arg in node.args]
                 if args and all(arg for arg in args):
@@ -634,7 +634,7 @@ def _extra_dependencies(
         # analysis. We conservatively deal with this by including all of its supported attributes
         # if a standalone reference is found.
         if all_model_attrs:
-            dependencies.model_attrs = {DBT_ALL_MODEL_ATTRS}
+            dependencies.model_attrs.all_attrs = True
 
         return dependencies
 

tests/dbt/test_manifest.py

@@ -6,6 +6,7 @@
 
 from sqlmesh.core.config import ModelDefaultsConfig
 from sqlmesh.dbt.basemodel import Dependencies
+from sqlmesh.dbt.common import ModelAttrs
 from sqlmesh.dbt.context import DbtContext
 from sqlmesh.dbt.manifest import ManifestHelper, _convert_jinja_test_to_macro
 from sqlmesh.dbt.profile import Profile
@@ -33,7 +34,7 @@ def test_manifest_helper(caplog):
     assert models["top_waiters"].dependencies == Dependencies(
         refs={"sushi.waiter_revenue_by_day", "waiter_revenue_by_day"},
         variables={"top_waiters:revenue", "top_waiters:limit"},
-        model_attrs={"columns", "config"},
+        model_attrs=ModelAttrs(attrs={"columns", "config"}),
         macros=[
             MacroReference(name="get_top_waiters_limit"),
             MacroReference(name="ref"),