Edit on GitHub

sqlmesh.dbt.loader

 1from __future__ import annotations
 2
 3import typing as t
 4from pathlib import Path
 5
 6from sqlmesh.core.audit import Audit
 7from sqlmesh.core.config import Config
 8from sqlmesh.core.hooks import HookRegistry
 9from sqlmesh.core.loader import Loader
10from sqlmesh.core.macros import MacroRegistry
11from sqlmesh.core.model import Model
12from sqlmesh.dbt.basemodel import BaseModelConfig
13from sqlmesh.dbt.common import DbtContext
14from sqlmesh.dbt.profile import Profile
15from sqlmesh.dbt.project import Project
16from sqlmesh.utils import UniqueKeyDict
17
18
19def sqlmesh_config(project_root: t.Optional[Path] = None, **kwargs: t.Any) -> Config:
20    project_root = project_root or Path()
21    context = DbtContext(project_root=project_root)
22    profile = Profile.load(context)
23
24    return Config(
25        default_connection=profile.default_target,
26        connections=profile.to_sqlmesh(),
27        loader=DbtLoader,
28        **kwargs,
29    )
30
31
32class DbtLoader(Loader):
33    def _load_scripts(self) -> t.Tuple[MacroRegistry, HookRegistry]:
34        macro_files = list(Path(self._context.path, "macros").glob("**/*.sql"))
35        for file in macro_files:
36            self._track_file(file)
37
38        return (
39            UniqueKeyDict("macros"),
40            UniqueKeyDict("hooks"),
41        )
42
43    def _load_models(
44        self,
45        macros: MacroRegistry,
46        hooks: HookRegistry,
47    ) -> UniqueKeyDict[str, Model]:
48        models: UniqueKeyDict = UniqueKeyDict("models")
49
50        project = Project.load(
51            DbtContext(project_root=self._context.path, target_name=self._context.connection)
52        )
53        for path in project.project_files:
54            self._track_file(path)
55
56        context = project.context.copy()
57
58        for package_name, package in project.packages.items():
59            context.add_sources(package.sources)
60            context.add_seeds(package.seeds)
61            context.add_models(package.models)
62            context.jinja_macros.add_macros(
63                package.macros,
64                package=package_name if package_name != context.project_name else None,
65            )
66
67        # First render all the config and discover dependencies
68        for package in project.packages.values():
69            context.variables = package.variables
70
71            package.sources = {k: v.render_config(context) for k, v in package.sources.items()}
72            package.seeds = {k: v.render_config(context) for k, v in package.seeds.items()}
73            package.models = {k: v.render_config(context) for k, v in package.models.items()}
74
75            context.add_sources(package.sources)
76            context.add_seeds(package.seeds)
77            context.add_models(package.models)
78
79        # Now that config is rendered, create the sqlmesh models
80        for package in project.packages.values():
81            context.variables = package.variables
82            package_models: t.Dict[str, BaseModelConfig] = {**package.models, **package.seeds}
83            models.update(
84                {model.model_name: model.to_sqlmesh(context) for model in package_models.values()}
85            )
86
87        return models
88
89    def _load_audits(self) -> UniqueKeyDict[str, Audit]:
90        return UniqueKeyDict("audits")
def sqlmesh_config( project_root: Optional[pathlib.Path] = None, **kwargs: Any) -> sqlmesh.core.config.root.Config:
20def sqlmesh_config(project_root: t.Optional[Path] = None, **kwargs: t.Any) -> Config:
21    project_root = project_root or Path()
22    context = DbtContext(project_root=project_root)
23    profile = Profile.load(context)
24
25    return Config(
26        default_connection=profile.default_target,
27        connections=profile.to_sqlmesh(),
28        loader=DbtLoader,
29        **kwargs,
30    )
class DbtLoader(sqlmesh.core.loader.Loader):
33class DbtLoader(Loader):
34    def _load_scripts(self) -> t.Tuple[MacroRegistry, HookRegistry]:
35        macro_files = list(Path(self._context.path, "macros").glob("**/*.sql"))
36        for file in macro_files:
37            self._track_file(file)
38
39        return (
40            UniqueKeyDict("macros"),
41            UniqueKeyDict("hooks"),
42        )
43
44    def _load_models(
45        self,
46        macros: MacroRegistry,
47        hooks: HookRegistry,
48    ) -> UniqueKeyDict[str, Model]:
49        models: UniqueKeyDict = UniqueKeyDict("models")
50
51        project = Project.load(
52            DbtContext(project_root=self._context.path, target_name=self._context.connection)
53        )
54        for path in project.project_files:
55            self._track_file(path)
56
57        context = project.context.copy()
58
59        for package_name, package in project.packages.items():
60            context.add_sources(package.sources)
61            context.add_seeds(package.seeds)
62            context.add_models(package.models)
63            context.jinja_macros.add_macros(
64                package.macros,
65                package=package_name if package_name != context.project_name else None,
66            )
67
68        # First render all the config and discover dependencies
69        for package in project.packages.values():
70            context.variables = package.variables
71
72            package.sources = {k: v.render_config(context) for k, v in package.sources.items()}
73            package.seeds = {k: v.render_config(context) for k, v in package.seeds.items()}
74            package.models = {k: v.render_config(context) for k, v in package.models.items()}
75
76            context.add_sources(package.sources)
77            context.add_seeds(package.seeds)
78            context.add_models(package.models)
79
80        # Now that config is rendered, create the sqlmesh models
81        for package in project.packages.values():
82            context.variables = package.variables
83            package_models: t.Dict[str, BaseModelConfig] = {**package.models, **package.seeds}
84            models.update(
85                {model.model_name: model.to_sqlmesh(context) for model in package_models.values()}
86            )
87
88        return models
89
90    def _load_audits(self) -> UniqueKeyDict[str, Audit]:
91        return UniqueKeyDict("audits")

Abstract base class to load macros and models for a context