sqlmesh 0.213.1.dev1__py3-none-any.whl → 0.227.2.dev4__py3-none-any.whl

sqlmesh/dbt/loader.py

@@ -5,11 +5,13 @@ import sys
 import typing as t
 import sqlmesh.core.dialect as d
 from pathlib import Path
+from collections import defaultdict
 from sqlmesh.core.config import (
     Config,
     ConnectionConfig,
     GatewayConfig,
     ModelDefaultsConfig,
+    DbtConfig as RootDbtConfig,
 )
 from sqlmesh.core.environment import EnvironmentStatements
 from sqlmesh.core.loader import CacheBase, LoadedProject, Loader
@@ -48,11 +50,21 @@ def sqlmesh_config(
     dbt_profile_name: t.Optional[str] = None,
     dbt_target_name: t.Optional[str] = None,
     variables: t.Optional[t.Dict[str, t.Any]] = None,
+    threads: t.Optional[int] = None,
     register_comments: t.Optional[bool] = None,
+    infer_state_schema_name: bool = False,
+    profiles_dir: t.Optional[Path] = None,
     **kwargs: t.Any,
 ) -> Config:
     project_root = project_root or Path()
-    context = DbtContext(project_root=project_root, profile_name=dbt_profile_name)
+    context = DbtContext(
+        project_root=project_root, profiles_dir=profiles_dir, profile_name=dbt_profile_name
+    )
+
+    # note: Profile.load() is called twice with different DbtContext's:
+    # - once here with the above DbtContext (to determine connnection / gateway config which has to be set up before everything else)
+    # - again on the SQLMesh side via GenericContext.load() -> DbtLoader._load_projects() -> Project.load() which constructs a fresh DbtContext and ignores the above one
+    # it's important to ensure that the DbtContext created within the DbtLoader uses the same project root / profiles dir that we use here
     profile = Profile.load(context, target_name=dbt_target_name)
     model_defaults = kwargs.pop("model_defaults", ModelDefaultsConfig())
     if model_defaults.dialect is None:
