modify_site.py

"""A modification workflow for a site."""

from functools import partial
from typing import Annotated

from orchestrator.forms import FormPage
from orchestrator.targets import Target
from orchestrator.types import FormGenerator, State, SubscriptionLifecycle, UUIDstr
from orchestrator.workflow import StepList, begin, done, step, workflow
from orchestrator.workflows.steps import (
    resync,
    set_status,
    store_process_subscription,
    unsync,
)
from orchestrator.workflows.utils import wrap_modify_initial_input_form
from pydantic import AfterValidator, ConfigDict
from pydantic_forms.validators import ReadOnlyField

from gso.products.product_blocks.site import SiteTier
from gso.products.product_types.site import Site
from gso.utils.types.coordinates import LatitudeCoordinate, LongitudeCoordinate
from gso.utils.types.ip_address import IPAddress
from gso.utils.types.unique_field import validate_field_is_unique
from gso.workflows.shared import modify_summary_form


def initial_input_form_generator(subscription_id: UUIDstr) -> FormGenerator:
    """Gather input from the operator on what to change about the selected site subscription."""
    subscription = Site.from_subscription(subscription_id)

    class ModifySiteForm(FormPage):
        model_config = ConfigDict(title="Modify Site")

        site_name: ReadOnlyField(subscription.site.site_name, default_type=str)  # type: ignore[valid-type]
        site_city: str = subscription.site.site_city
        site_country: ReadOnlyField(subscription.site.site_country, default_type=str)  # type: ignore[valid-type]
        site_country_code: ReadOnlyField(subscription.site.site_country_code, default_type=str)  # type: ignore[valid-type]
        site_latitude: LatitudeCoordinate = subscription.site.site_latitude
        site_longitude: LongitudeCoordinate = subscription.site.site_longitude
        site_bgp_community_id: Annotated[int, AfterValidator(partial(validate_field_is_unique, subscription_id))] = (
            subscription.site.site_bgp_community_id
        )
        site_internal_id: Annotated[int, AfterValidator(partial(validate_field_is_unique, subscription_id))] = (
            subscription.site.site_internal_id
        )
        site_tier: ReadOnlyField(subscription.site.site_tier, default_type=SiteTier)  # type: ignore[valid-type]
        site_ts_address: (
            Annotated[IPAddress, AfterValidator(partial(validate_field_is_unique, subscription_id))] | None
        ) = subscription.site.site_ts_address

    user_input = yield ModifySiteForm
    user_input = user_input.model_dump()
    summary_fields = [
        "site_bgp_community_id",
        "site_internal_id",
        "site_ts_address",
        "site_city",
        "site_latitude",
        "site_longitude",
    ]
    yield from modify_summary_form(user_input, subscription.site, summary_fields)

    return user_input


@step("Modify subscription")
def modify_site_subscription(
    subscription: Site,
    site_city: str,
    site_latitude: LatitudeCoordinate,
    site_longitude: LongitudeCoordinate,
    site_bgp_community_id: int,
    site_internal_id: int,
    site_ts_address: IPAddress,
) -> State:
    """Update the subscription model in the service database."""
    subscription.site.site_city = site_city
    subscription.site.site_latitude = site_latitude
    subscription.site.site_longitude = site_longitude
    subscription.site.site_bgp_community_id = site_bgp_community_id
    subscription.site.site_internal_id = site_internal_id
    subscription.site.site_ts_address = site_ts_address

    subscription.description = f"Site in {site_city}, {subscription.site.site_country}"

    return {"subscription": subscription}


@workflow(
    "Modify Site",
    initial_input_form=wrap_modify_initial_input_form(initial_input_form_generator),
    target=Target.MODIFY,
)
def modify_site() -> StepList:
    """Modify a site subscription.

    * Update the subscription model in the service database
    """
    return (
        begin
        >> store_process_subscription(Target.MODIFY)
        >> unsync
        >> modify_site_subscription
        >> set_status(SubscriptionLifecycle.ACTIVE)
        >> resync
        >> done
    )