diff --git a/CHANGELOG.md b/CHANGELOG.md index f8fe94d4b..62426485b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -71,6 +71,9 @@ to include examples, links to docs, or any other relevant information. - Marked system Nexus envelope payloads so nested payloads can be detected and visited after the envelope is already stored as a payload. +- Suppressed intentional passthrough notifications for the initial workflow + module import while retaining configured notifications for that module's + import-time dependencies. ### Security diff --git a/temporalio/worker/workflow_sandbox/_importer.py b/temporalio/worker/workflow_sandbox/_importer.py index 1ab0a1dd6..78c52615d 100644 --- a/temporalio/worker/workflow_sandbox/_importer.py +++ b/temporalio/worker/workflow_sandbox/_importer.py @@ -66,6 +66,7 @@ def __init__( "__main__": types.ModuleType("__main__"), } self.modules_checked_for_restrictions: set[str] = set() + self._initial_workflow_module_import: str | None = None self.import_func = self._import if not LOG_TRACE else self._traced_import # Pre-collect restricted builtins self.restricted_builtins: list[tuple[str, _ThreadLocalCallable, Callable]] = [] @@ -159,6 +160,16 @@ def applied(self) -> Iterator[None]: finally: Importer._thread_local_current.importer = orig_importer + @contextmanager + def initial_workflow_module_import(self, module_name: str) -> Iterator[None]: + """Suppress notifications only for the initial workflow module import.""" + previous_module = self._initial_workflow_module_import + self._initial_workflow_module_import = module_name + try: + yield None + finally: + self._initial_workflow_module_import = previous_module + @contextmanager def _unapplied(self) -> Iterator[None]: orig_importer = Importer.current_importer() @@ -307,19 +318,30 @@ def _is_import_notification_policy_applied( return policy in self.restrictions.import_notification_policy def _maybe_passthrough_module(self, name: str) -> types.ModuleType | None: + # The workflow module itself must be loaded into every sandbox instance. + # That import is intentional, but its import-time dependencies still need + # to follow the configured notification policy. + is_initial_workflow_module = name == self._initial_workflow_module_import + # If imports not passed through and all modules are not passed through # and name not in passthrough modules, check parents if ( not temporalio.workflow.unsafe.is_imports_passed_through() and not self.module_configured_passthrough(name) ): - if self._is_import_notification_policy_applied( - temporalio.workflow.SandboxImportNotificationPolicy.RAISE_ON_UNINTENTIONAL_PASSTHROUGH + if ( + not is_initial_workflow_module + and self._is_import_notification_policy_applied( + temporalio.workflow.SandboxImportNotificationPolicy.RAISE_ON_UNINTENTIONAL_PASSTHROUGH + ) ): raise UnintentionalPassthroughError(name) - if self._is_import_notification_policy_applied( - temporalio.workflow.SandboxImportNotificationPolicy.WARN_ON_UNINTENTIONAL_PASSTHROUGH + if ( + not is_initial_workflow_module + and self._is_import_notification_policy_applied( + temporalio.workflow.SandboxImportNotificationPolicy.WARN_ON_UNINTENTIONAL_PASSTHROUGH + ) ): warnings.warn( f"Module {name} was not intentionally passed through to the sandbox." diff --git a/temporalio/worker/workflow_sandbox/_runner.py b/temporalio/worker/workflow_sandbox/_runner.py index 7f06bfcd6..98c442018 100644 --- a/temporalio/worker/workflow_sandbox/_runner.py +++ b/temporalio/worker/workflow_sandbox/_runner.py @@ -135,14 +135,17 @@ def _create_instance(self) -> None: if module_name == "__main__": module_name = "__temporal_main__" try: - # Import user code - self._run_code( - "with __temporal_importer.applied():\n" - # Import the workflow code - f" from {module_name} import {self.instance_details.defn.cls.__name__} as __temporal_workflow_class\n" - f" from {self.runner_class.__module__} import {self.runner_class.__name__} as __temporal_runner_class\n", - __temporal_importer=self.importer, - ) + # Import user code. The workflow module is intentionally loaded + # into every sandbox instance, while its import-time dependencies + # must retain the configured warning/error policy. + with self.importer.initial_workflow_module_import(module_name): + self._run_code( + "with __temporal_importer.applied():\n" + # Import the workflow code + f" from {module_name} import {self.instance_details.defn.cls.__name__} as __temporal_workflow_class\n" + f" from {self.runner_class.__module__} import {self.runner_class.__name__} as __temporal_runner_class\n", + __temporal_importer=self.importer, + ) # Set context as in runtime self.importer.restriction_context.is_runtime = True diff --git a/tests/worker/workflow_sandbox/test_runner.py b/tests/worker/workflow_sandbox/test_runner.py index 288da0861..f1e97bddb 100644 --- a/tests/worker/workflow_sandbox/test_runner.py +++ b/tests/worker/workflow_sandbox/test_runner.py @@ -525,6 +525,59 @@ async def run(self) -> None: ) from err +async def test_workflow_sandbox_initial_workflow_import_does_not_warn() -> None: + restrictions = dataclasses.replace( + SandboxRestrictions.default, + import_notification_policy=SandboxImportNotificationPolicy.WARN_ON_UNINTENTIONAL_PASSTHROUGH, + ) + + with warnings.catch_warnings(record=True) as recorder: + warnings.simplefilter("always") + definition = workflow._Definition.from_class(LazyImportWorkflow) + assert definition is not None + SandboxedWorkflowRunner(restrictions).prepare_workflow(definition) + + assert ( + "Module tests.worker.workflow_sandbox.test_runner was not intentionally " + "passed through to the sandbox." + not in {str(warning.message) for warning in recorder} + ) + + +async def test_workflow_sandbox_initial_workflow_import_warns_for_dependencies() -> ( + None +): + from tests.worker.workflow_sandbox.testmodules.initial_import_warning_dependency import ( + InitialImportWarningDependencyWorkflow, + ) + + restrictions = dataclasses.replace( + SandboxRestrictions.default, + import_notification_policy=SandboxImportNotificationPolicy.WARN_ON_UNINTENTIONAL_PASSTHROUGH, + ) + + with warnings.catch_warnings(record=True) as recorder: + warnings.simplefilter("always") + definition = workflow._Definition.from_class( + InitialImportWarningDependencyWorkflow + ) + assert definition is not None + SandboxedWorkflowRunner(restrictions).prepare_workflow(definition) + + assert ( + "Module tests.worker.workflow_sandbox.testmodules.lazy_module was not " + "intentionally passed through to the sandbox." + in {str(warning.message) for warning in recorder} + ) + + raising_restrictions = dataclasses.replace( + restrictions, + import_notification_policy=SandboxImportNotificationPolicy.RAISE_ON_UNINTENTIONAL_PASSTHROUGH, + ) + with pytest.raises(UnintentionalPassthroughError, match="lazy_module"): + SandboxedWorkflowRunner(raising_restrictions).prepare_workflow(definition) + + async def test_workflow_sandbox_import_default_warnings(client: Client): restrictions = dataclasses.replace( SandboxRestrictions.default, diff --git a/tests/worker/workflow_sandbox/testmodules/initial_import_warning_dependency.py b/tests/worker/workflow_sandbox/testmodules/initial_import_warning_dependency.py new file mode 100644 index 000000000..bcc28753c --- /dev/null +++ b/tests/worker/workflow_sandbox/testmodules/initial_import_warning_dependency.py @@ -0,0 +1,9 @@ +import tests.worker.workflow_sandbox.testmodules.lazy_module # noqa: F401 +from temporalio import workflow + + +@workflow.defn +class InitialImportWarningDependencyWorkflow: + @workflow.run + async def run(self) -> None: + pass