Jw/snow 2036324 include dbt profile (#2196)

* feat: [SNOW-2036324] validate dbt profile

* refactor: [SNOW-2036324] PR fixes

Author: sfc-gh-jwilkowski

src/snowflake/cli/_plugins/dbt/manager.py

@@ -14,11 +14,12 @@
 
 from __future__ import annotations
 
-from click import ClickException
+import yaml
 from snowflake.cli._plugins.object.manager import ObjectManager
 from snowflake.cli._plugins.stage.manager import StageManager
 from snowflake.cli.api.console import cli_console
-from snowflake.cli.api.constants import ObjectType
+from snowflake.cli.api.constants import DEFAULT_SIZE_LIMIT_MB, ObjectType
+from snowflake.cli.api.exceptions import CliError
 from snowflake.cli.api.identifiers import FQN
 from snowflake.cli.api.secure_path import SecurePath
 from snowflake.cli.api.sql_execution import SqlExecutionMixin
@@ -44,12 +45,30 @@ def deploy(
     ) -> SnowflakeCursor:
         dbt_project_path = path / "dbt_project.yml"
         if not dbt_project_path.exists():
-            raise ClickException(
+            raise CliError(
                 f"dbt_project.yml does not exist in directory {path.path.absolute()}."
             )
 
+        with dbt_project_path.open(read_file_limit_mb=DEFAULT_SIZE_LIMIT_MB) as fd:
+            dbt_project = yaml.safe_load(fd)
+            try:
+                profile = dbt_project["profile"]
+            except KeyError:
+                raise CliError("`profile` is not defined in dbt_project.yml")
+
+        dbt_profiles_path = path / "profiles.yml"
+        if not dbt_profiles_path.exists():
+            raise CliError(
+                f"profiles.yml does not exist in directory {path.path.absolute()}."
+            )
+
+        with dbt_profiles_path.open(read_file_limit_mb=DEFAULT_SIZE_LIMIT_MB) as fd:
+            profiles = yaml.safe_load(fd)
+            if profile not in profiles:
+                raise CliError(f"profile {profile} is not defined in profiles.yml")
+
         if self.exists(name=name) and force is not True:
-            raise ClickException(
+            raise CliError(
                 f"DBT project {name} already exists. Use --force flag to overwrite"
             )
 

tests/dbt/test_dbt_commands.py

@@ -17,6 +17,7 @@
 from unittest import mock
 
 import pytest
+import yaml
 from snowflake.cli._plugins.dbt.constants import OUTPUT_COLUMN_NAME, RESULT_COLUMN_NAME
 from snowflake.cli.api.identifiers import FQN
 
@@ -56,8 +57,20 @@ class TestDBTDeploy:
     @pytest.fixture
     def dbt_project_path(self, tmp_path_factory):
         source_path = tmp_path_factory.mktemp("dbt_project")
-        dbt_file = source_path / "dbt_project.yml"
-        dbt_file.touch()
+        dbt_project_file = source_path / "dbt_project.yml"
+        dbt_project_file.write_text(yaml.dump({"profile": "dev"}))
+        dbt_profiles_file = source_path / "profiles.yml"
+        dbt_profiles_file.write_text(
+            yaml.dump(
+                {
+                    "dev": {
+                        "outputs": {
+                            "local": {"account": "test_account", "database": "testdb"}
+                        }
+                    }
+                },
+            )
+        )
         yield source_path
 
     @pytest.fixture
@@ -105,68 +118,78 @@ def test_deploys_project_from_source(
             dbt_project_path, "@MockDatabase.MockSchema.dbt_TEST_PIPELINE_stage"
         )
 
+    @pytest.mark.parametrize("exists", (True, False))
     @mock.patch("snowflake.cli._plugins.dbt.manager.StageManager.put_recursive")
     @mock.patch("snowflake.cli._plugins.dbt.manager.StageManager.create")
-    def test_dbt_version_from_option_has_precedence_over_file(
+    def test_force_flag_uses_create_or_replace(
         self,
         _mock_create,
         _mock_put_recursive,
+        exists,
         mock_connect,
         runner,
         dbt_project_path,
         mock_exists,
     ):
+        mock_exists.return_value = exists
+
         result = runner.invoke(
             [
                 "dbt",
                 "deploy",
                 "TEST_PIPELINE",
                 f"--source={dbt_project_path}",
+                "--force",
             ]
         )
 
         assert result.exit_code == 0, result.output
-        assert (
-            mock_connect.mocked_ctx.get_query()
-            == """CREATE DBT PROJECT TEST_PIPELINE
-FROM @MockDatabase.MockSchema.dbt_TEST_PIPELINE_stage"""
+        assert mock_connect.mocked_ctx.get_query().startswith(
+            "CREATE OR REPLACE DBT PROJECT"
         )
 
-    @pytest.mark.parametrize("exists", (True, False))
-    @mock.patch("snowflake.cli._plugins.dbt.manager.StageManager.put_recursive")
-    @mock.patch("snowflake.cli._plugins.dbt.manager.StageManager.create")
-    def test_force_flag_uses_create_or_replace(
-        self,
-        _mock_create,
-        _mock_put_recursive,
-        exists,
-        mock_connect,
-        runner,
-        dbt_project_path,
-        mock_exists,
+    def test_raises_when_dbt_project_yml_is_not_available(
+        self, dbt_project_path, mock_connect, runner
     ):
-        mock_exists.return_value = exists
+        dbt_file = dbt_project_path / "dbt_project.yml"
+        dbt_file.unlink()
 
         result = runner.invoke(
             [
                 "dbt",
                 "deploy",
                 "TEST_PIPELINE",
                 f"--source={dbt_project_path}",
-                "--force",
-            ]
+            ],
         )
 
-        assert result.exit_code == 0, result.output
-        assert mock_connect.mocked_ctx.get_query().startswith(
-            "CREATE OR REPLACE DBT PROJECT"
+        assert result.exit_code == 1, result.output
+        assert f"dbt_project.yml does not exist in directory" in result.output
+        assert mock_connect.mocked_ctx.get_query() == ""
+
+    def test_raises_when_dbt_project_yml_does_not_specify_profile(
+        self, dbt_project_path, mock_connect, runner
+    ):
+        with open((dbt_project_path / "dbt_project.yml"), "w") as f:
+            yaml.dump({}, f)
+
+        result = runner.invoke(
+            [
+                "dbt",
+                "deploy",
+                "TEST_PIPELINE",
+                f"--source={dbt_project_path}",
+            ],
         )
 
-    def test_raises_when_dbt_project_is_not_available(
+        assert result.exit_code == 1, result.output
+        assert "`profile` is not defined in dbt_project.yml" in result.output
+        assert mock_connect.mocked_ctx.get_query() == ""
+
+    def test_raises_when_profiles_yml_is_not_available(
         self, dbt_project_path, mock_connect, runner
     ):
-        dbt_file = dbt_project_path / "dbt_project.yml"
-        dbt_file.unlink()
+        (dbt_project_path / "profiles.yml").unlink()
 
         result = runner.invoke(
             [
@@ -178,10 +201,29 @@ def test_raises_when_dbt_project_is_not_available(
         )
 
         assert result.exit_code == 1, result.output
-        assert f"dbt_project.yml does not exist in directory" in result.output
+        assert f"profiles.yml does not exist in directory" in result.output
+        assert mock_connect.mocked_ctx.get_query() == ""
+
+    def test_raises_when_profiles_yml_does_not_contain_selected_profile(
+        self, dbt_project_path, mock_connect, runner
+    ):
+        with open((dbt_project_path / "profiles.yml"), "w") as f:
+            yaml.dump({}, f)
+
+        result = runner.invoke(
+            [
+                "dbt",
+                "deploy",
+                "TEST_PIPELINE",
+                f"--source={dbt_project_path}",
+            ],
+        )
+
+        assert result.exit_code == 1, result.output
+        assert "profile dev is not defined in profiles.yml" in result.output
         assert mock_connect.mocked_ctx.get_query() == ""
 
-    def test_raises_when_dbt_project_exists_and_is_not_force(
+    def test_raises_when_dbt_object_exists_and_is_not_force(
         self, dbt_project_path, mock_connect, runner, mock_exists
     ):
         mock_exists.return_value = True

tests_integration/test_data/projects/dbt_project/profiles.yml

@@ -0,0 +1,5 @@
+dbt_integration_project:
+  target: dev
+  outputs: 
+    dev:
+      type: snowflake

tests_integration/test_dbt.py

@@ -14,20 +14,22 @@
 import datetime
 
 import pytest
+import yaml
 
 
 @pytest.mark.integration
 @pytest.mark.qa_only
-def test_dbt_deploy(
+def test_dbt(
     runner,
     snowflake_session,
     test_database,
     project_directory,
 ):
-    with project_directory("dbt_project"):
+    with project_directory("dbt_project") as root_dir:
         # Given a local dbt project
         ts = int(datetime.datetime.now().timestamp())
         name = f"dbt_project_{ts}"
+        _setup_dbt_profile(root_dir, snowflake_session)
 
         # When it's deployed
         result = runner.invoke_with_connection_json(["dbt", "deploy", name])
@@ -65,3 +67,16 @@ def test_dbt_deploy(
         )
         assert len(result.json) == 1, result.json
         assert result.json[0]["COUNT"] == 1, result.json[0]
+
+
+def _setup_dbt_profile(root_dir, snowflake_session):
+    with open((root_dir / "profiles.yml"), "r") as f:
+        profiles = yaml.safe_load(f)
+    dev_profile = profiles["dbt_integration_project"]["outputs"]["dev"]
+    dev_profile["database"] = snowflake_session.database
+    dev_profile["account"] = snowflake_session.account
+    dev_profile["user"] = snowflake_session.user
+    dev_profile["role"] = snowflake_session.role
+    dev_profile["warehouse"] = snowflake_session.warehouse
+    dev_profile["schema"] = snowflake_session.schema
+    (root_dir / "profiles.yml").write_text(yaml.dump(profiles))