Make it easier to modify sandbox restrictions

Author: tconley1428

temporalio/worker/workflow_sandbox/_runner.py

@@ -7,6 +7,7 @@
 from __future__ import annotations
 
 import threading
+from dataclasses import dataclass
 from datetime import datetime, timedelta, timezone
 from typing import Any, Optional, Sequence, Type
 
@@ -52,9 +53,12 @@
 )
 
 
+@dataclass
 class SandboxedWorkflowRunner(WorkflowRunner):
     """Runner for workflows in a sandbox."""
 
+    restrictions: SandboxRestrictions = SandboxRestrictions.default
+
     def __init__(
         self,
         *,
@@ -70,8 +74,8 @@ def __init__(
                 re-imported and instantiated for *each* workflow run.
         """
         super().__init__()
+        self.restrictions = restrictions
         self._runner_class = runner_class
-        self._restrictions = restrictions
         self._worker_level_failure_exception_types: Sequence[type[BaseException]] = []
 
     def prepare_workflow(self, defn: temporalio.workflow._Definition) -> None:
@@ -94,7 +98,7 @@ def prepare_workflow(self, defn: temporalio.workflow._Definition) -> None:
 
     def create_instance(self, det: WorkflowInstanceDetails) -> WorkflowInstance:
         """Implements :py:meth:`WorkflowRunner.create_instance`."""
-        return _Instance(det, self._runner_class, self._restrictions)
+        return _Instance(det, self._runner_class, self.restrictions)
 
     def set_worker_level_failure_exception_types(
         self, types: Sequence[type[BaseException]]

tests/test_plugins.py

@@ -1,4 +1,5 @@
 import warnings
+from typing import cast
 
 import pytest
 
@@ -7,6 +8,7 @@
 from temporalio.client import Client, ClientConfig, OutboundInterceptor
 from temporalio.testing import WorkflowEnvironment
 from temporalio.worker import Worker, WorkerConfig
+from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner
 from tests.worker.test_worker import never_run_activity
 
 
@@ -64,6 +66,9 @@ def configure_worker(self, config: WorkerConfig) -> WorkerConfig:
 class MyWorkerPlugin(temporalio.worker.Plugin):
     def configure_worker(self, config: WorkerConfig) -> WorkerConfig:
         config["task_queue"] = "replaced_queue"
+        runner = config.get("workflow_runner")
+        if isinstance(runner, SandboxedWorkflowRunner):
+            runner.restrictions.passthrough_modules.add("my_module")
         return super().configure_worker(config)
 
     async def run_worker(self, worker: Worker) -> None:
@@ -111,3 +116,19 @@ async def test_worker_duplicated_plugin(client: Client) -> None:
 
     assert len(warning_list) == 1
     assert "The same plugin type" in str(warning_list[0].message)
+
+
+async def test_worker_sandbox_restrictions(client: Client) -> None:
+    with warnings.catch_warnings(record=True) as warning_list:
+        worker = Worker(
+            client,
+            task_queue="queue",
+            activities=[never_run_activity],
+            plugins=[MyWorkerPlugin()],
+        )
+    assert (
+        "my_module"
+        in cast(
+            SandboxedWorkflowRunner, worker.config().get("workflow_runner")
+        ).restrictions.passthrough_modules
+    )