|
| 1 | +from __future__ import annotations |
| 2 | + |
| 3 | +import logging |
| 4 | +import re |
| 5 | +import urllib.parse |
| 6 | +from typing import ( |
| 7 | + Any, |
| 8 | + Optional, |
| 9 | +) |
| 10 | + |
| 11 | +import nexusrpc |
| 12 | + |
| 13 | +import temporalio.api.common.v1 |
| 14 | +import temporalio.api.enums.v1 |
| 15 | +import temporalio.client |
| 16 | + |
| 17 | +logger = logging.getLogger(__name__) |
| 18 | + |
| 19 | +_LINK_URL_PATH_REGEX = re.compile( |
| 20 | + r"^/namespaces/(?P<namespace>[^/]+)/workflows/(?P<workflow_id>[^/]+)/(?P<run_id>[^/]+)/history$" |
| 21 | +) |
| 22 | +LINK_EVENT_ID_PARAM_NAME = "eventID" |
| 23 | +LINK_EVENT_TYPE_PARAM_NAME = "eventType" |
| 24 | + |
| 25 | + |
| 26 | +def workflow_handle_to_workflow_execution_started_event_link( |
| 27 | + handle: temporalio.client.WorkflowHandle[Any, Any], |
| 28 | +) -> temporalio.api.common.v1.Link.WorkflowEvent: |
| 29 | + """Create a WorkflowEvent link corresponding to a started workflow""" |
| 30 | + if handle.first_execution_run_id is None: |
| 31 | + raise ValueError( |
| 32 | + f"Workflow handle {handle} has no first execution run ID. " |
| 33 | + f"Cannot create WorkflowExecutionStarted event link." |
| 34 | + ) |
| 35 | + return temporalio.api.common.v1.Link.WorkflowEvent( |
| 36 | + namespace=handle._client.namespace, |
| 37 | + workflow_id=handle.id, |
| 38 | + run_id=handle.first_execution_run_id, |
| 39 | + event_ref=temporalio.api.common.v1.Link.WorkflowEvent.EventReference( |
| 40 | + event_id=1, |
| 41 | + event_type=temporalio.api.enums.v1.EventType.EVENT_TYPE_WORKFLOW_EXECUTION_STARTED, |
| 42 | + ), |
| 43 | + # TODO(nexus-preview): RequestIdReference |
| 44 | + ) |
| 45 | + |
| 46 | + |
| 47 | +def workflow_event_to_nexus_link( |
| 48 | + workflow_event: temporalio.api.common.v1.Link.WorkflowEvent, |
| 49 | +) -> nexusrpc.Link: |
| 50 | + """Convert a WorkflowEvent link into a nexusrpc link |
| 51 | +
|
| 52 | + Used when propagating links from a StartWorkflow response to a Nexus start operation |
| 53 | + response. |
| 54 | + """ |
| 55 | + scheme = "temporal" |
| 56 | + namespace = urllib.parse.quote(workflow_event.namespace) |
| 57 | + workflow_id = urllib.parse.quote(workflow_event.workflow_id) |
| 58 | + run_id = urllib.parse.quote(workflow_event.run_id) |
| 59 | + path = f"/namespaces/{namespace}/workflows/{workflow_id}/{run_id}/history" |
| 60 | + query_params = _event_reference_to_query_params(workflow_event.event_ref) |
| 61 | + return nexusrpc.Link( |
| 62 | + url=urllib.parse.urlunparse((scheme, "", path, "", query_params, "")), |
| 63 | + type=workflow_event.DESCRIPTOR.full_name, |
| 64 | + ) |
| 65 | + |
| 66 | + |
| 67 | +def nexus_link_to_workflow_event( |
| 68 | + link: nexusrpc.Link, |
| 69 | +) -> Optional[temporalio.api.common.v1.Link.WorkflowEvent]: |
| 70 | + """Convert a nexus link into a WorkflowEvent link |
| 71 | +
|
| 72 | + This is used when propagating links from a Nexus start operation request to a |
| 73 | + StartWorklow request. |
| 74 | + """ |
| 75 | + url = urllib.parse.urlparse(link.url) |
| 76 | + match = _LINK_URL_PATH_REGEX.match(url.path) |
| 77 | + if not match: |
| 78 | + logger.warning( |
| 79 | + f"Invalid Nexus link: {link}. Expected path to match {_LINK_URL_PATH_REGEX.pattern}" |
| 80 | + ) |
| 81 | + return None |
| 82 | + try: |
| 83 | + event_ref = _query_params_to_event_reference(url.query) |
| 84 | + except ValueError as err: |
| 85 | + logger.warning( |
| 86 | + f"Failed to parse event reference from Nexus link URL query parameters: {link} ({err})" |
| 87 | + ) |
| 88 | + return None |
| 89 | + |
| 90 | + groups = match.groupdict() |
| 91 | + return temporalio.api.common.v1.Link.WorkflowEvent( |
| 92 | + namespace=urllib.parse.unquote(groups["namespace"]), |
| 93 | + workflow_id=urllib.parse.unquote(groups["workflow_id"]), |
| 94 | + run_id=urllib.parse.unquote(groups["run_id"]), |
| 95 | + event_ref=event_ref, |
| 96 | + ) |
| 97 | + |
| 98 | + |
| 99 | +def _event_reference_to_query_params( |
| 100 | + event_ref: temporalio.api.common.v1.Link.WorkflowEvent.EventReference, |
| 101 | +) -> str: |
| 102 | + event_type_name = temporalio.api.enums.v1.EventType.Name(event_ref.event_type) |
| 103 | + if event_type_name.startswith("EVENT_TYPE_"): |
| 104 | + event_type_name = _event_type_constant_case_to_pascal_case( |
| 105 | + event_type_name.removeprefix("EVENT_TYPE_") |
| 106 | + ) |
| 107 | + return urllib.parse.urlencode( |
| 108 | + { |
| 109 | + "eventID": event_ref.event_id, |
| 110 | + "eventType": event_type_name, |
| 111 | + "referenceType": "EventReference", |
| 112 | + } |
| 113 | + ) |
| 114 | + |
| 115 | + |
| 116 | +def _query_params_to_event_reference( |
| 117 | + raw_query_params: str, |
| 118 | +) -> temporalio.api.common.v1.Link.WorkflowEvent.EventReference: |
| 119 | + """Return an EventReference from the query params or raise ValueError.""" |
| 120 | + query_params = urllib.parse.parse_qs(raw_query_params) |
| 121 | + |
| 122 | + [reference_type] = query_params.get("referenceType") or [""] |
| 123 | + if reference_type != "EventReference": |
| 124 | + raise ValueError( |
| 125 | + f"Expected Nexus link URL query parameter referenceType to be EventReference but got: {reference_type}" |
| 126 | + ) |
| 127 | + # event type |
| 128 | + [raw_event_type_name] = query_params.get(LINK_EVENT_TYPE_PARAM_NAME) or [""] |
| 129 | + if not raw_event_type_name: |
| 130 | + raise ValueError(f"query params do not contain event type: {query_params}") |
| 131 | + if raw_event_type_name.startswith("EVENT_TYPE_"): |
| 132 | + event_type_name = raw_event_type_name |
| 133 | + elif re.match("[A-Z][a-z]", raw_event_type_name): |
| 134 | + event_type_name = "EVENT_TYPE_" + _event_type_pascal_case_to_constant_case( |
| 135 | + raw_event_type_name |
| 136 | + ) |
| 137 | + else: |
| 138 | + raise ValueError(f"Invalid event type name: {raw_event_type_name}") |
| 139 | + |
| 140 | + # event id |
| 141 | + event_id = 0 |
| 142 | + [raw_event_id] = query_params.get(LINK_EVENT_ID_PARAM_NAME) or [""] |
| 143 | + if raw_event_id: |
| 144 | + try: |
| 145 | + event_id = int(raw_event_id) |
| 146 | + except ValueError: |
| 147 | + raise ValueError(f"Query params contain invalid event id: {raw_event_id}") |
| 148 | + |
| 149 | + return temporalio.api.common.v1.Link.WorkflowEvent.EventReference( |
| 150 | + event_type=temporalio.api.enums.v1.EventType.Value(event_type_name), |
| 151 | + event_id=event_id, |
| 152 | + ) |
| 153 | + |
| 154 | + |
| 155 | +def _event_type_constant_case_to_pascal_case(s: str) -> str: |
| 156 | + """Convert a CONSTANT_CASE string to PascalCase. |
| 157 | +
|
| 158 | + >>> _event_type_constant_case_to_pascal_case("NEXUS_OPERATION_SCHEDULED") |
| 159 | + "NexusOperationScheduled" |
| 160 | + """ |
| 161 | + return re.sub(r"(\b|_)([a-z])", lambda m: m.groups()[1].upper(), s.lower()) |
| 162 | + |
| 163 | + |
| 164 | +def _event_type_pascal_case_to_constant_case(s: str) -> str: |
| 165 | + """Convert a PascalCase string to CONSTANT_CASE. |
| 166 | +
|
| 167 | + >>> _event_type_pascal_case_to_constant_case("NexusOperationScheduled") |
| 168 | + "NEXUS_OPERATION_SCHEDULED" |
| 169 | + """ |
| 170 | + return re.sub(r"([A-Z])", r"_\1", s).lstrip("_").upper() |
0 commit comments