From 28471f638737aa3f5d345928ecb143e40cda5cf2 Mon Sep 17 00:00:00 2001
From: wckdman <petrov.still@gmail.com>
Date: Wed, 12 Apr 2023 14:05:46 +0200
Subject: [PATCH] add activate-launchplan command to pyflyte

Signed-off-by: Artem Petrov <58334441+wckdman@users.noreply.github.com>
---
 .../sdk_in_container/activate_launchplan.py   | 61 +++++++++++++++++++
 flytekit/clis/sdk_in_container/pyflyte.py     |  2 +
 .../cli/pyflyte/test_activate_launchplan.py   | 25 ++++++++
 3 files changed, 88 insertions(+)
 create mode 100644 flytekit/clis/sdk_in_container/activate_launchplan.py
 create mode 100644 tests/flytekit/unit/cli/pyflyte/test_activate_launchplan.py

diff --git a/flytekit/clis/sdk_in_container/activate_launchplan.py b/flytekit/clis/sdk_in_container/activate_launchplan.py
new file mode 100644
index 0000000000..757d0d3f51
--- /dev/null
+++ b/flytekit/clis/sdk_in_container/activate_launchplan.py
@@ -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")
diff --git a/flytekit/clis/sdk_in_container/pyflyte.py b/flytekit/clis/sdk_in_container/pyflyte.py
index 5e1136d14c..0b88815046 100644
--- a/flytekit/clis/sdk_in_container/pyflyte.py
+++ b/flytekit/clis/sdk_in_container/pyflyte.py
@@ -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
@@ -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__":
diff --git a/tests/flytekit/unit/cli/pyflyte/test_activate_launchplan.py b/tests/flytekit/unit/cli/pyflyte/test_activate_launchplan.py
new file mode 100644
index 0000000000..4f449444c6
--- /dev/null
+++ b/tests/flytekit/unit/cli/pyflyte/test_activate_launchplan.py
@@ -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