diff --git a/CHANGELOG.md b/CHANGELOG.md index 483c73d1d..0813f46f4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,9 @@ to include examples, links to docs, or any other relevant information. ### Added +- Added Nexus operation link propagation for Workflow Queries issued from operation handlers. The + queried Workflow link returned by the server is attached to the caller's Nexus operation event. + ### Changed ### Deprecated diff --git a/temporalio/client/_impl.py b/temporalio/client/_impl.py index 9595351f2..e689fd693 100644 --- a/temporalio/client/_impl.py +++ b/temporalio/client/_impl.py @@ -472,6 +472,9 @@ async def query_workflow(self, input: QueryWorkflowInput) -> Any: raise WorkflowQueryFailedError(err.message) else: raise + nexus_ctx = temporalio.nexus._operation_context._try_start_operation_context() + if nexus_ctx is not None and resp.HasField("link"): + nexus_ctx._add_response_link(resp.link) if resp.HasField("query_rejected"): raise WorkflowQueryRejectedError( WorkflowExecutionStatus(resp.query_rejected.status) diff --git a/temporalio/nexus/_operation_context.py b/temporalio/nexus/_operation_context.py index 06ffd9b0f..9d86e1616 100644 --- a/temporalio/nexus/_operation_context.py +++ b/temporalio/nexus/_operation_context.py @@ -300,9 +300,9 @@ def _add_start_workflow_response_link( def _add_response_link(self, link: temporalio.api.common.v1.Link | None) -> None: """Append a response link returned by an RPC the operation handler issued. - ``link`` is the ``common.v1.Link`` returned on a signal, signal-with-start, or start - response (or ``None`` against a server that did not return one). When present, it is - converted to a Nexus link and added to the operation's outbound links. + ``link`` is the ``common.v1.Link`` returned by a Temporal RPC (or ``None`` against a + server that did not return one). When present, it is converted to a Nexus link and added + to the operation's outbound links. This is only safe to call from the single thread/task that runs the operation handler. """ diff --git a/tests/nexus/test_link_propagation.py b/tests/nexus/test_link_propagation.py index 2c9f6eec7..fb95d08f8 100644 --- a/tests/nexus/test_link_propagation.py +++ b/tests/nexus/test_link_propagation.py @@ -1,8 +1,9 @@ """Unit tests for Nexus link propagation. -These exercise link propagation when a Nexus operation handler signals or starts a -workflow or activity against a mocked workflow service. End-to-end signal backlinks -require a server with EnableCHASMSignalBacklinks enabled and are not covered here. +These exercise link propagation when a Nexus operation handler queries, signals, or +starts a workflow or activity against a mocked workflow service. End-to-end signal +backlinks require a server with EnableCHASMSignalBacklinks enabled and are not covered +here. """ from __future__ import annotations @@ -35,6 +36,7 @@ import temporalio.nexus._token from temporalio.client._impl import _ClientImpl from temporalio.client._interceptor import ( + QueryWorkflowInput, SignalWorkflowInput, StartActivityInput, StartWorkflowInput, @@ -63,6 +65,19 @@ def _workflow_event_link( ) +def _workflow_link( + workflow_id: str, run_id: str, *, reason: str +) -> temporalio.api.common.v1.Link: + return temporalio.api.common.v1.Link( + workflow=temporalio.api.common.v1.Link.Workflow( + namespace=NAMESPACE, + workflow_id=workflow_id, + run_id=run_id, + reason=reason, + ) + ) + + def _inbound_nexus_link() -> temporalio.api.common.v1.Link: return _workflow_event_link( "caller-wf", @@ -132,6 +147,20 @@ def _signal_input() -> SignalWorkflowInput: ) +def _query_input() -> QueryWorkflowInput: + return QueryWorkflowInput( + id=WORKFLOW_ID, + run_id=None, + query="query-done", + args=[], + reject_condition=None, + headers={}, + ret_type=bool, + rpc_metadata={}, + rpc_timeout=None, + ) + + def _start_input(start_signal: str | None = None) -> StartWorkflowInput: return StartWorkflowInput( workflow="TestWorkflow", @@ -193,9 +222,56 @@ def _outbound_link_urls(ctx: Any) -> list[str]: return [link.url for link in ctx.nexus_context.outbound_links] +def test_response_link_captures_workflow_link( + nexus_ctx: _TemporalStartOperationContext, +) -> None: + nexus_ctx._add_response_link( + _workflow_link(WORKFLOW_ID, "target-run", reason="Query processed") + ) + + assert nexus_ctx.nexus_context.outbound_links == [ + nexusrpc.Link( + type=temporalio.api.common.v1.Link.Workflow.DESCRIPTOR.full_name, + url=( + "temporal:///namespaces/test-namespace/workflows/" + "wf-target/target-run?reason=Query+processed" + ), + ) + ] + + # ── signal ──────────────────────────────────────────────────────────────────────────────── +# Query responses differ from Signal responses by linking to the Workflow rather than an event. +async def test_query_captures_response_workflow_link( + nexus_ctx: _TemporalStartOperationContext, +) -> None: + payloads = await temporalio.converter.DataConverter.default.encode([False]) + workflow_service = mock.MagicMock() + workflow_service.query_workflow = mock.AsyncMock( + return_value=temporalio.api.workflowservice.v1.QueryWorkflowResponse( + query_result=temporalio.api.common.v1.Payloads(payloads=payloads), + link=_workflow_link(WORKFLOW_ID, "target-run", reason="Query processed"), + ) + ) + impl = _make_client_impl(workflow_service) + + result = await impl.query_workflow(_query_input()) + + assert result is False + assert nexus_ctx.nexus_context.outbound_links == [ + nexusrpc.Link( + type=temporalio.api.common.v1.Link.Workflow.DESCRIPTOR.full_name, + url=( + "temporal:///namespaces/test-namespace/workflows/" + "wf-target/target-run?reason=Query+processed" + ), + ) + ] + + +# Signal responses link to the event that accepted the Signal. async def test_signal_forwards_inbound_links_and_captures_response_backlink( nexus_ctx: _TemporalStartOperationContext, ) -> None: diff --git a/tests/nexus/test_temporal_operation.py b/tests/nexus/test_temporal_operation.py index d966c5cc6..37b0f664f 100644 --- a/tests/nexus/test_temporal_operation.py +++ b/tests/nexus/test_temporal_operation.py @@ -111,6 +111,7 @@ class TestService: sync_result: Operation[Input, str] custom_cancel: Operation[str, None] update_op: Operation[Input, str] + query_op: Operation[str, bool] echo_activity: Operation[Input, str] error_activity: Operation[Input, None] blocking_activity: Operation[str, None] @@ -287,6 +288,17 @@ async def update_op( update_id=input.update_id, ) + @nexus.temporal_operation + async def query_op( + self, + _ctx: nexus.TemporalStartOperationContext, + client: nexus.TemporalNexusClient, + input: str, + ) -> nexus.TemporalOperationResult[bool]: + handle = client.client.get_workflow_handle(input) + result = await handle.query(BlockingWorkflow.query_done) + return nexus.TemporalOperationResult.sync(result) + @nexus.temporal_operation async def echo_activity( self, @@ -817,6 +829,74 @@ async def run(self) -> None: async def unblock(self): self.done = True + @workflow.query + def query_done(self) -> bool: + return self.done + + +@workflow.defn +class QueryWorkflowCaller: + @workflow.run + async def run(self, input: Input) -> bool: + client = workflow.create_nexus_client( + service=TestService, endpoint=make_nexus_endpoint_name(input.task_queue) + ) + return await client.execute_operation(TestService.query_op, input.value) + + +async def test_temporal_operation_query_workflow( + client: Client, env: WorkflowEnvironment +) -> None: + task_queue = str(uuid.uuid4()) + endpoint_name = make_nexus_endpoint_name(task_queue) + await env.create_nexus_endpoint(endpoint_name, task_queue) + target_workflow_id = f"query-target-{uuid.uuid4()}" + + async with Worker( + env.client, + task_queue=task_queue, + nexus_service_handlers=[TestServiceHandler()], + workflows=[BlockingWorkflow, QueryWorkflowCaller], + ): + target_handle = await client.start_workflow( + BlockingWorkflow.run, + id=target_workflow_id, + task_queue=task_queue, + ) + caller_handle = await client.start_workflow( + QueryWorkflowCaller.run, + Input(value=target_workflow_id, task_queue=task_queue), + id=f"query-caller-{uuid.uuid4()}", + task_queue=task_queue, + ) + + try: + assert not await caller_handle.result() + + caller_history = await caller_handle.fetch_history() + completed_event = next( + event + for event in caller_history.events + if event.event_type == EventType.EVENT_TYPE_NEXUS_OPERATION_COMPLETED + ) + + target_history = await target_handle.fetch_history() + assert not any(event.links for event in target_history.events) + + if not completed_event.links: + pytest.skip("server did not return a Workflow Query response link") + assert target_handle.result_run_id is not None + assert Link( + workflow=Link.Workflow( + namespace=client.namespace, + workflow_id=target_workflow_id, + run_id=target_handle.result_run_id, + reason="Query processed", + ) + ) in list(completed_event.links) + finally: + await target_handle.cancel() + @workflow.defn class CancelBlockingWorkflowCaller: