Feat(sqlmesh_dbt): Add support for --project-dir and --profiles-dir (#5524)

Author: erindru

sqlmesh/core/config/loader.py

@@ -178,6 +178,7 @@ def load_config_from_paths(
 
             dbt_python_config = sqlmesh_config(
                 project_root=dbt_project_file.parent,
+                profiles_dir=kwargs.pop("profiles_dir", None),
                 dbt_profile_name=kwargs.pop("profile", None),
                 dbt_target_name=kwargs.pop("target", None),
                 variables=variables,

sqlmesh/dbt/context.py

@@ -37,6 +37,8 @@ class DbtContext:
     """Context for DBT environment"""
 
     project_root: Path = Path()
+    profiles_dir: t.Optional[Path] = None
+    """Optional override to specify the directory where profiles.yml is located, if not at the :project_root"""
     target_name: t.Optional[str] = None
     profile_name: t.Optional[str] = None
     project_schema: t.Optional[str] = None

sqlmesh/dbt/loader.py

@@ -53,10 +53,18 @@ def sqlmesh_config(
     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:
@@ -98,6 +106,7 @@ def sqlmesh_config(
 
     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),
@@ -116,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:
@@ -225,6 +237,7 @@ def _load_projects(self) -> t.List[Project]:
             project = Project.load(
                 DbtContext(
                     project_root=self.config_path,
+                    profiles_dir=self._profiles_dir,
                     target_name=target_name,
                     sqlmesh_config=self.config,
                 ),

sqlmesh/dbt/profile.py

@@ -60,16 +60,16 @@ def load(cls, context: DbtContext, target_name: t.Optional[str] = None) -> Profi
             if not context.profile_name:
                 raise ConfigError(f"{project_file.stem} must include project name.")
 
-        profile_filepath = cls._find_profile(context.project_root)
+        profile_filepath = cls._find_profile(context.project_root, context.profiles_dir)
         if not profile_filepath:
             raise ConfigError(f"{cls.PROFILE_FILE} not found.")
 
         target_name, target = cls._read_profile(profile_filepath, context, target_name)
         return Profile(profile_filepath, target_name, target)
 
     @classmethod
-    def _find_profile(cls, project_root: Path) -> t.Optional[Path]:
-        dir = os.environ.get("DBT_PROFILES_DIR", "")
+    def _find_profile(cls, project_root: Path, profiles_dir: t.Optional[Path]) -> t.Optional[Path]:
+        dir = os.environ.get("DBT_PROFILES_DIR", profiles_dir or "")
         path = Path(project_root, dir, cls.PROFILE_FILE)
         if path.exists():
             return path

sqlmesh_dbt/cli.py

@@ -84,6 +84,16 @@ def _cleanup() -> None:
     type=click.Choice(["debug", "info", "warn", "error", "none"]),
     help="Specify the minimum severity of events that are logged to the console and the log file.",
 )
+@click.option(
+    "--profiles-dir",
+    type=click.Path(exists=True, file_okay=False, path_type=Path),
+    help="Which directory to look in for the profiles.yml file. If not set, dbt will look in the current working directory first, then HOME/.dbt/",
+)
+@click.option(
+    "--project-dir",
+    type=click.Path(exists=True, file_okay=False, path_type=Path),
+    help="Which directory to look in for the dbt_project.yml file. Default is the current working directory and its parents.",
+)
 @click.pass_context
 @cli_global_error_handler
 def dbt(
@@ -92,6 +102,8 @@ def dbt(
     target: t.Optional[str] = None,
     debug: bool = False,
     log_level: t.Optional[str] = None,
+    profiles_dir: t.Optional[Path] = None,
+    project_dir: t.Optional[Path] = None,
 ) -> None:
     """
     An ELT tool for managing your SQL transformations and data models, powered by the SQLMesh engine.
@@ -105,7 +117,8 @@ def dbt(
     # that need to be known before we attempt to load the project
     ctx.obj = functools.partial(
         create,
-        project_dir=Path.cwd(),
+        project_dir=project_dir,
+        profiles_dir=profiles_dir,
         profile=profile,
         target=target,
         debug=debug,

sqlmesh_dbt/operations.py

@@ -232,6 +232,7 @@ def close(self) -> None:
 
 def create(
     project_dir: t.Optional[Path] = None,
+    profiles_dir: t.Optional[Path] = None,
     profile: t.Optional[str] = None,
     target: t.Optional[str] = None,
     vars: t.Optional[t.Dict[str, t.Any]] = None,
@@ -268,7 +269,11 @@ def create(
         sqlmesh_context = Context(
             paths=[project_dir],
             config_loader_kwargs=dict(
-                profile=profile, target=target, variables=vars, threads=threads
+                profile=profile,
+                target=target,
+                variables=vars,
+                threads=threads,
+                profiles_dir=profiles_dir,
             ),
             load=True,
             # DbtSelector selects based on dbt model fqn's rather than SQLMesh model names

tests/dbt/cli/test_global_flags.py

@@ -107,3 +107,77 @@ def test_log_level(invoke_cli: t.Callable[..., Result], create_empty_project: Em
     result = invoke_cli(["--log-level", "debug", "list"])
     assert result.exit_code == 0
     assert logging.getLogger("sqlmesh").getEffectiveLevel() == logging.DEBUG
+
+
+def test_profiles_dir(
+    invoke_cli: t.Callable[..., Result], create_empty_project: EmptyProjectCreator, tmp_path: Path
+):
+    project_dir, _ = create_empty_project(project_name="test_profiles_dir")
+
+    orig_profiles_yml = project_dir / "profiles.yml"
+    assert orig_profiles_yml.exists()
+
+    new_profiles_yml = tmp_path / "some_other_place" / "profiles.yml"
+    new_profiles_yml.parent.mkdir(parents=True)
+
+    orig_profiles_yml.rename(new_profiles_yml)
+    assert not orig_profiles_yml.exists()
+    assert new_profiles_yml.exists()
+
+    # should fail if we don't specify --profiles-dir
+    result = invoke_cli(["list"])
+    assert result.exit_code > 0, result.output
+    assert "profiles.yml not found" in result.output
+
+    # should pass if we specify --profiles-dir
+    result = invoke_cli(["--profiles-dir", str(new_profiles_yml.parent), "list"])
+    assert result.exit_code == 0, result.output
+    assert "Models in project" in result.output
+
+
+def test_project_dir(
+    invoke_cli: t.Callable[..., Result], create_empty_project: EmptyProjectCreator
+):
+    orig_project_dir, _ = create_empty_project(project_name="test_project_dir")
+
+    orig_project_yml = orig_project_dir / "dbt_project.yml"
+    assert orig_project_yml.exists()
+
+    new_project_yml = orig_project_dir / "nested" / "dbt_project.yml"
+    new_project_yml.parent.mkdir(parents=True)
+
+    orig_project_yml.rename(new_project_yml)
+    assert not orig_project_yml.exists()
+    assert new_project_yml.exists()
+
+    # should fail if we don't specify --project-dir
+    result = invoke_cli(["list"])
+    assert result.exit_code != 0, result.output
+    assert "Error:" in result.output
+
+    # should fail if the profiles.yml also doesnt exist at that --project-dir
+    result = invoke_cli(["--project-dir", str(new_project_yml.parent), "list"])
+    assert result.exit_code != 0, result.output
+    assert "profiles.yml not found" in result.output
+
+    # should pass if it can find both files, either because we specified --profiles-dir explicitly or the profiles.yml was found in --project-dir
+    result = invoke_cli(
+        [
+            "--project-dir",
+            str(new_project_yml.parent),
+            "--profiles-dir",
+            str(orig_project_dir),
+            "list",
+        ]
+    )
+    assert result.exit_code == 0, result.output
+    assert "Models in project" in result.output
+
+    orig_profiles_yml = orig_project_dir / "profiles.yml"
+    new_profiles_yml = new_project_yml.parent / "profiles.yml"
+    assert orig_profiles_yml.exists()
+    orig_profiles_yml.rename(new_profiles_yml)
+
+    result = invoke_cli(["--project-dir", str(new_project_yml.parent), "list"])
+    assert result.exit_code == 0, result.output
+    assert "Models in project" in result.output