diff --git a/CHANGELOG.md b/CHANGELOG.md index be3444177..48563102b 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/temporalio/client/__init__.py b/temporalio/client/__init__.py index cdb34b860..86030b1a8 100644 --- a/temporalio/client/__init__.py +++ b/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", diff --git a/temporalio/converter/__init__.py b/temporalio/converter/__init__.py index 68857049b..99e55a775 100644 --- a/temporalio/converter/__init__.py +++ b/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", diff --git a/temporalio/converter/_extstore.py b/temporalio/converter/_extstore.py index a946b2c0f..ad6fa2e4b 100644 --- a/temporalio/converter/_extstore.py +++ b/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)) diff --git a/tests/test_client_exports.py b/tests/test_client_exports.py index 5317e53c4..08e0bd3da 100644 --- a/tests/test_client_exports.py +++ b/tests/test_client_exports.py @@ -114,6 +114,7 @@ "StartWorkflowUpdateInput", "StartWorkflowUpdateWithStartInput", "StorageDriverActivityInfo", + "StorageDriverSelectContext", "StorageDriverStoreContext", "StorageDriverWorkflowInfo", "TLSConfig", diff --git a/tests/test_extstore.py b/tests/test_extstore.py index 9a058c582..4a52c65c6 100644 --- a/tests/test_extstore.py +++ b/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