diff --git a/flytekit/clis/sdk_in_container/launchplan.py b/flytekit/clis/sdk_in_container/launchplan.py new file mode 100644 index 0000000000..7a8bbb6926 --- /dev/null +++ b/flytekit/clis/sdk_in_container/launchplan.py @@ -0,0 +1,74 @@ +import click + +from flytekit.clis.sdk_in_container.helpers import get_and_save_remote_with_click_context +from flytekit.models.launch_plan import LaunchPlanState + +_launchplan_help = """ +The launchplan command activates or deactivates a specified or the latest version of the launchplan. +If ``--activate`` is chosen then the previous version of the launchplan will be deactivated. + +- ``launchplan`` refers to the name of the Launchplan +- ``launchplan_version`` is optional and should be a valid version for a Launchplan version. If not specified the latest will be used. +""" + + +@click.command("launchplan", help=_launchplan_help) +@click.option( + "-p", + "--project", + required=False, + type=str, + default="flytesnacks", + help="Fecth launchplan from this project", +) +@click.option( + "-d", + "--domain", + required=False, + type=str, + default="development", + help="Fetch launchplan from this domain", +) +@click.option( + "--activate/--deactivate", + required=True, + type=bool, + is_flag=True, + help="Activate or Deactivate the launchplan", +) +@click.argument( + "launchplan", + required=True, + type=str, +) +@click.argument( + "launchplan-version", + required=False, + type=str, + default=None, +) +@click.pass_context +def launchplan( + ctx: click.Context, + project: str, + domain: str, + activate: bool, + launchplan: str, + launchplan_version: str, +): + remote = get_and_save_remote_with_click_context(ctx, project, domain) + try: + launchplan = remote.fetch_launch_plan( + project=project, + domain=domain, + name=launchplan, + version=launchplan_version, + ) + state = LaunchPlanState.ACTIVE if activate else LaunchPlanState.INACTIVE + remote.client.update_launch_plan(id=launchplan.id, state=state) + click.secho( + f"\n Launchplan was set to {LaunchPlanState.enum_to_string(state)}: {launchplan.name}:{launchplan.id.version}", + fg="green", + ) + except StopIteration as e: + click.secho(f"{e.value}", fg="red") diff --git a/flytekit/clis/sdk_in_container/pyflyte.py b/flytekit/clis/sdk_in_container/pyflyte.py index 68726e5106..94c6cbd075 100644 --- a/flytekit/clis/sdk_in_container/pyflyte.py +++ b/flytekit/clis/sdk_in_container/pyflyte.py @@ -9,6 +9,7 @@ from flytekit.clis.sdk_in_container.build import build from flytekit.clis.sdk_in_container.constants import CTX_CONFIG_FILE, CTX_PACKAGES, CTX_VERBOSE from flytekit.clis.sdk_in_container.init import init +from flytekit.clis.sdk_in_container.launchplan import launchplan from flytekit.clis.sdk_in_container.local_cache import local_cache from flytekit.clis.sdk_in_container.package import package from flytekit.clis.sdk_in_container.register import register @@ -134,6 +135,7 @@ def main(ctx, pkgs: typing.List[str], config: str, verbose: bool): main.add_command(register) main.add_command(backfill) main.add_command(build) +main.add_command(launchplan) main.epilog if __name__ == "__main__": diff --git a/tests/flytekit/unit/cli/pyflyte/test_launchplan.py b/tests/flytekit/unit/cli/pyflyte/test_launchplan.py new file mode 100644 index 0000000000..1a461bfd35 --- /dev/null +++ b/tests/flytekit/unit/cli/pyflyte/test_launchplan.py @@ -0,0 +1,34 @@ +import pytest +from click.testing import CliRunner +from mock import mock + +from flytekit.clis.sdk_in_container import pyflyte +from flytekit.remote import FlyteRemote + + +@mock.patch("flytekit.clis.sdk_in_container.helpers.FlyteRemote", spec=FlyteRemote) +@pytest.mark.parametrize( + ("action", "expected_state"), + [ + ("activate", "ACTIVE"), + ("deactivate", "INACTIVE"), + ], +) +def test_pyflyte_launchplan(mock_remote, action, expected_state): + mock_remote.generate_console_url.return_value = "ex" + runner = CliRunner() + with runner.isolated_filesystem(): + result = runner.invoke( + pyflyte.main, + [ + "launchplan", + f"--{action}", + "-p", + "flytesnacks", + "-d", + "development", + "daily", + ], + ) + assert result.exit_code == 0 + assert f"Launchplan was set to {expected_state}: " in result.output