Move WorkflowRunOperationContext to _operation_context module

Author: dandavison

temporalio/nexus/__init__.py

@@ -8,6 +8,9 @@
 
 from ._decorators import workflow_run_operation as workflow_run_operation
 from ._operation_context import Info as Info
+from ._operation_context import (
+    WorkflowRunOperationContext as WorkflowRunOperationContext,
+)
 from ._operation_context import (
     _temporal_operation_context as _temporal_operation_context,
 )
@@ -18,7 +21,6 @@
 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 WorkflowRunOperationContext as WorkflowRunOperationContext
 
 
 class LoggerAdapter(logging.LoggerAdapter):

temporalio/nexus/_decorators.py

@@ -15,6 +15,7 @@
 )
 from nexusrpc.types import InputT, OutputT, ServiceHandlerT
 
+from temporalio.nexus._operation_context import WorkflowRunOperationContext
 from temporalio.nexus._operation_handlers import (
     WorkflowRunOperationHandler,
 )
@@ -25,7 +26,6 @@
     get_callable_name,
     get_workflow_run_start_method_input_and_output_type_annotations,
 )
-from temporalio.nexus._workflow import WorkflowRunOperationContext
 
 
 @overload

temporalio/nexus/_operation_context.py

@@ -5,10 +5,13 @@
 import urllib.parse
 from contextvars import ContextVar
 from dataclasses import dataclass
+from datetime import timedelta
 from typing import (
     Any,
     Callable,
+    Mapping,
     Optional,
+    Sequence,
     Union,
 )
 
@@ -19,6 +22,13 @@
 import temporalio.api.enums.v1
 import temporalio.client
 import temporalio.common
+from temporalio.nexus._token import WorkflowHandle
+from temporalio.types import (
+    MethodAsyncSingleParam,
+    ParamType,
+    ReturnType,
+    SelfType,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -92,6 +102,123 @@ def _temporal_cancel_operation_context(
         return _TemporalCancelOperationContext(ctx)
 
 
+@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,
+        )
+
+        start_operation_context.add_outbound_links(wf_handle)
+
+        return WorkflowHandle[ReturnType]._unsafe_from_client_workflow_handle(wf_handle)
+
+
 @dataclass
 class _TemporalStartOperationContext:
     nexus_operation_context: StartOperationContext

temporalio/nexus/_util.py

@@ -19,7 +19,7 @@
     ServiceHandlerT,
 )
 
-from temporalio.nexus._workflow import WorkflowRunOperationContext
+from temporalio.nexus._operation_context import WorkflowRunOperationContext
 
 from ._token import (
     WorkflowHandle as WorkflowHandle,

temporalio/nexus/_workflow.py

@@ -38,8 +38,7 @@
 )
 from temporalio.common import WorkflowIDConflictPolicy
 from temporalio.exceptions import CancelledError, NexusHandlerError, NexusOperationError
-from temporalio.nexus import workflow_run_operation
-from temporalio.nexus._workflow import WorkflowRunOperationContext
+from temporalio.nexus import WorkflowRunOperationContext, workflow_run_operation
 from temporalio.service import RPCError, RPCStatusCode
 from temporalio.worker import UnsandboxedWorkflowRunner, Worker
 from tests.helpers.nexus import create_nexus_endpoint, make_nexus_endpoint_name

tests/nexus/test_workflow_caller.py

@@ -13,8 +13,8 @@
 from nexusrpc.handler._decorators import operation_handler
 
 from temporalio import workflow
+from temporalio.nexus import WorkflowRunOperationContext
 from temporalio.nexus._operation_handlers import WorkflowRunOperationHandler
-from temporalio.nexus._workflow import WorkflowRunOperationContext
 from temporalio.testing import WorkflowEnvironment
 from temporalio.worker import Worker
 from tests.helpers.nexus import (