feat: [SNOW-1890085] add run_async option

Author: sfc-gh-jwilkowski

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

@@ -26,7 +26,7 @@
 from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
 from snowflake.cli.api.feature_flags import FeatureFlag
 from snowflake.cli.api.identifiers import FQN
-from snowflake.cli.api.output.types import CommandResult, QueryResult
+from snowflake.cli.api.output.types import CommandResult, MessageResult, QueryResult
 
 app = SnowTyperFactory(
     name="dbt",
@@ -112,6 +112,9 @@ def deploy_dbt(
 @global_options_with_connection
 def before_callback(
     name: FQN = DBTNameArgument,
+    run_async: Optional[bool] = typer.Option(
+        False, help="Run dbt command asynchronously and check it's result later."
+    ),
     **options,
 ):
     """Handles global options passed before the command and takes pipeline name to be accessed through child context later"""
@@ -134,4 +137,10 @@ def _dbt_execute(
         dbt_cli_args = ctx.args
         dbt_command = ctx.command.name
         name = ctx.parent.params["name"]
-        return QueryResult(DBTManager().execute(dbt_command, name, *dbt_cli_args))
+        run_async = ctx.parent.params["run_async"]
+        result = DBTManager().execute(dbt_command, name, run_async, *dbt_cli_args)
+        if not run_async:
+            return QueryResult(result)
+        return MessageResult(
+            f"Command submitted. You can check the result with `snow sql -q \"select execution_status from table(information_schema.query_history_by_user()) where query_id in ('{result.sfqid}');\"`"
+        )

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

@@ -75,8 +75,8 @@ def deploy(
                 query += f"\nWAREHOUSE='{execute_in_warehouse}'"
             return self.execute_query(query)
 
-    def execute(self, dbt_command: str, name: str, *dbt_cli_args):
+    def execute(self, dbt_command: str, name: str, run_async: bool, *dbt_cli_args):
         if dbt_cli_args:
             dbt_command = dbt_command + " " + " ".join([arg for arg in dbt_cli_args])
         query = f"EXECUTE DBT PROJECT {name} args='{dbt_command.strip()}'"
-        return self.execute_query(query)
+        return self.execute_query(query, _exec_async=run_async)

tests/dbt/test_dbt_commands.py

@@ -267,4 +267,24 @@ 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.kwargs[0]["_exec_async"] is False
         assert mock_connect.mocked_ctx.get_query() == expected_query
+
+    def test_execute_async(self, mock_connect, runner):
+        result = runner.invoke(
+            [
+                "dbt",
+                "execute",
+                "--run-async",
+                "pipeline_name",
+                "compile",
+            ]
+        )
+
+        assert result.exit_code == 0, result.output
+        assert result.output.startswith("Command submitted")
+        assert mock_connect.mocked_ctx.kwargs[0]["_exec_async"] is True
+        assert (
+            mock_connect.mocked_ctx.get_query()
+            == "EXECUTE DBT PROJECT pipeline_name args='compile'"
+        )