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 )
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