@@ -66,16 +78,45 @@ def sqlmesh_config(
     if not issubclass(loader, DbtLoader):
         raise ConfigError("The loader must be a DbtLoader.")
 
+    if threads is not None:
+        # the to_sqlmesh() function on TargetConfig maps self.threads -> concurrent_tasks
+        profile.target.threads = threads
+
+    gateway_kwargs = {}
+    if infer_state_schema_name:
+        profile_name = context.profile_name
+
+        # Note: we deliberately isolate state based on the target *schema* and not the target name.
+        # It is assumed that the project will define a target, eg 'dev', and then in each users own ~/.dbt/profiles.yml the schema
+        # for the 'dev' target is overriden to something user-specific, rather than making the target name itself user-specific.
+        # This means that the schema name is the indicator of isolated state, not the target name which may be re-used across multiple schemas.
+        target_schema = profile.target.schema_
+
+        # dbt-core doesnt allow schema to be undefined, but it does allow an empty string, and then just
+        # fails at runtime when `CREATE SCHEMA ""` doesnt work
+        if not target_schema:
+            raise ConfigError(
+                f"Target '{profile.target_name}' does not specify a schema.\n"
+                "A schema is required in order to infer where to store SQLMesh state"
+            )
+
+        inferred_state_schema_name = f"sqlmesh_state_{profile_name}_{target_schema}"
+        logger.info("Inferring state schema: %s", inferred_state_schema_name)
+        gateway_kwargs["state_schema"] = inferred_state_schema_name
+
     return Config(
         loader=loader,
+        loader_kwargs=dict(profiles_dir=profiles_dir),
         model_defaults=model_defaults,
         variables=variables or {},
+        dbt=RootDbtConfig(infer_state_schema_name=infer_state_schema_name),
         **{
             "default_gateway": profile.target_name if "gateways" not in kwargs else "",
             "gateways": {
                 profile.target_name: GatewayConfig(
                     connection=profile.target.to_sqlmesh(**target_to_sqlmesh_args),
                     state_connection=state_connection,
+                    **gateway_kwargs,
                 )
             },  # type: ignore
             **kwargs,
@@ -84,9 +125,12 @@ def sqlmesh_config(
 
 
 class DbtLoader(Loader):
-    def __init__(self, context: GenericContext, path: Path) -> None:
+    def __init__(
+        self, context: GenericContext, path: Path, profiles_dir: t.Optional[Path] = None
+    ) -> None:
         self._projects: t.List[Project] = []
         self._macros_max_mtime: t.Optional[float] = None
+        self._profiles_dir = profiles_dir
         super().__init__(context, path)
 
     def load(self) -> LoadedProject:
@@ -137,16 +181,22 @@ class DbtLoader(Loader):
                 package_context.set_and_render_variables(package.variables, package.name)
                 package_models: t.Dict[str, BaseModelConfig] = {**package.models, **package.seeds}
 
+                package_models_by_path: t.Dict[Path, t.List[BaseModelConfig]] = defaultdict(list)
                 for model in package_models.values():
-                    if isinstance(model, ModelConfig) and not model.sql_no_config:
+                    if isinstance(model, ModelConfig) and not model.sql.strip():
                         logger.info(f"Skipping empty model '{model.name}' at path '{model.path}'.")
                         continue
+                    package_models_by_path[model.path].append(model)
 
-                    sqlmesh_model = cache.get_or_load_models(
-                        model.path, loader=lambda: [_to_sqlmesh(model, package_context)]
-                    )[0]
-
-                    models[sqlmesh_model.fqn] = sqlmesh_model
+                for path, path_models in package_models_by_path.items():
+                    sqlmesh_models = cache.get_or_load_models(
+                        path,
+                        loader=lambda: [
+                            _to_sqlmesh(model, package_context) for model in path_models
+                        ],
+                    )
+                    for sqlmesh_model in sqlmesh_models:
+                        models[sqlmesh_model.fqn] = sqlmesh_model
 
             models.update(self._load_external_models(audits, cache))
 
@@ -165,7 +215,8 @@ class DbtLoader(Loader):
                 for test in package.tests.values():
                     logger.debug("Converting '%s' to sqlmesh format", test.name)
                     try:
-                        audits[test.name] = test.to_sqlmesh(package_context)
+                        audits[test.canonical_name] = test.to_sqlmesh(package_context)
+
                     except BaseMissingReferenceError as e:
                         ref_type = "model" if isinstance(e, MissingModelError) else "source"
                         logger.warning(
@@ -186,6 +237,7 @@ class DbtLoader(Loader):
             project = Project.load(
                 DbtContext(
                     project_root=self.config_path,
+                    profiles_dir=self._profiles_dir,
                     target_name=target_name,
                     sqlmesh_config=self.config,
                 ),

sqlmesh/dbt/manifest.py

@@ -11,7 +11,7 @@ from collections import defaultdict
 from functools import cached_property
 from pathlib import Path
 
-from dbt import constants as dbt_constants, flags
+from dbt import flags
 
 from sqlmesh.dbt.util import DBT_VERSION
 from sqlmesh.utils.conversions import make_serializable
@@ -19,6 +19,8 @@ from sqlmesh.utils.conversions import make_serializable
 # Override the file name to prevent dbt commands from invalidating the cache.
 
 if DBT_VERSION >= (1, 6, 0):
+    from dbt import constants as dbt_constants
+
     dbt_constants.PARTIAL_PARSE_FILE_NAME = "sqlmesh_partial_parse.msgpack"  # type: ignore
 else:
     from dbt.parser import manifest as dbt_manifest  # type: ignore
@@ -44,10 +46,10 @@ from dbt.tracking import do_not_track
 from sqlmesh.core import constants as c
 from sqlmesh.utils.errors import SQLMeshError
 from sqlmesh.core.config import ModelDefaultsConfig
-from sqlmesh.dbt.basemodel import Dependencies
 from sqlmesh.dbt.builtin import BUILTIN_FILTERS, BUILTIN_GLOBALS, OVERRIDDEN_MACROS
+from sqlmesh.dbt.common import Dependencies
 from sqlmesh.dbt.model import ModelConfig
-from sqlmesh.dbt.package import HookConfig, MacroConfig
+from sqlmesh.dbt.package import HookConfig, MacroConfig, MaterializationConfig
 from sqlmesh.dbt.seed import SeedConfig
 from sqlmesh.dbt.source import SourceConfig
 from sqlmesh.dbt.target import TargetConfig
@@ -61,6 +63,7 @@ from sqlmesh.utils.jinja import (
     extract_call_names,
     jinja_call_arg_name,
 )
+from sqlglot.helper import ensure_list
 
 if t.TYPE_CHECKING:
     from dbt.contracts.graph.manifest import Macro, Manifest
@@ -75,6 +78,7 @@ SeedConfigs = t.Dict[str, SeedConfig]
 SourceConfigs = t.Dict[str, SourceConfig]
 MacroConfigs = t.Dict[str, MacroConfig]
 HookConfigs = t.Dict[str, HookConfig]
+MaterializationConfigs = t.Dict[str, MaterializationConfig]
 
 
 IGNORED_PACKAGES = {"elementary"}
@@ -135,6 +139,7 @@ class ManifestHelper:
 
         self._on_run_start_per_package: t.Dict[str, HookConfigs] = defaultdict(dict)
         self._on_run_end_per_package: t.Dict[str, HookConfigs] = defaultdict(dict)
+        self._materializations: MaterializationConfigs = {}
 
     def tests(self, package_name: t.Optional[str] = None) -> TestConfigs:
         self._load_all()
@@ -164,6 +169,10 @@ class ManifestHelper:
         self._load_all()
         return self._on_run_end_per_package[package_name or self._project_name]
 
+    def materializations(self) -> MaterializationConfigs:
+        self._load_all()
+        return self._materializations
+
     @property
     def all_macros(self) -> t.Dict[str, t.Dict[str, MacroInfo]]:
         self._load_all()
@@ -213,6 +222,7 @@ class ManifestHelper:
         self._calls = {k: (v, False) for k, v in (self._call_cache.get("") or {}).items()}
 
         self._load_macros()
+        self._load_materializations()
         self._load_sources()
         self._load_tests()
         self._load_models_and_seeds()
@@ -250,11 +260,14 @@ class ManifestHelper:
 
     def _load_macros(self) -> None:
         for macro in self._manifest.macros.values():
+            if macro.name.startswith("materialization_"):
+                continue
+
             if macro.name.startswith("test_"):
                 macro.macro_sql = _convert_jinja_test_to_macro(macro.macro_sql)
 
             dependencies = Dependencies(macros=_macro_references(self._manifest, macro))
-            if not macro.name.startswith("materialization_") and not macro.name.startswith("test_"):
+            if not macro.name.startswith("test_"):
                 dependencies = dependencies.union(
                     self._extra_dependencies(macro.macro_sql, macro.package_name)
                 )
@@ -281,6 +294,32 @@ class ManifestHelper:
                 if pos > 0 and name[pos + 2 :] in adapter_macro_names:
                     macro_config.info.is_top_level = True
 
+    def _load_materializations(self) -> None:
+        for macro in self._manifest.macros.values():
+            if macro.name.startswith("materialization_"):
+                # Extract name and adapter ( "materialization_{name}_{adapter}" or "materialization_{name}_default")
+                name_parts = macro.name.split("_")
+                if len(name_parts) >= 3:
+                    mat_name = "_".join(name_parts[1:-1])
+                    adapter = name_parts[-1]
+
+                    dependencies = Dependencies(macros=_macro_references(self._manifest, macro))
+                    macro.macro_sql = _strip_jinja_materialization_tags(macro.macro_sql)
+                    dependencies = dependencies.union(
+                        self._extra_dependencies(macro.macro_sql, macro.package_name)
+                    )
+
+                    materialization_config = MaterializationConfig(
+                        name=mat_name,
+                        adapter=adapter,
+                        definition=macro.macro_sql,
+                        dependencies=dependencies,
+                        path=Path(macro.original_file_path),
+                    )
+
+                    key = f"{mat_name}_{adapter}"
+                    self._materializations[key] = materialization_config
+
     def _load_tests(self) -> None:
         for node in self._manifest.nodes.values():
             if node.resource_type != "test":
@@ -317,15 +356,17 @@ class ManifestHelper:
             )
 
             test_model = _test_model(node)
+            node_config = _node_base_config(node)
+            node_config["name"] = _build_test_name(node, dependencies)
 
             test = TestConfig(
                 sql=sql,
                 model_name=test_model,
                 test_kwargs=node.test_metadata.kwargs if hasattr(node, "test_metadata") else {},
                 dependencies=dependencies,
-                **_node_base_config(node),
+                **node_config,
             )
-            self._tests_per_package[node.package_name][node.name.lower()] = test
+            self._tests_per_package[node.package_name][node.unique_id] = test
             if test_model:
                 self._tests_by_owner[test_model].append(test)
 
@@ -338,10 +379,12 @@ class ManifestHelper:
                 continue
 
             macro_references = _macro_references(self._manifest, node)
-            tests = (
+            all_tests = (
                 self._tests_by_owner[node.name]
                 + self._tests_by_owner[f"{node.package_name}.{node.name}"]
             )
+            # Only include non-standalone tests (tests that don't reference other models)
+            tests = [test for test in all_tests if not test.is_standalone]
             node_config = _node_base_config(node)
 
             node_name = node.name
@@ -354,7 +397,15 @@ class ManifestHelper:
                 dependencies = Dependencies(
                     macros=macro_references, refs=_refs(node), sources=_sources(node)
                 )
-                dependencies = dependencies.union(self._extra_dependencies(sql, node.package_name))
+                dependencies = dependencies.union(
+                    self._extra_dependencies(sql, node.package_name, track_all_model_attrs=True)
+                )
+                for hook in [*node_config.get("pre-hook", []), *node_config.get("post-hook", [])]:
+                    dependencies = dependencies.union(
+                        self._extra_dependencies(
+                            hook["sql"], node.package_name, track_all_model_attrs=True
+                        )
+                    )
                 dependencies = dependencies.union(
                     self._flatten_dependencies_from_macros(dependencies.macros, node.package_name)
                 )
@@ -552,17 +603,37 @@ class ManifestHelper:
             dependencies = dependencies.union(macro_dependencies)
         return dependencies
 
-    def _extra_dependencies(self, target: str, package: str) -> Dependencies:
-        # We sometimes observe that the manifest doesn't capture all macros, refs, and sources within a macro.
-        # This behavior has been observed with macros like dbt.current_timestamp(), dbt_utils.slugify(), and source().
-        # Here we apply our custom extractor to make a best effort to supplement references captured in the manifest.
+    def _extra_dependencies(
+        self,
+        target: str,
+        package: str,
+        track_all_model_attrs: bool = False,
+    ) -> Dependencies:
+        """
+        We sometimes observe that the manifest doesn't capture all macros, refs, and sources within a macro.
+        This behavior has been observed with macros like dbt.current_timestamp(), dbt_utils.slugify(), and source().
+        Here we apply our custom extractor to make a best effort to supplement references captured in the manifest.
+        """
         dependencies = Dependencies()
+
+        # Whether all `model` attributes (e.g., `model.config`) should be included in the dependencies
+        all_model_attrs = False
+
         for call_name, node in extract_call_names(target, cache=self._calls):
             if call_name[0] == "config":
                 continue
-            elif isinstance(node, jinja2.nodes.Getattr):
+
+            if (
+                track_all_model_attrs
+                and not all_model_attrs
+                and isinstance(node, jinja2.nodes.Call)
+                and any(isinstance(a, jinja2.nodes.Name) and a.name == "model" for a in node.args)
+            ):
+                all_model_attrs = True
+
+            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):
@@ -606,6 +677,14 @@ class ManifestHelper:
                         call_name[0], call_name[1], dependencies.macros.append
                     )
 
+        # When `model` is referenced as-is, e.g. it's passed as an argument to a macro call like
+        # `{{ foo(model) }}`, we can't easily track the attributes that are actually used, because
+        # it may be aliased and hence tracking actual uses of `model` requires a proper data flow
+        # 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.all_attrs = True
+
         return dependencies
 
 
@@ -629,7 +708,7 @@ def _macro_references(
         return result
 
     for macro_node_id in node.depends_on.macros:
-        if not macro_node_id:
+        if not macro_node_id or macro_node_id == "None":
             continue
 
         macro_node = manifest.macros[macro_node_id]
@@ -667,7 +746,12 @@ def _test_model(node: ManifestNode) -> t.Optional[str]:
     attached_node = getattr(node, "attached_node", None)
     if attached_node:
         pieces = attached_node.split(".")
-        return pieces[-1] if pieces[0] in ["model", "seed"] else None
+        if pieces[0] in ["model", "seed"]:
+            # versioned models have format "model.package.model_name.v1" (4 parts)
+            if len(pieces) == 4:
+                return f"{pieces[2]}_{pieces[3]}"
+            return pieces[-1]
+        return None
 
     key_name = getattr(node, "file_key_name", None)
     if key_name:
@@ -700,3 +784,77 @@ def _convert_jinja_test_to_macro(test_jinja: str) -> str:
     macro = macro_tag + test_jinja[match.span()[-1] :]
 
     return re.sub(ENDTEST_REGEX, lambda m: m.group(0).replace("endtest", "endmacro"), macro)
+
+
+def _strip_jinja_materialization_tags(materialization_jinja: str) -> str:
+    MATERIALIZATION_TAG_REGEX = r"\s*{%-?\s*materialization\s+[^%]*%}\s*\n?"
+    ENDMATERIALIZATION_REGEX = r"{%-?\s*endmaterialization\s*-?%}\s*\n?"
+
+    if not re.match(MATERIALIZATION_TAG_REGEX, materialization_jinja):
+        return materialization_jinja
+
+    materialization_jinja = re.sub(
+        MATERIALIZATION_TAG_REGEX,
+        "",
+        materialization_jinja,
+        flags=re.IGNORECASE,
+    )
+
+    materialization_jinja = re.sub(
+        ENDMATERIALIZATION_REGEX,
+        "",
+        materialization_jinja,
+        flags=re.IGNORECASE,
+    )
+
+    return materialization_jinja.strip()
+
+
+def _build_test_name(node: ManifestNode, dependencies: Dependencies) -> str:
+    """
+    Build a user-friendly test name that includes the test's model/source, column,
+    and args for tests with custom user names. Needed because dbt only generates these
+    names for tests that do not specify the "name" field in their YAML definition.
+
+    Name structure
+    - Model test:  [namespace]_[test name]_[model name]_[column name]__[arg values]
+    - Source test: [namespace]_source_[test name]_[source name]_[table name]_[column name]__[arg values]
+    """
+    # standalone test
+    if not hasattr(node, "test_metadata"):
+        return node.name
+
+    model_name = _test_model(node)
+    source_name = None
+    if not model_name and dependencies.sources:
+        # extract source and table names
+        source_parts = list(dependencies.sources)[0].split(".")
+        source_name = "_".join(source_parts) if len(source_parts) == 2 else source_parts[-1]
+    entity_name = model_name or source_name or ""
+    entity_name = f"_{entity_name}" if entity_name else ""
+
+    name_prefix = ""
+    if namespace := getattr(node.test_metadata, "namespace", None):
+        name_prefix += f"{namespace}_"
+    if source_name and not model_name:
+        name_prefix += "source_"
+
+    metadata_kwargs = node.test_metadata.kwargs
+    arg_val_parts = []
+    for arg, val in sorted(metadata_kwargs.items()):
+        if arg == "model":
+            continue
+        if isinstance(val, dict):
+            val = list(val.values())
+        val = [re.sub("[^0-9a-zA-Z_]+", "_", str(v)) for v in ensure_list(val)]
+        arg_val_parts.extend(val)
+    unique_args = "__".join(arg_val_parts) if arg_val_parts else ""
+    unique_args = f"_{unique_args}" if unique_args else ""
+
+    auto_name = f"{name_prefix}{node.test_metadata.name}{entity_name}{unique_args}"
+
+    if node.name == auto_name:
+        return node.name
+
+    custom_prefix = name_prefix if source_name and not model_name else ""
+    return f"{custom_prefix}{node.name}{entity_name}{unique_args}"