Expose client as activity.client()

Author: dandavison

temporalio/activity.py

@@ -19,6 +19,7 @@
 from dataclasses import dataclass
 from datetime import datetime, timedelta
 from typing import (
+    TYPE_CHECKING,
     Any,
     Callable,
     Iterator,
@@ -42,6 +43,9 @@
 
 from .types import CallableType
 
+if TYPE_CHECKING:
+    from temporalio.client import Client
+
 
 @overload
 def defn(fn: CallableType) -> CallableType: ...
@@ -179,6 +183,7 @@ class _Context:
         temporalio.converter.PayloadConverter,
     ]
     runtime_metric_meter: Optional[temporalio.common.MetricMeter]
+    client: Optional[Client]
     cancellation_details: _ActivityCancellationDetailsHolder
     _logger_details: Optional[Mapping[str, Any]] = None
     _payload_converter: Optional[temporalio.converter.PayloadConverter] = None
@@ -271,13 +276,32 @@ def wait_sync(self, timeout: Optional[float] = None) -> None:
         self.thread_event.wait(timeout)
 
 
+def client() -> Client:
+    """Return a Temporal Client for use in the current activity.
+
+    Returns:
+        :py:class:`temporalio.client.Client` for use in the current activity.
+
+    Raises:
+        RuntimeError: When not in an activity.
+    """
+    client = _Context.current().client
+    if not client:
+        raise RuntimeError(
+            "No client available. The client is only available in async "
+            "(i.e. `async def`) activities; not in sync (i.e. `def`) activities. "
+            "In tests you can pass a client when creating ActivityEnvironment."
+        )
+    return client
+
+
 def in_activity() -> bool:
     """Whether the current code is inside an activity.
 
     Returns:
         True if in an activity, False otherwise.
     """
-    return not _current_context.get(None) is None
+    return _current_context.get(None) is not None
 
 
 def info() -> Info:
@@ -574,8 +598,10 @@ def _apply_to_callable(
                 fn=fn,
                 # iscoroutinefunction does not return true for async __call__
                 # TODO(cretz): Why can't MyPy handle this?
-                is_async=inspect.iscoroutinefunction(fn)
-                or inspect.iscoroutinefunction(fn.__call__),  # type: ignore
+                is_async=(
+                    inspect.iscoroutinefunction(fn)
+                    or inspect.iscoroutinefunction(fn.__call__)  # type: ignore
+                ),
                 no_thread_cancel_exception=no_thread_cancel_exception,
             ),
         )

temporalio/testing/_activity.py

@@ -16,6 +16,7 @@
 import temporalio.converter
 import temporalio.exceptions
 import temporalio.worker._activity
+from temporalio.client import Client
 
 _Params = ParamSpec("_Params")
 _Return = TypeVar("_Return")
@@ -63,7 +64,7 @@ class ActivityEnvironment:
             take effect. Default is noop.
     """
 
-    def __init__(self) -> None:
+    def __init__(self, client: Optional[Client] = None) -> None:
         """Create an ActivityEnvironment for running activity code."""
         self.info = _default_info
         self.on_heartbeat: Callable[..., None] = lambda *args: None
@@ -74,6 +75,7 @@ def __init__(self) -> None:
         self._cancelled = False
         self._worker_shutdown = False
         self._activities: Set[_Activity] = set()
+        self._client = client
         self._cancellation_details = (
             temporalio.activity._ActivityCancellationDetailsHolder()
         )
@@ -128,18 +130,21 @@ def run(
             The callable's result.
         """
         # Create an activity and run it
-        return _Activity(self, fn).run(*args, **kwargs)
+        return _Activity(self, fn, self._client).run(*args, **kwargs)
 
 
 class _Activity:
     def __init__(
         self,
         env: ActivityEnvironment,
         fn: Callable,
+        client: Optional[Client],
     ) -> None:
         self.env = env
         self.fn = fn
-        self.is_async = inspect.iscoroutinefunction(fn)
+        self.is_async = inspect.iscoroutinefunction(fn) or inspect.iscoroutinefunction(
+            fn.__call__  # type: ignore
+        )
         self.cancel_thread_raiser: Optional[
             temporalio.worker._activity._ThreadExceptionRaiser
         ] = None
@@ -163,11 +168,14 @@ def __init__(
                 thread_event=threading.Event(),
                 async_event=asyncio.Event() if self.is_async else None,
             ),
-            shield_thread_cancel_exception=None
-            if not self.cancel_thread_raiser
-            else self.cancel_thread_raiser.shielded,
+            shield_thread_cancel_exception=(
+                None
+                if not self.cancel_thread_raiser
+                else self.cancel_thread_raiser.shielded
+            ),
             payload_converter_class_or_instance=env.payload_converter,
             runtime_metric_meter=env.metric_meter,
+            client=client if self.is_async else None,
             cancellation_details=env._cancellation_details,
         )
         self.task: Optional[asyncio.Task] = None

temporalio/worker/_activity.py

