|
| 1 | +"""Stream calls outside Workflow code, which Activity and client handles share.""" |
| 2 | + |
| 3 | +from __future__ import annotations |
| 4 | + |
| 5 | +from collections.abc import AsyncGenerator, Awaitable, Callable, Sequence |
| 6 | +from dataclasses import dataclass |
| 7 | +from datetime import timedelta |
| 8 | +from typing import Any, TypeVar |
| 9 | + |
| 10 | +import temporalio.converter |
| 11 | +from temporalio.common._streams import _convert, _reader |
| 12 | +from temporalio.common._streams._backend import BackendCalls, CallOptions |
| 13 | +from temporalio.common._streams._defn import TopicDefinition |
| 14 | +from temporalio.common._streams._handle import StreamHandle, TopicBase, workflow_owner |
| 15 | +from temporalio.common._streams._items import End, close_status |
| 16 | +from temporalio.common._streams._reader import Item, Position, Subscription |
| 17 | +from temporalio.common._streams._writer import StreamWriter |
| 18 | +from temporalio.common._streams.system import models |
| 19 | + |
| 20 | +T = TypeVar("T") |
| 21 | + |
| 22 | + |
| 23 | +@dataclass |
| 24 | +class Access: |
| 25 | + """The calls of one Stream handle, bound to a backend.""" |
| 26 | + |
| 27 | + backend: BackendCalls |
| 28 | + converter: temporalio.converter.DataConverter |
| 29 | + namespace: str |
| 30 | + stream_id: Callable[[], Awaitable[str]] |
| 31 | + writer: Callable[[], StreamWriter] |
| 32 | + # A close sends these, and the Store checks them against the Stream's owner. |
| 33 | + workflow_id: str | None = None |
| 34 | + first_execution_run_id: str | None = None |
| 35 | + |
| 36 | + async def publish(self, topic: str, values: Sequence[Any]) -> int: |
| 37 | + """Write the values as Records now, and return the offset of the last one.""" |
| 38 | + if not values: |
| 39 | + raise TypeError("publish takes one or more values") |
| 40 | + return await self.writer().write([(topic, value) for value in values]) |
| 41 | + |
| 42 | + async def publish_records(self, records: Sequence[tuple[str, Any]]) -> int: |
| 43 | + """Write ``(topic, value)`` pairs as Records now, in order, and return the offset of |
| 44 | + the last one. |
| 45 | + """ |
| 46 | + return await self.writer().write(records) |
| 47 | + |
| 48 | + async def items( |
| 49 | + self, read: ReadSpec, position: Position |
| 50 | + ) -> AsyncGenerator[Item, None]: |
| 51 | + """Read the Stream's items in order, and move ``position`` past each page.""" |
| 52 | + items = _reader.read( |
| 53 | + self.backend.service(self.namespace, _convert.raw(self.converter)), |
| 54 | + self.converter, |
| 55 | + await self.stream_id(), |
| 56 | + position=position, |
| 57 | + topics=read.topics, |
| 58 | + topic_prefixes=read.topic_prefixes, |
| 59 | + types=read.types, |
| 60 | + default_type=read.default_type, |
| 61 | + from_offset=read.offset, |
| 62 | + last=read.last, |
| 63 | + wait_new_records=read.wait_new_records, |
| 64 | + create_wait=read.create_wait, |
| 65 | + ) |
| 66 | + try: |
| 67 | + async for item in items: |
| 68 | + yield item |
| 69 | + finally: |
| 70 | + await items.aclose() |
| 71 | + |
| 72 | + async def get_offset(self) -> int: |
| 73 | + """The next offset of the Stream, the newest offset plus one, or 0 if it has no Record. |
| 74 | +
|
| 75 | + Offsets are sparse and above zero, so ``offset=0`` reads from the start. |
| 76 | + """ |
| 77 | + service = self.backend.service(self.namespace) |
| 78 | + response = await ( |
| 79 | + await service.describe_stream( |
| 80 | + stream=models.StreamRef(id=await self.stream_id()) |
| 81 | + ) |
| 82 | + ) |
| 83 | + return response.last_offset + 1 if response.last_offset else 0 |
| 84 | + |
| 85 | + async def close(self, result: Any | None, failure: BaseException | None) -> None: |
| 86 | + """Write this Producer's Records, then close the Stream.""" |
| 87 | + writer = self.writer() |
| 88 | + await writer.close() |
| 89 | + producer = models._ProducerTransferTypeConverter().from_transfer_type( |
| 90 | + writer.producer, models.Producer |
| 91 | + ) |
| 92 | + status = close_status(failure) |
| 93 | + service = self.backend.service(self.namespace, self.converter) |
| 94 | + await ( |
| 95 | + await service.close_stream( |
| 96 | + stream=models.StreamRef( |
| 97 | + id=await self.stream_id(), |
| 98 | + owner=workflow_owner(self.workflow_id, self.first_execution_run_id), |
| 99 | + ), |
| 100 | + producer=producer, |
| 101 | + result=result, |
| 102 | + failure=failure, |
| 103 | + status=models.StreamCloseStatus(status) if status is not None else None, |
| 104 | + ) |
| 105 | + ) |
| 106 | + |
| 107 | + |
| 108 | +@dataclass |
| 109 | +class ReadSpec: |
| 110 | + """What one read asks for.""" |
| 111 | + |
| 112 | + topics: list[str] | None |
| 113 | + topic_prefixes: Sequence[str] |
| 114 | + types: dict[str, Any] |
| 115 | + default_type: Any |
| 116 | + offset: int | End | None |
| 117 | + last: int | None |
| 118 | + wait_new_records: bool |
| 119 | + create_wait: timedelta | None |
| 120 | + |
| 121 | + |
| 122 | +def read_spec( |
| 123 | + handle: StreamHandle[Any], |
| 124 | + topics: Sequence[ |
| 125 | + str | TopicDefinition[Any] | _AsyncTopic[Any] | SyncTopicBase[Any] |
| 126 | + ], |
| 127 | + *, |
| 128 | + topic_prefixes: Sequence[str], |
| 129 | + default_type: Any, |
| 130 | + offset: int | End | None, |
| 131 | + last: int | None, |
| 132 | + wait_new_records: bool, |
| 133 | + create_wait: timedelta | None, |
| 134 | +) -> ReadSpec: |
| 135 | + """The Topic names of a read, the decode type of each known Topic, and its options. |
| 136 | +
|
| 137 | + Raises: |
| 138 | + ValueError: ``last`` is not above zero, or the call also gives ``offset``. |
| 139 | + """ |
| 140 | + if last is not None and (last < 1 or offset is not None): |
| 141 | + raise ValueError("last takes a count above zero and no offset") |
| 142 | + types = handle._known_types() |
| 143 | + names = [] |
| 144 | + for topic in topics: |
| 145 | + if isinstance(topic, (_AsyncTopic, SyncTopicBase)): |
| 146 | + name, topic_type = topic.name, topic.type |
| 147 | + elif isinstance(topic, TopicDefinition): |
| 148 | + name, topic_type = handle._topic_spec(topic, None) |
| 149 | + else: |
| 150 | + # A Topic name that the definition doesn't have gets no type here, so its values |
| 151 | + # decode with the read's default type. |
| 152 | + name, topic_type = topic, types.get(topic) |
| 153 | + names.append(name) |
| 154 | + if topic_type is not None: |
| 155 | + types[name] = topic_type |
| 156 | + return ReadSpec( |
| 157 | + topics=names or None, |
| 158 | + topic_prefixes=topic_prefixes, |
| 159 | + types=types, |
| 160 | + default_type=default_type, |
| 161 | + offset=offset, |
| 162 | + last=last, |
| 163 | + wait_new_records=wait_new_records, |
| 164 | + create_wait=create_wait, |
| 165 | + ) |
| 166 | + |
| 167 | + |
| 168 | +def subscription( |
| 169 | + items: Callable[[Position], AsyncGenerator[Item, None]], |
| 170 | + read: ReadSpec, |
| 171 | + options: CallOptions | None = None, |
| 172 | +) -> Subscription[Any]: |
| 173 | + """A read over the items that ``items`` gives.""" |
| 174 | + # A read from END or of the last Records has no resume offset until its first page. |
| 175 | + if read.last is not None or read.offset is End.END: |
| 176 | + position = Position(None) |
| 177 | + else: |
| 178 | + position = Position(read.offset or 0) |
| 179 | + return Subscription(items(position), position, options) |
| 180 | + |
| 181 | + |
| 182 | +class _AsyncTopic(TopicBase[T]): |
| 183 | + # A Topic of an Activity or client handle. Each publish writes before it returns. |
| 184 | + |
| 185 | + async def publish(self, *values: T) -> int: |
| 186 | + """Write one Record for each value, in order, and return the last Record's offset. |
| 187 | +
|
| 188 | + A list is one value. |
| 189 | + """ |
| 190 | + return await self._stream._access().publish(self._name, values) |
| 191 | + |
| 192 | + |
| 193 | +class SyncTopicBase(TopicBase[T]): |
| 194 | + """Base class of a Topic whose calls block, for sync Activities.""" |
0 commit comments