From 77513f01fbf88313537cb1e0fdc901695b83eef8 Mon Sep 17 00:00:00 2001
From: Artem <58334441+wckdman@users.noreply.github.com>
Date: Tue, 18 Apr 2023 22:09:24 +0200
Subject: [PATCH] Feature: add activate-launchplan command to pyflyte (#1588)

* add activate-launchplan command to pyflyte

Signed-off-by: Artem Petrov <58334441+wckdman@users.noreply.github.com>

* fix linting

Signed-off-by: Artem Petrov <58334441+wckdman@users.noreply.github.com>

* rework launchplan command

Signed-off-by: Artem Petrov <58334441+wckdman@users.noreply.github.com>

* fix linting

Signed-off-by: Artem Petrov <58334441+wckdman@users.noreply.github.com>

---------

Signed-off-by: Artem Petrov <58334441+wckdman@users.noreply.github.com>
Signed-off-by: Fabio Graetz <fabiograetz@googlemail.com>
---
 flytekit/clis/sdk_in_container/launchplan.py  | 74 +++++++++++++++++++
 flytekit/clis/sdk_in_container/pyflyte.py     |  2 +
 .../unit/cli/pyflyte/test_launchplan.py       | 34 +++++++++
 3 files changed, 110 insertions(+)
 create mode 100644 flytekit/clis/sdk_in_container/launchplan.py
 create mode 100644 tests/flytekit/unit/cli/pyflyte/test_launchplan.py

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