feat: [SNOW-1890085] dbt deploy: add execute_in_warehouse flag

Author: sfc-gh-jwilkowski

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

@@ -73,6 +73,9 @@ def deploy_dbt(
     dbt_adapter_version: str = typer.Option(
         help="dbt-snowflake adapter version to be used",
     ),
+    execute_in_warehouse: Optional[str] = typer.Option(
+        None, help="Warehouse to use when running `dbt execute` commands"
+    ),
     **options,
 ) -> CommandResult:
     """
@@ -85,7 +88,12 @@ def deploy_dbt(
         path = Path(source)
     return QueryResult(
         DBTManager().deploy(
-            path.resolve(), name, dbt_version, dbt_adapter_version, force=force
+            path.resolve(),
+            name,
+            dbt_version,
+            dbt_adapter_version,
+            execute_in_warehouse,
+            force=force,
         )
     )
 

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

@@ -50,6 +50,7 @@ def deploy(
         name: FQN,
         dbt_version: Optional[str],
         dbt_adapter_version: str,
+        execute_in_warehouse: Optional[str],
         force: bool,
     ) -> SnowflakeCursor:
         dbt_project_path = path.joinpath("dbt_project.yml")
@@ -81,6 +82,8 @@ def deploy(
             query = f"""{'CREATE OR REPLACE' if force is True else 'CREATE'} DBT PROJECT {name}
 FROM {stage_name} MAIN_FILE='{staged_dbt_project_path}'
 DBT_VERSION='{dbt_version}' DBT_ADAPTER_VERSION='{dbt_adapter_version}'"""
+            if execute_in_warehouse:
+                query += f" WAREHOUSE='{execute_in_warehouse}'"
             return self.execute_query(query)
 
     def execute(self, dbt_command: str, name: str, *dbt_cli_args):

tests/dbt/test_dbt_commands.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+from textwrap import dedent
 from unittest import mock
 
 import pytest
@@ -128,6 +129,30 @@ def test_force_flag_uses_create_or_replace(
             "CREATE OR REPLACE DBT PROJECT"
         )
 
+    @mock.patch("snowflake.cli._plugins.dbt.manager.StageManager.put_recursive")
+    @mock.patch("snowflake.cli._plugins.dbt.manager.StageManager.create")
+    def test_execute_in_warehouse(
+        self, _mock_create, _mock_put_recursive, mock_connect, runner, dbt_project_path
+    ):
+
+        result = runner.invoke(
+            [
+                "dbt",
+                "deploy",
+                "TEST_PIPELINE",
+                f"--source={dbt_project_path}",
+                "--dbt-adapter-version=3.4.5",
+                "--execute-in-warehouse=XL",
+            ]
+        )
+
+        assert result.exit_code == 0, result.output
+        assert mock_connect.mocked_ctx.get_query() == dedent(
+            """CREATE DBT PROJECT TEST_PIPELINE
+FROM @MockDatabase.MockSchema.dbt_TEST_PIPELINE_stage MAIN_FILE='@MockDatabase.MockSchema.dbt_TEST_PIPELINE_stage/dbt_project.yml'
+DBT_VERSION='1.2.3' DBT_ADAPTER_VERSION='3.4.5' WAREHOUSE='XL'"""
+        )
+
     def test_raises_when_dbt_project_is_not_available(
         self, dbt_project_path, mock_connect, runner
     ):