Sandbox and type fixes (#202)

Fixes #200
Fixes #201
Fixes #199
Fixes #198

Author: cretz

temporalio/bridge/runtime.py

@@ -147,5 +147,5 @@ class TelemetryConfig:
     logging: Optional[LoggingConfig] = LoggingConfig.default
     """Logging configuration."""
 
-    metrics: Optional[PrometheusConfig] = None
+    metrics: Optional[MetricsConfig] = None
     """Metrics configuration."""

temporalio/contrib/opentelemetry.py

@@ -35,6 +35,14 @@
 import temporalio.worker
 import temporalio.workflow
 
+# OpenTelemetry dynamically, lazily chooses its context implementation at
+# runtime. When first accessed, they use pkg_resources.iter_entry_points + load.
+# The load uses built-in open() which we don't allow in sandbox mode at runtime,
+# only import time. Therefore if the first use of a OTel context is inside the
+# sandbox, which it may be for a workflow worker, this will fail. So instead we
+# eagerly reference it here to force loading at import time instead of lazily.
+opentelemetry.context.get_current()
+
 default_text_map_propagator = opentelemetry.propagators.composite.CompositePropagator(
     [
         opentelemetry.trace.propagation.tracecontext.TraceContextTextMapPropagator(),

temporalio/worker/workflow_sandbox/_in_sandbox.py

@@ -11,6 +11,7 @@
 import temporalio.bridge.proto.workflow_activation
 import temporalio.bridge.proto.workflow_completion
 import temporalio.worker._workflow_instance
+import temporalio.workflow
 
 logger = logging.getLogger(__name__)
 
@@ -34,22 +35,10 @@ def __init__(
     ) -> None:
         """Create in-sandbox instance."""
         _trace("Initializing workflow %s in sandbox", workflow_class)
-        # We have to replace the given instance instance details with new one
-        # replacing references to the workflow class
-        old_defn = instance_details.defn
-        new_defn = dataclasses.replace(
-            old_defn,
-            cls=workflow_class,
-            run_fn=getattr(workflow_class, old_defn.run_fn.__name__),
-            signals={
-                k: dataclasses.replace(v, fn=getattr(workflow_class, v.fn.__name__))
-                for k, v in old_defn.signals.items()
-            },
-            queries={
-                k: dataclasses.replace(v, fn=getattr(workflow_class, v.fn.__name__))
-                for k, v in old_defn.queries.items()
-            },
-        )
+        # We expect to be able to get the workflow definition back off the
+        # class. We can't use the definition that was given to us because it has
+        # type hints and references to outside-of-sandbox types.
+        new_defn = temporalio.workflow._Definition.must_from_class(workflow_class)
         new_instance_details = dataclasses.replace(instance_details, defn=new_defn)
 
         # Instantiate the runner and the instance

temporalio/worker/workflow_sandbox/_restrictions.py

@@ -437,7 +437,9 @@ def _public_callables(parent: Any, *, exclude: Set[str] = set()) -> Set[str]:
         # TODO(cretz): Fix issues with class extensions on restricted proxy
         # "argparse": SandboxMatcher.all_uses_runtime,
         "bz2": SandboxMatcher(use={"open"}),
-        "concurrent": SandboxMatcher.all_uses_runtime,
+        "concurrent": SandboxMatcher(
+            children={"futures": SandboxMatcher.all_uses_runtime}
+        ),
         # Python's own re lib registers itself. This is mostly ok to not
         # restrict since it's just global picklers that people may want to
         # register globally.

tests/worker/workflow_sandbox/test_runner.py

@@ -7,13 +7,14 @@
 import time
 import uuid
 from dataclasses import dataclass
-from datetime import date
+from datetime import date, timedelta
+from enum import IntEnum
 from typing import Callable, Dict, List, Optional, Sequence, Type
 
 import pytest
 
 import temporalio.worker.workflow_sandbox._restrictions
-from temporalio import workflow
+from temporalio import activity, workflow
 from temporalio.client import Client, WorkflowFailureError, WorkflowHandle
 from temporalio.exceptions import ApplicationError
 from temporalio.worker import Worker
@@ -205,6 +206,8 @@ async def test_workflow_sandbox_restrictions(client: Client):
             # General library allowed calls
             "import datetime\ndatetime.date(2001, 1, 1)",
             "import uuid\nuuid.uuid5(uuid.NAMESPACE_DNS, 'example.com')",
+            # Other imports we want to allow
+            "from concurrent.futures import ThreadPoolExecutor",
         ]
         for code in valid_code_to_check:
             await client.execute_workflow(
@@ -310,6 +313,53 @@ async def test_workflow_sandbox_access_stack(client: Client):
         )
 
 
+class InstanceCheckEnum(IntEnum):
+    FOO = 1
+    BAR = 2
+
+
+@dataclass
+class InstanceCheckData:
+    some_enum: InstanceCheckEnum
+
+
+@activity.defn
+async def instance_check_activity(param: InstanceCheckData) -> InstanceCheckData:
+    assert isinstance(param, InstanceCheckData)
+    assert param.some_enum is InstanceCheckEnum.BAR
+    return param
+
+
+@workflow.defn
+class InstanceCheckWorkflow:
+    @workflow.run
+    async def run(self, param: InstanceCheckData) -> InstanceCheckData:
+        assert isinstance(param, InstanceCheckData)
+        assert param.some_enum is InstanceCheckEnum.BAR
+        # Exec child if not a child, otherwise exec activity
+        if workflow.info().parent is None:
+            return await workflow.execute_child_workflow(
+                InstanceCheckWorkflow.run, param
+            )
+        return await workflow.execute_activity(
+            instance_check_activity,
+            param,
+            schedule_to_close_timeout=timedelta(minutes=1),
+        )
+
+
+async def test_workflow_sandbox_instance_check(client: Client):
+    async with new_worker(
+        client, InstanceCheckWorkflow, activities=[instance_check_activity]
+    ) as worker:
+        await client.execute_workflow(
+            InstanceCheckWorkflow.run,
+            InstanceCheckData(some_enum=InstanceCheckEnum.BAR),
+            id=f"workflow-{uuid.uuid4()}",
+            task_queue=worker.task_queue,
+        )
+
+
 def new_worker(
     client: Client,
     *workflows: Type,