diff --git a/tests/contrib/streams/examples/README.md b/tests/contrib/streams/examples/README.md index 9564002ee..1b001933d 100644 --- a/tests/contrib/streams/examples/README.md +++ b/tests/contrib/streams/examples/README.md @@ -13,3 +13,13 @@ Each file is a complete program for one section of the proposal. `test_examples. | Worker setup | [`worker_setup.py`](worker_setup.py) | The proposal's Nexus examples are a later stage, so they have no file here. + +## Jetty v2 + +These add Workflow reads, the native backend, and two backends in one Worker to the same API. `test_examples.py` runs each one on three lanes. The memory and Redis lanes use a memory Store with the backend name `native` in place of the server. The native lane uses `NativeStreamPlugin` against a server with the Stream service, and it runs only when `TEMPORAL_STREAMS_NATIVE_TARGET` is set. + +| Section | File | +|---|---| +| One Worker with Redis and native, one Stream on each | [`two_backends.py`](two_backends.py) | +| A Workflow reads another Workflow's Stream | [`workflow_reader.py`](workflow_reader.py) | +| Old Streams stay on Redis, new ones go to native | [`migrate_redis_to_native.py`](migrate_redis_to_native.py) | diff --git a/tests/contrib/streams/examples/migrate_redis_to_native.py b/tests/contrib/streams/examples/migrate_redis_to_native.py new file mode 100644 index 000000000..e0d9f0726 --- /dev/null +++ b/tests/contrib/streams/examples/migrate_redis_to_native.py @@ -0,0 +1,103 @@ +"""Jetty v2, "Move from Redis to native Streams". + +Old Streams stay on Redis, and new ones go to the server. Step 1 runs the old code on a Worker +with only ``RedisStreamPlugin``. Step 2 deploys new code on a Worker with both plugins. The new +code picks ``native`` behind ``workflow.patched``, so a run that started on the old code replays +the same and goes on with its Redis Stream. Each handle carries its backend, so the client reads +each Stream where it is. Redis stays the default, since a change of the default changes the +commands of a replayed create. +""" + +from __future__ import annotations + +import uuid +from datetime import timedelta +from typing import Any + +from temporalio import workflow +from temporalio.client import Client, WorkflowHandle +from temporalio.contrib import streams +from temporalio.contrib.streams import workflow as stream_workflow +from temporalio.contrib.streams.native import NativeStreamPlugin +from temporalio.contrib.streams.redis import RedisStreamPlugin +from temporalio.worker import Worker + + +@streams.defn(name="events") +class Events: + tokens = streams.topic(str) + + +class _Chat: + def __init__(self) -> None: + self._more: list[str] = [] + + @workflow.signal + def more(self, text: str) -> None: + self._more.append(text) + + async def chat(self, stream: stream_workflow.WorkflowStream[Events]) -> None: + await stream.topic(Events.tokens).publish("hello") + await workflow.wait_condition(lambda: bool(self._more)) + await stream.topic(Events.tokens).publish(*self._more) + + +@workflow.defn(name="Chat") +class ChatBeforeTheMove(_Chat): + @workflow.run + async def run(self) -> None: + await self.chat(await stream_workflow.create_stream(Events)) + + +@workflow.defn(name="Chat") +class Chat(_Chat): + @workflow.run + async def run(self) -> None: + # A run from before the move has no patch marker, so it keeps its Redis Stream. + backend = "native" if workflow.patched("streams-on-native") else None + await self.chat(await stream_workflow.create_stream(Events, backend=backend)) + + +async def connect_before(target_host: str, redis_url: str) -> Client: + return await Client.connect(target_host, plugins=[RedisStreamPlugin(redis_url)]) + + +async def connect_during(target_host: str, redis_url: str) -> Client: + return await Client.connect( + target_host, + plugins=[RedisStreamPlugin(redis_url, default=True), NativeStreamPlugin()], + ) + + +async def read(client: Client, handle: WorkflowHandle[Any, None]) -> list[str]: + await handle.result() + # A UI that reads chats from both sides of the move looks on native first, then Redis. + stream = streams.StreamClient(client).workflow_stream( + handle, Events, backend=("native", "redis") + ) + tokens = [ + item.data + async for item in stream.subscribe() + if isinstance(item, streams.Record) + ] + return [f"{stream.backend}: {token}" for token in tokens] + + +async def main(before: Client, during: Client, task_queue: str) -> list[str]: + # Step 1: the old code starts a chat on Redis. + async with Worker(before, task_queue=task_queue, workflows=[ChatBeforeTheMove]): + old = await before.start_workflow( + ChatBeforeTheMove.run, id=f"chat-{uuid.uuid4()}", task_queue=task_queue + ) + stream = streams.StreamClient(before).workflow_stream(old, Events) + async for _ in stream.subscribe(create_wait=timedelta(seconds=30)): + # The old chat wrote its first Record, so it runs on the old code. + break + # Step 2: the new code goes on with the old chat, and starts a new one on the server. + async with Worker(during, task_queue=task_queue, workflows=[Chat]): + await old.signal(ChatBeforeTheMove.more, "old chat goes on") + new = await during.start_workflow( + Chat.run, id=f"chat-{uuid.uuid4()}", task_queue=task_queue + ) + await new.signal(Chat.more, "new chat") + return [*await read(during, old), *await read(during, new)] diff --git a/tests/contrib/streams/examples/test_examples.py b/tests/contrib/streams/examples/test_examples.py index 4b0a6d69c..b9cf92d02 100644 --- a/tests/contrib/streams/examples/test_examples.py +++ b/tests/contrib/streams/examples/test_examples.py @@ -2,16 +2,25 @@ import uuid +import pytest + from temporalio.client import Client from temporalio.contrib import streams +from temporalio.contrib.streams.native import NativeStreamPlugin +from temporalio.contrib.streams.redis import RedisStreamPlugin +from temporalio.contrib.streams.testing import MemoryStreamPlugin +from tests.contrib.streams.conftest import NATIVE_NAMESPACE, NATIVE_TARGET from tests.contrib.streams.examples import ( activity_streams_tokens, agent_publishes_turns, client_creates_stream, + migrate_redis_to_native, resume_from_offset, turn_scoped_stream, + two_backends, ui_backend_reads, worker_setup, + workflow_reader, ) @@ -90,3 +99,82 @@ async def test_worker_setup_on_redis(client: Client, redis_url: str): ] migrating = await worker_setup.connect_while_migrating(target, redis_url, redis_url) assert await worker_setup.create_on_second_backend(migrating) == "redis-b" + + +# The Jetty v2 examples run with a native stand-in, a memory Store with the backend name +# `native`, and on a server with the Stream service when TEMPORAL_STREAMS_NATIVE_TARGET is set. +LANES = ["memory", "redis", "native"] + + +def lane_plugins( + lane: str, request: pytest.FixtureRequest +) -> tuple[streams.StreamPlugin, streams.StreamPlugin]: + """The plugins of a lane: one named `redis` and one named `native`.""" + if lane == "memory": + return MemoryStreamPlugin(name="redis", default=True), MemoryStreamPlugin( + name="native" + ) + if lane == "native" and not NATIVE_TARGET: + pytest.skip( + "TEMPORAL_STREAMS_NATIVE_TARGET names no server with the Stream service" + ) + redis_url = request.getfixturevalue("redis_url") + redis = RedisStreamPlugin(redis_url, key_prefix=f"t-{uuid.uuid4()}", default=True) + if lane == "redis": + return redis, MemoryStreamPlugin(name="native") + return redis, NativeStreamPlugin() + + +async def lane_client( + lane: str, client: Client, plugins: list[streams.StreamPlugin] +) -> Client: + if lane != "native": + config = client.config() + config["plugins"] = plugins + return Client(**config) + if not NATIVE_TARGET: + pytest.skip( + "TEMPORAL_STREAMS_NATIVE_TARGET names no server with the Stream service" + ) + return await Client.connect( + NATIVE_TARGET, namespace=NATIVE_NAMESPACE, plugins=plugins + ) + + +@pytest.mark.parametrize("lane", LANES) +async def test_two_backends(lane: str, client: Client, request: pytest.FixtureRequest): + two = await lane_client(lane, client, list(lane_plugins(lane, request))) + assert await two_backends.main(two, task_queue()) == { + "audit": ["started", "summarized"], + "live": ["hello", "there", "summary"], + } + + +@pytest.mark.parametrize("lane", LANES) +async def test_workflow_reader( + lane: str, client: Client, request: pytest.FixtureRequest +): + # Each lane reads from one backend: the memory Store, Redis or the server. + plugin = {"memory": 0, "redis": 0, "native": 1}[lane] + one = await lane_client(lane, client, [lane_plugins(lane, request)[plugin]]) + assert await workflow_reader.main(one, task_queue()) == [ + "find", + "the", + "answer", + "agent said FIND THE ANSWER", + ] + + +@pytest.mark.parametrize("lane", LANES) +async def test_migrate_redis_to_native( + lane: str, client: Client, request: pytest.FixtureRequest +): + redis, native = lane_plugins(lane, request) + before = await lane_client(lane, client, [redis]) + during = await lane_client(lane, client, [redis, native]) + assert await migrate_redis_to_native.main(before, during, task_queue()) == [ + "redis: hello", + "redis: old chat goes on", + "native: hello", + "native: new chat", + ] diff --git a/tests/contrib/streams/examples/two_backends.py b/tests/contrib/streams/examples/two_backends.py new file mode 100644 index 000000000..c60e55b84 --- /dev/null +++ b/tests/contrib/streams/examples/two_backends.py @@ -0,0 +1,94 @@ +"""Jetty v2, "Two backends in one Worker". + +The client has ``RedisStreamPlugin`` and ``NativeStreamPlugin``, and the Worker takes both from +it. Each Stream handle carries its backend name, so every call goes to the backend that keeps +the Stream. A create with no ``backend=`` goes to the default one, here Redis. Tests swap a +plugin for ``MemoryStreamPlugin(name="native")`` with no change to Workflow code. +""" + +from __future__ import annotations + +import uuid +from dataclasses import dataclass +from datetime import timedelta + +from temporalio import activity, workflow +from temporalio.client import Client +from temporalio.contrib import streams +from temporalio.contrib.streams import workflow as stream_workflow +from temporalio.contrib.streams.native import NativeStreamPlugin +from temporalio.contrib.streams.redis import RedisStreamPlugin +from temporalio.worker import Worker + + +@streams.defn(name="events") +class Events: + tokens = streams.topic(str) + + +# Two Streams of one definition, under their own names. +Audit = streams.with_name(Events, "audit") +Live = streams.with_name(Events, "live") + + +@dataclass +class SummarizeInput: + audit: streams.ActivityStream[Events] + live: streams.ActivityStream[Events] + + +@activity.defn +async def summarize(input: SummarizeInput) -> str: + # One Activity writes to both Streams. Each write goes to its own backend. + await input.live.topic(Events.tokens).publish("summary") + await input.audit.topic(Events.tokens).publish("summarized") + return "done" + + +@workflow.defn +class TwoStreams: + @workflow.run + async def run(self) -> str: + # The audit log stays on the default backend, Redis. The live view is on the server. + audit = await stream_workflow.create_stream(Audit) + live = await stream_workflow.create_stream(Live, backend="native") + await audit.topic(Events.tokens).publish("started") + await live.topic(Events.tokens).publish("hello", "there") + return await workflow.execute_activity( + summarize, + SummarizeInput( + streams.ActivityStream.from_handle(audit), + streams.ActivityStream.from_handle(live), + ), + start_to_close_timeout=timedelta(seconds=30), + ) + + +async def connect(target_host: str, redis_url: str) -> Client: + return await Client.connect( + target_host, + plugins=[RedisStreamPlugin(redis_url, default=True), NativeStreamPlugin()], + ) + + +async def main(client: Client, task_queue: str) -> dict[str, list[str]]: + async with Worker( + client, task_queue=task_queue, workflows=[TwoStreams], activities=[summarize] + ): + handle = await client.start_workflow( + TwoStreams.run, id=f"two-{uuid.uuid4()}", task_queue=task_queue + ) + await handle.result() + stream_client = streams.StreamClient(client) + read: dict[str, list[str]] = {} + for name, definition, backend in ( + ("audit", Audit, "redis"), + ("live", Live, "native"), + ): + stream = stream_client.workflow_stream(handle, definition, backend=backend) + read[name] = [ + item.data + async for item in stream.subscribe() + if isinstance(item, streams.Record) + ] + return read diff --git a/tests/contrib/streams/examples/workflow_reader.py b/tests/contrib/streams/examples/workflow_reader.py new file mode 100644 index 000000000..f5a70b2b1 --- /dev/null +++ b/tests/contrib/streams/examples/workflow_reader.py @@ -0,0 +1,68 @@ +"""Jetty v2, "A Workflow reads another Workflow's Stream". + +An orchestrator Workflow follows the turns of an agent Workflow. It gets the agent's Stream as +an argument and reads it with ``subscribe``, an async iterator in Workflow code. Each page is in +History, so a replay reads the same items. On a native Stream a read at the end waits on the +server. On Redis it waits on a Local Activity, then on a durable timer. The Workflow code is +the same on both. +""" + +from __future__ import annotations + +import uuid +from datetime import timedelta + +from temporalio import workflow +from temporalio.client import Client +from temporalio.contrib import streams +from temporalio.contrib.streams import workflow as stream_workflow +from temporalio.worker import Worker + + +@streams.defn(name="events") +class Events: + tokens = streams.topic(str) + + +@workflow.defn +class Agent: + @workflow.run + async def run(self, prompt: str) -> str: + stream = await stream_workflow.create_stream(Events) + for word in prompt.split(): + # Each word goes out in its own Workflow Task, so the follower waits for some. + await workflow.sleep(timedelta(milliseconds=200)) + await stream.topic(Events.tokens).publish(word) + return prompt.upper() + + +@workflow.defn +class Follower: + @workflow.run + async def run(self, agent: stream_workflow.WorkflowStream[Events]) -> list[str]: + seen = [] + async for item in agent.subscribe(create_wait=timedelta(seconds=30)): + if isinstance(item, streams.Record): + seen.append(item.data) + elif isinstance(item, streams.Closed): + # The agent's Stream closes with the agent's result. + seen.append(f"agent said {item.result}") + return seen + + +async def main(client: Client, task_queue: str) -> list[str]: + async with Worker(client, task_queue=task_queue, workflows=[Agent, Follower]): + agent = await client.start_workflow( + Agent.run, + "find the answer", + id=f"agent-{uuid.uuid4()}", + task_queue=task_queue, + ) + # The handle has the first run id, so the client builds the Stream with no call. + stream = streams.StreamClient(client).workflow_stream(agent, Events) + return await client.execute_workflow( + Follower.run, + stream_workflow.WorkflowStream.from_handle(stream), + id=f"follower-{uuid.uuid4()}", + task_queue=task_queue, + )