feat: [SNOW-1890085] implement dbt deploy command

Author: sfc-gh-jwilkowski

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

@@ -15,13 +15,16 @@
 from __future__ import annotations
 
 import logging
-from typing import Annotated
+from pathlib import Path
+from typing import Optional
 
 import typer
 from snowflake.cli._plugins.dbt.constants import DBT_COMMANDS
 from snowflake.cli._plugins.dbt.manager import DBTManager
 from snowflake.cli.api.commands.decorators import global_options_with_connection
+from snowflake.cli.api.commands.flags import identifier_argument
 from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
+from snowflake.cli.api.identifiers import FQN
 from snowflake.cli.api.output.types import CommandResult, QueryResult
 
 app = SnowTyperFactory(
@@ -32,6 +35,9 @@
 log = logging.getLogger(__name__)
 
 
+DBTNameArgument = identifier_argument(sf_object="DBT Object", example="my_pipeline")
+
+
 @app.command(
     "list",
     requires_connection=True,
@@ -45,6 +51,34 @@ def list_dbts(
     return QueryResult(DBTManager().list())
 
 
+@app.command(
+    "deploy",
+    requires_connection=True,
+)
+def deploy_dbt(
+    name: FQN = DBTNameArgument,
+    source: Optional[str] = typer.Option(
+        help="Path to directory containing dbt files to deploy. Defaults to current working directory.",
+        show_default=False,
+        default=None,
+    ),
+    force: Optional[bool] = typer.Option(
+        False,
+        help="Overwrites conflicting files in the project, if any.",
+    ),
+    **options,
+) -> CommandResult:
+    """
+    Copy dbt files and create or update dbt on Snowflake project.
+    """
+    # TODO: options for DBT version?
+    if source is None:
+        path = Path.cwd()
+    else:
+        path = Path(source)
+    return QueryResult(DBTManager().deploy(path.resolve(), name, force=force))
+
+
 # `execute` is a pass through command group, meaning that all params after command should be passed over as they are,
 # suppressing usual CLI behaviour for displaying help or formatting options.
 dbt_execute_app = SnowTyperFactory(
@@ -57,9 +91,7 @@ def list_dbts(
 @dbt_execute_app.callback()
 @global_options_with_connection
 def before_callback(
-    name: Annotated[
-        str, typer.Argument(help="Name of the dbt object to execute command on.")
-    ],
+    name: FQN = DBTNameArgument,
     **options,
 ):
     """Handles global options passed before the command and takes pipeline name to be accessed through child context later"""

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

@@ -14,6 +14,12 @@
 
 from __future__ import annotations
 
+from pathlib import Path
+
+from click import ClickException
+from snowflake.cli._plugins.stage.manager import StageManager
+from snowflake.cli.api.console import cli_console
+from snowflake.cli.api.identifiers import FQN
 from snowflake.cli.api.sql_execution import SqlExecutionMixin
 from snowflake.connector.cursor import SnowflakeCursor
 
@@ -23,6 +29,25 @@ def list(self) -> SnowflakeCursor:  # noqa: A003
         query = "SHOW DBT"
         return self.execute_query(query)
 
+    def deploy(self, path: Path, name: FQN, force: bool) -> SnowflakeCursor:
+        # TODO: what to do with force?
+        if not path.joinpath("dbt_project.yml").exists():
+            raise ClickException(f"dbt_project.yml does not exist in provided path.")
+
+        with cli_console.phase("Creating temporary stage"):
+            stage_manager = StageManager()
+            stage_fqn = FQN.from_string(f"dbt_{name}_stage").using_context()
+            stage_name = stage_manager.get_standard_stage_prefix(stage_fqn)
+            stage_manager.create(stage_fqn, temporary=True)
+
+        with cli_console.phase("Copying project files to stage"):
+            results = list(stage_manager.put_recursive(path, stage_name))
+            cli_console.step(f"Copied {len(results)} files")
+
+        with cli_console.phase("Creating DBT project"):
+            query = f"CREATE OR REPLACE DBT {name} FROM {stage_name}"
+            return self.execute_query(query)
+
     def execute(self, dbt_command: str, name: str, *dbt_cli_args):
         query = f"EXECUTE DBT {name} {dbt_command}"
         if dbt_cli_args:

tests/dbt/test_dbt_commands.py

@@ -12,9 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 from unittest import mock
 
 import pytest
+from snowflake.cli.api.identifiers import FQN
 
 
 @pytest.fixture
@@ -26,76 +29,127 @@ def mock_connect(mock_ctx):
         yield _fixture
 
 
-def test_dbt_list(mock_connect, runner):
-
-    result = runner.invoke(["dbt", "list"])
-
-    assert result.exit_code == 0, result.output
-    assert mock_connect.mocked_ctx.get_query() == "SHOW DBT"
-
-
-@pytest.mark.parametrize(
-    "args,expected_query",
-    [
-        pytest.param(
-            [
-                "dbt",
-                "execute",
-                "pipeline_name",
-                "test",
-            ],
-            "EXECUTE DBT pipeline_name test",
-            id="simple-command",
-        ),
-        pytest.param(
-            [
-                "dbt",
-                "execute",
-                "pipeline_name",
-                "run",
-                "-f",
-                "--select @source:snowplow,tag:nightly models/export",
-            ],
-            "EXECUTE DBT pipeline_name run -f --select @source:snowplow,tag:nightly models/export",
-            id="with-dbt-options",
-        ),
-        pytest.param(
-            ["dbt", "execute", "pipeline_name", "compile", "--vars '{foo:bar}'"],
-            "EXECUTE DBT pipeline_name compile --vars '{foo:bar}'",
-            id="with-dbt-vars",
-        ),
-        pytest.param(
-            [
-                "dbt",
-                "execute",
-                "pipeline_name",
-                "compile",
-                "--format=TXT",  # collision with CLI's option; unsupported option
-                "-v",  # collision with CLI's option
-                "-h",
-                "--debug",
-                "--info",
-                "--config-file=/",
-            ],
-            "EXECUTE DBT pipeline_name compile --format=TXT -v -h --debug --info --config-file=/",
-            id="with-dbt-conflicting-options",
-        ),
-        pytest.param(
-            [
-                "dbt",
-                "execute",
-                "--format=JSON",
-                "pipeline_name",
-                "compile",
-            ],
-            "EXECUTE DBT pipeline_name compile",
-            id="with-cli-flag",
-        ),
-    ],
-)
-def test_dbt_execute(mock_connect, runner, args, expected_query):
-
-    result = runner.invoke(args)
-
-    assert result.exit_code == 0, result.output
-    assert mock_connect.mocked_ctx.get_query() == expected_query
+class TestDBTList:
+    def test_dbt_list(self, mock_connect, runner):
+
+        result = runner.invoke(["dbt", "list"])
+
+        assert result.exit_code == 0, result.output
+        assert mock_connect.mocked_ctx.get_query() == "SHOW DBT"
+
+
+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()
+        yield source_path
+
+    @pytest.fixture
+    def mock_cli_console(self):
+        with mock.patch("snowflake.cli.api.console") as _fixture:
+            yield _fixture
+
+    @mock.patch("snowflake.cli._plugins.dbt.manager.StageManager.put_recursive")
+    @mock.patch("snowflake.cli._plugins.dbt.manager.StageManager.create")
+    def test_deploys_project_from_source(
+        self, mock_create, mock_put_recursive, mock_connect, runner, dbt_project_path
+    ):
+
+        result = runner.invoke(
+            ["dbt", "deploy", "TEST_PIPELINE", f"--source={dbt_project_path}"]
+        )
+
+        assert result.exit_code == 0, result.output
+        assert (
+            mock_connect.mocked_ctx.get_query()
+            == "CREATE OR REPLACE DBT TEST_PIPELINE FROM @MockDatabase.MockSchema.dbt_TEST_PIPELINE_stage"
+        )
+        stage_fqn = FQN.from_string(f"dbt_TEST_PIPELINE_stage").using_context()
+        mock_create.assert_called_once_with(stage_fqn, temporary=True)
+        mock_put_recursive.assert_called_once_with(
+            dbt_project_path, "@MockDatabase.MockSchema.dbt_TEST_PIPELINE_stage"
+        )
+
+    def test_raises_when_dbt_project_is_not_available(
+        self, dbt_project_path, mock_connect, runner
+    ):
+        dbt_file = dbt_project_path / "dbt_project.yml"
+        dbt_file.unlink()
+
+        result = runner.invoke(
+            ["dbt", "deploy", "TEST_PIPELINE", f"--source={dbt_project_path}"]
+        )
+
+        assert result.exit_code == 1, result.output
+        assert "dbt_project.yml does not exist in provided path." in result.output
+        assert mock_connect.mocked_ctx.get_query() == ""
+
+
+class TestDBTExecute:
+    @pytest.mark.parametrize(
+        "args,expected_query",
+        [
+            pytest.param(
+                [
+                    "dbt",
+                    "execute",
+                    "pipeline_name",
+                    "test",
+                ],
+                "EXECUTE DBT pipeline_name test",
+                id="simple-command",
+            ),
+            pytest.param(
+                [
+                    "dbt",
+                    "execute",
+                    "pipeline_name",
+                    "run",
+                    "-f",
+                    "--select @source:snowplow,tag:nightly models/export",
+                ],
+                "EXECUTE DBT pipeline_name run -f --select @source:snowplow,tag:nightly models/export",
+                id="with-dbt-options",
+            ),
+            pytest.param(
+                ["dbt", "execute", "pipeline_name", "compile", "--vars '{foo:bar}'"],
+                "EXECUTE DBT pipeline_name compile --vars '{foo:bar}'",
+                id="with-dbt-vars",
+            ),
+            pytest.param(
+                [
+                    "dbt",
+                    "execute",
+                    "pipeline_name",
+                    "compile",
+                    "--format=TXT",  # collision with CLI's option; unsupported option
+                    "-v",  # collision with CLI's option
+                    "-h",
+                    "--debug",
+                    "--info",
+                    "--config-file=/",
+                ],
+                "EXECUTE DBT pipeline_name compile --format=TXT -v -h --debug --info --config-file=/",
+                id="with-dbt-conflicting-options",
+            ),
+            pytest.param(
+                [
+                    "dbt",
+                    "execute",
+                    "--format=JSON",
+                    "pipeline_name",
+                    "compile",
+                ],
+                "EXECUTE DBT pipeline_name compile",
+                id="with-cli-flag",
+            ),
+        ],
+    )
+    def test_dbt_execute(self, mock_connect, runner, args, expected_query):
+
+        result = runner.invoke(args)
+
+        assert result.exit_code == 0, result.output
+        assert mock_connect.mocked_ctx.get_query() == expected_query