add start and terminate workflow tasks for routers

gso/migrations/versions/2025-04-29_e86c16d809ab_add_moodi_run_and_terminate_workflow_.py

+"""add moodi run and terminate workflow tasks for routers.
+
+Revision ID: e86c16d809ab
+Revises: a3177c5f9641
+Create Date: 2025-04-29 17:14:51.627317
+
+"""
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = 'e86c16d809ab'
+down_revision = 'a3177c5f9641'
+branch_labels = None
+depends_on = None
+
+
+from orchestrator.migrations.helpers import create_workflow, delete_workflow
+
+new_workflows = [
+    {
+        "name": "run_moodi",
+        "target": "SYSTEM",
+        "description": "Run Moodi Service",
+        "product_type": "Router"
+    },
+    {
+        "name": "terminate_moodi",
+        "target": "SYSTEM",
+        "description": "Terminate Moodi Service",
+        "product_type": "Router"
+    }
+]
+
+
+def upgrade() -> None:
+    conn = op.get_bind()
+    for workflow in new_workflows:
+        create_workflow(conn, workflow)
+
+
+def downgrade() -> None:
+    conn = op.get_bind()
+    for workflow in new_workflows:
+        delete_workflow(conn, workflow["name"])

gso/workflows/__init__.py

@@ -54,6 +54,8 @@ LazyWorkflowInstance("gso.workflows.router.create_imported_router", "create_impo
 LazyWorkflowInstance("gso.workflows.router.validate_router", "validate_router")
 LazyWorkflowInstance("gso.workflows.router.promote_p_to_pe", "promote_p_to_pe")
 LazyWorkflowInstance("gso.workflows.router.modify_kentik_license", "modify_router_kentik_license")
+LazyWorkflowInstance("gso.workflows.router.run_moodi", "run_moodi")
+LazyWorkflowInstance("gso.workflows.router.terminate_moodi", "terminate_moodi")
 
 #  Switch workflows
 LazyWorkflowInstance("gso.workflows.switch.create_switch", "create_switch")

gso/workflows/router/run_moodi.py

+"""Run Moodi monitoring service."""
+
+from orchestrator.targets import Target
+from orchestrator.workflow import StepList, begin, conditional, done, step, workflow
+from orchestrator.workflows.steps import store_process_subscription
+from orchestrator.workflows.utils import wrap_modify_initial_input_form
+from pydantic_forms.types import State, UUIDstr
+
+from gso.products.product_types.router import Router
+from gso.utils.shared_enums import Vendor
+from gso.utils.workflow_steps import start_moodi
+
+
+@step("Prepare required keys in state")
+def prepare_state(subscription_id: UUIDstr) -> State:
+    """Add required keys to the state for the workflow to run successfully."""
+    router = Router.from_subscription(subscription_id)
+
+    return {"subscription": router}
+
+
+@workflow("Run Moodi Service", target=Target.SYSTEM, initial_input_form=wrap_modify_initial_input_form(None))
+def run_moodi() -> StepList:
+    """Run Moodi monitoring service."""
+    is_juniper_router = conditional(lambda state: state["subscription"]["router"]["vendor"] == Vendor.JUNIPER)
+
+    return (
+        begin
+        >> store_process_subscription(Target.SYSTEM)
+        >> prepare_state
+        >> is_juniper_router(done)
+        >> start_moodi()
+        >> done
+    )

gso/workflows/router/terminate_moodi.py

+"""Terminate Moodi monitoring service."""
+
+from orchestrator.targets import Target
+from orchestrator.workflow import StepList, begin, conditional, done, step, workflow
+from orchestrator.workflows.steps import store_process_subscription
+from orchestrator.workflows.utils import wrap_modify_initial_input_form
+from pydantic_forms.types import State, UUIDstr
+
+from gso.products.product_types.router import Router
+from gso.utils.shared_enums import Vendor
+from gso.utils.workflow_steps import stop_moodi
+
+
+@step("Prepare required keys in state")
+def prepare_state(subscription_id: UUIDstr) -> State:
+    """Add required keys to the state for the workflow to run successfully."""
+    router = Router.from_subscription(subscription_id)
+
+    return {"subscription": router}
+
+
+@workflow("Terminate Moodi Service", target=Target.SYSTEM, initial_input_form=wrap_modify_initial_input_form(None))
+def terminate_moodi() -> StepList:
+    """Terminate Moodi monitoring service."""
+    is_juniper_router = conditional(lambda state: state["subscription"]["router"]["vendor"] == Vendor.JUNIPER)
+
+    return (
+        begin
+        >> store_process_subscription(Target.SYSTEM)
+        >> prepare_state
+        >> is_juniper_router(done)
+        >> stop_moodi()
+        >> done
+    )