Wire context for signal, query, and update

Author: dandavison

temporalio/client.py

@@ -62,6 +62,7 @@
 import temporalio.service
 import temporalio.workflow
 from temporalio.activity import ActivityCancellationDetails
+from temporalio.converter import WorkflowSerializationContext
 from temporalio.service import (
     HttpConnectProxyConfig,
     KeepAliveConfig,
@@ -6062,8 +6063,14 @@ async def query_workflow(self, input: QueryWorkflowInput) -> Any:
             )
         req.query.query_type = input.query
         if input.args:
+            context = WorkflowSerializationContext(
+                namespace=self._client.namespace,
+                workflow_id=input.id,
+            )
             req.query.query_args.payloads.extend(
-                await self._client.data_converter.encode(input.args)
+                await self._client.data_converter._with_context(context).encode(
+                    input.args
+                )
             )
         if input.headers is not None:
             await self._apply_headers(input.headers, req.query.header.fields)
@@ -6087,7 +6094,11 @@ async def query_workflow(self, input: QueryWorkflowInput) -> Any:
         if not resp.query_result.payloads:
             return None
         type_hints = [input.ret_type] if input.ret_type else None
-        results = await self._client.data_converter.decode(
+        context = WorkflowSerializationContext(
+            namespace=self._client.namespace,
+            workflow_id=input.id,
+        )
+        results = await self._client.data_converter._with_context(context).decode(
             resp.query_result.payloads, type_hints
         )
         if not results:
@@ -6108,8 +6119,14 @@ async def signal_workflow(self, input: SignalWorkflowInput) -> None:
             request_id=str(uuid.uuid4()),
         )
         if input.args:
+            context = temporalio.converter.WorkflowSerializationContext(
+                namespace=self._client.namespace,
+                workflow_id=input.id,
+            )
             req.input.payloads.extend(
-                await self._client.data_converter.encode(input.args)
+                await self._client.data_converter._with_context(context).encode(
+                    input.args
+                )
             )
         if input.headers is not None:
             await self._apply_headers(input.headers, req.header.fields)

temporalio/worker/_workflow_instance.py

@@ -576,12 +576,17 @@ async def run_update() -> None:
                 self._in_progress_updates[job.id] = HandlerExecution(
                     job.name, defn.unfinished_policy, job.id
                 )
+                context = temporalio.converter.WorkflowSerializationContext(
+                    namespace=self._info.namespace,
+                    workflow_id=self._info.workflow_id,
+                )
                 args = self._process_handler_args(
                     job.name,
                     job.input,
                     defn.name,
                     defn.arg_types,
                     defn.dynamic_vararg,
+                    context,
                 )
                 handler_input = HandleUpdateInput(
                     id=job.id,
@@ -600,6 +605,7 @@ async def run_update() -> None:
                             defn.name,
                             defn.arg_types,
                             defn.dynamic_vararg,
+                            context,
                         )
                         handler_input.args = args
 
@@ -692,12 +698,17 @@ async def run_query() -> None:
                         )
 
                     # Create input
+                    context = temporalio.converter.WorkflowSerializationContext(
+                        namespace=self._info.namespace,
+                        workflow_id=self._info.workflow_id,
+                    )
                     args = self._process_handler_args(
                         job.query_type,
                         job.arguments,
                         defn.name,
                         defn.arg_types,
                         defn.dynamic_vararg,
+                        context,
                     )
                     input = HandleQueryInput(
                         id=job.query_id,
@@ -706,7 +717,12 @@ async def run_query() -> None:
                         headers=job.headers,
                     )
                     success = await self._inbound.handle_query(input)
-                    result_payloads = self._payload_converter.to_payloads([success])
+                    converter = self._payload_converter
+                    if isinstance(
+                        converter, temporalio.converter.WithSerializationContext
+                    ):
+                        converter = converter.with_context(context)
+                    result_payloads = converter.to_payloads([success])
                     if len(result_payloads) != 1:
                         raise ValueError(
                             f"Expected 1 result payload, got {len(result_payloads)}"
@@ -716,11 +732,20 @@ async def run_query() -> None:
                 command.respond_to_query.succeeded.response.CopyFrom(result_payloads[0])
             except Exception as err:
                 try:
+                    context = temporalio.converter.WorkflowSerializationContext(
+                        namespace=self._info.namespace,
+                        workflow_id=self._info.workflow_id,
+                    )
+                    converter = self._payload_converter
+                    if isinstance(
+                        converter, temporalio.converter.WithSerializationContext
+                    ):
+                        converter = converter.with_context(context)
                     command = self._add_command()
                     command.respond_to_query.query_id = job.query_id
                     self._failure_converter.to_failure(
                         err,
-                        self._payload_converter,
+                        converter,
                         command.respond_to_query.failed,
                     )
                 except Exception as inner_err:
@@ -2106,35 +2131,41 @@ def _process_handler_args(
         defn_name: Optional[str],
         defn_arg_types: Optional[List[Type]],
         defn_dynamic_vararg: bool,
+        context: temporalio.converter.SerializationContext,
     ) -> List[Any]:
         # If dynamic old-style vararg, args become name + varargs of given arg
         # types. If dynamic new-style raw value sequence, args become name +
         # seq of raw values.
         if not defn_name and defn_dynamic_vararg:
             # Take off the string type hint for conversion
             arg_types = defn_arg_types[1:] if defn_arg_types else None
-            return [job_name] + self._convert_payloads(job_input, arg_types)
+            return [job_name] + self._convert_payloads(job_input, arg_types, context)
         if not defn_name:
             return [
                 job_name,
                 self._convert_payloads(
-                    job_input, [temporalio.common.RawValue] * len(job_input)
+                    job_input, [temporalio.common.RawValue] * len(job_input), context
                 ),
             ]
-        return self._convert_payloads(job_input, defn_arg_types)
+        return self._convert_payloads(job_input, defn_arg_types, context)
 
     def _process_signal_job(
         self,
         defn: temporalio.workflow._SignalDefinition,
         job: temporalio.bridge.proto.workflow_activation.SignalWorkflow,
     ) -> None:
         try:
+            context = temporalio.converter.WorkflowSerializationContext(
+                namespace=self._info.namespace,
+                workflow_id=self._info.workflow_id,
+            )
             args = self._process_handler_args(
                 job.signal_name,
                 job.input,
                 defn.name,
                 defn.arg_types,
                 defn.dynamic_vararg,
+                context,
             )
         except Exception:
             logger.exception(