Fix: Use dbt source_name for source method and add support package arg for ref method. (#935)

Author: crericha

sqlmesh/dbt/basemodel.py

@@ -22,7 +22,6 @@
     GeneralConfig,
     QuotingConfig,
     SqlStr,
-    context_for_dependencies,
 )
 from sqlmesh.dbt.test import TestConfig
 from sqlmesh.utils import AttributeDict
@@ -66,6 +65,8 @@ class BaseModelConfig(GeneralConfig):
             (eg. 'parquet')
         path: The file path of the model
         dependencies: The macro, source, var, and ref dependencies used to execute the model and its hooks
+        name: Name of the model.
+        package_name: Name of the package that defines the model.
         database: Database the model is stored in
         schema: Custom schema name added to the model schema name
         alias: Relation identifier for this model instead of the filename
@@ -85,6 +86,8 @@ class BaseModelConfig(GeneralConfig):
     dependencies: Dependencies = Dependencies()
 
     # DBT configuration fields
+    name: str = ""
+    package_name: str = ""
     schema_: str = Field("", alias="schema")
     database: t.Optional[str] = None
     alias: t.Optional[str] = None
@@ -163,7 +166,14 @@ def table_name(self) -> str:
         return self.alias or self.path.stem
 
     @property
-    def model_name(self) -> str:
+    def config_name(self) -> str:
+        """
+        Get the model's config name (package_name.name)
+        """
+        return f"{self.package_name}.{self.name}"
+
+    @property
+    def sql_name(self) -> str:
         """
         Get the sqlmesh model name
 
@@ -206,7 +216,7 @@ def model_function(self) -> AttributeDict[str, t.Any]:
 
     def sqlmesh_model_kwargs(self, context: DbtContext) -> t.Dict[str, t.Any]:
         """Get common sqlmesh model parameters"""
-        model_context = context_for_dependencies(context, self.dependencies)
+        model_context = context.context_for_dependencies(self.dependencies)
         jinja_macros = model_context.jinja_macros.trim(self.dependencies.macros)
         jinja_macros.global_objs.update(
             {
@@ -228,8 +238,8 @@ def sqlmesh_model_kwargs(self, context: DbtContext) -> t.Dict[str, t.Any]:
             "audits": [(test.name, {}) for test in self.tests],
             "columns": column_types_to_sqlmesh(self.columns) or None,
             "column_descriptions_": column_descriptions_to_sqlmesh(self.columns) or None,
-            "depends_on": {context.refs[ref] for ref in self.dependencies.refs}.union(
-                {context.sources[source].source_name for source in self.dependencies.sources}
+            "depends_on": {model.sql_name for model in model_context.refs.values()}.union(
+                {source.sql_name for source in model_context.sources.values()}
             ),
             "jinja_macros": jinja_macros,
             "path": self.path,

sqlmesh/dbt/builtin.py

@@ -152,11 +152,9 @@ def var(name: str, default: t.Optional[str] = None) -> str:
 
 
 def generate_ref(refs: t.Dict[str, t.Any]) -> t.Callable:
-
-    # TODO suport package name
     def ref(package: str, name: t.Optional[str] = None) -> t.Optional[BaseRelation]:
-        name = name or package
-        relation_info = refs.get(name)
+        ref_name = f"{package}.{name}" if name else package
+        relation_info = refs.get(ref_name)
         if relation_info is None:
             return None
 

sqlmesh/dbt/common.py

@@ -18,9 +18,6 @@
 
 T = t.TypeVar("T", bound="GeneralConfig")
 
-if t.TYPE_CHECKING:
-    from sqlmesh.dbt.context import DbtContext
-
 PROJECT_FILENAME = "dbt_project.yml"
 
 JINJA_ONLY = {
@@ -179,34 +176,6 @@ def dict(self, *args: t.Any, **kwargs: t.Any) -> t.Dict[str, t.Any]:
         return out
 
 
-def context_for_dependencies(context: DbtContext, dependencies: Dependencies) -> DbtContext:
-    dependency_context = context.copy()
-
-    models = {}
-    seeds = {}
-    sources = {}
-
-    for ref in dependencies.refs:
-        if ref in context.seeds:
-            seeds[ref] = context.seeds[ref]
-        elif ref in context.models:
-            models[ref] = context.models[ref]
-        else:
-            raise ConfigError(f"Model '{ref}' was not found.")
-
-    for source in dependencies.sources:
-        if source in context.sources:
-            sources[source] = context.sources[source]
-        else:
-            raise ConfigError(f"Source '{source}' was not found.")
-
-    dependency_context.sources = sources
-    dependency_context.seeds = seeds
-    dependency_context.models = models
-
-    return dependency_context
-
-
 def extract_jinja_config(input: str) -> t.Tuple[str, str]:
     def jinja_end(sql: str, start: int) -> int:
         cursor = start

sqlmesh/dbt/context.py

@@ -14,6 +14,7 @@
 if t.TYPE_CHECKING:
     from jinja2 import Environment
 
+    from sqlmesh.dbt.basemodel import Dependencies
     from sqlmesh.dbt.model import ModelConfig
     from sqlmesh.dbt.seed import SeedConfig
     from sqlmesh.dbt.source import SourceConfig
@@ -40,7 +41,7 @@ class DbtContext:
     _models: t.Dict[str, ModelConfig] = field(default_factory=dict)
     _seeds: t.Dict[str, SeedConfig] = field(default_factory=dict)
     _sources: t.Dict[str, SourceConfig] = field(default_factory=dict)
-    _refs: t.Dict[str, str] = field(default_factory=dict)
+    _refs: t.Dict[str, t.Union[ModelConfig, SeedConfig]] = field(default_factory=dict)
 
     _target: t.Optional[TargetConfig] = None
 
@@ -128,9 +129,12 @@ def add_sources(self, sources: t.Dict[str, SourceConfig]) -> None:
         self._jinja_environment = None
 
     @property
-    def refs(self) -> t.Dict[str, str]:
+    def refs(self) -> t.Dict[str, t.Union[ModelConfig, SeedConfig]]:
         if not self._refs:
-            self._refs = {k: v.model_name for k, v in {**self._seeds, **self._models}.items()}  # type: ignore
+            # Refs can be called with or without package name.
+            for model in {**self._seeds, **self._models}.values():  # type: ignore
+                self._refs[model.name] = model
+                self._refs[model.config_name] = model
         return self._refs
 
     @property
@@ -162,14 +166,46 @@ def jinja_environment(self) -> Environment:
 
     @property
     def jinja_globals(self) -> t.Dict[str, JinjaGlobalAttribute]:
-        refs: t.Dict[str, t.Union[ModelConfig, SeedConfig]] = {**self.models, **self.seeds}
         output: t.Dict[str, JinjaGlobalAttribute] = {
             "vars": AttributeDict(self.variables),
-            "refs": AttributeDict({k: v.relation_info for k, v in refs.items()}),
+            "refs": AttributeDict({k: v.relation_info for k, v in self.refs.items()}),
             "sources": AttributeDict({k: v.relation_info for k, v in self.sources.items()}),
         }
         if self.project_name is not None:
             output["project_name"] = self.project_name
         if self._target is not None:
             output["target"] = self._target.attribute_dict()
         return output
+
+    def context_for_dependencies(self, dependencies: Dependencies) -> DbtContext:
+        from sqlmesh.dbt.model import ModelConfig
+        from sqlmesh.dbt.seed import SeedConfig
+
+        dependency_context = self.copy()
+
+        models = {}
+        seeds = {}
+        sources = {}
+
+        for ref in dependencies.refs:
+            model = self.refs.get(ref)
+            if model:
+                if isinstance(model, SeedConfig):
+                    seeds[ref] = t.cast(SeedConfig, model)
+                else:
+                    models[ref] = t.cast(ModelConfig, model)
+            else:
+                raise ConfigError(f"Model '{ref}' was not found.")
+
+        for source in dependencies.sources:
+            if source in self.sources:
+                sources[source] = self.sources[source]
+            else:
+                raise ConfigError(f"Source '{source}' was not found.")
+
+        dependency_context.sources = sources
+        dependency_context.seeds = seeds
+        dependency_context.models = models
+        dependency_context._refs = {**dependency_context._seeds, **dependency_context._models}  # type: ignore
+
+        return dependency_context

sqlmesh/dbt/loader.py

@@ -99,7 +99,7 @@ def _load_models(
 
             models.update(
                 {
-                    model.model_name: cache.get_or_load_model(
+                    model.sql_name: cache.get_or_load_model(
                         model.path, lambda: self._to_sqlmesh(model, context)
                     )
                     for model in package_models.values()
@@ -139,7 +139,7 @@ def _load_project(self) -> Project:
 
     @classmethod
     def _to_sqlmesh(cls, config: BMC, context: DbtContext) -> Model:
-        logger.debug("Converting '%s' to sqlmesh format", config.model_name)
+        logger.debug("Converting '%s' to sqlmesh format", config.sql_name)
         return config.to_sqlmesh(context)
 
     def _compute_yaml_max_mtime_per_subfolder(self, root: Path) -> t.Dict[Path, float]:

sqlmesh/dbt/manifest.py

@@ -61,6 +61,8 @@ def __init__(
         self._sources_per_package: t.Dict[str, SourceConfigs] = defaultdict(dict)
         self._macros_per_package: t.Dict[str, MacroConfigs] = defaultdict(dict)
 
+        self._tests_by_owner: t.Dict[str, t.List[TestConfig]] = defaultdict(list)
+
     def tests(self, package_name: t.Optional[str] = None) -> TestConfigs:
         self._load_all()
         return self._tests_per_package[package_name or self._project_name]
@@ -109,7 +111,7 @@ def _load_sources(self) -> None:
                 **source_dict,
             )
             self._sources_per_package[source.package_name][
-                source_config.source_name
+                source_config.config_name
             ] = source_config
 
     def _load_macros(self) -> None:
@@ -162,20 +164,26 @@ def _load_tests(self) -> None:
                 )
                 continue
 
-            self._tests_per_package[node.package_name][node.name.lower()] = TestConfig(
+            test = TestConfig(
                 sql=node.raw_code if DBT_VERSION >= (1, 3) else node.raw_sql,  # type: ignore
                 owner=test_owner,
                 test_kwargs=node.test_metadata.kwargs if hasattr(node, "test_metadata") else {},
                 dependencies=dependencies,
                 **_node_base_config(node),
             )
+            self._tests_per_package[node.package_name][node.name.lower()] = test
+            self._tests_by_owner[test_owner].append(test)
 
     def _load_models_and_seeds(self) -> None:
         for node in self._manifest.nodes.values():
             if node.resource_type not in ("model", "seed"):
                 continue
 
             macro_references = _macro_references(self._manifest, node)
+            tests = (
+                self._tests_by_owner[node.name]
+                + self._tests_by_owner[f"{node.package_name}.{node.name}"]
+            )
 
             if node.resource_type == "model":
                 self._models_per_package[node.package_name][node.name] = ModelConfig(
@@ -185,13 +193,13 @@ def _load_models_and_seeds(self) -> None:
                         refs=_refs(node),
                         sources=_sources(node),
                     ),
-                    tests=_tests_for_node(node, self._tests_per_package[node.package_name]),
+                    tests=tests,
                     **_node_base_config(node),
                 )
             else:
                 self._seeds_per_package[node.package_name][node.name] = SeedConfig(
                     dependencies=Dependencies(macros=macro_references),
-                    tests=_tests_for_node(node, self._tests_per_package[node.package_name]),
+                    tests=tests,
                     **_node_base_config(node),
                 )
 
@@ -259,9 +267,9 @@ def _macro_references(
 
 def _refs(node: ManifestNode) -> t.Set[str]:
     if DBT_VERSION >= (1, 5):
-        return {r.name for r in node.refs}  # type: ignore
+        return {f"{r.package}.{r.name}" if r.package else r.name for r in node.refs}  # type: ignore
     else:
-        return {r[1] if len(r) > 1 else r[0] for r in node.refs}  # type: ignore
+        return {".".join(r) for r in node.refs}  # type: ignore
 
 
 def _sources(node: ManifestNode) -> t.Set[str]:
@@ -309,10 +317,6 @@ def _node_base_config(node: ManifestNode) -> t.Dict[str, t.Any]:
     }
 
 
-def _tests_for_node(node: ManifestNode, tests: t.Dict[str, TestConfig]) -> t.List[TestConfig]:
-    return [test for test in tests.values() if test.owner == node.name]
-
-
 def _convert_jinja_test_to_macro(test_jinja: str) -> str:
     TEST_TAG_REGEX = "\s*{%\s*test\s+"
     ENDTEST_REGEX = "{%\s*endtest\s*%}"

sqlmesh/dbt/model.py

@@ -152,13 +152,13 @@ def model_kind(self, target: TargetConfig) -> ModelKind:
                     and strategy not in INCREMENTAL_BY_UNIQUE_KEY_STRATEGIES
                 ):
                     raise ConfigError(
-                        f"{self.model_name}: SQLMesh IncrementalByUniqueKey is not compatible with '{strategy}'"
+                        f"{self.sql_name}: SQLMesh IncrementalByUniqueKey is not compatible with '{strategy}'"
                         f" incremental strategy. Supported strategies include {collection_to_str(INCREMENTAL_BY_UNIQUE_KEY_STRATEGIES)}."
                     )
                 return IncrementalByUniqueKeyKind(unique_key=self.unique_key, **incremental_kwargs)
 
             raise ConfigError(
-                f"{self.model_name}: Incremental materialization requires either a "
+                f"{self.sql_name}: Incremental materialization requires either a "
                 f"time_column {collection_to_str(INCREMENTAL_BY_TIME_STRATEGIES)}) or a "
                 f"unique_key ({collection_to_str(INCREMENTAL_BY_UNIQUE_KEY_STRATEGIES.union(['none']))}) configuration."
             )
@@ -205,10 +205,10 @@ def to_sqlmesh(self, context: DbtContext) -> Model:
                 optional_kwargs[field] = field_val
 
         if not context.target:
-            raise ConfigError(f"Target required to load '{self.model_name}' into SQLMesh.")
+            raise ConfigError(f"Target required to load '{self.sql_name}' into SQLMesh.")
 
         return create_sql_model(
-            self.model_name,
+            self.sql_name,
             expressions[0],
             dialect=dialect,
             kind=self.model_kind(context.target),

sqlmesh/dbt/seed.py

@@ -21,7 +21,7 @@ class SeedConfig(BaseModelConfig):
     def to_sqlmesh(self, context: DbtContext) -> Model:
         """Converts the dbt seed into a SQLMesh model."""
         return create_seed_model(
-            self.model_name,
+            self.sql_name,
             SeedKind(path=self.path.absolute()),
             **self.sqlmesh_model_kwargs(context),
         )

sqlmesh/dbt/source.py

@@ -14,8 +14,8 @@
 class SourceConfig(GeneralConfig):
     """
     Args:
-        config_name: The schema.table_name names declared in source config
-        name: The name of the source or table
+        name: The name of the table
+        source_name: The name of the source that defines the table
         database: Name of the database where the table is stored. By default, the project's target database is used.
         schema: The scehma name as stored in the database. If not specified, the source name is used.
         identifier: The table name as stored in the database. If not specified, the source table name is used
@@ -28,11 +28,9 @@ class SourceConfig(GeneralConfig):
         columns: Columns within the source
     """
 
-    # sqlmesh fields
-    config_name: str = ""
-
     # DBT configuration fields
-    name: t.Optional[str] = None
+    name: str = ""
+    source_name_: str = Field("", alias="source_name")
     database: t.Optional[str] = None
     schema_: t.Optional[str] = Field(None, alias="schema")
     identifier: t.Optional[str] = None
@@ -54,7 +52,11 @@ def table_name(self) -> t.Optional[str]:
         return self.identifier or self.name
 
     @property
-    def source_name(self) -> str:
+    def config_name(self) -> str:
+        return f"{self.source_name_}.{self.name}"
+
+    @property
+    def sql_name(self) -> str:
         return ".".join(part for part in (self.schema_, self.table_name) if part)
 
     @property