Skip to content

Commit

Permalink
add activate-launchplan command to pyflyte
Browse files Browse the repository at this point in the history
Signed-off-by: Artem Petrov <[email protected]>
  • Loading branch information
wckdman committed Apr 13, 2023
1 parent e3cee83 commit 28471f6
Show file tree
Hide file tree
Showing 3 changed files with 88 additions and 0 deletions.
61 changes: 61 additions & 0 deletions flytekit/clis/sdk_in_container/activate_launchplan.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,61 @@
import click
from flytekit.models.launch_plan import LaunchPlanState

from flytekit.clis.sdk_in_container.helpers import get_and_save_remote_with_click_context

_activate_launchplan_help = """
The activate-launchplan command activates a specified or the latest version of the launchplan and disables a previous version.
- ``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("activate-launchplan", help=_activate_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.argument(
"launchplan",
required=True,
type=str,
)
@click.argument(
"launchplan-version",
required=False,
type=str,
default=None,
)
@click.pass_context
def activate_launchplan(
ctx: click.Context,
project: str,
domain: str,
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,
)
remote.client.update_launch_plan(id=launchplan.id, state=LaunchPlanState.ACTIVE)
click.secho(f"\n Launchplan was activated: {launchplan.name}:{launchplan.id.version}", fg="green")
except StopIteration as e:
click.secho(f"{e.value}", fg="red")
2 changes: 2 additions & 0 deletions flytekit/clis/sdk_in_container/pyflyte.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
from google.protobuf.json_format import MessageToJson

from flytekit import configuration
from flytekit.clis.sdk_in_container.activate_launchplan import activate_launchplan
from flytekit.clis.sdk_in_container.backfill import backfill
from flytekit.clis.sdk_in_container.constants import CTX_CONFIG_FILE, CTX_PACKAGES, CTX_VERBOSE
from flytekit.clis.sdk_in_container.init import init
Expand Down Expand Up @@ -132,6 +133,7 @@ def main(ctx, pkgs: typing.List[str], config: str, verbose: bool):
main.add_command(run)
main.add_command(register)
main.add_command(backfill)
main.add_command(activate_launchplan)
main.epilog

if __name__ == "__main__":
Expand Down
25 changes: 25 additions & 0 deletions tests/flytekit/unit/cli/pyflyte/test_activate_launchplan.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
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)
def test_pyflyte_activate_launchplan(mock_remote):
mock_remote.generate_console_url.return_value = "ex"
runner = CliRunner()
with runner.isolated_filesystem():
result = runner.invoke(
pyflyte.main,
[
"activate-launchplan",
"-p",
"flytesnacks",
"-d",
"development",
"daily",
],
)
assert result.exit_code == 0
assert "Launchplan was activated: " in result.output

0 comments on commit 28471f6

Please sign in to comment.