💥 feat(extstore): pass a dedicated select context to the driver selector

CHANGELOG.md

@@ -26,6 +26,11 @@ to include examples, links to docs, or any other relevant information.
 
 ### :boom: Breaking Changes
 
+- Experimental external storage: `ExternalStorage.driver_selector` is now called with a
+  `StorageDriverSelectContext` instead of a `StorageDriverStoreContext`. Update the annotation;
+  the new type carries the same `target` field. Since selectors are plain callables, a stale
+  annotation fails type checking rather than at runtime.
+
 ### Fixed
 
 ### Security

temporalio/client/__init__.py

@@ -8,6 +8,7 @@
     DataConverter,
     SerializationContext,
     StorageDriverActivityInfo,
+    StorageDriverSelectContext,
     StorageDriverStoreContext,
     StorageDriverWorkflowInfo,
     WithSerializationContext,
@@ -351,6 +352,7 @@
     "DataConverter",
     "SerializationContext",
     "StorageDriverActivityInfo",
+    "StorageDriverSelectContext",
     "StorageDriverStoreContext",
     "StorageDriverWorkflowInfo",
     "WithSerializationContext",

temporalio/converter/__init__.py

@@ -10,6 +10,7 @@
     StorageDriverActivityInfo,
     StorageDriverClaim,
     StorageDriverRetrieveContext,
+    StorageDriverSelectContext,
     StorageDriverStoreContext,
     StorageDriverWorkflowInfo,
     StorageWarning,
@@ -61,6 +62,7 @@
     "StorageDriverActivityInfo",
     "StorageDriverClaim",
     "StorageDriverRetrieveContext",
+    "StorageDriverSelectContext",
     "StorageDriverStoreContext",
     "StorageDriverWorkflowInfo",
     "StorageWarning",

temporalio/converter/_extstore.py

@@ -143,7 +143,25 @@ class StorageDriverActivityInfo:
 
 @dataclass(frozen=True)
 class StorageDriverStoreContext:
-    """Context passed to :meth:`StorageDriver.store` and ``driver_selector`` calls.
+    """Context passed to :meth:`StorageDriver.store` calls.
+
+    .. warning::
+        This API is experimental.
+    """
+
+    target: StorageDriverActivityInfo | StorageDriverWorkflowInfo | None = None
+    """The workflow or activity for which this payload is being stored.
+
+    For payloads being stored on behalf of an explicit target (e.g. a child
+    workflow being started, an activity being scheduled, an external workflow
+    being signaled), this is that target's identity.  When no explicit target
+    exists the current execution context (workflow or activity) is used as the
+    target instead."""
+
+
+@dataclass(frozen=True)
+class StorageDriverSelectContext:
+    """Context passed to :attr:`ExternalStorage.driver_selector` calls.
 
     .. warning::
         This API is experimental.
@@ -257,7 +275,7 @@ class ExternalStorage:
     """
 
     driver_selector: (
-        Callable[[StorageDriverStoreContext, Payload], StorageDriver | None] | None
+        Callable[[StorageDriverSelectContext, Payload], StorageDriver | None] | None
     ) = None
     """Controls which driver stores a given payload. A callable that returns the
     driver instance to use, or ``None`` to leave the payload stored inline.
@@ -288,6 +306,14 @@ class ExternalStorage:
     )
     """Store context bound to this instance via :meth:`_with_store_context`."""
 
