Add --release-tag option to pull CLI command (#3343)

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: [email protected] <[email protected]>

Author: devin-ai-integration[bot]

ee/vellum_cli/__init__.py

@@ -208,13 +208,19 @@ def get_command(self, ctx, cmd_name):
     help="""Directory to pull the workflow into. If not specified, \
 the workflow will be pulled into the current working directory.""",
 )
+@click.option(
+    "--release-tag",
+    type=str,
+    help="""Release tag to use when pulling from a deployment.""",
+)
 def pull(
     ctx: click.Context,
     include_json: Optional[bool],
     exclude_code: Optional[bool],
     strict: Optional[bool],
     include_sandbox: Optional[bool],
     target_directory: Optional[str],
+    release_tag: Optional[str],
 ) -> None:
     """Pull Resources from Vellum"""
 
@@ -225,6 +231,7 @@ def pull(
             strict=strict,
             include_sandbox=include_sandbox,
             target_directory=target_directory,
+            release_tag=release_tag,
         )
 
 
@@ -271,6 +278,11 @@ def pull(
     type=str,
     help="The specific Workspace config to use when pulling",
 )
+@click.option(
+    "--release-tag",
+    type=str,
+    help="""Release tag to use when pulling from a deployment.""",
+)
 def workflows_pull(
     module: Optional[str],
     include_json: Optional[bool],
@@ -281,6 +293,7 @@ def workflows_pull(
     include_sandbox: Optional[bool],
     target_directory: Optional[str],
     workspace: Optional[str],
+    release_tag: Optional[str],
 ) -> None:
     """
     Pull Workflows from Vellum. If a module is provided, only the Workflow for that module will be pulled.
@@ -297,6 +310,7 @@ def workflows_pull(
         include_sandbox=include_sandbox,
         target_directory=target_directory,
         workspace=workspace,
+        release_tag=release_tag,
     )
 
 
@@ -332,13 +346,19 @@ def workflows_pull(
     help="""Directory to pull the workflow into. If not specified, \
 the workflow will be pulled into the current working directory.""",
 )
+@click.option(
+    "--release-tag",
+    type=str,
+    help="""Release tag to use when pulling from a deployment.""",
+)
 def pull_module(
     ctx: click.Context,
     include_json: Optional[bool],
     exclude_code: Optional[bool],
     strict: Optional[bool],
     include_sandbox: Optional[bool],
     target_directory: Optional[str],
+    release_tag: Optional[str],
 ) -> None:
     """Pull a specific module from Vellum"""
 
@@ -350,6 +370,7 @@ def pull_module(
             strict=strict,
             include_sandbox=include_sandbox,
             target_directory=target_directory,
+            release_tag=release_tag,
         )
 
 

ee/vellum_cli/pull.py

@@ -145,6 +145,7 @@ def pull_command(
     include_sandbox: Optional[bool] = None,
     target_directory: Optional[str] = None,
     workspace: Optional[str] = None,
+    release_tag: Optional[str] = None,
 ) -> None:
     load_dotenv(dotenv_path=os.path.join(os.getcwd(), ".env"))
     logger = load_cli_logger()
@@ -198,6 +199,7 @@ def pull_command(
 
     response = client.workflows.pull(
         pk,
+        release_tag=release_tag,
         request_options={"additional_query_parameters": query_parameters},
     )
 

ee/vellum_cli/tests/test_pull.py

@@ -1337,3 +1337,33 @@ def test_pull__workspace_option__uses_different_api_key(mock_module, vellum_clie
     assert os.path.exists(workflow_py)
     with open(workflow_py) as f:
         assert f.read() == "print('hello')"
+
+
+@pytest.mark.parametrize(
+    "base_command",
+    [
+        ["pull"],
+        ["workflows", "pull"],
+    ],
+    ids=["pull", "workflows_pull"],
+)
+def test_pull__release_tag(vellum_client, mock_module, base_command):
+    """Tests that the --release-tag option is passed to the API."""
+
+    # GIVEN a module on the user's filesystem
+    module = mock_module.module
+
+    # AND the workflow pull API call returns a zip file
+    vellum_client.workflows.pull.return_value = iter([zip_file_map({"workflow.py": "print('hello')"})])
+
+    # WHEN the user runs the pull command with --release-tag
+    runner = CliRunner()
+    result = runner.invoke(cli_main, base_command + [module, "--release-tag", "my-release-tag"])
+
+    # THEN the command returns successfully
+    assert result.exit_code == 0
+
+    # AND the pull api is called with release_tag="my-release-tag"
+    vellum_client.workflows.pull.assert_called_once()
+    call_args = vellum_client.workflows.pull.call_args.kwargs
+    assert call_args["release_tag"] == "my-release-tag"