|
3 | 3 | import dataclasses |
4 | 4 | import logging |
5 | 5 | from collections.abc import Awaitable, Mapping, MutableMapping, Sequence |
| 6 | +from contextlib import contextmanager |
6 | 7 | from contextvars import ContextVar |
7 | 8 | from dataclasses import dataclass |
8 | 9 | from datetime import timedelta |
9 | 10 | from typing import ( |
10 | 11 | TYPE_CHECKING, |
11 | 12 | Any, |
12 | 13 | Callable, |
| 14 | + Generator, |
13 | 15 | Optional, |
14 | 16 | Union, |
15 | 17 | overload, |
|
47 | 49 | ContextVar("temporal-cancel-operation-context") |
48 | 50 | ) |
49 | 51 |
|
| 52 | +_temporal_nexus_backing_workflow_start_context: ContextVar[bool] = ContextVar( |
| 53 | + "temporal-nexus-backing-workflow-start-context" |
| 54 | +) |
| 55 | + |
50 | 56 |
|
51 | 57 | @dataclass(frozen=True) |
52 | 58 | class Info: |
@@ -96,6 +102,19 @@ def _try_temporal_context() -> ( |
96 | 102 | return start_ctx or cancel_ctx |
97 | 103 |
|
98 | 104 |
|
| 105 | +@contextmanager |
| 106 | +def _nexus_backing_workflow_start_context() -> Generator[None, None, None]: |
| 107 | + token = _temporal_nexus_backing_workflow_start_context.set(True) |
| 108 | + try: |
| 109 | + yield |
| 110 | + finally: |
| 111 | + _temporal_nexus_backing_workflow_start_context.reset(token) |
| 112 | + |
| 113 | + |
| 114 | +def _in_nexus_backing_workflow_start_context() -> bool: |
| 115 | + return _temporal_nexus_backing_workflow_start_context.get(False) |
| 116 | + |
| 117 | + |
99 | 118 | @dataclass |
100 | 119 | class _TemporalStartOperationContext: |
101 | 120 | """Context for a Nexus start operation being handled by a Temporal Nexus Worker.""" |
@@ -396,48 +415,40 @@ async def start_workflow( |
396 | 415 | Nexus caller is itself a workflow, this means that the workflow in the caller |
397 | 416 | namespace web UI will contain links to the started workflow, and vice versa. |
398 | 417 | """ |
399 | | - # TODO(nexus-preview): When sdk-python supports on_conflict_options, Typescript does this: |
400 | | - # if (workflowOptions.workflowIdConflictPolicy === 'USE_EXISTING') { |
401 | | - # internalOptions.onConflictOptions = { |
402 | | - # attachLinks: true, |
403 | | - # attachCompletionCallbacks: true, |
404 | | - # attachRequestId: true, |
405 | | - # }; |
406 | | - # } |
407 | | - |
408 | 418 | # We must pass nexus_completion_callbacks, workflow_event_links, and request_id, |
409 | 419 | # but these are deliberately not exposed in overloads, hence the type-check |
410 | 420 | # violation. |
411 | | - wf_handle = await self._temporal_context.client.start_workflow( # type: ignore |
412 | | - workflow=workflow, |
413 | | - arg=arg, |
414 | | - args=args, |
415 | | - id=id, |
416 | | - task_queue=task_queue or self._temporal_context.info().task_queue, |
417 | | - result_type=result_type, |
418 | | - execution_timeout=execution_timeout, |
419 | | - run_timeout=run_timeout, |
420 | | - task_timeout=task_timeout, |
421 | | - id_reuse_policy=id_reuse_policy, |
422 | | - id_conflict_policy=id_conflict_policy, |
423 | | - retry_policy=retry_policy, |
424 | | - cron_schedule=cron_schedule, |
425 | | - memo=memo, |
426 | | - search_attributes=search_attributes, |
427 | | - static_summary=static_summary, |
428 | | - static_details=static_details, |
429 | | - start_delay=start_delay, |
430 | | - start_signal=start_signal, |
431 | | - start_signal_args=start_signal_args, |
432 | | - rpc_metadata=rpc_metadata, |
433 | | - rpc_timeout=rpc_timeout, |
434 | | - request_eager_start=request_eager_start, |
435 | | - priority=priority, |
436 | | - versioning_override=versioning_override, |
437 | | - callbacks=self._temporal_context._get_callbacks(), |
438 | | - workflow_event_links=self._temporal_context._get_workflow_event_links(), |
439 | | - request_id=self._temporal_context.nexus_context.request_id, |
440 | | - ) |
| 421 | + with _nexus_backing_workflow_start_context(): |
| 422 | + wf_handle = await self._temporal_context.client.start_workflow( # type: ignore |
| 423 | + workflow=workflow, |
| 424 | + arg=arg, |
| 425 | + args=args, |
| 426 | + id=id, |
| 427 | + task_queue=task_queue or self._temporal_context.info().task_queue, |
| 428 | + result_type=result_type, |
| 429 | + execution_timeout=execution_timeout, |
| 430 | + run_timeout=run_timeout, |
| 431 | + task_timeout=task_timeout, |
| 432 | + id_reuse_policy=id_reuse_policy, |
| 433 | + id_conflict_policy=id_conflict_policy, |
| 434 | + retry_policy=retry_policy, |
| 435 | + cron_schedule=cron_schedule, |
| 436 | + memo=memo, |
| 437 | + search_attributes=search_attributes, |
| 438 | + static_summary=static_summary, |
| 439 | + static_details=static_details, |
| 440 | + start_delay=start_delay, |
| 441 | + start_signal=start_signal, |
| 442 | + start_signal_args=start_signal_args, |
| 443 | + rpc_metadata=rpc_metadata, |
| 444 | + rpc_timeout=rpc_timeout, |
| 445 | + request_eager_start=request_eager_start, |
| 446 | + priority=priority, |
| 447 | + versioning_override=versioning_override, |
| 448 | + callbacks=self._temporal_context._get_callbacks(), |
| 449 | + workflow_event_links=self._temporal_context._get_workflow_event_links(), |
| 450 | + request_id=self._temporal_context.nexus_context.request_id, |
| 451 | + ) |
441 | 452 |
|
442 | 453 | self._temporal_context._add_outbound_links(wf_handle) |
443 | 454 |
|
|
0 commit comments