@@ -69,6 +69,7 @@ def __init__(
         data_converter: temporalio.converter.DataConverter,
         interceptors: Sequence[Interceptor],
         metric_meter: temporalio.common.MetricMeter,
+        client: temporalio.client.Client,
         encode_headers: bool,
     ) -> None:
         self._bridge_worker = bridge_worker
@@ -86,6 +87,7 @@ def __init__(
             None
         )
         self._seen_sync_activity = False
+        self._client = client
 
         # Validate and build activity dict
         self._activities: Dict[str, temporalio.activity._Definition] = {}
@@ -569,11 +571,14 @@ async def _execute_activity(
                 heartbeat=None,
                 cancelled_event=running_activity.cancelled_event,
                 worker_shutdown_event=self._worker_shutdown_event,
-                shield_thread_cancel_exception=None
-                if not running_activity.cancel_thread_raiser
-                else running_activity.cancel_thread_raiser.shielded,
+                shield_thread_cancel_exception=(
+                    None
+                    if not running_activity.cancel_thread_raiser
+                    else running_activity.cancel_thread_raiser.shielded
+                ),
                 payload_converter_class_or_instance=self._data_converter.payload_converter,
                 runtime_metric_meter=None if sync_non_threaded else self._metric_meter,
+                client=self._client if not running_activity.sync else None,
                 cancellation_details=running_activity.cancellation_details,
             )
         )
@@ -838,11 +843,12 @@ def _execute_sync_activity(
             worker_shutdown_event=temporalio.activity._CompositeEvent(
                 thread_event=worker_shutdown_event, async_event=None
             ),
-            shield_thread_cancel_exception=None
-            if not cancel_thread_raiser
-            else cancel_thread_raiser.shielded,
+            shield_thread_cancel_exception=(
+                None if not cancel_thread_raiser else cancel_thread_raiser.shielded
+            ),
             payload_converter_class_or_instance=payload_converter_class_or_instance,
             runtime_metric_meter=runtime_metric_meter,
+            client=None,
             cancellation_details=cancellation_details,
         )
     )

temporalio/worker/_worker.py

@@ -411,8 +411,10 @@ def __init__(
                 data_converter=client_config["data_converter"],
                 interceptors=interceptors,
                 metric_meter=self._runtime.metric_meter,
-                encode_headers=client_config["header_codec_behavior"]
-                == HeaderCodecBehavior.CODEC,
+                client=client,
+                encode_headers=(
+                    client_config["header_codec_behavior"] == HeaderCodecBehavior.CODEC
+                ),
             )
         self._nexus_worker: Optional[_NexusWorker] = None
         if nexus_service_handlers:

tests/testing/test_activity.py

@@ -2,8 +2,12 @@
 import threading
 import time
 from contextvars import copy_context
+from unittest.mock import Mock
+
+import pytest
 
 from temporalio import activity
+from temporalio.client import Client
 from temporalio.exceptions import CancelledError
 from temporalio.testing import ActivityEnvironment
 
@@ -122,3 +126,30 @@ async def assert_equals(a: str, b: str) -> None:
 
     assert type(expected_err) == type(actual_err)
     assert str(expected_err) == str(actual_err)
+
+
+async def test_error_on_access_client_in_activity_environment_without_client():
+    saw_error: bool = False
+
+    def my_activity() -> None:
+        with pytest.raises(RuntimeError, match="No client available"):
+            activity.client()
+        nonlocal saw_error
+        saw_error = True
+
+    env = ActivityEnvironment()
+    env.run(my_activity)
+    assert saw_error
+
+
+async def test_access_client_in_activity_environment_with_client():
+    got_client: bool = False
+
+    def my_activity() -> None:
+        nonlocal got_client
+        if activity.client():
+            got_client = True
+
+    env = ActivityEnvironment(client=Mock(spec=Client))
+    env.run(my_activity)
+    assert got_client

tests/worker/test_activity.py

@@ -94,6 +94,50 @@ async def get_name(name: str) -> str:
     assert result.result == "Name: my custom activity name!"
 
 
+async def test_client_available_in_async_activities(
+    client: Client, worker: ExternalWorker
+):
+    with pytest.raises(RuntimeError) as err:
+        activity.client()
+    assert str(err.value) == "Not in activity context"
+
+    captured_client: Optional[Client] = None
+
+    @activity.defn
+    async def capture_client() -> None:
+        nonlocal captured_client
+        captured_client = activity.client()
+
+    await _execute_workflow_with_activity(client, worker, capture_client)
+    assert captured_client is client
+
+
+async def test_client_not_available_in_sync_activities(
+    client: Client, worker: ExternalWorker
+):
+    saw_error = False
+
+    @activity.defn
+    def some_activity() -> None:
+        with pytest.raises(
+            RuntimeError, match="The client is only available in async"
+        ) as err:
+            activity.client()
+        nonlocal saw_error
+        saw_error = True
+
+    await _execute_workflow_with_activity(
+        client,
+        worker,
+        some_activity,
+        worker_config={
+            "activity_executor": concurrent.futures.ThreadPoolExecutor(1),
+            "max_concurrent_activities": 1,
+        },
+    )
+    assert saw_error
+
+
 async def test_activity_info(
     client: Client, worker: ExternalWorker, env: WorkflowEnvironment
 ):