+    _select_context: StorageDriverSelectContext = dataclasses.field(
+        default=StorageDriverSelectContext(target=None),
+        init=False,
+        repr=False,
+        compare=False,
+    )
+    """Selector context derived from :attr:`_store_context`."""
+
     _claim_converter: ClassVar[JSONProtoPayloadConverter] = JSONProtoPayloadConverter()
     _legacy_claim_converter: ClassVar[JSONPlainPayloadConverter] = (
         JSONPlainPayloadConverter(encoding=_REFERENCE_ENCODING.decode())
@@ -325,7 +351,7 @@ def __post_init__(self) -> None:
         object.__setattr__(self, "_driver_map", driver_map)
 
     def _select_driver(
-        self, context: StorageDriverStoreContext, payload: Payload
+        self, context: StorageDriverSelectContext, payload: Payload
     ) -> StorageDriver | None:
         """Returns the driver to use for this payload, or None to pass through."""
         if payload.ByteSize() < self.payload_size_threshold:
@@ -354,12 +380,15 @@ def _with_store_context(self, ctx: StorageDriverStoreContext) -> ExternalStorage
         """Return a copy of this instance with ``ctx`` bound as the store context."""
         result = dataclasses.replace(self)
         object.__setattr__(result, "_store_context", ctx)
+        object.__setattr__(
+            result, "_select_context", StorageDriverSelectContext(target=ctx.target)
+        )
         return result
 
     async def _store_payload(self, payload: Payload) -> Payload:
         start_time = time.monotonic()
 
-        driver = self._select_driver(self._store_context, payload)
+        driver = self._select_driver(self._select_context, payload)
         if driver is None:
             return payload
 
@@ -401,7 +430,7 @@ async def _store_payload_sequence(
 
         to_store: list[tuple[int, Payload, StorageDriver]] = []
         for index, payload in enumerate(payloads):
-            driver = self._select_driver(self._store_context, payload)
+            driver = self._select_driver(self._select_context, payload)
             if driver is None:
                 continue
             to_store.append((index, payload, driver))

tests/test_client_exports.py

@@ -114,6 +114,7 @@
     "StartWorkflowUpdateInput",
     "StartWorkflowUpdateWithStartInput",
     "StorageDriverActivityInfo",
+    "StorageDriverSelectContext",
     "StorageDriverStoreContext",
     "StorageDriverWorkflowInfo",
     "TLSConfig",

tests/test_extstore.py

@@ -15,7 +15,9 @@
     StorageDriver,
     StorageDriverClaim,
     StorageDriverRetrieveContext,
+    StorageDriverSelectContext,
     StorageDriverStoreContext,
+    StorageDriverWorkflowInfo,
 )
 from temporalio.converter._extstore import _REFERENCE_ENCODING, _StorageReference
 from temporalio.converter._payload_converter import JSONProtoPayloadConverter
@@ -49,6 +51,7 @@ def __init__(
         self._storage: dict[str, bytes] = {}
         self._store_calls = 0
         self._retrieve_calls = 0
+        self._store_contexts: list[StorageDriverStoreContext] = []
 
     def name(self) -> str:
         return self._driver_name
@@ -59,6 +62,7 @@ async def store(
         payloads: Sequence[Payload],
     ) -> list[StorageDriverClaim]:
         self._store_calls += 1
+        self._store_contexts.append(context)
         start_index = len(self._storage)
 
         entries = [
@@ -537,6 +541,35 @@ async def test_no_selector_second_driver_is_retrieve_only(self):
         assert driver_a._retrieve_calls == 0  # never consulted
         assert driver_b._retrieve_calls == 1
 
+    async def test_selector_receives_select_context_with_target(self):
+        """The selector is handed a StorageDriverSelectContext -- not the
+        StorageDriverStoreContext the driver receives -- carrying the same
+        target."""
+        driver = InMemoryTestDriver("test-driver")
+        seen: list[object] = []
+
+        def selector(context: object, _payload: Payload) -> StorageDriver:
+            seen.append(context)
+            return driver
+
+        target = StorageDriverWorkflowInfo(
+            namespace="ns", id="wf-id", type="MyWorkflow", run_id="run-id"
+        )
+        storage = ExternalStorage(
+            drivers=[driver],
+            driver_selector=selector,
+            payload_size_threshold=50,
+        )._with_store_context(StorageDriverStoreContext(target=target))
+
+        converter = DataConverter(external_storage=storage)
+        await converter.encode(["x" * 200])
+
+        assert len(seen) == 1
+        assert isinstance(seen[0], StorageDriverSelectContext)
+        assert seen[0].target == target
+        assert isinstance(driver._store_contexts[0], StorageDriverStoreContext)
+        assert driver._store_contexts[0].target == target
+
     async def test_selector_routes_payloads_to_different_drivers_in_single_batch(self):
         """When a selector routes different payloads to different drivers, a
         single encode([v1, v2, ...]) call batches payloads per driver so each