From 5f0d59e9bd53888e4ba4a6a11f18af3e65d6fc66 Mon Sep 17 00:00:00 2001
From: Karel van Klink <karel.vanklink@geant.org>
Date: Fri, 2 Aug 2024 15:10:06 +0200
Subject: [PATCH] Add modification workflow for routers to change Kentik
license
---
...e_add_router_modification_workflow_for_.py | 39 ++++++++++++++
gso/workflows/__init__.py | 1 +
gso/workflows/router/modify_kentik_license.py | 54 +++++++++++++++++++
3 files changed, 94 insertions(+)
create mode 100644 gso/migrations/versions/2024-08-02_87a05eddee3e_add_router_modification_workflow_for_.py
create mode 100644 gso/workflows/router/modify_kentik_license.py
diff --git a/gso/migrations/versions/2024-08-02_87a05eddee3e_add_router_modification_workflow_for_.py b/gso/migrations/versions/2024-08-02_87a05eddee3e_add_router_modification_workflow_for_.py
new file mode 100644
index 00000000..13162210
--- /dev/null
+++ b/gso/migrations/versions/2024-08-02_87a05eddee3e_add_router_modification_workflow_for_.py
@@ -0,0 +1,39 @@
+"""Add router modification workflow for kentik licenses.
+
+Revision ID: 87a05eddee3e
+Revises: 41fd1ae225aq
+Create Date: 2024-08-02 15:09:42.597063
+
+"""
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = '87a05eddee3e'
+down_revision = '41fd1ae225aq'
+branch_labels = None
+depends_on = None
+
+
+from orchestrator.migrations.helpers import create_workflow, delete_workflow
+
+new_workflows = [
+ {
+ "name": "modify_router_kentik_license",
+ "target": "MODIFY",
+ "description": "Modify Kentik license",
+ "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"])
diff --git a/gso/workflows/__init__.py b/gso/workflows/__init__.py
index 291a9e25..96638b85 100644
--- a/gso/workflows/__init__.py
+++ b/gso/workflows/__init__.py
@@ -46,6 +46,7 @@ LazyWorkflowInstance("gso.workflows.router.import_router", "import_router")
LazyWorkflowInstance("gso.workflows.router.create_imported_router", "create_imported_router")
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")
# Site workflows
LazyWorkflowInstance("gso.workflows.site.create_site", "create_site")
diff --git a/gso/workflows/router/modify_kentik_license.py b/gso/workflows/router/modify_kentik_license.py
new file mode 100644
index 00000000..3d6fd825
--- /dev/null
+++ b/gso/workflows/router/modify_kentik_license.py
@@ -0,0 +1,54 @@
+"""A workflow that modifies the Kentik license of a router."""
+
+from orchestrator.forms import FormPage
+from orchestrator.targets import Target
+from orchestrator.types import FormGenerator, State, UUIDstr
+from orchestrator.utils.errors import ProcessFailureError
+from orchestrator.workflow import StepList, begin, done, step, workflow
+from orchestrator.workflows.steps import resync, store_process_subscription, unsync
+from orchestrator.workflows.utils import wrap_modify_initial_input_form
+from pydantic_forms.validators import Choice
+
+from gso.products.product_types.router import Router
+from gso.services.kentik_client import KentikClient
+
+
+def _initial_input_form(subscription_id: UUIDstr) -> FormGenerator:
+ router = Router.from_subscription(subscription_id)
+
+ class ModifyKentikLicenseForm(FormPage):
+ @staticmethod
+ def _get_available_plans() -> list[dict[str, str]]:
+ kentik_plans = KentikClient().get_plans()
+ return [{plan["name"]: plan["id"]} for plan in kentik_plans]
+
+ new_plan_id = Choice("Select a Kentik license", _get_available_plans())
+
+ user_input = yield ModifyKentikLicenseForm
+
+ return user_input.model_dump() | {"subscription": router}
+
+
+@step("Update device license in Kentik")
+def update_kentik_license(subscription: Router, new_plan_id: int) -> State:
+ """Update a Kentik device with a new license attached to it."""
+ kentik_client = KentikClient()
+ kentik_device = kentik_client.get_device_by_name(subscription.router.router_fqdn)
+ if "id" not in kentik_device:
+ msg = "Failed to find Kentik device by name"
+ raise ProcessFailureError(msg, details=kentik_device)
+
+ updated_kentik_device = kentik_client.update_device(kentik_device["id"], {"plan_id": new_plan_id})
+ updated_kentik_device.pop("custom_column_data", None)
+
+ return {"kentik_device": updated_kentik_device}
+
+
+@workflow(
+ "Modify Kentik license",
+ initial_input_form=wrap_modify_initial_input_form(_initial_input_form),
+ target=Target.MODIFY,
+)
+def modify_router_kentik_license() -> StepList:
+ """Apply a selected Kentik license on an existing PE router."""
+ return begin >> store_process_subscription(Target.MODIFY) >> unsync >> update_kentik_license >> resync >> done
--
GitLab