Move start_workflow to WorkflowRunOperationContext

Author: dandavison

temporalio/nexus/__init__.py

@@ -18,7 +18,7 @@
 from ._operation_context import info as info
 from ._operation_handlers import cancel_operation as cancel_operation
 from ._token import WorkflowHandle as WorkflowHandle
-from ._workflow import start_workflow as start_workflow
+from ._workflow import WorkflowRunOperationContext as WorkflowRunOperationContext
 
 
 class LoggerAdapter(logging.LoggerAdapter):

temporalio/nexus/_decorators.py

@@ -25,16 +25,17 @@
     get_callable_name,
     get_workflow_run_start_method_input_and_output_type_annotations,
 )
+from temporalio.nexus._workflow import WorkflowRunOperationContext
 
 
 @overload
 def workflow_run_operation(
     start: Callable[
-        [ServiceHandlerT, StartOperationContext, InputT],
+        [ServiceHandlerT, WorkflowRunOperationContext, InputT],
         Awaitable[WorkflowHandle[OutputT]],
     ],
 ) -> Callable[
-    [ServiceHandlerT, StartOperationContext, InputT],
+    [ServiceHandlerT, WorkflowRunOperationContext, InputT],
     Awaitable[WorkflowHandle[OutputT]],
 ]: ...
 
@@ -46,12 +47,12 @@ def workflow_run_operation(
 ) -> Callable[
     [
         Callable[
-            [ServiceHandlerT, StartOperationContext, InputT],
+            [ServiceHandlerT, WorkflowRunOperationContext, InputT],
             Awaitable[WorkflowHandle[OutputT]],
         ]
     ],
     Callable[
-        [ServiceHandlerT, StartOperationContext, InputT],
+        [ServiceHandlerT, WorkflowRunOperationContext, InputT],
         Awaitable[WorkflowHandle[OutputT]],
     ],
 ]: ...
@@ -60,26 +61,26 @@ def workflow_run_operation(
 def workflow_run_operation(
     start: Optional[
         Callable[
-            [ServiceHandlerT, StartOperationContext, InputT],
+            [ServiceHandlerT, WorkflowRunOperationContext, InputT],
             Awaitable[WorkflowHandle[OutputT]],
         ]
     ] = None,
     *,
     name: Optional[str] = None,
 ) -> Union[
     Callable[
-        [ServiceHandlerT, StartOperationContext, InputT],
+        [ServiceHandlerT, WorkflowRunOperationContext, InputT],
         Awaitable[WorkflowHandle[OutputT]],
     ],
     Callable[
         [
             Callable[
-                [ServiceHandlerT, StartOperationContext, InputT],
+                [ServiceHandlerT, WorkflowRunOperationContext, InputT],
                 Awaitable[WorkflowHandle[OutputT]],
             ]
         ],
         Callable[
-            [ServiceHandlerT, StartOperationContext, InputT],
+            [ServiceHandlerT, WorkflowRunOperationContext, InputT],
             Awaitable[WorkflowHandle[OutputT]],
         ],
     ],
@@ -90,11 +91,11 @@ def workflow_run_operation(
 
     def decorator(
         start: Callable[
-            [ServiceHandlerT, StartOperationContext, InputT],
+            [ServiceHandlerT, WorkflowRunOperationContext, InputT],
             Awaitable[WorkflowHandle[OutputT]],
         ],
     ) -> Callable[
-        [ServiceHandlerT, StartOperationContext, InputT],
+        [ServiceHandlerT, WorkflowRunOperationContext, InputT],
         Awaitable[WorkflowHandle[OutputT]],
     ]:
         (
@@ -108,7 +109,7 @@ def operation_handler_factory(
             async def _start(
                 ctx: StartOperationContext, input: InputT
             ) -> WorkflowHandle[OutputT]:
-                return await start(self, ctx, input)
+                return await start(self, WorkflowRunOperationContext(ctx), input)
 
             _start.__doc__ = start.__doc__
             return WorkflowRunOperationHandler(_start, input_type, output_type)

temporalio/nexus/_operation_handlers.py

@@ -41,7 +41,8 @@ class WorkflowRunOperationHandler(OperationHandler[InputT, OutputT]):
 
     Use this class to create an operation handler that starts a workflow by passing your
     ``start`` method to the constructor. Your ``start`` method must use
-    :py:func:`temporalio.nexus.start_workflow` to start the workflow.
+    :py:func:`temporalio.nexus.WorkflowRunOperationContext.start_workflow` to start the
+    workflow.
     """
 
     def __init__(

temporalio/nexus/_util.py

@@ -13,21 +13,22 @@
     Union,
 )
 
-from nexusrpc.handler import StartOperationContext
 from nexusrpc.types import (
     InputT,
     OutputT,
     ServiceHandlerT,
 )
 
+from temporalio.nexus._workflow import WorkflowRunOperationContext
+
 from ._token import (
     WorkflowHandle as WorkflowHandle,
 )
 
 
 def get_workflow_run_start_method_input_and_output_type_annotations(
     start: Callable[
-        [ServiceHandlerT, StartOperationContext, InputT],
+        [ServiceHandlerT, WorkflowRunOperationContext, InputT],
         Awaitable[WorkflowHandle[OutputT]],
     ],
 ) -> tuple[

temporalio/nexus/_workflow.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from dataclasses import dataclass
 from datetime import timedelta
 from typing import (
     Any,
@@ -9,6 +10,8 @@
     Union,
 )
 
+from nexusrpc.handler import StartOperationContext
+
 import temporalio.api.common.v1
 import temporalio.api.enums.v1
 import temporalio.common
@@ -22,116 +25,118 @@
 )
 
 
-# Overload for single-param workflow
-# TODO(nexus-prerelease): bring over other overloads
-async def start_workflow(
-    workflow: MethodAsyncSingleParam[SelfType, ParamType, ReturnType],
-    arg: ParamType,
-    *,
-    id: str,
-    task_queue: Optional[str] = None,
-    execution_timeout: Optional[timedelta] = None,
-    run_timeout: Optional[timedelta] = None,
-    task_timeout: Optional[timedelta] = None,
-    id_reuse_policy: temporalio.common.WorkflowIDReusePolicy = temporalio.common.WorkflowIDReusePolicy.ALLOW_DUPLICATE,
-    id_conflict_policy: temporalio.common.WorkflowIDConflictPolicy = temporalio.common.WorkflowIDConflictPolicy.UNSPECIFIED,
-    retry_policy: Optional[temporalio.common.RetryPolicy] = None,
-    cron_schedule: str = "",
-    memo: Optional[Mapping[str, Any]] = None,
-    search_attributes: Optional[
-        Union[
-            temporalio.common.TypedSearchAttributes,
-            temporalio.common.SearchAttributes,
-        ]
-    ] = None,
-    static_summary: Optional[str] = None,
-    static_details: Optional[str] = None,
-    start_delay: Optional[timedelta] = None,
-    start_signal: Optional[str] = None,
-    start_signal_args: Sequence[Any] = [],
-    rpc_metadata: Mapping[str, str] = {},
-    rpc_timeout: Optional[timedelta] = None,
-    request_eager_start: bool = False,
-    priority: temporalio.common.Priority = temporalio.common.Priority.default,
-    versioning_override: Optional[temporalio.common.VersioningOverride] = None,
-) -> WorkflowHandle[ReturnType]:
-    """Start a workflow that will deliver the result of the Nexus operation.
-
-    The workflow will be started in the same namespace as the Nexus worker, using
-    the same client as the worker. If task queue is not specified, the worker's task
-    queue will be used.
-
-    See :py:meth:`temporalio.client.Client.start_workflow` for all arguments.
-
-    The return value is :py:class:`temporalio.nexus.WorkflowOperationToken`.
-    Use :py:meth:`temporalio.nexus.WorkflowOperationToken.to_workflow_handle`
-    to get a :py:class:`temporalio.client.WorkflowHandle` for interacting with the
-    workflow.
-
-    The workflow will be started as usual, with the following modifications:
-
-    - On workflow completion, Temporal server will deliver the workflow result to
-        the Nexus operation caller, using the callback from the Nexus operation start
-        request.
-
-    - The request ID from the Nexus operation start request will be used as the
-        request ID for the start workflow request.
-
-    - Inbound links to the caller that were submitted in the Nexus start operation
-        request will be attached to the started workflow and, outbound links to the
-        started workflow will be added to the Nexus start operation response. If the
-        Nexus caller is itself a workflow, this means that the workflow in the caller
-        namespace web UI will contain links to the started workflow, and vice versa.
-    """
-    tctx = _TemporalNexusOperationContext.get()
-    start_operation_context = tctx._temporal_start_operation_context
-    if not start_operation_context:
-        raise RuntimeError(
-            "temporalio.nexus.start_workflow() must be called from "
-            "within a Nexus start operation context"
+@dataclass
+class WorkflowRunOperationContext:
+    start_operation_context: StartOperationContext
+
+    # Overload for single-param workflow
+    # TODO(nexus-prerelease): bring over other overloads
+    async def start_workflow(
+        self,
+        workflow: MethodAsyncSingleParam[SelfType, ParamType, ReturnType],
+        arg: ParamType,
+        *,
+        id: str,
+        task_queue: Optional[str] = None,
+        execution_timeout: Optional[timedelta] = None,
+        run_timeout: Optional[timedelta] = None,
+        task_timeout: Optional[timedelta] = None,
+        id_reuse_policy: temporalio.common.WorkflowIDReusePolicy = temporalio.common.WorkflowIDReusePolicy.ALLOW_DUPLICATE,
+        id_conflict_policy: temporalio.common.WorkflowIDConflictPolicy = temporalio.common.WorkflowIDConflictPolicy.UNSPECIFIED,
+        retry_policy: Optional[temporalio.common.RetryPolicy] = None,
+        cron_schedule: str = "",
+        memo: Optional[Mapping[str, Any]] = None,
+        search_attributes: Optional[
+            Union[
+                temporalio.common.TypedSearchAttributes,
+                temporalio.common.SearchAttributes,
+            ]
+        ] = None,
+        static_summary: Optional[str] = None,
+        static_details: Optional[str] = None,
+        start_delay: Optional[timedelta] = None,
+        start_signal: Optional[str] = None,
+        start_signal_args: Sequence[Any] = [],
+        rpc_metadata: Mapping[str, str] = {},
+        rpc_timeout: Optional[timedelta] = None,
+        request_eager_start: bool = False,
+        priority: temporalio.common.Priority = temporalio.common.Priority.default,
+        versioning_override: Optional[temporalio.common.VersioningOverride] = None,
+    ) -> WorkflowHandle[ReturnType]:
+        """Start a workflow that will deliver the result of the Nexus operation.
+
+        The workflow will be started in the same namespace as the Nexus worker, using
+        the same client as the worker. If task queue is not specified, the worker's task
+        queue will be used.
+
+        See :py:meth:`temporalio.client.Client.start_workflow` for all arguments.
+
+        The return value is :py:class:`temporalio.nexus.WorkflowHandle`.
+
+        The workflow will be started as usual, with the following modifications:
+
+        - On workflow completion, Temporal server will deliver the workflow result to
+            the Nexus operation caller, using the callback from the Nexus operation start
+            request.
+
+        - The request ID from the Nexus operation start request will be used as the
+            request ID for the start workflow request.
+
+        - Inbound links to the caller that were submitted in the Nexus start operation
+            request will be attached to the started workflow and, outbound links to the
+            started workflow will be added to the Nexus start operation response. If the
+            Nexus caller is itself a workflow, this means that the workflow in the caller
+            namespace web UI will contain links to the started workflow, and vice versa.
+        """
+        tctx = _TemporalNexusOperationContext.get()
+        start_operation_context = tctx._temporal_start_operation_context
+        if not start_operation_context:
+            raise RuntimeError(
+                "WorkflowRunOperationContext.start_workflow() must be called from "
+                "within a Nexus start operation context"
+            )
+
+        # TODO(nexus-preview): When sdk-python supports on_conflict_options, Typescript does this:
+        # if (workflowOptions.workflowIdConflictPolicy === 'USE_EXISTING') {
+        #     internalOptions.onConflictOptions = {
+        #     attachLinks: true,
+        #     attachCompletionCallbacks: true,
+        #     attachRequestId: true,
+        #     };
+        # }
+
+        # We must pass nexus_completion_callbacks, workflow_event_links, and request_id,
+        # but these are deliberately not exposed in overloads, hence the type-check
+        # violation.
+        wf_handle = await tctx.client.start_workflow(  # type: ignore
+            workflow=workflow,
+            arg=arg,
+            id=id,
+            task_queue=task_queue or tctx.info().task_queue,
+            execution_timeout=execution_timeout,
+            run_timeout=run_timeout,
+            task_timeout=task_timeout,
+            id_reuse_policy=id_reuse_policy,
+            id_conflict_policy=id_conflict_policy,
+            retry_policy=retry_policy,
+            cron_schedule=cron_schedule,
+            memo=memo,
+            search_attributes=search_attributes,
+            static_summary=static_summary,
+            static_details=static_details,
+            start_delay=start_delay,
+            start_signal=start_signal,
+            start_signal_args=start_signal_args,
+            rpc_metadata=rpc_metadata,
+            rpc_timeout=rpc_timeout,
+            request_eager_start=request_eager_start,
+            priority=priority,
+            versioning_override=versioning_override,
+            nexus_completion_callbacks=start_operation_context.get_completion_callbacks(),
+            workflow_event_links=start_operation_context.get_workflow_event_links(),
+            request_id=start_operation_context.nexus_operation_context.request_id,
         )
 
-    # TODO(nexus-preview): When sdk-python supports on_conflict_options, Typescript does this:
-    # if (workflowOptions.workflowIdConflictPolicy === 'USE_EXISTING') {
-    #     internalOptions.onConflictOptions = {
-    #     attachLinks: true,
-    #     attachCompletionCallbacks: true,
-    #     attachRequestId: true,
-    #     };
-    # }
-
-    # We must pass nexus_completion_callbacks, workflow_event_links, and request_id,
-    # but these are deliberately not exposed in overloads, hence the type-check
-    # violation.
-    wf_handle = await tctx.client.start_workflow(  # type: ignore
-        workflow=workflow,
-        arg=arg,
-        id=id,
-        task_queue=task_queue or tctx.info().task_queue,
-        execution_timeout=execution_timeout,
-        run_timeout=run_timeout,
-        task_timeout=task_timeout,
-        id_reuse_policy=id_reuse_policy,
-        id_conflict_policy=id_conflict_policy,
-        retry_policy=retry_policy,
-        cron_schedule=cron_schedule,
-        memo=memo,
-        search_attributes=search_attributes,
-        static_summary=static_summary,
-        static_details=static_details,
-        start_delay=start_delay,
-        start_signal=start_signal,
-        start_signal_args=start_signal_args,
-        rpc_metadata=rpc_metadata,
-        rpc_timeout=rpc_timeout,
-        request_eager_start=request_eager_start,
-        priority=priority,
-        versioning_override=versioning_override,
-        nexus_completion_callbacks=start_operation_context.get_completion_callbacks(),
-        workflow_event_links=start_operation_context.get_workflow_event_links(),
-        request_id=start_operation_context.nexus_operation_context.request_id,
-    )
-
-    start_operation_context.add_outbound_links(wf_handle)
-
-    return WorkflowHandle[ReturnType]._unsafe_from_client_workflow_handle(wf_handle)
+        start_operation_context.add_outbound_links(wf_handle)
+
+        return WorkflowHandle[ReturnType]._unsafe_from_client_workflow_handle(wf_handle)